#ifndef CAFFE2_OPERATORS_SOFTMAX_OP_H_ #define CAFFE2_OPERATORS_SOFTMAX_OP_H_ #include "caffe2/core/context.h" #include "caffe2/core/logging.h" #include "caffe2/core/operator.h" #include "caffe2/utils/math.h" namespace caffe2 { template class SoftmaxOp final : public Operator { public: template explicit SoftmaxOp(Args&&... args) : Operator(std::forward(args)...), axis_(this->template GetSingleArgument("axis", 1)) {} USE_OPERATOR_CONTEXT_FUNCTIONS; bool RunOnDevice() override; protected: int axis_; Tensor scale_; Tensor rowmax_; Tensor sum_multiplier_; }; template class SoftmaxGradientOp final : public Operator { public: template explicit SoftmaxGradientOp(Args&&... args) : Operator(std::forward(args)...), axis_(this->template GetSingleArgument("axis", 1)) {} USE_OPERATOR_CONTEXT_FUNCTIONS; bool RunOnDevice() override; protected: int axis_; Tensor scale_; Tensor sum_multiplier_; }; } // namespace caffe2 #endif // CAFFE2_OPERATORS_SOFTMAX_OP_H_