From 7cc5dec2f6f29a536e204ff9fedb2a5d86510d2d Mon Sep 17 00:00:00 2001
From: Wang Yao <kevinthesunwy@gmail.com>
Date: Tue, 31 Mar 2020 16:20:25 -0700
Subject: [PATCH 1/6] Support TF Frontend Static TensorArray

---
 python/tvm/relay/frontend/common.py           |  54 ++--
 python/tvm/relay/frontend/tensorflow.py       | 281 ++++++++++++++++--
 .../frontend/tensorflow/test_forward.py       | 142 +++++----
 topi/python/topi/util.py                      |   4 +-
 4 files changed, 384 insertions(+), 97 deletions(-)

diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
index 5465e50bba43..76d4eb6d3c37 100644
--- a/python/tvm/relay/frontend/common.py
+++ b/python/tvm/relay/frontend/common.py
@@ -21,7 +21,7 @@
 import numpy as np
 
 import tvm
-from tvm.ir import IRModule
+from tvm.ir import IRModule, TypeCall
 from topi.util import get_const_tuple
 
 from .. import expr as _expr
@@ -456,12 +456,19 @@ def get_name(node):
 
 def infer_type(node, mod=None):
     """A method to infer the type of an intermediate node in the relay graph."""
-    new_mod = IRModule.from_expr(node)
-    if mod is not None:
-        new_mod.update(mod)
-    new_mod = _transform.InferType()(new_mod)
-    entry = new_mod["main"]
-    return entry if isinstance(node, _function.Function) else entry.body
+    if isinstance(mod, IRModule):
+        mod["main"] = _function.Function([], node)
+        mod = _transform.InferType()(mod)
+        entry = mod["main"]
+        return entry.body
+    else:
+        new_mod = IRModule.from_expr(node)
+        if mod is not None:
+            new_mod.update(mod)
+            new_mod = _transform.InferType()(new_mod)
+        entry = new_mod["main"]
+        return entry if isinstance(node, _function.Function) else entry.body
+
 
 def infer_shape(inputs, mod=None):
     """A method to get the output type of an intermediate node in the graph."""
@@ -473,15 +480,6 @@ def infer_shape(inputs, mod=None):
     # The return type is not a tensor, for example List
     return checked_type
 
-def infer_channels(inputs, transpose=False):
-    """A hack for getting 'channels' or 'units' since caffe2 does not provide
-    these attributes. We check the shape of weights provided to get the number.
-    """
-    out_type = infer_type(inputs)
-    out_shapes = [get_const_tuple(out_type.checked_type.shape)]
-    channels = out_shapes[0][0] if not transpose else out_shapes[0][1]
-    return channels
-
 
 def infer_value(input_val, params, mod=None):
     """A hack for getting the value of an expression by evaluating a
@@ -505,7 +503,7 @@ def infer_value(input_val, params, mod=None):
         return m.get_output(0)
     except Exception:
         if isinstance(mod, IRModule):
-            mod["main"] = _expr.Function(analysis.free_vars(input_val), input_val)
+            mod["main"] = _function.Function(analysis.free_vars(input_val), input_val)
         else:
             mod = IRModule.from_expr(input_val)
         exc = tvm.relay.create_executor("debug", mod=mod, ctx=tvm.cpu(), target="llvm")
@@ -548,6 +546,28 @@ def new_var(name_hint,
     return _expr.var(name_hint, type_annotation, shape, dtype)
 
 
+def check_tensor_array_shape(expr, dtype, mod):
+    """Check whether a tensor array has fixed rank shape.
+    Return its shape if yes. Otherwise, return None.
+    """
+    checked_type = infer_shape(expr, mod)
+    assert isinstance(checked_type, TypeCall), "Input must be a tensor array."
+    ta_type_str = checked_type.args[0].func.name_hint
+    static_ta_ty_start = "static_tensor_{}".format(dtype)
+    if ta_type_str.startswith(static_ta_ty_start):
+        shape_str = ta_type_str.replace("{}_".format(static_ta_ty_start), '') \
+            .replace("_t", '')
+        shape = []
+        if "scalar" not in shape_str:
+            for dim_str in shape_str.split("_"):
+                if dim_str == "?":
+                    shape.append(Any())
+                else:
+                    shape.append(int(dim_str))
+        return tuple(shape)
+    return None
+
+
 class Renamer(object):
     """A simply renamer for operators.
 
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 84318c384418..0cea94be0563 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -26,7 +26,7 @@
 import tvm
 
 from tvm.ir import IRModule
-from tvm.relay.prelude import Prelude
+from tvm.relay.prelude import Prelude, StaticTensorArrayOps
 from tvm.ir import structural_hash as s_hash
 
 from .. import analysis
@@ -40,6 +40,7 @@
 from .common import infer_channels as _infer_channels
 from .common import infer_value as _infer_value
 from .common import infer_value_simulated as _infer_value_simulated
+from .common import check_tensor_array_shape as _check_tensor_array_shape
 
 __all__ = ['from_tensorflow']
 
@@ -259,8 +260,6 @@ def _impl(inputs, attr, params, mod):
         if opname == 'conv_transpose' and attr['data_format'] == 'NHWC':
             # transform to NCHW for TVM backend compatible and set 'flip_layout'
             # to have output flip back to NHWC
-            tmp_shape = _infer_shape(inputs[2], mod)
-            tmp_shape = [tmp_shape[ii] for ii in (0, 3, 1, 2)]
             inputs[2] = _op.transpose(inputs[2], axes=(0, 3, 1, 2))
             attr['strides'][1], attr['strides'][2], attr['strides'][3] = \
                 attr['strides'][3], attr['strides'][1], attr['strides'][2]
@@ -790,24 +789,124 @@ def _impl(inputs, attr, params, mod):
 def _tensor_array():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr.get('dtype').name
-        tensor_array_constructor = prelude.get_var('tensor_array', dtype_str)
-        return tensor_array_constructor(_op.take(inputs[0], tvm.relay.const(0)))
+        assert not attr["dynamic_size"], "Dynamic size tensor array is " \
+                                         "not supported in TVM yet."
+        if attr['identical_element_shapes']:
+            # For static tensor array, we need to create with static constructor.
+            # However, at this point we might not know the input shape yet.
+            # We create a static tensor array with dummy shape and record it in
+            # _static_tensor_array_map. Later when creating other tensor array ops
+            # which uses this tensor array, we reconstruct this tensor array with
+            # actual shape.
+            dummy_shape = ()
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           dummy_shape)
+            static_tensor_array_ops.register()
+            tensor_array_constructor = prelude.get_var_static('tensor_array',
+                                                              dtype_str,
+                                                              dummy_shape)
+            tensor_array =  tensor_array_constructor(inputs[0])
+            _static_tensor_array_map[tensor_array] = None
+        else:
+            tensor_array_constructor = prelude.get_var('tensor_array', dtype_str)
+            tensor_array =  tensor_array_constructor(inputs[0])
+        return tensor_array
     return _impl
 
 def _tensor_array_scatter():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr.get('T').name
