[Relay] Fix a bug in tensor_array_scatter (#6890)

* [Relay] Fix a bug in tensor_array_scatter

  tensor_array_scatter constructs helper functions according to dtype
  and shape of element. When there are multiple scatter operations with
  same dtype and element shape but different indicies_shape, there will
  be name conflict in prelude.

* Refine get_name

Author: lixiaoquan

python/tvm/relay/prelude.py

@@ -73,7 +73,7 @@ def get_tensor_array_shape(expr, dtype, prelude):
     return None
 
 
-def _get_name_static(canonical, dtype, shape, batch_dim=None):
+def _get_name_static(canonical, dtype, shape, batch_dim=None, extra_shapes=None):
     """Get name for static shape tensor array op
 
     By design, static ADT tensor in TVM has type name in the format
@@ -100,14 +100,12 @@ def _get_name_static(canonical, dtype, shape, batch_dim=None):
     name : String
         The tensor array op name
     """
-    dim_names = []
-    for dim in shape:
-        if isinstance(dim, Any):
-            dim_names.append("any")
-        else:
-            dim_names.append(str(dim))
+    shape_str = _to_str(shape)
 
-    shape_str = "_".join(dim_names)
+    if extra_shapes is not None:
+        for n, s in extra_shapes.items():
+            extra_shape_str = "_{}_{}".format(n, _to_str(s))
+            shape_str += extra_shape_str
 
     if len(shape_str) == 0:
         shape_str = "scalar"
@@ -120,6 +118,16 @@ def _get_name_static(canonical, dtype, shape, batch_dim=None):
     return "{}_{}_batch{}_{}".format(canonical, dtype, str(batch_dim), shape_str)
 
 
+def _to_str(shape):
+    dim_names = []
+    for dim in shape:
+        if isinstance(dim, Any):
+            dim_names.append("any")
+        else:
+            dim_names.append(str(dim))
+    return "_".join(dim_names)
+
+
 class StaticTensorArrayOps(object):
     """Contains tensor array related ops for fixed rank tensor array"""
 
@@ -131,9 +139,9 @@ def __init__(self, prelude, dtype, shape, batch_dim=None):
         self.batch_dim = batch_dim
         self.list, self.cons, self.nil = self.prelude.mod.get_type("List")
 
-    def get_name(self, canonical):
+    def get_name(self, canonical, extra_shapes=None):
         """Get name corresponding to the canonical name"""
-        return _get_name_static(canonical, self.dtype, self.shape, self.batch_dim)
+        return _get_name_static(canonical, self.dtype, self.shape, self.batch_dim, extra_shapes)
 
     def get_global_var(self, canonical):
         """Get global corresponding to the canonical name"""
@@ -408,11 +416,16 @@ def define_tensor_array_scatter(self, indices_shape=None, force_update=False):
         # When this operator has already been registered, only update
         # when force_update is set. This should be used only when we need to
         # redefine this op for static indices shape.
-        tensor_array_scatter_name = self.get_name("tensor_array_scatter")
+
+        extra_shapes = {"indices": indices_shape} if indices_shape is not None else None
+        tensor_array_scatter_name = self.get_name("tensor_array_scatter", extra_shapes)
         if hasattr(self.prelude, tensor_array_scatter_name) and not force_update:
             return
 
-        tensor_array_scatter_helper_name = self.get_name("tensor_array_scatter_helper")
+        tensor_array_scatter_helper_name = self.get_name(
+            "tensor_array_scatter_helper", extra_shapes
+        )
+
         tensor_array_scatter_helper_var = self._create_global_var(tensor_array_scatter_helper_name)
         ta = Var("ta", self.list(self.tensor_type_var()))
         current = Var("current", scalar_type("int32"))

tests/python/frontend/tensorflow/test_forward.py

@@ -1537,19 +1537,29 @@ def run(dtype_str, infer_shape):
                 element_shape = tf.TensorShape([tf.Dimension(None)])
             else:
                 element_shape = None
-            t = tf.constant(np.array([[1.0], [2.0], [3.0]]).astype(dtype_str), dtype=dtype)
-            indices = tf.constant([2, 1, 0])
-            ta1 = tf.TensorArray(
-                dtype=dtype, size=3, infer_shape=infer_shape, element_shape=element_shape
-            )
-            ta2 = ta1.scatter(indices, t)
-            out0 = ta2.read(0)
-            out1 = ta2.read(1)
-            out2 = ta2.read(2)
+            ta0 = _construct_scatter(dtype, dtype_str, element_shape, infer_shape, 3)
+            out0 = ta0.read(0)
+            out1 = ta0.read(1)
+            out2 = ta0.read(2)
+            ta1 = _construct_scatter(dtype, dtype_str, element_shape, infer_shape, 4)
+            out4 = ta1.read(0)
             g = tf.get_default_graph()
             compare_tf_with_tvm([], [], ["TensorArrayReadV3:0"], mode="vm")
             compare_tf_with_tvm([], [], ["TensorArrayReadV3_1:0"], mode="vm")
             compare_tf_with_tvm([], [], ["TensorArrayReadV3_2:0"], mode="vm")
+            compare_tf_with_tvm([], [], ["TensorArrayReadV3_2:0", out4.name], mode="vm")
+
+    def _construct_scatter(dtype, dtype_str, element_shape, infer_shape, size):
+        arr = [[float(i)] for i in range(size)]
+        indices_arr = [i for i in range(size - 1, -1, -1)]
+
+        t = tf.constant(np.array(arr).astype(dtype_str), dtype=dtype)
+        indices = tf.constant(indices_arr)
+        ta1 = tf.TensorArray(
+            dtype=dtype, size=size, infer_shape=infer_shape, element_shape=element_shape
+        )
+        ta2 = ta1.scatter(indices, t)
+        return ta2
 
     for dtype in ["float32", "int8"]:
         run(dtype, False)