review feedback

Author: guschmue

tests/test_graph.py

@@ -365,7 +365,7 @@ def print_handler(ctx, node, name, args):
     def test_custom_op(self):
         """Custom op test."""
 
-        @tf_op("Print", type_map={"Print": "Identity"})
+        @tf_op("Print", onnx_op="Identity")
         class Print:
             @classmethod
             def version_1(cls, ctx, node, **kwargs):

tf2onnx/custom_opsets/__init__.py

@@ -3,4 +3,3 @@
 """ custom tf2onnx mapping functions. """
 
 from . import ms
-from .. import constants

tf2onnx/handler.py

@@ -16,10 +16,18 @@
 # pylint: disable=unused-argument,missing-docstring,invalid-name
 
 class tf_op:
+    """Class to implement the decorator to register handlers that map tf to onnx."""
+
     _OPSETS = collections.OrderedDict()
     _MAPPING = None
 
     def __init__(self, name, domain=constants.ONNX_DOMAIN, **kwargs):
+        """Called decorator from decorator.
+
+        :param name: The name of the tensorflow operator.
+        :param domain: The domain the operator belongs to, defaults to onnx.
+        :param kwargs: Dictionary that are passed to the handler. A key 'onnx_op' will change the operator name.
+        """
         if not isinstance(name, list):
             name = [name]
         self.name = name
@@ -39,8 +47,15 @@ def __call__(self, func):
                 opset_dict = opset[version]
                 for name in self.name:
                     opset_dict[name] = (v, self.kwargs)
+        return func
 
     def register_compat_handler(self, func, version):
+        """Register old style custom handler.
+
+        :param func: The handler.
+        :param version: The domain the operator belongs to, defaults to onnx.
+        :param version: The version of the handler.
+        """
         opset = tf_op._OPSETS.get(self.domain)
         if not opset:
             opset = []
@@ -55,8 +70,13 @@ def get_opsets():
         return tf_op._OPSETS
 
     @staticmethod
-    def create_mapping(max_opset, extra_opsets):
-        mapping = {constants.ONNX_DOMAIN: max_opset}
+    def create_mapping(max_onnx_opset_version, extra_opsets):
+        """Create the final mapping dictionary by stacking domains and opset versions.
+
+        :param max_onnx_opset_version: The highest onnx opset the resulting graph may use.
+        :param extra_opsets: Extra opsets the resulting graph may use.
+        """
+        mapping = {constants.ONNX_DOMAIN: max_onnx_opset_version}
         if extra_opsets:
             for extra_opset in extra_opsets:
                 mapping[extra_opset.domain] = extra_opset.version
@@ -72,7 +92,14 @@ def create_mapping(max_opset, extra_opsets):
         return ops_mapping
 
     @staticmethod
-    def find_op(name):
+    def find_effective_op(name):
+        """Find the effective version of an op create_mapping.
+           This is used if we need to compose ops from other ops where we'd need to find the
+           op that is doing to be used in the final graph, for example there is a custom op
+           that overrides a onnx op ...
+
+        :param name: The operator name.
+        """
         map_info = tf_op._MAPPING.get(name)
         if map_info is None:
             return None

tf2onnx/onnx_opset/logical.py

@@ -39,16 +39,16 @@ def logical_compare_op(ctx, node, **kwargs):
                 ctx.set_dtype(inp_cast.output[0], target_dtype)
 
 
-@tf_op(["LogicalNot", "NotEqual"],
-       type_map={"LogicalNot": "Not"})
+@tf_op(["LogicalNot", "NotEqual"], onnx_op="Not")
 class DirectOp:
     @classmethod
     def version_4(cls, ctx, node, **kwargs):
         pass
 
 
-@tf_op(["Equal", "Greater", "Less", "LogicalAnd", "LogicalOr", "LogicalAnd", "LogicalOr"],
-       type_map={"LogicalAnd": "And", "LogicalOr": "Or"})
+@tf_op(["Equal", "Greater", "Less"])
+@tf_op("LogicalAnd", onnx_op="And")
+@tf_op("LogicalOr", onnx_op="Or")
 class BroadcastOp(common.BroadcastOp):
     pass
 
@@ -64,8 +64,8 @@ def version_7(cls, ctx, node, **kwargs):
         logical_compare_op(ctx, node, **kwargs)
 
 
-@tf_op(["GreaterEqual", "LessEqual"],
-       type_map={"GreaterEqual": "Less", "LessEqual": "Greater"})
+@tf_op("GreaterEqual", onnx_op="Less")
+@tf_op("LessEqual", onnx_op="Greater")
 class GreaterLessEqual:
     @classmethod
     def version_7(cls, ctx, node, **kwargs):