-        values_rank = len(inputs[2].type_annotation.shape)
-        unstack_name = "tensor_array_unstack_tensor{}".format(values_rank)
-        unstack_function = prelude.get_var(unstack_name, dtype_str)
-        values = unstack_function(inputs[2])
-        tensor_array_scatter_func = prelude.get_var('tensor_array_scatter', dtype_str)
-        return tensor_array_scatter_func(inputs[0], inputs[1], values)
+        input_ta = inputs[0]
+        input_shape = _check_tensor_array_shape(input_ta, dtype_str, prelude.mod)
+        values_shape = _infer_shape(inputs[2], prelude.mod)
+        input_t_shape = values_shape[1:]
+        indices_shape = _infer_shape(inputs[1], prelude.mod)
+
+        if input_shape is None:
+            values_rank = len(values_shape)
+            unstack_name = "tensor_array_unstack_tensor{}".format(values_rank)
+            unstack_function = prelude.get_var(unstack_name, dtype_str)
+            values = unstack_function(inputs[2])
+            tensor_array_scatter_func = prelude.get_var('tensor_array_scatter', dtype_str)
+        else:
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           input_t_shape)
+            static_tensor_array_ops.register()
+            # For scatter operation, it is possible to write to a newly create
+            # tensor array. We need to check and recreate its input tensor array.
+            if input_ta in _static_tensor_array_map and \
+                    _static_tensor_array_map[input_ta] is None:
+                ta_constructor = prelude.get_var_static('tensor_array',
+                                                        dtype_str,
+                                                        input_t_shape)
+                new_ta =  ta_constructor(input_ta.args[0])
+                _static_tensor_array_map[input_ta] = new_ta
+                input_ta = new_ta
+
+            # Register static indices shape
+            if isinstance(indices_shape[0], int):
+                static_tensor_array_ops.define_tensor_array_scatter(indices_shape, True)
+            tensor_array_scatter_func = prelude.get_var_static('tensor_array_scatter',
+                                                               dtype_str,
+                                                               input_t_shape)
+
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           values_shape)
+            static_tensor_array_ops.register()
+            unstack_function = prelude.get_var_static('tensor_array_unstack',
+                                                      dtype_str,
+                                                      values_shape)
+            values = unstack_function(inputs[2])
+        ret = tensor_array_scatter_func(input_ta, inputs[1], values)
+        return ret
     return _impl
 
 def _tensor_array_gather():
     def _impl(inputs, attr, params, prelude):
-        return prelude.tensor_array_gather(inputs[2], inputs[1])
+        dtype_str = attr.get('dtype').name
+        input_shape = _check_tensor_array_shape(inputs[2], dtype_str, prelude.mod)
+        indices_shape = _infer_shape(inputs[1], prelude.mod)
+
+        if input_shape is None:
+            gather_func = prelude.get_var('tensor_array_gather', dtype_str)
+            return gather_func(inputs[2], inputs[1])
+        else:
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           input_shape)
+            static_tensor_array_ops.register()
+            if not isinstance(indices_shape[0], int):
+                gather_function = prelude.get_var_static('tensor_array_gather',
+                                                         dtype_str,
+                                                         input_shape)
+                out_tensor_t = gather_function(inputs[2], inputs[1])
+
+                # Output shape is (indices_shape[0],) + input_shape
+                static_tensor_array_ops.define_tensor_get_data((indices_shape[0],) + input_shape)
+                get_data_func = prelude.get_var_static('tensor_get_data',
+                                                       dtype_str,
+                                                       input_shape)
+                return get_data_func(out_tensor_t)
+            else:
+                # For fixed length indices, directly generate static shape output
+                read_func = prelude.get_var_static('tensor_array_read',
+                                                   dtype_str,
+                                                   input_shape)
+                static_tensor_array_ops.define_tensor_get_data(input_shape)
+                get_data_func = prelude.get_var_static('tensor_get_data',
+                                                       dtype_str,
+                                                       input_shape)
+                tensor_list = []
+                for i in range(indices_shape[0]):
+                    index = _op.take(inputs[1], tvm.relay.const(i))
+                    out_tensor = get_data_func(read_func(inputs[2], index))
+                    tensor_list.append(_op.expand_dims(out_tensor, axis=0))
+
+                return _op.concatenate(tensor_list, axis=0)
+
     return _impl
 
 def _tensor_array_size():
@@ -817,37 +916,160 @@ def _impl(inputs, attr, params, prelude):
 
 def _tensor_array_write():
     def _impl(inputs, attr, params, prelude):
-        input_rank = len(inputs[2].type_annotation.shape)
-        dtype = attr.get('T').name
+        dtype_str = attr.get('T').name
+        input_ta = inputs[3]
+        input_ta_shape = _check_tensor_array_shape(input_ta, dtype_str, prelude.mod)
+        input_t_shape = _infer_shape(inputs[2], prelude.mod)
+        input_rank = len(input_t_shape)
+
+        if input_ta_shape is None:
+            tensor_name = 'tensor{}'.format(input_rank)
+            tensor_func = prelude.get_var(tensor_name, dtype_str)
+            v = tensor_func(inputs[2])
+            write_func = prelude.get_var('tensor_array_write', dtype_str)
+        else:
+            # For write operation, it is possible to write to a newly create
+            # tensor array. We need to check and recreate its input tensor array.
+            if input_ta in _static_tensor_array_map and \
+                    _static_tensor_array_map[input_ta] is None:
+                static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                               dtype_str,
+                                                               input_t_shape)
+                static_tensor_array_ops.register()
+                ta_constructor = prelude.get_var_static('tensor_array',
+                                                        dtype_str,
+                                                        input_t_shape)
+                new_ta =  ta_constructor(input_ta.args[0])
+                _static_tensor_array_map[input_ta] = new_ta
+                input_ta = new_ta
+                input_ta_shape = input_t_shape
+            else:
+                input_ta_rank = len(input_ta_shape)
+                assert input_ta_rank == input_rank, "Shape rank mismatch: {} vs {}". \
+                    format(input_ta_rank, input_rank)
+                static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                               dtype_str,
+                                                               input_ta_shape)
+                static_tensor_array_ops.register()
 
