public:
explicit SoftmaxOp(int axis) : OpDescription(1, 1), _axis(axis) {}
- int getAxis() const { return _axis; }
+ int getAxis() const
+ {
+ if (_axis < 0)
+ {
+ // Negative axis is used to index starting from the last element of the shape
+ // -1 means last element, -2 means second from end, like in python
+ int res = _axis + getInputShape(0).rank();
+ assert(res >= 0);
+ return res;
+ }
+ return _axis;
+ }
private:
int _axis;
std::vector<INode::Ref> OpCreator::createSoftmax(InputOps inputs, InputParams params,
const SoftmaxOptions *opts)
{
- // TODO: here assuming that softmax is applied to a 1-d tensor
- return createOp<ops::SoftmaxOp>(inputs, ActivationFunctionType_NONE, 1);
+ // -1 represents last one dimension
+ return createOp<ops::SoftmaxOp>(inputs, ActivationFunctionType_NONE, -1);
}
std::vector<INode::Ref> OpCreator::createReshape(InputOps inputs, InputParams params,
#include "nnc/core/IR/model/operations/operation.h"
+#include "nnc/core/IR/model/operations/softmax_op.h"
#include <gtest/gtest.h>
ASSERT_EQ(inShape, op.getInputShape(0));
ASSERT_EQ(outShape, op.getOutputShape(0));
}
+
+TEST(OpDescription, SoftmaxAxisTest) {
+ Shape inShape{1,2,3};
+
+ ops::SoftmaxOp op_1(1);
+ op_1.setInputShape(0, inShape);
+ ASSERT_EQ(op_1.getAxis(), 1);
+
+ ops::SoftmaxOp op_n1(-1);
+ op_n1.setInputShape(0, inShape);
+ ASSERT_EQ(op_n1.getAxis(), 2);
+
+ ops::SoftmaxOp op_n3(-3);
+ op_n3.setInputShape(0, inShape);
+ ASSERT_EQ(op_n3.getAxis(), 0);
+}
\ No newline at end of file