support for old custom op registration, fix graph unit test

Author: guschmue

tests/test_graph.py

@@ -22,10 +22,12 @@
 from tf2onnx.graph import GraphUtil
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.tfonnx import process_tf_graph
+from tf2onnx.handler import tf_op
+
 from common import get_test_config, unittest_main
 
 
-# pylint: disable=missing-docstring
+# pylint: disable=missing-docstring,unused-argument,unused-variable
 
 def onnx_to_graphviz(g, include_attrs=False):
     """Return dot for graph."""
@@ -331,10 +333,10 @@ def rewrite_test(g, ops):
                 'output [op_type=Identity] input1:0 -> Add input1:0 -> Add Add:0 -> output }',
                 onnx_to_graphviz(g))
 
-    def test_custom_op(self):
-        """Custom op test."""
+    def test_custom_op_depreciated(self):
+        """Custom op test using old depreciated api."""
 
-        def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
+        def print_handler(ctx, node, name, args):
             # replace tf.Print() with Identity
             #   T output = Print(T input, data, @list(type) U, @string message, @int first_n, @int summarize)
             # becomes:
@@ -360,6 +362,32 @@ def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
             self.assertEqual(g.opset, self.config.opset)
             self.assertEqual(g.extra_opset, [constants.TENSORFLOW_OPSET])
 
+    def test_custom_op(self):
+        """Custom op test."""
+
+        @tf_op("Print", type_map={"Print": "Identity"})
+        class Print:
+            @classmethod
+            def version_1(cls, ctx, node, **kwargs):
+                self.assertEqual(node.type, "Identity")
+                node.domain = constants.TENSORFLOW_OPSET.domain
+                del node.input[1:]
+                return node
+
+        with tf.Session() as sess:
+            x = tf.placeholder(tf.float32, [2, 3], name="input1")
+            x_ = tf.Print(x, [x], "hello")
+            _ = tf.identity(x_, name="output")
+            g = process_tf_graph(sess.graph,
+                                 opset=self.config.opset,
+                                 extra_opset=[constants.TENSORFLOW_OPSET])
+            self.assertEqual(
+                'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Print [domain="ai.onnx.converters.tensorflow" '
+                'op_type=Identity] output [op_type=Identity] input1:0 -> Print Print:0 -> output }',
+                onnx_to_graphviz(g))
+            self.assertEqual(g.opset, self.config.opset)
+            self.assertEqual(g.extra_opset, [constants.TENSORFLOW_OPSET])
+
     def test_extra_opset(self):
         extra_opset = [
             utils.make_opsetid(constants.MICROSOFT_DOMAIN, 1),

tf2onnx/handler.py

@@ -38,11 +38,20 @@ def __call__(self, func):
                 for name in self.name:
                     opset_dict[name] = (v, self.kwargs)
 
+    def register_compat_handler(self, func, version):
+        opset = tf_op._OPSETS.get(self.domain)
+        if not opset:
+            opset = []
+            tf_op._OPSETS[self.domain] = opset
+            while version >= len(opset):
+                opset.append({})
+            opset_dict = opset[version]
+            opset_dict[self.name[0]] = (func, self.kwargs)
+
     @staticmethod
     def get_opsets():
         return tf_op._OPSETS
 
-
     @staticmethod
     def create_mapping(max_opset, extra_opsets):
         mapping = {"onnx": max_opset}
@@ -60,7 +69,6 @@ def create_mapping(max_opset, extra_opsets):
         tf_op._MAPPING = ops_mapping
         return ops_mapping
 
-
     @staticmethod
     def find_op(name):
         map_info = tf_op._MAPPING.get(name)

tf2onnx/tfonnx.py

@@ -717,7 +717,34 @@ def process_tf_graph(tf_graph, continue_on_error=False, verbose=False, target=No
     # apply custom ops on top of the assembled opset. We can either complement the opset
     # or override existing ops with a custom op.
     if custom_op_handlers is not None:
-        custom_opset = {k: v for k, v in custom_op_handlers.items()}
+        # below is a bit tricky since there are a few api's:
+        # 1. the future way we want custom ops to be registered with the @tf_op decorator. THose handlers will be
+        #     registered via the decorator on load of the module ... nothing is required here.
+        # 2. the old custom op api: a dictionary of {name: (func, args[])
+        #     We deal with this by using a compat_handler that wraps to old handler with a new style handler.
+        #     This is tempoary to give people give to move to the new api and after tf2onnx-1.5 we want to remove this
+        custom_opset = {}
+        for k, v in custom_op_handlers.items():
+            # FIXME: remove this after tf2onnx-1.5
+            def compat_handler(ctx, node, **kwargs):
+                # wrap old handler
+                name = node.name
+                args = kwargs["args"]
+                func = kwargs["func"]
+                return func(ctx, node, name, args)
+
+            args = v[1]
+            kwargs = {"func": v[0]}
+            if args:
+                type_map = {k: args[0]}
+                kwargs["type_map"] = type_map
+                args = args[1:]
+            kwargs["args"] = args
+            new_handler = handler.tf_op(k,
+                                        domain=constants.TENSORFLOW_OPSET.domain,
+                                        kwargs=kwargs)
+            new_handler.register_compat_handler(compat_handler, 1)
+            custom_opset[k] = (compat_handler, kwargs)
         ops_mapping.update(custom_opset)
 
     infer_shape_for_graph(g)