-        tensor_name = 'tensor{}'.format(input_rank)
-        tensor_func = prelude.get_var(tensor_name, dtype)
-        v = tensor_func(inputs[2])
-        write_func = prelude.get_var('tensor_array_write', dtype)
+            tensor_func = prelude.get_var_static("tensor_constructor",
+                                                 dtype_str,
+                                                 input_ta_shape)
+            v = tensor_func(inputs[2])
+            write_func = prelude.get_var_static('tensor_array_write',
+                                                dtype_str,
+                                                input_ta_shape)
 
-        return write_func(inputs[3], _op.take(inputs[1], tvm.relay.const(0)), v)
+        return write_func(input_ta, _op.take(inputs[1], tvm.relay.const(0)), v)
     return _impl
 
 def _tensor_array_read():
     def _impl(inputs, attr, params, prelude):
-        read_func = prelude.get_var('tensor_array_read', attr.get('dtype').name)
-        return read_func(inputs[2], _op.take(inputs[1], tvm.relay.const(0)))
+        dtype_str = attr['dtype'].name
+        input_shape = _check_tensor_array_shape(inputs[2], dtype_str, prelude.mod)
+
+        if input_shape is None:
+            read_func = prelude.get_var('tensor_array_read', dtype_str)
+            return read_func(inputs[2], _op.take(inputs[1], tvm.relay.const(0)))
+        else:
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           input_shape)
+            static_tensor_array_ops.register()
+            static_tensor_array_ops.define_tensor_get_data(input_shape)
+            read_func = prelude.get_var_static("tensor_array_read", dtype_str, input_shape)
+            out_tensor = read_func(inputs[2], _op.take(inputs[1], tvm.relay.const(0)))
+            get_data_func = prelude.get_var_static('tensor_get_data',
+                                                   dtype_str,
+                                                   input_shape)
+            return get_data_func(out_tensor)
     return _impl
 
 def _tensor_array_split():
     def _impl(inputs, attr, params, prelude):
-        input_rank = len(inputs[1].type_annotation.shape)
         dtype_str = attr.get('T').name
-        v = prelude.get_var("tensor{}".format(input_rank), dtype_str)(inputs[1])
+        input_ta = inputs[0]
+        input_ta_shape = _check_tensor_array_shape(input_ta, dtype_str, prelude.mod)
+        input_t_shape = _infer_shape(inputs[1], prelude.mod)
+        input_rank = len(input_t_shape)
         lengths = _op.cast(inputs[2], 'int32')
-        split_var = prelude.get_var('tensor_array_split', dtype_str)
-        return split_var(inputs[0], v, lengths)
+        lengths_shape = _infer_shape(lengths, prelude.mod)
+        value_shape = _infer_shape(inputs[1], prelude.mod)
+
+        if input_ta_shape is None:
+            v = prelude.get_var("tensor{}".format(input_rank), dtype_str)(inputs[1])
+            split_func = prelude.get_var('tensor_array_split', dtype_str)
+        else:
+            # For split operation, it is possible to write to a newly create
+            # tensor array. We need to check and recreate its input tensor array.
+            if input_ta in _static_tensor_array_map and \
+                    _static_tensor_array_map[input_ta] is None:
+                input_ta_shape = (Any(),) + input_t_shape[1:]
+                static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                               dtype_str,
+                                                               input_ta_shape)
+                static_tensor_array_ops.register()
+                ta_constructor = prelude.get_var_static('tensor_array',
+                                                        dtype_str,
+                                                        input_ta_shape)
+                new_ta =  ta_constructor(input_ta.args[0])
+                _static_tensor_array_map[input_ta] = new_ta
+                input_ta = new_ta
+            else:
+                input_ta_rank = len(input_ta_shape)
+                assert input_ta_rank == input_rank, "Shape rank mismatch: {} vs {}". \
+                    format(input_ta_rank, input_rank)
+                static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                               dtype_str,
+                                                               input_ta_shape)
+                static_tensor_array_ops.register()
+
+            # Check static value/indices shape
+            if isinstance(value_shape[0], int) or isinstance(lengths_shape[0], int):
+                static_tensor_array_ops.define_tensor_array_split(value_shape,
+                                                                  lengths_shape,
+                                                                  True)
+
+            tensor_func_name = prelude.get_name_static("tensor_constructor",
+                                                       dtype_str,
+                                                       value_shape)
+            if not hasattr(prelude, tensor_func_name):
+                static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                               dtype_str,
+                                                               value_shape)
+                static_tensor_array_ops.register()
+            tensor_func = prelude.get_var_static("tensor_constructor",
+                                                 dtype_str,
+                                                 value_shape)
+            v = tensor_func(inputs[1])
+            split_func = prelude.get_var_static('tensor_array_split',
+                                                dtype_str,
+                                                input_ta_shape)
+
+        return split_func(input_ta, v, lengths)
     return _impl
 
 def _tensor_array_concat():
     def _impl(inputs, attr, params, prelude):
-        concat_func = prelude.get_var('tensor_array_concat', attr['dtype'].name)
-        return concat_func(inputs[1])
+        dtype_str = attr['dtype'].name
+        input_shape = _check_tensor_array_shape(inputs[1], dtype_str, prelude.mod)
+
+        if input_shape is None:
+            concat_func = prelude.get_var('tensor_array_concat', dtype_str)
+            return concat_func(inputs[1])
+        else:
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           input_shape)
+            static_tensor_array_ops.register()
+            concat_func = prelude.get_var_static("tensor_array_concat", dtype_str, input_shape)
+            out_tensor = concat_func(inputs[1])
+            static_tensor_array_ops.define_tensor_get_data((Any(),) + input_shape[1:])
+            get_data_func = prelude.get_var_static('tensor_get_data',
+                                                   dtype_str,
+                                                   input_shape)
+            return get_data_func(out_tensor)
+
     return _impl
 
 def _tile():
