Merge pull request #936 from onnx/gs/add-is_finite

add support for tf.math.is_finite

Author: guschmue

tests/test_backend.py

@@ -3190,12 +3190,13 @@ def func(X, K):
                 k_val = np.array(raw_k).astype(np.int32)
                 self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: x_val, _INPUT1: k_val})
 
-    @check_opset_min_version(12)
-    def test_inverse(self):
+    @test_ms_domain()
+    def test_inverse(self, extra_opset):
+        # this depends on onnx Inverse which was removed from opset-12 but does exists in the ms-domain
         x_val = np.random.random([5, 5]).astype(np.float32)
         def func(x):
             return tf.linalg.inv(x, name=_TFOUTPUT)
-        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val}, process_args={"extra_opset": [extra_opset]})
 
     @check_opset_min_version(12)
     def test_squared_distance(self):
@@ -3224,6 +3225,15 @@ def func(x, y):
                    tf.math.greater_equal(x, y, name=_TFOUTPUT1)
         self._run_test_case(func, [_OUTPUT, _OUTPUT1], {_INPUT: x_val, _INPUT1: y_val})
 
+    @check_tf_min_version("1.14", "required for tf.math.is_finite")
+    @check_opset_min_version(10)
+    def test_is_finite(self):
+        x_val = np.array([5.0, 4.8, 6.8, np.inf, np.nan], dtype=np.float32)
+        def func(x):
+            y = tf.math.is_finite(x)
+            return tf.identity(y, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
+
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/custom_opsets/ms.py

@@ -96,7 +96,7 @@ def version_1(cls, ctx, node, **kwargs):
 @tf_op("CropAndResize", domain=constants.MICROSOFT_DOMAIN)
 class CropAndResize:
     @classmethod
-    def version_11(cls, ctx, node, **kwargs):
+    def version_1(cls, ctx, node, **kwargs):
         """ utilize contrib cropandresize """
         node.attr['method'].name = 'mode'
         node.domain = constants.MICROSOFT_DOMAIN
@@ -107,7 +107,7 @@ def version_11(cls, ctx, node, **kwargs):
 @tf_op("MatrixInverse", domain=constants.MICROSOFT_DOMAIN, onnx_op="Inverse")
 class Inverse:
     @classmethod
-    def version_12(cls, ctx, node, **kwargs):
+    def version_1(cls, ctx, node, **kwargs):
         utils.make_sure(node.get_attr('adjoint').i == 0, "adjoint must be false")
         del node.attr["adjoint"]
         node.domain = constants.MICROSOFT_DOMAIN

tf2onnx/onnx_opset/math.py

@@ -566,10 +566,23 @@ def version_12(cls, ctx, node, **kwargs):
         del node.attr["N"]
 
 
-@tf_op("MatrixInverse", onnx_op="Inverse")
-class Inverse:
+@tf_op("IsFinite")
+class IsFinite:
     @classmethod
-    def version_12(cls, ctx, node, **kwargs):
-        utils.make_sure(node.get_attr('adjoint').i == 0, "adjoint must be false")
-        del node.attr["adjoint"]
-        node.domain = constants.MICROSOFT_DOMAIN
+    def version_10(cls, ctx, node, **kwargs):
+        # map to onnx as:
+        # not (isinf(x) or isnan(x))
+
+        shapes = node.output_shapes
+        dtypes = [onnx_pb.TensorProto.BOOL] * len(node.output_dtypes)
+
+        ctx.remove_node(node.name)
+
+        inf_node = ctx.make_node("IsInf", inputs=node.input, name=utils.make_name(node.name),
+                                 shapes=shapes, dtypes=dtypes)
+        nan_node = ctx.make_node("IsNaN", inputs=node.input, name=utils.make_name(node.name),
+                                 shapes=shapes, dtypes=dtypes)
+        or_node = ctx.make_node("Or", inputs=[inf_node.output[0], nan_node.output[0]], name=utils.make_name(node.name),
+                                shapes=shapes, dtypes=dtypes)
+        _ = ctx.make_node("Not", inputs=or_node.output, name=node.name,
+                          shapes=shapes, dtypes=dtypes)

tf2onnx/onnx_opset/tensor.py

@@ -625,6 +625,16 @@ def version_7(cls, ctx, node, **kwargs):
             return
         raise ValueError("non-const dim is not supported")
 
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        dim_node = node.inputs[1]
+        if dim_node.is_const():
+            node.type = "Unsqueeze"
+            dim = dim_node.get_tensor_value()
+            node.set_attr("axes", [dim])
+            ctx.remove_input(node, node.input[1])
+            return
+        raise ValueError("non-const dim is not supported")
 
 @tf_op("StridedSlice")
 class StridedSlice:

tools/dump-onnx.py

@@ -14,6 +14,7 @@
 import collections
 import re
 
+import onnx
 from onnx import ModelProto
 from onnx import helper, shape_inference