From 0dbc8a98caed9fc5e35b7dd5946eab7d69751dc5 Mon Sep 17 00:00:00 2001
From: Lianmin Zheng <mercy_zheng@sjtu.edu.cn>
Date: Sat, 6 Oct 2018 18:03:14 -0700
Subject: [PATCH] [RELAY] Add softmax (#1841)

---
 docs/langref/relay_op.rst            |  1 +
 include/tvm/relay/attrs/nn.h         | 10 +++++++
 python/tvm/relay/op/nn/nn.py         | 20 +++++++++++++
 src/relay/op/nn/convolution.cc       |  6 ++--
 src/relay/op/nn/nn.cc                | 43 ++++++++++++++++++++++++++++
 tests/python/relay/test_op_level1.py | 14 +++++++++
 6 files changed, 91 insertions(+), 3 deletions(-)
 create mode 100644 src/relay/op/nn/nn.cc

diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst
index 8566404561b2..c7db9364a72e 100644
--- a/docs/langref/relay_op.rst
+++ b/docs/langref/relay_op.rst
@@ -28,6 +28,7 @@ This level enables fully connected multi-layer perceptron.
    tvm.relay.sigmoid
    tvm.relay.add
    tvm.relay.expand_dims
+   tvm.relay.nn.softmax
 
 **Level 2: Convolutions**
 
diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h
index b364079f06fc..0de0164a562f 100644
--- a/include/tvm/relay/attrs/nn.h
+++ b/include/tvm/relay/attrs/nn.h
@@ -67,6 +67,16 @@ struct ConvAttrs : public tvm::AttrsNode<ConvAttrs> {
   }
 };
 
+/*! \brief Attributes used in softmax operators */
+struct SoftmaxAttrs : public tvm::AttrsNode<SoftmaxAttrs> {
+  int axis;
+
+  TVM_DECLARE_ATTRS(SoftmaxAttrs, "relay.attrs.SoftmaxAttrs") {
+      TVM_ATTR_FIELD(axis).set_default(1)
+          .describe("The axis to sum over when computing softmax.");
+  }
+};
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_NN_H_
diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py
index f2d60d48eaad..3b168c6fce21 100644
--- a/python/tvm/relay/op/nn/nn.py
+++ b/python/tvm/relay/op/nn/nn.py
@@ -86,3 +86,23 @@ def conv2d(data,
     return _make.conv2d(data, weight, strides, padding, dilation,
                         groups, channels, kernel_size, data_layout,
                         weight_layout, out_layout, out_dtype)
+
+
+def softmax(data, axis):
+    r"""Computes softmax.
+
+    .. math:: \text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}
+
+    .. note::
+        This operator can be optimized away for inference.
+
+    Parameters
+    ----------
+    data: relay.Expr
+        The input data to the operator.
+
+    axis: int
+        The axis to sum over when computing softmax
+    """
+
+    return _make.softmax(data, axis)
diff --git a/src/relay/op/nn/convolution.cc b/src/relay/op/nn/convolution.cc
index 920fc68d51e8..ba424128640c 100644
--- a/src/relay/op/nn/convolution.cc
+++ b/src/relay/op/nn/convolution.cc
@@ -49,9 +49,9 @@ bool Conv2DRel(const Array<Type>& types,
     CHECK_EQ(param->dilation.size(), 2);
     std::vector<IndexExpr> wshape(
         {param->channels / param->groups,
-              data->shape[1] / param->groups,
-              param->kernel_size[0],
-              param->kernel_size[1]});
+         data->shape[1] / param->groups,
+         param->kernel_size[0],
+         param->kernel_size[1]});
     wshape = ConvertLayout(wshape, kOIHW, kernel_layout);
     wshape[kernel_layout.indexof('O')] *= param->groups;
     channels = param->channels;
diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc
new file mode 100644
index 000000000000..b34d248d1704
--- /dev/null
+++ b/src/relay/op/nn/nn.cc
@@ -0,0 +1,43 @@
+/*!
+ *  Copyright (c) 2018 by Contributors
+ * \file nn.cc
+ * \brief Property def of nn operators.
+ */
+
+#include <tvm/relay/op.h>
+#include <tvm/relay/attrs/nn.h>
+#include "../type_relations.h"
+
+namespace tvm {
+namespace relay {
+
+
+TVM_REGISTER_API("relay.op.nn._make.softmax")
+.set_body([](const TVMArgs& args, TVMRetValue* rv) {
+  auto make_func = [](Expr data, int axis) {
+    auto attrs = make_node<SoftmaxAttrs>();
+    attrs->axis = axis;
+    static const Op& op = Op::Get("nn.softmax");
+    return CallNode::make(op, {data}, Attrs(attrs), {});
+  };
+
+  runtime::detail::unpack_call<Expr, 2>(make_func, args, rv);
+});
+
+RELAY_REGISTER_OP("nn.softmax")
+    .describe(R"code(Softmax layer.
+
+.. math:: \text{softmax}(x)_i = \frac{exp(x_i)}{\sum_j exp(x_j)}
+
+.. note::
+    This operator can be optimized away for inference.
+
+- **data**: The input data
+)code" TVM_ADD_FILELINE)
+.set_num_inputs(1)
+.add_argument("data", "Tensor", "The input tensor.")
+.set_support_level(1)
+.add_type_rel("Identity", IdentityRel);
+
+}  // namespace relay
+}  // namespace tvm
diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py
index 621d40e79b64..654e184e8e23 100644
--- a/tests/python/relay/test_op_level1.py
+++ b/tests/python/relay/test_op_level1.py
@@ -16,6 +16,19 @@ def test_expand_dims_infer_type():
         (n, t, 1, 100), "float32")
 
 
+def test_softmax():
+    ib = relay.ir_builder.IRBuilder()
+    n, d = tvm.var("n"), tvm.var("d")
+    x = ib.param("x", relay.ty.TensorType((n, d), "float32"))
+    with ib.function(x) as func:
+        ib.ret(relay.nn.softmax(x, axis=1))
+    ib.ret(func)
+
+    func = relay.ir_pass.infer_type(ib.env, func.to_func())
+    ftype = func.checked_type()
+    assert ftype.ret_type == relay.ty.TensorType((n, d), "float32")
+
+
 def test_unary_op():
     for op in [relay.exp,
                relay.log,
@@ -34,3 +47,4 @@ def test_unary_op():
 if __name__ == "__main__":
     test_expand_dims_infer_type()
     test_unary_op()
+    test_softmax()