clean up

apache · Nov 21, 2018 · 1747a21 · 1747a21
1 parent 844e006
commit 1747a21
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py
@@ -231,6 +231,7 @@ def zeros_compute(attrs, inputs, output_type, target):
     return [topi.full(output_type.shape, output_type.dtype, 0.0)]
 
 register_schedule("zeros", schedule_broadcast)
+register_pattern("zeros", OpPattern.ELEMWISE)
 
 # zeros_like
 @register_compute("zeros_like")
@@ -247,6 +248,7 @@ def ones_compute(attrs, inputs, output_type, target):
     return [topi.full(output_type.shape, output_type.dtype, 1.0)]
 
 register_schedule("ones", schedule_broadcast)
+register_pattern("ones", OpPattern.ELEMWISE)
 
 # ones_like
 @register_compute("ones_like")
@@ -263,11 +265,3 @@ def clip_compute(attrs, inputs, output_type, target):
     return [topi.clip(inputs[0], attrs.a_min, attrs.a_max)]
 
 register_schedule("clip", schedule_elemwise)
-
-# concatenate
-@register_compute("concatenate")
-def concatenate_compute(attrs, inputs, output_type, target):
-    return [topi.concatenate(inputs, axis=attrs.axis)]
-
-register_schedule("concatenate", schedule_injective)
-register_pattern("concatenate", OpPattern.INJECTIVE)
diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py
@@ -1,6 +1,7 @@
 #pylint: disable=invalid-name, unused-argument
 """Backend compiler related feature registration"""
 from __future__ import absolute_import
+import topi
 from . import op as _reg
 from .op import schedule_injective, OpPattern
 
@@ -14,3 +15,11 @@
 # layout_transform
 _reg.register_schedule("layout_transform", schedule_injective)
 _reg.register_pattern("layout_transform", OpPattern.INJECTIVE)
+
+# concatenate
+@_reg.register_compute("concatenate")
+def concatenate_compute(attrs, inputs, output_type, target):
+    return [topi.concatenate(inputs, axis=attrs.axis)]
+
+_reg.register_schedule("concatenate", schedule_injective)
+_reg.register_pattern("concatenate", OpPattern.INJECTIVE)
diff --git a/src/relay/op/tensor/unary.cc b/src/relay/op/tensor/unary.cc
@@ -31,7 +31,7 @@ RELAY_REGISTER_UNARY_OP("exp")
 .set_support_level(1);
 
 RELAY_REGISTER_UNARY_OP("sqrt")
-.describe(R"code(Returns the sqrt input array, computed element-wise.
+.describe(R"code(Returns the rsqrt input array, computed element-wise.
 
 .. math::
    sqrt(x)
@@ -42,12 +42,12 @@ RELAY_REGISTER_UNARY_OP("sqrt")
 RELAY_REGISTER_UNARY_OP("zeros_like")
 .describe(R"code(Returns an array of zeros, with same type and shape as the input.
 )code" TVM_ADD_FILELINE)
-.set_support_level(1);
+.set_support_level(4);
 
 RELAY_REGISTER_UNARY_OP("ones_like")
 .describe(R"code(Returns an array of ones, with same type and shape as the input.
 )code" TVM_ADD_FILELINE)
-.set_support_level(1);
+.set_support_level(4);
 
 RELAY_REGISTER_UNARY_OP("sigmoid")
 .describe(R"code(Returns the sigmoid input array, computed element-wise.