Merge pull request #1142 from onnx/tom/TensorScatterNdUpdate

Added support for TensorScatterUpdate

Author: TomWildenhain-Microsoft

tests/test_backend.py

@@ -3157,6 +3157,30 @@ def func(x):
             return tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case(func, [_OUTPUT], {_INPUT: x_val})
 
+    @check_tf_min_version("1.14", "tensor_scatter_nd_update needs tf 1.14")
+    @check_opset_min_version(11, "ScatterND")
+    def test_tensor_scatter_update(self):
+        x_val = np.array([10, 20, 30, 40], dtype=np.int32).reshape((4))
+        y_val = np.array([0, 2], dtype=np.int64).reshape((2, 1))
+        z_val = np.array([8, 11], dtype=np.int32).reshape((2))
+
+        def func(x, y, z):
+            x_ = tf.tensor_scatter_nd_update(x, y, z)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+
+    @check_tf_min_version("1.14", "tensor_scatter_nd_update needs tf 1.14")
+    @check_opset_min_version(11, "ScatterND")
+    def test_tensor_scatter_update_cast_indices(self):
+        x_val = np.array([10, 20, 30, 40], dtype=np.int32).reshape((4))
+        y_val = np.array([0, 2], dtype=np.int32).reshape((2, 1))
+        z_val = np.array([8, 11], dtype=np.int32).reshape((2))
+
+        def func(x, y, z):
+            x_ = tf.tensor_scatter_nd_update(x, y, z)
+            return tf.identity(x_, name=_TFOUTPUT)
+        self._run_test_case(func, [_OUTPUT], {_INPUT: x_val, _INPUT1: y_val, _INPUT2: z_val})
+
     @check_opset_min_version(11, "ScatterND")
     def test_scatternd_1d(self):
         x_val = np.array([4, 3, 1, 7], dtype=np.int32).reshape((4, 1))

tf2onnx/onnx_opset/tensor.py

@@ -539,6 +539,14 @@ def version_11(cls, ctx, node, **kwargs):
         ctx.replace_inputs(node, [node.input[2], node.input[0], node.input[1]])
 
 
+@tf_op("TensorScatterUpdate", onnx_op="ScatterND")
+class TensorScatterUpdate:
+    @classmethod
+    def version_11(cls, ctx, node, **kwargs):
+        if ctx.get_dtype(node.input[1]) != TensorProto.INT64:
+            ctx.insert_new_node_on_input(node, "Cast", node.input[1], to=TensorProto.INT64)
+
+
 @tf_op("Split")
 class Split:
     @classmethod