File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 1+ #ifndef INFINI_OPS_BASE_SOFTMAX_H_
2+ #define INFINI_OPS_BASE_SOFTMAX_H_
3+
4+ #include " operator.h"
5+
6+ namespace infini ::ops {
7+
8+ class Softmax : public Operator <Softmax> {
9+ public:
10+ Softmax (const Tensor input, const int64_t dim, Tensor out)
11+ : input_shape_{input.shape ()},
12+ input_strides_{input.strides ()},
13+ input_type_{input.dtype ()},
14+ out_shape_{out.shape ()},
15+ out_strides_{out.strides ()},
16+ out_type_{out.dtype ()},
17+ dim_{dim},
18+ device_index_{out.device ().index ()} {}
19+
20+ virtual void operator ()(const Tensor input, const int64_t dim,
21+ Tensor out) const = 0;
22+
23+ protected:
24+ Tensor::Shape input_shape_;
25+
26+ Tensor::Strides input_strides_;
27+
28+ DataType input_type_;
29+
30+ Tensor::Shape out_shape_;
31+
32+ Tensor::Strides out_strides_;
33+
34+ DataType out_type_;
35+
36+ int64_t dim_{};
37+
38+ int device_index_{0 };
39+ };
40+
41+ } // namespace infini::ops
42+
43+ #endif
You can’t perform that action at this time.
0 commit comments