[GitHub] [tvm] zackcquic commented on a change in pull request #8079: Complete register op from python

[GitBox <gi...@apache.org>]

zackcquic commented on a change in pull request #8079:
URL: https://github.com/apache/tvm/pull/8079#discussion_r635750074



##########
File path: python/tvm/ir/op.py
##########
@@ -85,17 +85,76 @@ def reset_attr(self, attr_name):
         """
         _ffi_api.OpResetAttr(self, attr_name)
 
+    def add_type_rel(self, rel_name, type_rel_func=None):

Review comment:
       Could you provide example and document how to rewrite ```type_rel_func```? Since it seems different with c++ API.
   
   Also type_rel_func is used in ```InferType``` pass, could you provide a unittest test that applies ```InferType``` on module with custom ops?
   
   

##########
File path: tests/python/relay/test_ir_op.py
##########
@@ -103,11 +104,23 @@ def test_op_register():
     """Tests register_op functionality."""
     op_name = "custom_op"
 
-    tvm.ir.register_op(op_name)
-    tvm.ir.register_op_attr(op_name, "num_inputs", 2, 256)
+    def custom_op_rel(x):
+        return x + 1
 
-    assert tvm.ir.Op.get(op_name).name == op_name
-    assert tvm.ir.Op.get(op_name).num_inputs == 2
+    _op.register(op_name, r"code(Add two tensor with inner broadcasting.)code")
+    _op.get(op_name).set_num_inputs(2)
+    _op.get(op_name).add_argument("data_0", "Tensor", "The input data tensor.")
+    _op.get(op_name).add_argument("data_1", "Tensor", "The input data tensor.")
+    _op.get(op_name).add_type_rel("CustomOp", custom_op_rel)
+    _op.get(op_name).set_support_level(1)
+    _op.register_pattern(op_name, _op.OpPattern.ELEMWISE)
+    _op.register_stateful(op_name, False)
+
+    assert _op.get(op_name).name == op_name
+    assert _op.get(op_name).num_inputs == 2
+    assert _op.get(op_name).get_attr("CustomOp")(1) == 2

Review comment:
       ```add_type_rel``` doesn't add attribute to op but register type relations.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org