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

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

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



##########
File path: python/tvm/ir/op.py
##########
@@ -85,17 +85,87 @@ def reset_attr(self, attr_name):
         """
         _ffi_api.OpResetAttr(self, attr_name)
 
+    def add_type_rel(self, rel_name, type_rel_func=None):
+        """Attach the type function corresponding to the return type.
 
-def register_op(op_name):
-    """Register an operator by name
+        Parameters
+        ----------
+        rel_name : str
+            The type relation name to register.
+
+        type_rel_func: function (args: List[Type], attrs: Attrs) -> Type
+            The backing relation function which can solve an arbitrary relation on variables.
+            Differences with type_rel_func in C++:
+            1, when type_rel_func is not None:
+               1) OpAddTypeRel on C++ side will adjust type_rel_func with TypeReporter to
+                  calling convention of relay type system.
+               2) type_rel_func returns output argument's type, return None means can't
+                  infer output's type.
+               3) only support single output operators for now, the last argument is output tensor.

Review comment:
       We can support when we need, add num_of_output in type_rel_func, the last num_of_output is output,  so we can reporter->Assign outputs.




-- 
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