Added make_custom_op_from_attrs_str function.

Author: JoeyChou-SiMa-ai

python/tvm/relay/op/nn/custom_operation.py

@@ -20,17 +20,33 @@
 from tvm.ir import Attrs
 
 from . import _make
-from typing import List, Tuple
-from tvm.custom_operation_config import (
-    CustomOpConfigInfo, CustomOperationConfig
-)
+from typing import List
+from tvm.custom_operation_config import CustomOperationConfig
 import tvm._ffi
 import json
 
 
 MAX_TENSOR_INPUTS = 5
 
 
+def make_custom_op_from_attrs_str(tensor_inputs: List[expr.ExprWithOp],
+                                  custom_op_attr_str: str):
+
+    if len(tensor_inputs) == 1:
+        return _make.custom_op_1(*tensor_inputs, custom_op_attr_str)
+    elif len(tensor_inputs) == 2:
+        return _make.custom_op_2(*tensor_inputs, custom_op_attr_str)
+    elif len(tensor_inputs) == 3:
+        return _make.custom_op_3(*tensor_inputs, custom_op_attr_str)
+    elif len(tensor_inputs) == 4:
+        return _make.custom_op_4(*tensor_inputs, custom_op_attr_str)
+    elif len(tensor_inputs) == 5:
+        return _make.custom_op_5(*tensor_inputs, custom_op_attr_str)
+    else:
+        msg = "Unsupported number of input tensor arguments (%d)." % (len(tensor_inputs))
+        raise AssertionError(msg)
+
+
 def custom_op(inputs, input_types, name, code, func_name, datatype, compiler_flags):
     """
     Create a Relay IR node for the custom operation.  Specifically, a
@@ -82,20 +98,7 @@ def custom_op(inputs, input_types, name, code, func_name, datatype, compiler_fla
     }
 
     custom_op_attr_str = json.dumps(custom_op_attrs)
-
-    if len(tensor_inputs) == 1:
-        return _make.custom_op_1(*tensor_inputs, custom_op_attr_str)
-    elif len(tensor_inputs) == 2:
-        return _make.custom_op_2(*tensor_inputs, custom_op_attr_str)
-    elif len(tensor_inputs) == 3:
-        return _make.custom_op_3(*tensor_inputs, custom_op_attr_str)
-    elif len(tensor_inputs) == 4:
-        return _make.custom_op_4(*tensor_inputs, custom_op_attr_str)
-    elif len(tensor_inputs) == 5:
-        return _make.custom_op_5(*tensor_inputs, custom_op_attr_str)
-    else:
-        msg = "Unsupported number of input tensor arguments (%d)." % (len(tensor_inputs))
-        raise AssertionError(msg)
+    return make_custom_op_from_attrs_str(tensor_inputs, custom_op_attr_str)
 
 
 def is_valid_attribute(input):
@@ -110,7 +113,7 @@ def is_valid_attribute(input):
 
     if input_type == list and type(input[0]) in [int, float]:
         for elem in input:
-            if type(elem) != type(input[0]):
+            if not isinstance(elem, type(input[0])):
                 return False
         return True