Merge pull request #477 from mindest/implement_isinf

implement IsInf for opset 10

Author: nbcsm

tests/test_backend.py

@@ -2103,6 +2103,19 @@ def test_space_to_batchnd(self):
         _ = tf.space_to_batch_nd(input_x, block_size, pad, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: input_val})
 
+    @check_opset_min_version(10, "is_inf")
+    def test_isinf(self):
+        x_types = [np.float32, np.float64]
+        for x_type in x_types:
+            x_val1 = np.array([1.0, -2.0, 3.0, -4.0], dtype=x_type)
+            x_val2 = np.array([np.inf, np.inf, np.inf, np.inf], dtype=x_type).reshape((2, 2))
+            x_val3 = np.array([1.0, np.inf, -3.0, np.inf, 5.0, np.inf, -7.0, np.inf], dtype=x_type).reshape((2, 2, 2))
+            for x_val in [x_val1, x_val2, x_val3]:
+                x = tf.placeholder(x_type, x_val.shape, name=_TFINPUT)
+                x_ = tf.is_inf(x)
+                _ = tf.identity(x_, name=_TFOUTPUT)
+                self._run_test_case([_OUTPUT], {_INPUT: x_val})
+                tf.reset_default_graph()
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/tensor.py

@@ -922,3 +922,13 @@ def version_4(cls, ctx, node, **kwargs):
         reorganize_node = ctx.make_node(node.type, trans1.output, attr={"blocksize": blocksize[0]})
         ctx.make_node("Transpose", reorganize_node.output, {"perm": [1, 2, 3, 0]}, name=node.name, outputs=node.output,
                       shapes=shapes, dtypes=dtypes)
+
+
+@tf_op("IsInf", onnx_op="IsInf")
+class IsInf:
+    @classmethod
+    def version_10(cls, ctx, node, **kwargs):
+        node_dtype = ctx.get_dtype(node.input[0])
+        utils.make_sure(node_dtype, "Dtype of {} is None".format(node.name))
+        if node_dtype not in [onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.DOUBLE]:
+            raise ValueError("dtype " + str(node_dtype) + " is not supported in onnx for now")