@@ -1360,7 +1582,7 @@ def _impl(inputs, attr, params, mod):
 
         return AttrCvt(
             op_name="arange",
-            ignores=['Tidx'],
+            ignores=['Tidx', '_class'],
             extras={'start': start,
                     'stop': limit,
                     'step': delta,
@@ -2073,6 +2295,9 @@ def _get_abs_layer_name(node):
 # 1.x.
 _control_flow_nodes = ['Merge', 'Switch', 'NextIteration', 'Exit', 'Enter', 'LoopCond']
 
+# A map to record tensor array with fixed rank shape
+_static_tensor_array_map = {}
+
 class RewriteSubgraph(ExprMutator):
     """
     A helper class to rewrite expr in while loop function to variable
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 5c0391f6588d..bc51bbb0c725 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -839,63 +839,73 @@ def test_forward_squeeze():
     _test_squeeze(np.arange(6).reshape((1, 2, 1, 3, 1)), [-3, -5, -1])
 
 
-def test_tensor_array_constructor():
-    def run(dtype_str):
+#######################################################################
+# TensorArray
+# -----------
+def test_tensor_array_write_read():
+    def run(dtype_str, infer_shape):
         with tf.Graph().as_default():
             dtype = tf_dtypes[dtype_str]
-            t = tf.constant(np.array([[1.0, 2.0], [3.0, 4.0]]).astype(dtype_str), dtype=dtype)
-            t2 = tf.constant(np.array([[1.0, 2.0], [3.0, 4.0]]).astype(dtype_str), dtype=dtype)
-            ta1 = tf.TensorArray(dtype=dtype, size=2, infer_shape=False, dynamic_size=False)
-            ta2 = ta1.write(0, t)
+            np_data = np.array([[1.0, 2.0], [3.0, 4.0]]).astype(dtype_str)
+            in_data = [np_data, np_data]
+            t1 = tf.constant(np_data, dtype=dtype)
+            t2 = tf.constant(np_data, dtype=dtype)
+            ta1 = tf.TensorArray(dtype=dtype, size=2, infer_shape=infer_shape)
+            ta2 = ta1.write(0, t1)
             ta3 = ta2.write(1, t2)
             out = ta3.read(0)
             g = tf.get_default_graph()
-            compare_tf_with_tvm([], [], 'TensorArrayReadV3:0', mode='debug')
-    for dtype in tf_dtypes.keys():
-        run(dtype)
+            compare_tf_with_tvm([], [], 'TensorArrayReadV3:0', mode='vm')
+
+    for dtype in ["float32", "int8"]:
+        run(dtype, False)
+        run(dtype, True)
 
 
 def test_tensor_array_scatter():
-    def run(dtype_str):
+    def run(dtype_str, infer_shape):
         with tf.Graph().as_default():
             dtype =  tf_dtypes[dtype_str]
             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=False, dynamic_size=False)
+                                 infer_shape=infer_shape)
             ta2 = ta1.scatter(indices, t)
             out0 = ta2.read(0)
             out1 = ta2.read(1)
             out2 = ta2.read(2)
             g = tf.get_default_graph()
-            compare_tf_with_tvm([], [], ['TensorArrayReadV3:0'], mode='debug')
-            compare_tf_with_tvm([], [], ['TensorArrayReadV3_1:0'], mode='debug')
-            compare_tf_with_tvm([], [], ['TensorArrayReadV3_2:0'], mode='debug')
-    for dtype in tf_dtypes.keys():
-        run(dtype)
-
-# TODO(wweic): Fix gather issue with PartialEvaluate
-# def test_tensor_array_gather():
-#     with tf.Graph().as_default():
-#         dtype = 'float32'
-#         t = tf.constant([[1.0], [2.0], [3.0]])
-#         scatter_indices = tf.constant([2, 1, 0])
-#         gather_indices = tf.constant([1, 2])
-#         ta1 = tf.TensorArray(dtype=tf.float32, size=3, infer_shape=False, dynamic_size=False)
-#         ta2 = ta1.scatter(scatter_indices, t)
-#         t1 = ta2.gather(gather_indices)
-#         g = tf.get_default_graph()
-#         compare_tf_with_tvm([], [], ['TensorArrayGatherV3:0'], mode='debug')
+            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')
+    for dtype in ["float32", "int8"]:
+        run(dtype, False)
+        run(dtype, True)
+
+
+def test_tensor_array_gather():
+    def run(dtype_str, infer_shape):
+        with tf.Graph().as_default():
+            dtype =  tf_dtypes[dtype_str]
+            t = tf.constant(np.array([[1.0], [2.0], [3.0]]).astype(dtype_str))
+            scatter_indices = tf.constant([2, 1, 0])
+            gather_indices = tf.constant([1, 2])
+            ta1 = tf.TensorArray(dtype=dtype, size=3, infer_shape=infer_shape)
+            ta2 = ta1.scatter(scatter_indices, t)
+            t1 = ta2.gather(gather_indices)
+            g = tf.get_default_graph()
+            compare_tf_with_tvm([], [], ['TensorArrayGatherV3:0'], mode='vm')
+    for dtype in ["float32", "int8"]:
+        run(dtype, True)
 
 
 def test_tensor_array_split():
-    def run(dtype_str):
+    def run(dtype_str, infer_shape):
         with tf.Graph().as_default():
             dtype =  tf_dtypes[dtype_str]
             t = tf.constant(np.array([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0], [8.0]]).astype(dtype_str), dtype=dtype)
             split_length = tf.constant([2, 2, 2, 2], dtype=tf.int32)
-            ta1 = tf.TensorArray(dtype=dtype, size=4,
-                                 infer_shape=False, dynamic_size=False)
+            ta1 = tf.TensorArray(dtype=dtype, size=4, infer_shape=infer_shape)
             ta2 = ta1.split(t, split_length)
             out0 = ta2.read(0)
             out1 = ta2.read(1)
@@ -906,56 +916,76 @@ def run(dtype_str):
             compare_tf_with_tvm([], [], ['TensorArrayReadV3_1:0'], mode='debug')
             compare_tf_with_tvm([], [], ['TensorArrayReadV3_2:0'], mode='debug')
             compare_tf_with_tvm([], [], ['TensorArrayReadV3_3:0'], mode='debug')
-    for dtype in tf_dtypes.keys():
-        run(dtype)
+    for dtype in ["float32", "int8"]:
+        run(dtype, False)
+        run(dtype, True)
 
 
 def test_tensor_array_concat():
-    def run(dtype_str):
+    def run(dtype_str, infer_shape):
         with tf.Graph().as_default():
             dtype = tf_dtypes[dtype_str]
             t = tf.constant(np.array([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0], [8.0]]).astype(dtype_str), dtype=dtype)
             split_length = tf.constant([2, 2, 2, 2], dtype=tf.int32)
             ta1 = tf.TensorArray(dtype=dtype, size=4,
-                                 infer_shape=False, dynamic_size=False)
+                                 infer_shape=infer_shape)
             ta2 = ta1.split(t, split_length)
             t = ta2.concat()
             out = tf.identity(t)
             compare_tf_with_tvm([], [], ['Identity:0'], mode='debug')
-    for dtype in tf_dtypes.keys():
-        run(dtype)
+    for dtype in ["float32", "int8"]:
+        run(dtype, False)
+        run(dtype, True)
 
 
 def test_tensor_array_size():
-    def run(dtype_str):
+    def run(dtype_str, infer_shape):
         with tf.Graph().as_default():
             dtype =  tf_dtypes[dtype_str]
-            ta1 = tf.TensorArray(dtype=dtype, size=2, infer_shape=False, dynamic_size=False)
+            ta1 = tf.TensorArray(dtype=dtype, size=2, infer_shape=infer_shape)
             out = ta1.size()
             g = tf.get_default_graph()
             compare_tf_with_tvm([], [], 'TensorArraySizeV3:0', mode='debug')
-    for dtype in tf_dtypes.keys():
-        run(dtype)
+    for dtype in ["float32", "int8"]:
+        run(dtype, False)
+        run(dtype, True)
+
+
+def test_tensor_array_stack():
+    def run(dtype_str, infer_shape):
+        with tf.Graph().as_default():
+            dtype =  tf_dtypes[dtype_str]
+            t = tf.constant(np.array([[1.0], [2.0], [3.0]]).astype(dtype_str))
+            scatter_indices = tf.constant([2, 1, 0])
+            ta1 = tf.TensorArray(dtype=dtype, size=3, infer_shape=infer_shape)
+            ta2 = ta1.scatter(scatter_indices, t)
+            t1 = ta2.stack()
+            print(t1)
+            g = tf.get_default_graph()
+
+            compare_tf_with_tvm([], [], ['TensorArrayStack/TensorArrayGatherV3:0'], mode='vm')
+    for dtype in ["float32", "int8"]:
+        run(dtype, True)
+
 
 def test_tensor_array_unstack():
-    def run(dtype_str, input_shape):
+    def run(dtype_str, input_shape, infer_shape):
         with tf.Graph().as_default():
             dtype = tf_dtypes[dtype_str]
             t = tf.constant(np.random.choice([0, 1, 2, 3],
                                              size=input_shape).astype(dtype.name))
-            ta1 = tf.TensorArray(dtype=dtype, infer_shape=False, size=input_shape[0])
+            ta1 = tf.TensorArray(dtype=dtype, infer_shape=infer_shape, size=input_shape[0])
             ta2 = ta1.unstack(t)
             out0 = ta2.size()
             out1 = ta2.read(0)
             compare_tf_with_tvm([], [], 'TensorArraySizeV3:0', mode='debug')
             compare_tf_with_tvm([], [], 'TensorArrayReadV3:0', mode='debug')
-    for dtype in tf_dtypes.keys():
-        run(dtype, (5,))
-        run(dtype, (5, 5))
-        run(dtype, (5, 5, 5))
-        run(dtype, (5, 5, 5, 5))
-        run(dtype, (5, 5, 5, 5, 5))
-        run(dtype, (5, 5, 5, 5, 5, 5))
+    for dtype in ["float32", "int8"]:
+        run(dtype, (5,), False)
+        run(dtype, (5, 5), True)
+        run(dtype, (5, 5, 5), False)
+        run(dtype, (5, 5, 5, 5), True)
+
 
 #######################################################################
 # ConcatV2
@@ -3300,6 +3330,16 @@ def test_forward_isfinite():
     test_forward_reduce_any()
     test_forward_reduce_min()
 
+    # TensorArray
+    test_tensor_array_write_read()
+    test_tensor_array_concat()
+    test_tensor_array_scatter()
+    test_tensor_array_gather()
+    test_tensor_array_size()
+    test_tensor_array_split()
+    test_tensor_array_stack()
+    test_tensor_array_unstack()
+
     # General
     test_forward_multi_input()
     test_forward_multi_output()
diff --git a/topi/python/topi/util.py b/topi/python/topi/util.py
index 681535761f83..50a6a36edc46 100644
--- a/topi/python/topi/util.py
+++ b/topi/python/topi/util.py
@@ -166,12 +166,14 @@ def get_const_tuple(in_tuple):
     """
     ret = []
     for elem in in_tuple:
-        if isinstance(elem, tvm.tir.Var):
+        if isinstance(elem, (tvm.tir.Var, tvm.tir.expr.Any)):
             ret.append(elem)
         elif not isinstance(elem, (tvm.tir.IntImm, int)):
             elem = tvm.tir.ir_pass.Simplify(elem)
             if not isinstance(elem, tvm.tir.IntImm):
                 ret.append(elem)
+            else:
+                ret.append(get_const_int(elem))
         else:
             ret.append(get_const_int(elem))
     return tuple(ret)

From 1eb561e9d7892daf3dd34cd93bf5f8a3b3105239 Mon Sep 17 00:00:00 2001
From: Wang Yao <kevinthesunwy@gmail.com>
Date: Sun, 5 Apr 2020 15:24:19 -0700
Subject: [PATCH 2/6] Fix pylint

---
 python/tvm/relay/frontend/common.py           | 10 +++
 python/tvm/relay/frontend/tensorflow.py       | 61 ++++++++++++++-----
 .../frontend/tensorflow/test_forward.py       | 10 +--
 3 files changed, 62 insertions(+), 19 deletions(-)

diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
index 76d4eb6d3c37..b71d4576209c 100644
--- a/python/tvm/relay/frontend/common.py
+++ b/python/tvm/relay/frontend/common.py
@@ -29,6 +29,7 @@
 from .. import transform as _transform
 from .. import op as _op
 from .. import analysis
+from ..ty import Any
 
 
 class RequiredAttr(object):
@@ -469,6 +470,15 @@ def infer_type(node, mod=None):
         entry = new_mod["main"]
         return entry if isinstance(node, _function.Function) else entry.body
 
+def infer_channels(inputs, transpose=False):
+    """A hack for getting 'channels' or 'units' since caffe2 does not provide
+    these attributes. We check the shape of weights provided to get the number.
+    """
+    out_type = infer_type(inputs)
+    out_shapes = [get_const_tuple(out_type.checked_type.shape)]
+    channels = out_shapes[0][0] if not transpose else out_shapes[0][1]
+    return channels
+
 
 def infer_shape(inputs, mod=None):
     """A method to get the output type of an intermediate node in the graph."""
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 0cea94be0563..e57b6eb6feb6 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -33,6 +33,7 @@
 from .. import expr as _expr
 from .. import function as _function
 from .. import op as _op
+from ..ty import Any
 from ..expr_functor import ExprMutator, ExprVisitor
 from .common import AttrCvt, get_relay_op
 from .common import infer_type as _infer_type
@@ -788,12 +789,38 @@ def _impl(inputs, attr, params, mod):
 
 def _tensor_array():
     def _impl(inputs, attr, params, prelude):
+        try:
+            from tensorflow.python.framework import tensor_util
+        except ImportError as e:
+            raise ImportError(
+                "Unable to import tensorflow which is required {}".format(e))
+
         dtype_str = attr.get('dtype').name
         assert not attr["dynamic_size"], "Dynamic size tensor array is " \
                                          "not supported in TVM yet."
-        if attr['identical_element_shapes']:
-            # For static tensor array, we need to create with static constructor.
-            # However, at this point we might not know the input shape yet.
+
+        raw_elem_shape = tensor_util.TensorShapeProtoToList(attr['element_shape'])
+        elem_shape = []
+        for dim in raw_elem_shape:
+            if dim < 0:
+                elem_shape.append(Any())
+            else:
+                elem_shape.append(dim)
+
+        if elem_shape:
+            # Element shape is specified.
+            # Directly create static tensor array with given shape.
+            static_tensor_array_ops = StaticTensorArrayOps(prelude,
+                                                           dtype_str,
+                                                           elem_shape)
+            static_tensor_array_ops.register()
+            tensor_array_constructor = prelude.get_var_static('tensor_array',
+                                                              dtype_str,
+                                                              elem_shape)
+            tensor_array = tensor_array_constructor(inputs[0])
+            _static_tensor_array_map[tensor_array] = tensor_array
+        elif attr['identical_element_shapes']:
+            # identical_element_shapes is set but element shape is not given.
             # We create a static tensor array with dummy shape and record it in
             # _static_tensor_array_map. Later when creating other tensor array ops
             # which uses this tensor array, we reconstruct this tensor array with
@@ -806,11 +833,11 @@ def _impl(inputs, attr, params, prelude):
             tensor_array_constructor = prelude.get_var_static('tensor_array',
                                                               dtype_str,
                                                               dummy_shape)
-            tensor_array =  tensor_array_constructor(inputs[0])
+            tensor_array = tensor_array_constructor(inputs[0])
             _static_tensor_array_map[tensor_array] = None
         else:
             tensor_array_constructor = prelude.get_var('tensor_array', dtype_str)
-            tensor_array =  tensor_array_constructor(inputs[0])
+            tensor_array = tensor_array_constructor(inputs[0])
         return tensor_array
     return _impl
 
@@ -841,7 +868,7 @@ def _impl(inputs, attr, params, prelude):
                 ta_constructor = prelude.get_var_static('tensor_array',
                                                         dtype_str,
                                                         input_t_shape)
-                new_ta =  ta_constructor(input_ta.args[0])
+                new_ta = ta_constructor(input_ta.args[0])
                 _static_tensor_array_map[input_ta] = new_ta
                 input_ta = new_ta
 
@@ -872,7 +899,7 @@ def _impl(inputs, attr, params, prelude):
 
         if input_shape is None:
             gather_func = prelude.get_var('tensor_array_gather', dtype_str)
-            return gather_func(inputs[2], inputs[1])
+            out = gather_func(inputs[2], inputs[1])
         else:
             static_tensor_array_ops = StaticTensorArrayOps(prelude,
                                                            dtype_str,
@@ -889,7 +916,7 @@ def _impl(inputs, attr, params, prelude):
                 get_data_func = prelude.get_var_static('tensor_get_data',
                                                        dtype_str,
                                                        input_shape)
-                return get_data_func(out_tensor_t)
+                out = get_data_func(out_tensor_t)
             else:
                 # For fixed length indices, directly generate static shape output
                 read_func = prelude.get_var_static('tensor_array_read',
@@ -905,8 +932,9 @@ def _impl(inputs, attr, params, prelude):
                     out_tensor = get_data_func(read_func(inputs[2], index))
                     tensor_list.append(_op.expand_dims(out_tensor, axis=0))
 
-                return _op.concatenate(tensor_list, axis=0)
+                out = _op.concatenate(tensor_list, axis=0)
 
+        return out
     return _impl
 
 def _tensor_array_size():
@@ -939,7 +967,7 @@ def _impl(inputs, attr, params, prelude):
                 ta_constructor = prelude.get_var_static('tensor_array',
                                                         dtype_str,
                                                         input_t_shape)
-                new_ta =  ta_constructor(input_ta.args[0])
+                new_ta = ta_constructor(input_ta.args[0])
                 _static_tensor_array_map[input_ta] = new_ta
                 input_ta = new_ta
                 input_ta_shape = input_t_shape
@@ -970,7 +998,7 @@ def _impl(inputs, attr, params, prelude):
 
         if input_shape is None:
             read_func = prelude.get_var('tensor_array_read', dtype_str)
-            return read_func(inputs[2], _op.take(inputs[1], tvm.relay.const(0)))
+            out = read_func(inputs[2], _op.take(inputs[1], tvm.relay.const(0)))
         else:
             static_tensor_array_ops = StaticTensorArrayOps(prelude,
                                                            dtype_str,
@@ -982,7 +1010,9 @@ def _impl(inputs, attr, params, prelude):
             get_data_func = prelude.get_var_static('tensor_get_data',
                                                    dtype_str,
                                                    input_shape)
-            return get_data_func(out_tensor)
+            out = get_data_func(out_tensor)
+
+        return out
     return _impl
 
 def _tensor_array_split():
@@ -1012,7 +1042,7 @@ def _impl(inputs, attr, params, prelude):
                 ta_constructor = prelude.get_var_static('tensor_array',
                                                         dtype_str,
                                                         input_ta_shape)
-                new_ta =  ta_constructor(input_ta.args[0])
+                new_ta = ta_constructor(input_ta.args[0])
                 _static_tensor_array_map[input_ta] = new_ta
                 input_ta = new_ta
             else:
@@ -1056,7 +1086,7 @@ def _impl(inputs, attr, params, prelude):
 
         if input_shape is None:
             concat_func = prelude.get_var('tensor_array_concat', dtype_str)
-            return concat_func(inputs[1])
+            out = concat_func(inputs[1])
         else:
             static_tensor_array_ops = StaticTensorArrayOps(prelude,
                                                            dtype_str,
@@ -1068,8 +1098,9 @@ def _impl(inputs, attr, params, prelude):
             get_data_func = prelude.get_var_static('tensor_get_data',
                                                    dtype_str,
                                                    input_shape)
-            return get_data_func(out_tensor)
+            out = get_data_func(out_tensor)
 
+        return out
     return _impl
 
 def _tile():
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index bc51bbb0c725..d73e900c01af 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -843,14 +843,15 @@ def test_forward_squeeze():
 # TensorArray
 # -----------
 def test_tensor_array_write_read():
-    def run(dtype_str, infer_shape):
+    def run(dtype_str, infer_shape, element_shape):
         with tf.Graph().as_default():
             dtype = tf_dtypes[dtype_str]
             np_data = np.array([[1.0, 2.0], [3.0, 4.0]]).astype(dtype_str)
             in_data = [np_data, np_data]
             t1 = tf.constant(np_data, dtype=dtype)
             t2 = tf.constant(np_data, dtype=dtype)
-            ta1 = tf.TensorArray(dtype=dtype, size=2, infer_shape=infer_shape)
+            ta1 = tf.TensorArray(dtype=dtype, size=2, infer_shape=infer_shape,
+                                 element_shape=element_shape)
             ta2 = ta1.write(0, t1)
             ta3 = ta2.write(1, t2)
             out = ta3.read(0)
@@ -858,8 +859,9 @@ def run(dtype_str, infer_shape):
             compare_tf_with_tvm([], [], 'TensorArrayReadV3:0', mode='vm')
 
     for dtype in ["float32", "int8"]:
-        run(dtype, False)
-        run(dtype, True)
+        run(dtype, False, None)
+        run(dtype, False, tf.TensorShape([None, 2]))
+        run(dtype, True, None)
 
 
 def test_tensor_array_scatter():

From 978656a96546840e9ed602850e73090f6317fea0 Mon Sep 17 00:00:00 2001
From: Yao Wang <kevinthesunwy@gmail.com>
Date: Sun, 5 Apr 2020 22:55:00 +0000
Subject: [PATCH 3/6] Fix lint

---
 python/tvm/relay/frontend/common.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
index b71d4576209c..38955f23efb2 100644
--- a/python/tvm/relay/frontend/common.py
+++ b/python/tvm/relay/frontend/common.py
@@ -461,14 +461,16 @@ def infer_type(node, mod=None):
         mod["main"] = _function.Function([], node)
         mod = _transform.InferType()(mod)
         entry = mod["main"]
-        return entry.body
+        ret = entry.body
     else:
         new_mod = IRModule.from_expr(node)
         if mod is not None:
             new_mod.update(mod)
             new_mod = _transform.InferType()(new_mod)
         entry = new_mod["main"]
-        return entry if isinstance(node, _function.Function) else entry.body
+        ret = entry if isinstance(node, _function.Function) else entry.body
+
+    return ret
 
 def infer_channels(inputs, transpose=False):
     """A hack for getting 'channels' or 'units' since caffe2 does not provide

From a8d092cd105ebccd98631c56c8b7078eeb45ab84 Mon Sep 17 00:00:00 2001
From: Yao Wang <kevinthesunwy@gmail.com>
Date: Mon, 6 Apr 2020 18:07:24 +0000
Subject: [PATCH 4/6] Move get_tensor_array_shape into prelude

---
 python/tvm/relay/frontend/common.py     | 23 ------------
 python/tvm/relay/frontend/tensorflow.py | 15 ++++----
 python/tvm/relay/prelude.py             | 49 +++++++++++++++++++++++--
 3 files changed, 53 insertions(+), 34 deletions(-)

diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
index 38955f23efb2..c62b50b7bb9c 100644
--- a/python/tvm/relay/frontend/common.py
+++ b/python/tvm/relay/frontend/common.py
@@ -29,7 +29,6 @@
 from .. import transform as _transform
 from .. import op as _op
 from .. import analysis
-from ..ty import Any
 
 
 class RequiredAttr(object):
@@ -558,28 +557,6 @@ def new_var(name_hint,
     return _expr.var(name_hint, type_annotation, shape, dtype)
 
 
-def check_tensor_array_shape(expr, dtype, mod):
-    """Check whether a tensor array has fixed rank shape.
-    Return its shape if yes. Otherwise, return None.
-    """
-    checked_type = infer_shape(expr, mod)
-    assert isinstance(checked_type, TypeCall), "Input must be a tensor array."
-    ta_type_str = checked_type.args[0].func.name_hint
-    static_ta_ty_start = "static_tensor_{}".format(dtype)
-    if ta_type_str.startswith(static_ta_ty_start):
-        shape_str = ta_type_str.replace("{}_".format(static_ta_ty_start), '') \
-            .replace("_t", '')
-        shape = []
-        if "scalar" not in shape_str:
-            for dim_str in shape_str.split("_"):
-                if dim_str == "?":
-                    shape.append(Any())
-                else:
-                    shape.append(int(dim_str))
-        return tuple(shape)
-    return None
-
-
 class Renamer(object):
     """A simply renamer for operators.
 
diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index e57b6eb6feb6..2214a50b2f3a 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -26,7 +26,7 @@
 import tvm
 
 from tvm.ir import IRModule
-from tvm.relay.prelude import Prelude, StaticTensorArrayOps
+from tvm.relay.prelude import Prelude, StaticTensorArrayOps, get_tensor_array_shape
 from tvm.ir import structural_hash as s_hash
 
 from .. import analysis
@@ -41,7 +41,6 @@
 from .common import infer_channels as _infer_channels
 from .common import infer_value as _infer_value
 from .common import infer_value_simulated as _infer_value_simulated
-from .common import check_tensor_array_shape as _check_tensor_array_shape
 
 __all__ = ['from_tensorflow']
 
@@ -845,7 +844,7 @@ def _tensor_array_scatter():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr.get('T').name
         input_ta = inputs[0]
-        input_shape = _check_tensor_array_shape(input_ta, dtype_str, prelude.mod)
+        input_shape = get_tensor_array_shape(input_ta, dtype_str, prelude)
         values_shape = _infer_shape(inputs[2], prelude.mod)
         input_t_shape = values_shape[1:]
         indices_shape = _infer_shape(inputs[1], prelude.mod)
@@ -894,7 +893,7 @@ def _impl(inputs, attr, params, prelude):
 def _tensor_array_gather():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr.get('dtype').name
-        input_shape = _check_tensor_array_shape(inputs[2], dtype_str, prelude.mod)
+        input_shape = get_tensor_array_shape(inputs[2], dtype_str, prelude)
         indices_shape = _infer_shape(inputs[1], prelude.mod)
 
         if input_shape is None:
@@ -946,7 +945,7 @@ def _tensor_array_write():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr.get('T').name
         input_ta = inputs[3]
-        input_ta_shape = _check_tensor_array_shape(input_ta, dtype_str, prelude.mod)
+        input_ta_shape = get_tensor_array_shape(input_ta, dtype_str, prelude)
         input_t_shape = _infer_shape(inputs[2], prelude.mod)
         input_rank = len(input_t_shape)
 
@@ -994,7 +993,7 @@ def _impl(inputs, attr, params, prelude):
 def _tensor_array_read():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr['dtype'].name
-        input_shape = _check_tensor_array_shape(inputs[2], dtype_str, prelude.mod)
+        input_shape = get_tensor_array_shape(inputs[2], dtype_str, prelude)
 
         if input_shape is None:
             read_func = prelude.get_var('tensor_array_read', dtype_str)
@@ -1019,7 +1018,7 @@ def _tensor_array_split():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr.get('T').name
         input_ta = inputs[0]
-        input_ta_shape = _check_tensor_array_shape(input_ta, dtype_str, prelude.mod)
+        input_ta_shape = get_tensor_array_shape(input_ta, dtype_str, prelude)
         input_t_shape = _infer_shape(inputs[1], prelude.mod)
         input_rank = len(input_t_shape)
         lengths = _op.cast(inputs[2], 'int32')
@@ -1082,7 +1081,7 @@ def _impl(inputs, attr, params, prelude):
 def _tensor_array_concat():
     def _impl(inputs, attr, params, prelude):
         dtype_str = attr['dtype'].name
-        input_shape = _check_tensor_array_shape(inputs[1], dtype_str, prelude.mod)
+        input_shape = get_tensor_array_shape(inputs[1], dtype_str, prelude)
 
         if input_shape is None:
             concat_func = prelude.get_var('tensor_array_concat', dtype_str)
diff --git a/python/tvm/relay/prelude.py b/python/tvm/relay/prelude.py
index 47c3ba7b43b0..243eace0fb94 100644
--- a/python/tvm/relay/prelude.py
+++ b/python/tvm/relay/prelude.py
@@ -16,7 +16,7 @@
 # under the License.
 # pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name
 """A prelude containing useful global functions and ADT definitions."""
-from tvm.ir import IRModule
+from tvm.ir import IRModule, TypeCall
 
 from .ty import GlobalTypeVar, TensorType, Any, scalar_type
 from .expr import Var, GlobalVar, If, const
@@ -24,8 +24,51 @@
 from .op.tensor import add, subtract, equal
 from .adt import Constructor, TypeData, Clause, Match
 from .adt import PatternConstructor, PatternVar, PatternWildcard
-from . import op
-
+from . import op, transform
+
+
+def get_tensor_array_shape(expr, dtype, prelude):
+    """Get the static shape of a tensor array if it has fixed rank shape.
+
+    By design, static ADT tensor in TVM has type name in the format
+    of static_tensor_dim0_dim1_..._dimN_t.
+
+    Parameters
+    ----------
+    expr : Relay Expr
+        Input expression.
+
+    dtype : str
+        Data type.
+
+    prelude : Prelude
+        Tensor array prelude
+
+    Returns
+    -------
+    shape : tuple of (int, Any) or None
+        The output shape. None if input tensor array
+        has dynamic shape.
+    """
+    mod = prelude.mod
+    mod["main"] = Function([], expr)
+    mod = transform.InferType()(mod)
+    checked_type = mod["main"].body.checked_type
+    assert isinstance(checked_type, TypeCall), "Input must be a tensor array."
+    ta_type_str = checked_type.args[0].func.name_hint
+    static_ta_ty_start = "static_tensor_{}".format(dtype)
+    if ta_type_str.startswith(static_ta_ty_start):
+        shape_str = ta_type_str.replace("{}_".format(static_ta_ty_start), '') \
+            .replace("_t", '')
+        shape = []
+        if "scalar" not in shape_str:
+            for dim_str in shape_str.split("_"):
+                if dim_str == "?":
+                    shape.append(Any())
+                else:
+                    shape.append(int(dim_str))
+        return tuple(shape)
+    return None
 
 def _get_name_static(canonical, dtype, shape):
     """Get name for static shape tensor array op corresponding

From bfc5a772c2ac92d85d819ec12f1f7d79ae4f8b9a Mon Sep 17 00:00:00 2001
From: Yao Wang <kevinthesunwy@gmail.com>
Date: Mon, 6 Apr 2020 18:09:35 +0000
Subject: [PATCH 5/6] Fix lint

---
 python/tvm/relay/frontend/common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
index c62b50b7bb9c..2790ba328799 100644
--- a/python/tvm/relay/frontend/common.py
+++ b/python/tvm/relay/frontend/common.py
@@ -21,7 +21,7 @@
 import numpy as np
 
 import tvm
-from tvm.ir import IRModule, TypeCall
+from tvm.ir import IRModule
 from topi.util import get_const_tuple
 
 from .. import expr as _expr

From 51e79727187f3c71e94f5018e8d1f873c6781459 Mon Sep 17 00:00:00 2001
From: Yao Wang <kevinthesunwy@gmail.com>
Date: Fri, 10 Apr 2020 19:03:50 +0000
Subject: [PATCH 6/6] Fix common

---
 python/tvm/relay/frontend/common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/tvm/relay/frontend/common.py b/python/tvm/relay/frontend/common.py
index 2790ba328799..e86890f3639a 100644
--- a/python/tvm/relay/frontend/common.py
+++ b/python/tvm/relay/frontend/common.py
@@ -487,7 +487,7 @@ def infer_shape(inputs, mod=None):
     checked_type = out_type.checked_type
     if hasattr(checked_type, 'shape'):
         # Regular operator that outputs tensors
-        return get_const_tuple(out_type.checked_type.shape)
+        return get_const_tuple(checked_type.shape)
     # The return type is not a tensor, for example List
     return checked_type