From eec3648aa967fffdd670132181a62a4eb1996407 Mon Sep 17 00:00:00 2001
From: Meghan Cowan <cowanmeg@cs.washington.edu>
Date: Fri, 5 Oct 2018 21:54:47 -0700
Subject: [PATCH] [RELAY] Add sigmoid relay operator (#1836)
---
docs/langref/relay_op.rst | 2 ++
python/tvm/relay/op/tensor.py | 16 ++++++++++++++++
src/relay/op/tensor/unary.cc | 11 +++++++++++
tests/python/relay/test_op_level1.py | 16 ++++++++++++++++
4 files changed, 45 insertions(+)
diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst
index bcaece0bf..a4728eaa7 100644
--- a/docs/langref/relay_op.rst
+++ b/docs/langref/relay_op.rst
@@ -25,6 +25,7 @@ This level enables fully connected multi-layer perceptron.
tvm.relay.log
tvm.relay.sqrt
tvm.relay.exp
+ tvm.relay.sigmoid
tvm.relay.add
tvm.relay.expand_dims
@@ -61,6 +62,7 @@ Level 1 Definitions
.. autofunction:: tvm.relay.log
.. autofunction:: tvm.relay.sqrt
.. autofunction:: tvm.relay.exp
+.. autofunction:: tvm.relay.sigmoid
.. autofunction:: tvm.relay.add
diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py
index c8c42c1a6..05538c401 100644
--- a/python/tvm/relay/op/tensor.py
+++ b/python/tvm/relay/op/tensor.py
@@ -60,6 +60,22 @@ def sqrt(data):
return _make.sqrt(data)
+def sigmoid(data):
+ """Compute elementwise sigmoid of data.
+
+ Parameters
+ ----------
+ data : relay.Expr
+ The input data
+
+ Returns
+ -------
+ result : relay.Expr
+ The computed result.
+ """
+ return _make.sigmoid(data)
+
+
def add(lhs, rhs):
"""Addition with numpy-style broadcasting.
diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc
index 798d4aa79..77bb08e14 100644
--- a/src/relay/op/tensor/unary.cc
+++ b/src/relay/op/tensor/unary.cc
@@ -65,6 +65,17 @@ RELAY_REGISTER_UNARY_OP("sqrt")
.add_type_rel("Identity", IdentityRel);
+RELAY_REGISTER_UNARY_OP("sigmoid")
+.describe(R"code(Returns the sigmoid input array, computed element-wise.
+
+.. math::
+ sigmoid(x)
+
+)code" TVM_ADD_FILELINE)
+.set_support_level(1)
+.add_type_rel("Identity", IdentityRel);
+
+
// Concat
TVM_REGISTER_API("relay.op._make.concat")
.set_body_typed<Expr(Expr)>([](Expr tuple) {
diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py
index c1c8b03c1..9cfca9630 100644
--- a/tests/python/relay/test_op_level1.py
+++ b/tests/python/relay/test_op_level1.py
@@ -16,5 +16,21 @@ def test_expand_dims_infer_type():
(n, t, 1, 100), "float32")
+def test_unary_op():
+ for op in [relay.exp,
+ relay.log,
+ relay.sqrt,
+ relay.sigmoid]:
+ ib = relay.ir_builder.IRBuilder()
+ x = ib.param("x", relay.TensorType((10, 4), "int32"))
+ with ib.function(x) as func:
+ ib.ret(op(x.var))
+ ib.ret(func)
+ func = relay.ir_pass.infer_type(ib.env, func.to_func())
+ ftype = func.checked_type()
+ assert ftype.ret_type == relay.TensorType((10, 4), "int32")
+
+
if __name__ == "__main__":
test_expand_dims_infer_type()
+ test_unary_op()
--
GitLab