tf2onnx/onnx_opset/math.py

@@ -27,8 +27,7 @@ class BroadcastOp(common.BroadcastOp):
     pass
 
 
-@tf_op(["RealDiv", "TruncateDiv"],
-       type_map={"RealDiv": "Div", "TruncateDiv": "Div"})
+@tf_op(["RealDiv", "TruncateDiv"], onnx_op="Div")
 class RealDiv(common.BroadcastOp):
     pass
 
@@ -55,8 +54,8 @@ def version_9(cls, ctx, node, **kwargs):
         pass
 
 
-@tf_op(["Minimum", "Maximum"],
-       type_map={"Minimum": "Min", "Maximum": "Max"})
+@tf_op("Minimum", onnx_op="Min")
+@tf_op("Maximum", onnx_op="Max")
 class MinMaxOp:
     @classmethod
     def version_4(cls, ctx, node, **kwargs):

tf2onnx/onnx_opset/nn.py

@@ -287,9 +287,8 @@ def version_4(cls, ctx, node, **kwargs):
         conv_convert_inputs(ctx, node, with_kernel=True, new_kernel_shape=new_kernel_shape)
 
 
-@tf_op(["AvgPool", "AvgPool3D", "MaxPool", "MaxPoolV2"],
-       type_map={"AvgPool": "AveragePool", "AvgPool3D": "AveragePool",
-                 "MaxPool": "MaxPool", "MaxPoolV2": "MaxPool"})
+@tf_op(["AvgPool", "AvgPool3D"], onnx_op="AveragePool")
+@tf_op(["MaxPool", "MaxPoolV2"], onnx_op="MaxPool")
 class PoolOp:
     @classmethod
     def version_4(cls, ctx, node, **kwargs):

tf2onnx/onnx_opset/reduction.py

@@ -22,9 +22,11 @@
 
 # pylint: disable=unused-argument,missing-docstring
 
-@tf_op(["Min", "Max", "Mean", "Sum", "Prod"],
-       type_map={"Min": "ReduceMin", "Max": "ReduceMax", "Mean": "ReduceMean",
-                 "Sum": "ReduceSum", "Prod": "ReduceProd"})
+@tf_op("Min", onnx_op="ReduceMin")
+@tf_op("Max", onnx_op="ReduceMax")
+@tf_op("Mean", onnx_op="ReduceMean")
+@tf_op("Sum", onnx_op="ReduceSum")
+@tf_op("Prod", onnx_op="ReduceProd")
 class ReduceOpBase:
     @classmethod
     def version_4(cls, ctx, node, **kwargs):

tf2onnx/onnx_opset/tensor.py

@@ -787,7 +787,7 @@ def version_4(cls, ctx, node, **kwargs):
         ctx.copy_shape(node.output[0], output_cast.output[0])
 
 
-@tf_op("IsNan", type_map={"IsNan": "IsNaN"})
+@tf_op("IsNan", onnx_op="IsNaN")
 class IsNan:
     @classmethod
     def version_9(cls, ctx, node, **kwargs):

tf2onnx/tfonnx.py

@@ -512,7 +512,7 @@ def rewrite_conv2d_with_pad(g, ops):
         g.replace_input(conv, conv.input[0], pad.input[0])
         # convert Conv2D
         conv.type = "Conv"
-        func, info = handler.tf_op.find_op("Conv2D")
+        func, info = handler.tf_op.find_effective_op("Conv2D")
         func(g, conv)
         conv.skip_conversion = True
         conv.set_attr("auto_pad", "NOTSET")
@@ -541,12 +541,10 @@ def tensorflow_onnx_mapping(g, continue_on_error, ops_mapping):
         mapped_op[op] += 1
         func, kwargs = map_info
         if kwargs:
-            # if there is a type_map key we'll map the old type to a new type
-            type_map = kwargs.get("type_map")
-            if type_map:
-                new_type = type_map.get(node.type)
-                if new_type:
-                    node.type = new_type
+            # if there is a onnx_op key we'll map the old type to a new type
+            onnx_op = kwargs.get("onnx_op")
+            if onnx_op:
+                node.type = onnx_op
         try:
             body_graphs = node.get_body_graphs()
             if body_graphs:
@@ -736,8 +734,8 @@ def compat_handler(ctx, node, **kwargs):
             args = v[1]
             kwargs = {"func": v[0]}
             if args:
-                type_map = {k: args[0]}
-                kwargs["type_map"] = type_map
+                onnx_op = args[0]
+                kwargs["onnx_op"] = onnx_op
                 args = args[1:]
             kwargs["args"] = args
             new_handler = handler.tf_op(k,