From 2d30945979532a7887016e8d8ff5f4181a28a3b5 Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Tue, 18 May 2021 23:55:04 +0000
Subject: [PATCH 01/10] added test infrastructure for frozen TF2 models

---
 tests/python/frontend/tensorflow2/common.py   | 111 ++++++
 .../tensorflow2/test_functional_models.py     | 364 ++++++++++++++++++
 .../tensorflow2/test_sequential_models.py     | 108 ++++++
 3 files changed, 583 insertions(+)
 create mode 100644 tests/python/frontend/tensorflow2/common.py
 create mode 100644 tests/python/frontend/tensorflow2/test_functional_models.py
 create mode 100644 tests/python/frontend/tensorflow2/test_sequential_models.py

diff --git a/tests/python/frontend/tensorflow2/common.py b/tests/python/frontend/tensorflow2/common.py
new file mode 100644
index 000000000000..338df00ca544
--- /dev/null
+++ b/tests/python/frontend/tensorflow2/common.py
@@ -0,0 +1,111 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-self, invalid-name, unused-argument, too-many-lines, len-as-condition, broad-except
+# pylint: disable=import-outside-toplevel, redefined-builtin
+"""TF2 to relay converter test utilities"""
+
+import tvm
+from tvm import relay
+
+from tvm.runtime.vm import VirtualMachine
+import tvm.contrib.graph_executor as runtime
+from tvm.relay.frontend.tensorflow import from_tensorflow
+import tvm.testing
+
+import tensorflow as tf
+from tensorflow.python.eager.def_function import Function
+
+def vmobj_to_list(o):
+    if isinstance(o, tvm.nd.NDArray):
+        out = o.asnumpy().tolist()
+    elif isinstance(o, tvm.runtime.container.ADT):
+        result = []
+        for f in o:
+            result.append(vmobj_to_list(f))
+        out = result
+    else:
+        raise RuntimeError("Unknown object type: %s" % type(o))
+    return out
+
+def run_tf_code(func, input_):
+    if type(func) is Function:
+        out = func(input_)
+        if isinstance(out, list):
+            a = [x.numpy() for x in out]
+        else:
+            a = out.numpy()
+    else:
+        a = func(tf.constant(input_))
+        if type(a) is dict:
+            a = [x.numpy() for x in a.values()]
+            if len(a) == 1:
+                a = a[0]
+        elif type(a) is list:
+            a = [x.numpy() for x in a]
+            if len(a) == 1:
+                a = a[0]
+        else:
+            a = a.numpy()
+    return a
+
+def compile_graph_runtime(mod, params, target = "llvm", target_host = "llvm",
+            opt_level=3, output_sig=None):
+    with tvm.transform.PassContext(opt_level):
+        lib = relay.build(mod, target=target, target_host=target_host, params=params)
+    return lib
+
+def compile_vm(mod, params, target = "llvm", target_host = "llvm",
+               opt_level=3, disabled_pass=None, output_sig=None):
+    with tvm.transform.PassContext(opt_level, disabled_pass=disabled_pass):
+        mod = relay.transform.InferType()(mod)
+        vm_exec = relay.vm.compile(mod, target, target_host, params=params)
+    return vm_exec
+
+def run_vm(vm_exec, input_, ctx = tvm.cpu(0)):
+    vm = VirtualMachine(vm_exec, ctx)
+    _out = vm.invoke("main", input_)
+    return vmobj_to_list(_out)
+
+def run_graph(lib, input_, ctx = tvm.cpu(0)):
+    mod = runtime.GraphModule(lib["default"](ctx))
+    mod.set_input(0, input_)
+    mod.run()
+    _out = mod.get_output(0).asnumpy()
+    return _out
+
+def compare_tf_tvm(gdef, input_, output_, vm=True, output_sig=None):
+    """ compare tf and tvm execution for the same input.
+
+    Parameters
+    ----------
+    func: tf function. can be from saved model or not. different ways to pass input
+        from saved model: <class 'tensorflow.python.saved_model.load._WrapperFunction'>
+        not from saved model:  <class 'tensorflow.python.eager.def_function.Function'>
+
+    mod: compiled relay module (vm or graph runtime). converted from tf func.
+
+    input_: a single numpy array object
+
+    """
+    mod, params = from_tensorflow(gdef, outputs=output_sig)
+    if vm:
+        exec_ = compile_vm(mod, params, output_sig=output_sig)
+        tvm_out = run_vm(exec_, input_)
+    else:
+        lib = compile_graph_runtime(mod, params, output_sig=output_sig)
+        tvm_out = run_graph(lib, input_)
+    tvm.testing.assert_allclose(output_, tvm_out, atol=1e-5)
diff --git a/tests/python/frontend/tensorflow2/test_functional_models.py b/tests/python/frontend/tensorflow2/test_functional_models.py
new file mode 100644
index 000000000000..ab0b9f56ae71
--- /dev/null
+++ b/tests/python/frontend/tensorflow2/test_functional_models.py
@@ -0,0 +1,364 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-self, invalid-name, unused-argument, too-many-lines, len-as-condition, broad-except
+# pylint: disable=import-outside-toplevel, redefined-builtin
+"""TF2 to relay converter test: tests basic examples"""
+
+import tempfile
+import tensorflow as tf
+import numpy as np
+import pytest
+from common import compare_tf_tvm
+from common import run_tf_code
+
+""" simple function to test x=x+1"""
+class AddOne(tf.Module):
+
+    def get_input(self):
+        return np.array(1.0, dtype='float32')
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
+    def func(self, x):
+        return x + 1
+
+class AddOne2D(AddOne):
+
+    def get_input(self):
+        return np.ones((2, 2), dtype='float32')
+
+    """2D array as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+    def func(self, x):
+        return x + 1
+
+class AddOne2DConstant(AddOne):
+
+    def get_input(self):
+        return np.ones((2, 2), dtype='float32')
+
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
+    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+    def func(self, x):
+        return x + np.ones((2, 2), dtype='float32')
+
+class SubOne2DConstant(tf.Module):
+
+    def get_input(self):
+        return np.ones((2, 2), dtype='float32')
+
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
+    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+    def func(self, x):
+        return x - np.ones((2, 2), dtype='float32')
+
+class MulOne2DConstant(tf.Module):
+
+    def get_input(self):
+        return np.ones((2, 2), dtype='float32')
+
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+    def func(self, x):
+        return x * np.ones((2, 2), dtype='float32')
+
+class DivOne2DConstant(tf.Module):
+
+    def get_input(self):
+        return np.ones((2, 2), dtype='float32')
+
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
+    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+    def func(self, x):
+        return x / np.ones((2, 2), dtype='float32')
+
+
+class StridedSlice(tf.Module):
+    def get_input(self):
+        return np.ones((3,2,3), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(3,2,3), dtype=tf.float32)])
+    def func(self, x):
+        return tf.strided_slice(x, [1, 0, 0], [2, 1, 3], [1, 1, 1])
+
+
+class Shape(tf.Module):
+    def get_input(self):
+        return np.ones((3,2,3), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
+    def func(self, x):
+        a = tf.ones_like(tf.raw_ops.Shape(input=x), dtype=tf.float32)
+        return a + x
+
+
+class Pack(tf.Module):
+    def get_input(self):
+        return np.ones((2,3), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3), dtype=tf.float32)])
+    def func(self, x):
+        return tf.raw_ops.Pack(values=[x, x], axis=0)
+
+
+class Split(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        a,b,c = tf.split(x, 3, axis=1)
+        return tf.raw_ops.Pack(values=[a,b,c], axis=1)
+
+
+class Maximum(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        a,b = tf.split(x, 2, axis=1)
+        return tf.math.maximum(a, b, name=None)
+
+
+class Less(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        a,b = tf.split(x, 2, axis=1)
+        return tf.math.less(a, b, name=None)
+
+
+class Equal(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        a,b = tf.split(x, 2, axis=1)
+        return tf.math.equal(a, b, name=None)
+
+
+
+class Cast(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.cast(x, tf.int32)
+
+
+class ExpandDims(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.expand_dims(x, axis=2)
+
+
+class Transpose(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        x = tf.expand_dims(x, axis=2)
+        return tf.transpose(x, perm=[0,2,1])
+
+class Reshape(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.reshape(x, (1,2,15))
+
+
+class Tanh(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.math.tanh(x)
+
+
+class Sigmoid(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.math.sigmoid(x)
+
+
+class Relu(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.nn.relu(x)
+
+
+
+class Floor(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        return tf.math.floor(x)
+
+
+class FloorMod(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        a, b = tf.split(x, 2, axis=1)
+        return tf.math.floormod(a, b)
+
+class ConcatV2(tf.Module):
+    def get_input(self):
+        return np.ones((1,30), dtype=np.float32)
+
+    """scalar as input"""
+    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+    def func(self, x):
+        a,b,c = tf.split(x, 3, axis=1)
+        return tf.raw_ops.ConcatV2(values=[a, b, c], axis=1)
+
+
+def _function_graph(TestClass):
+    f = TestClass().func
+    gdef = f.get_concrete_function().graph.as_graph_def()
+    gdef_ops = list(set([n.op for n in gdef.node]))
+    input_ = TestClass().get_input()
+    output = run_tf_code(f, input_)
+    return gdef, input_, output
+
+
+def _model_graph(TestClass):
+    model = TestClass()
+    with tempfile.TemporaryDirectory() as model_path:
+        tf.saved_model.save(model, model_path)
+        imported = tf.saved_model.load(model_path)
+
+    f = imported.signatures["serving_default"]
+    gdef = f.graph.as_graph_def(add_shapes=True)
+
+    input_ = model.get_input()
+    output = run_tf_code(f, input_)
+    return gdef, input_, output
+
+
+def run_func_graph(TestClass, use_vm=False):
+    compare_tf_tvm(*_function_graph(TestClass), vm=use_vm)
+
+def run_model_graph(TestClass, output_sig=None):
+    compare_tf_tvm(*_model_graph(TestClass), vm=True, output_sig=output_sig)
+
+def run_all(TestClass):
+    run_model_graph(TestClass)
+    for use_vm in [True, False]:
+        run_func_graph(TestClass, use_vm=use_vm)
+
+
+def test_basic_ops():
+    run_all(AddOne)
+    run_all(AddOne2D)
+    run_all(AddOne2DConstant)
+    run_all(SubOne2DConstant)
+    run_all(MulOne2DConstant)
+    run_all(DivOne2DConstant)
+
+def test_strided_slice():
+    run_all(StridedSlice)
+
+def test_shape():
+    run_all(Shape)
+
+def test_pack():
+    run_all(Pack)
+
+def test_split():
+    run_all(Split)
+
+def test_max():
+    run_all(Maximum)
+
+def test_less():
+    run_all(Less)
+
+def test_equal():
+    run_all(Equal)
+
+def test_floor():
+    run_all(Floor)
+    run_all(FloorMod)
+
+def test_concat_v2():
+    run_all(ConcatV2)
+
+def test_cast():
+    run_all(Cast)
+
+def test_expand_dims():
+    run_all(ExpandDims)
+
+def test_transpose():
+    run_all(Transpose)
+
+def test_reshape():
+    run_all(Reshape)
+    
+def test_tanh():
+    run_all(Tanh)
+
+def test_sigmoid():
+    run_all(Sigmoid)
+
+def test_relu():
+    run_all(Relu)
+
+if __name__ == "__main__":
+    pytest.main([__file__])
diff --git a/tests/python/frontend/tensorflow2/test_sequential_models.py b/tests/python/frontend/tensorflow2/test_sequential_models.py
new file mode 100644
index 000000000000..702b58c5b9b7
--- /dev/null
+++ b/tests/python/frontend/tensorflow2/test_sequential_models.py
@@ -0,0 +1,108 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=import-self, invalid-name, unused-argument, too-many-lines, len-as-condition, broad-except
+# pylint: disable=import-outside-toplevel, redefined-builtin
+"""TF2 to relay converter test: x+1 as a very basic example"""
+
+import tempfile
+import numpy as np
+import pytest
+import tensorflow as tf
+from tensorflow.python.framework.convert_to_constants \
+    import convert_variables_to_constants_v2
+
+from common import compare_tf_tvm
+from common import run_tf_code
+
+
+def run_sequential_model(model_fn, input_shape):
+    def get_input(shape):
+        _input = np.random.uniform(0,1,shape).astype(dtype='float32')
+        return _input
+
+    def save_and_reload(_model):
+        with tempfile.TemporaryDirectory() as model_path:
+            tf.saved_model.save(_model, model_path)
+            loaded = tf.saved_model.load(model_path)
+            func = loaded.signatures["serving_default"]
+            frozen_func = convert_variables_to_constants_v2(func)
+        return frozen_func
+
+    def model_graph(model, input_shape):
+        _input = get_input(input_shape)
+        f = save_and_reload(model(input_shape))
+        _output = run_tf_code(f, _input)
+        gdef = f.graph.as_graph_def(add_shapes=True)
+        return gdef, _input, _output
+
+    compare_tf_tvm(*model_graph(model_fn, input_shape), vm=True, output_sig=None)
+
+
+def dense_model(input_shape, num_units=128):
+    return tf.keras.Sequential(
+        [
+            tf.keras.layers.Flatten(input_shape=input_shape[1:]),
+            tf.keras.layers.Dense(num_units)
+        ]
+    )
+
+def mnist_model(input_shape):
+    return tf.keras.Sequential([
+        tf.keras.layers.Flatten(input_shape=input_shape[1:]),
+        tf.keras.layers.Dense(128, activation="relu"),
+        tf.keras.layers.Dense(10)
+    ])
+
+
+def conv2d_model(input_shape, kernel=(3,3), filters=16):
+    model = tf.keras.Sequential([
+        tf.keras.layers.Input(shape=input_shape[1:], batch_size=1),
+        tf.keras.layers.Conv2D(filters, kernel)
+    ])
+    return model
+
+def maxpool_model(input_shape, pool_size=(2, 2)):
+    model = tf.keras.Sequential([
+        tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:])
+    ])
+    return model
+
+def maxpool_batchnorm_model(input_shape, pool_size=(2,2)):
+    model = tf.keras.Sequential([
+        tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:]),
+        tf.keras.layers.BatchNormalization()
+    ])
+    return model
+
+
+def test_dense_model():
+    run_sequential_model(dense_model,  input_shape = (1, 28, 28))
+
+def test_mnist_model():
+    run_sequential_model(mnist_model,  input_shape = (1, 28, 28))
+
+def test_conv2d_model():
+    run_sequential_model(conv2d_model,  input_shape = (1, 32, 32, 3))
+
+def test_maxpool_model():
+    run_sequential_model(maxpool_model, input_shape= (1, 32, 32, 3))
+
+def test_maxpool_batchnorm_model():
+    run_sequential_model(maxpool_batchnorm_model, input_shape= (1, 32, 32, 3))
+
+if __name__ == "__main__":
+    pytest.main([__file__])
\ No newline at end of file

From edb62ca7017aa8058921598966145e729d79b12e Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Tue, 18 May 2021 23:57:24 +0000
Subject: [PATCH 02/10] linting with black

---
 tests/python/frontend/tensorflow2/common.py   |  23 +++-
 .../tensorflow2/test_functional_models.py     | 130 ++++++++++++------
 .../tensorflow2/test_sequential_models.py     |  72 +++++-----
 3 files changed, 142 insertions(+), 83 deletions(-)

diff --git a/tests/python/frontend/tensorflow2/common.py b/tests/python/frontend/tensorflow2/common.py
index 338df00ca544..0eda019062c2 100644
--- a/tests/python/frontend/tensorflow2/common.py
+++ b/tests/python/frontend/tensorflow2/common.py
@@ -29,6 +29,7 @@
 import tensorflow as tf
 from tensorflow.python.eager.def_function import Function
 
+
 def vmobj_to_list(o):
     if isinstance(o, tvm.nd.NDArray):
         out = o.asnumpy().tolist()
@@ -41,6 +42,7 @@ def vmobj_to_list(o):
         raise RuntimeError("Unknown object type: %s" % type(o))
     return out
 
+
 def run_tf_code(func, input_):
     if type(func) is Function:
         out = func(input_)
@@ -62,33 +64,40 @@ def run_tf_code(func, input_):
             a = a.numpy()
     return a
 
-def compile_graph_runtime(mod, params, target = "llvm", target_host = "llvm",
-            opt_level=3, output_sig=None):
+
+def compile_graph_runtime(
+    mod, params, target="llvm", target_host="llvm", opt_level=3, output_sig=None
+):
     with tvm.transform.PassContext(opt_level):
         lib = relay.build(mod, target=target, target_host=target_host, params=params)
     return lib
 
-def compile_vm(mod, params, target = "llvm", target_host = "llvm",
-               opt_level=3, disabled_pass=None, output_sig=None):
+
+def compile_vm(
+    mod, params, target="llvm", target_host="llvm", opt_level=3, disabled_pass=None, output_sig=None
+):
     with tvm.transform.PassContext(opt_level, disabled_pass=disabled_pass):
         mod = relay.transform.InferType()(mod)
         vm_exec = relay.vm.compile(mod, target, target_host, params=params)
     return vm_exec
 
-def run_vm(vm_exec, input_, ctx = tvm.cpu(0)):
+
+def run_vm(vm_exec, input_, ctx=tvm.cpu(0)):
     vm = VirtualMachine(vm_exec, ctx)
     _out = vm.invoke("main", input_)
     return vmobj_to_list(_out)
 
-def run_graph(lib, input_, ctx = tvm.cpu(0)):
+
+def run_graph(lib, input_, ctx=tvm.cpu(0)):
     mod = runtime.GraphModule(lib["default"](ctx))
     mod.set_input(0, input_)
     mod.run()
     _out = mod.get_output(0).asnumpy()
     return _out
 
+
 def compare_tf_tvm(gdef, input_, output_, vm=True, output_sig=None):
-    """ compare tf and tvm execution for the same input.
+    """compare tf and tvm execution for the same input.
 
     Parameters
     ----------
diff --git a/tests/python/frontend/tensorflow2/test_functional_models.py b/tests/python/frontend/tensorflow2/test_functional_models.py
index ab0b9f56ae71..c448e63dc2ee 100644
--- a/tests/python/frontend/tensorflow2/test_functional_models.py
+++ b/tests/python/frontend/tensorflow2/test_functional_models.py
@@ -26,85 +26,91 @@
 from common import run_tf_code
 
 """ simple function to test x=x+1"""
-class AddOne(tf.Module):
 
+
+class AddOne(tf.Module):
     def get_input(self):
-        return np.array(1.0, dtype='float32')
+        return np.array(1.0, dtype="float32")
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
     def func(self, x):
         return x + 1
 
-class AddOne2D(AddOne):
 
+class AddOne2D(AddOne):
     def get_input(self):
-        return np.ones((2, 2), dtype='float32')
+        return np.ones((2, 2), dtype="float32")
 
     """2D array as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
         return x + 1
 
-class AddOne2DConstant(AddOne):
 
+class AddOne2DConstant(AddOne):
     def get_input(self):
-        return np.ones((2, 2), dtype='float32')
+        return np.ones((2, 2), dtype="float32")
 
     """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
-        return x + np.ones((2, 2), dtype='float32')
+        return x + np.ones((2, 2), dtype="float32")
 
-class SubOne2DConstant(tf.Module):
 
+class SubOne2DConstant(tf.Module):
     def get_input(self):
-        return np.ones((2, 2), dtype='float32')
+        return np.ones((2, 2), dtype="float32")
 
     """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
-        return x - np.ones((2, 2), dtype='float32')
+        return x - np.ones((2, 2), dtype="float32")
 
-class MulOne2DConstant(tf.Module):
 
+class MulOne2DConstant(tf.Module):
     def get_input(self):
-        return np.ones((2, 2), dtype='float32')
+        return np.ones((2, 2), dtype="float32")
 
     """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
-        return x * np.ones((2, 2), dtype='float32')
+        return x * np.ones((2, 2), dtype="float32")
 
-class DivOne2DConstant(tf.Module):
 
+class DivOne2DConstant(tf.Module):
     def get_input(self):
-        return np.ones((2, 2), dtype='float32')
+        return np.ones((2, 2), dtype="float32")
 
     """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
-        return x / np.ones((2, 2), dtype='float32')
+        return x / np.ones((2, 2), dtype="float32")
 
 
 class StridedSlice(tf.Module):
     def get_input(self):
-        return np.ones((3,2,3), dtype=np.float32)
+        return np.ones((3, 2, 3), dtype=np.float32)
 
     """scalar as input"""
-    @tf.function(input_signature=[tf.TensorSpec(shape=(3,2,3), dtype=tf.float32)])
+
+    @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
     def func(self, x):
         return tf.strided_slice(x, [1, 0, 0], [2, 1, 3], [1, 1, 1])
 
 
 class Shape(tf.Module):
     def get_input(self):
-        return np.ones((3,2,3), dtype=np.float32)
+        return np.ones((3, 2, 3), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
     def func(self, x):
         a = tf.ones_like(tf.raw_ops.Shape(input=x), dtype=tf.float32)
@@ -113,9 +119,10 @@ def func(self, x):
 
 class Pack(tf.Module):
     def get_input(self):
-        return np.ones((2,3), dtype=np.float32)
+        return np.ones((2, 3), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3), dtype=tf.float32)])
     def func(self, x):
         return tf.raw_ops.Pack(values=[x, x], axis=0)
@@ -123,54 +130,58 @@ def func(self, x):
 
 class Split(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
-        a,b,c = tf.split(x, 3, axis=1)
-        return tf.raw_ops.Pack(values=[a,b,c], axis=1)
+        a, b, c = tf.split(x, 3, axis=1)
+        return tf.raw_ops.Pack(values=[a, b, c], axis=1)
 
 
 class Maximum(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
-        a,b = tf.split(x, 2, axis=1)
+        a, b = tf.split(x, 2, axis=1)
         return tf.math.maximum(a, b, name=None)
 
 
 class Less(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
-        a,b = tf.split(x, 2, axis=1)
+        a, b = tf.split(x, 2, axis=1)
         return tf.math.less(a, b, name=None)
 
 
 class Equal(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
-        a,b = tf.split(x, 2, axis=1)
+        a, b = tf.split(x, 2, axis=1)
         return tf.math.equal(a, b, name=None)
 
 
-
 class Cast(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.cast(x, tf.int32)
@@ -178,9 +189,10 @@ def func(self, x):
 
 class ExpandDims(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.expand_dims(x, axis=2)
@@ -188,29 +200,33 @@ def func(self, x):
 
 class Transpose(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         x = tf.expand_dims(x, axis=2)
-        return tf.transpose(x, perm=[0,2,1])
+        return tf.transpose(x, perm=[0, 2, 1])
+
 
 class Reshape(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
-        return tf.reshape(x, (1,2,15))
+        return tf.reshape(x, (1, 2, 15))
 
 
 class Tanh(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.math.tanh(x)
@@ -218,9 +234,10 @@ def func(self, x):
 
 class Sigmoid(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.math.sigmoid(x)
@@ -228,20 +245,21 @@ def func(self, x):
 
 class Relu(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.nn.relu(x)
 
 
-
 class Floor(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.math.floor(x)
@@ -249,22 +267,25 @@ def func(self, x):
 
 class FloorMod(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b = tf.split(x, 2, axis=1)
         return tf.math.floormod(a, b)
 
+
 class ConcatV2(tf.Module):
     def get_input(self):
-        return np.ones((1,30), dtype=np.float32)
+        return np.ones((1, 30), dtype=np.float32)
 
     """scalar as input"""
+
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
-        a,b,c = tf.split(x, 3, axis=1)
+        a, b, c = tf.split(x, 3, axis=1)
         return tf.raw_ops.ConcatV2(values=[a, b, c], axis=1)
 
 
@@ -294,9 +315,11 @@ def _model_graph(TestClass):
 def run_func_graph(TestClass, use_vm=False):
     compare_tf_tvm(*_function_graph(TestClass), vm=use_vm)
 
+
 def run_model_graph(TestClass, output_sig=None):
     compare_tf_tvm(*_model_graph(TestClass), vm=True, output_sig=output_sig)
 
+
 def run_all(TestClass):
     run_model_graph(TestClass)
     for use_vm in [True, False]:
@@ -311,54 +334,71 @@ def test_basic_ops():
     run_all(MulOne2DConstant)
     run_all(DivOne2DConstant)
 
+
 def test_strided_slice():
     run_all(StridedSlice)
 
+
 def test_shape():
     run_all(Shape)
 
+
 def test_pack():
     run_all(Pack)
 
+
 def test_split():
     run_all(Split)
 
+
 def test_max():
     run_all(Maximum)
 
+
 def test_less():
     run_all(Less)
 
+
 def test_equal():
     run_all(Equal)
 
+
 def test_floor():
     run_all(Floor)
     run_all(FloorMod)
 
+
 def test_concat_v2():
     run_all(ConcatV2)
 
+
 def test_cast():
     run_all(Cast)
 
+
 def test_expand_dims():
     run_all(ExpandDims)
 
+
 def test_transpose():
     run_all(Transpose)
 
+
 def test_reshape():
     run_all(Reshape)
-    
+
+
 def test_tanh():
     run_all(Tanh)
 
+
 def test_sigmoid():
     run_all(Sigmoid)
 
+
 def test_relu():
     run_all(Relu)
 
+
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/tests/python/frontend/tensorflow2/test_sequential_models.py b/tests/python/frontend/tensorflow2/test_sequential_models.py
index 702b58c5b9b7..e0869ef48bd9 100644
--- a/tests/python/frontend/tensorflow2/test_sequential_models.py
+++ b/tests/python/frontend/tensorflow2/test_sequential_models.py
@@ -22,8 +22,7 @@
 import numpy as np
 import pytest
 import tensorflow as tf
-from tensorflow.python.framework.convert_to_constants \
-    import convert_variables_to_constants_v2
+from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
 from common import compare_tf_tvm
 from common import run_tf_code
@@ -31,7 +30,7 @@
 
 def run_sequential_model(model_fn, input_shape):
     def get_input(shape):
-        _input = np.random.uniform(0,1,shape).astype(dtype='float32')
+        _input = np.random.uniform(0, 1, shape).astype(dtype="float32")
         return _input
 
     def save_and_reload(_model):
@@ -53,56 +52,67 @@ def model_graph(model, input_shape):
 
 
 def dense_model(input_shape, num_units=128):
+    return tf.keras.Sequential(
+        [tf.keras.layers.Flatten(input_shape=input_shape[1:]), tf.keras.layers.Dense(num_units)]
+    )
+
+
+def mnist_model(input_shape):
     return tf.keras.Sequential(
         [
             tf.keras.layers.Flatten(input_shape=input_shape[1:]),
-            tf.keras.layers.Dense(num_units)
+            tf.keras.layers.Dense(128, activation="relu"),
+            tf.keras.layers.Dense(10),
         ]
     )
 
-def mnist_model(input_shape):
-    return tf.keras.Sequential([
-        tf.keras.layers.Flatten(input_shape=input_shape[1:]),
-        tf.keras.layers.Dense(128, activation="relu"),
-        tf.keras.layers.Dense(10)
-    ])
-
-
-def conv2d_model(input_shape, kernel=(3,3), filters=16):
-    model = tf.keras.Sequential([
-        tf.keras.layers.Input(shape=input_shape[1:], batch_size=1),
-        tf.keras.layers.Conv2D(filters, kernel)
-    ])
+
+def conv2d_model(input_shape, kernel=(3, 3), filters=16):
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.Input(shape=input_shape[1:], batch_size=1),
+            tf.keras.layers.Conv2D(filters, kernel),
+        ]
+    )
     return model
 
+
 def maxpool_model(input_shape, pool_size=(2, 2)):
-    model = tf.keras.Sequential([
-        tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:])
-    ])
+    model = tf.keras.Sequential(
+        [tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:])]
+    )
     return model
 
-def maxpool_batchnorm_model(input_shape, pool_size=(2,2)):
-    model = tf.keras.Sequential([
-        tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:]),
-        tf.keras.layers.BatchNormalization()
-    ])
+
+def maxpool_batchnorm_model(input_shape, pool_size=(2, 2)):
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:]),
+            tf.keras.layers.BatchNormalization(),
+        ]
+    )
     return model
 
 
 def test_dense_model():
-    run_sequential_model(dense_model,  input_shape = (1, 28, 28))
+    run_sequential_model(dense_model, input_shape=(1, 28, 28))
+
 
 def test_mnist_model():
-    run_sequential_model(mnist_model,  input_shape = (1, 28, 28))
+    run_sequential_model(mnist_model, input_shape=(1, 28, 28))
+
 
 def test_conv2d_model():
-    run_sequential_model(conv2d_model,  input_shape = (1, 32, 32, 3))
+    run_sequential_model(conv2d_model, input_shape=(1, 32, 32, 3))
+
 
 def test_maxpool_model():
-    run_sequential_model(maxpool_model, input_shape= (1, 32, 32, 3))
+    run_sequential_model(maxpool_model, input_shape=(1, 32, 32, 3))
+
 
 def test_maxpool_batchnorm_model():
-    run_sequential_model(maxpool_batchnorm_model, input_shape= (1, 32, 32, 3))
+    run_sequential_model(maxpool_batchnorm_model, input_shape=(1, 32, 32, 3))
+
 
 if __name__ == "__main__":
-    pytest.main([__file__])
\ No newline at end of file
+    pytest.main([__file__])

From 79a73ec1167dc51c9a6109c6f8b12b9a5c076ebb Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Wed, 19 May 2021 00:23:04 +0000
Subject: [PATCH 03/10] removing some comments

---
 .../tensorflow2/test_functional_models.py     | 60 ++++---------------
 1 file changed, 12 insertions(+), 48 deletions(-)

diff --git a/tests/python/frontend/tensorflow2/test_functional_models.py b/tests/python/frontend/tensorflow2/test_functional_models.py
index c448e63dc2ee..db19f88fc65e 100644
--- a/tests/python/frontend/tensorflow2/test_functional_models.py
+++ b/tests/python/frontend/tensorflow2/test_functional_models.py
@@ -25,70 +25,68 @@
 from common import compare_tf_tvm
 from common import run_tf_code
 
-""" simple function to test x=x+1"""
-
 
 class AddOne(tf.Module):
+    """ simple function to test x=x+1; scalar as input"""
+
     def get_input(self):
         return np.array(1.0, dtype="float32")
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
     def func(self, x):
         return x + 1
 
 
 class AddOne2D(AddOne):
+    """2D array as input"""
+
     def get_input(self):
         return np.ones((2, 2), dtype="float32")
 
-    """2D array as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
         return x + 1
 
 
 class AddOne2DConstant(AddOne):
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
     def get_input(self):
         return np.ones((2, 2), dtype="float32")
 
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
         return x + np.ones((2, 2), dtype="float32")
 
 
 class SubOne2DConstant(tf.Module):
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
     def get_input(self):
         return np.ones((2, 2), dtype="float32")
 
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
         return x - np.ones((2, 2), dtype="float32")
 
 
 class MulOne2DConstant(tf.Module):
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
     def get_input(self):
         return np.ones((2, 2), dtype="float32")
 
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
         return x * np.ones((2, 2), dtype="float32")
 
 
 class DivOne2DConstant(tf.Module):
+    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+
     def get_input(self):
         return np.ones((2, 2), dtype="float32")
 
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
     def func(self, x):
         return x / np.ones((2, 2), dtype="float32")
@@ -98,8 +96,6 @@ class StridedSlice(tf.Module):
     def get_input(self):
         return np.ones((3, 2, 3), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
     def func(self, x):
         return tf.strided_slice(x, [1, 0, 0], [2, 1, 3], [1, 1, 1])
@@ -109,8 +105,6 @@ class Shape(tf.Module):
     def get_input(self):
         return np.ones((3, 2, 3), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
     def func(self, x):
         a = tf.ones_like(tf.raw_ops.Shape(input=x), dtype=tf.float32)
@@ -121,8 +115,6 @@ class Pack(tf.Module):
     def get_input(self):
         return np.ones((2, 3), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3), dtype=tf.float32)])
     def func(self, x):
         return tf.raw_ops.Pack(values=[x, x], axis=0)
@@ -132,8 +124,6 @@ class Split(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b, c = tf.split(x, 3, axis=1)
@@ -144,8 +134,6 @@ class Maximum(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b = tf.split(x, 2, axis=1)
@@ -156,8 +144,6 @@ class Less(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b = tf.split(x, 2, axis=1)
@@ -168,8 +154,6 @@ class Equal(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b = tf.split(x, 2, axis=1)
@@ -180,8 +164,6 @@ class Cast(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.cast(x, tf.int32)
@@ -191,8 +173,6 @@ class ExpandDims(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.expand_dims(x, axis=2)
@@ -202,8 +182,6 @@ class Transpose(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         x = tf.expand_dims(x, axis=2)
@@ -214,8 +192,6 @@ class Reshape(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.reshape(x, (1, 2, 15))
@@ -225,8 +201,6 @@ class Tanh(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.math.tanh(x)
@@ -236,8 +210,6 @@ class Sigmoid(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.math.sigmoid(x)
@@ -247,8 +219,6 @@ class Relu(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.nn.relu(x)
@@ -258,8 +228,6 @@ class Floor(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         return tf.math.floor(x)
@@ -269,8 +237,6 @@ class FloorMod(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b = tf.split(x, 2, axis=1)
@@ -281,8 +247,6 @@ class ConcatV2(tf.Module):
     def get_input(self):
         return np.ones((1, 30), dtype=np.float32)
 
-    """scalar as input"""
-
     @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
     def func(self, x):
         a, b, c = tf.split(x, 3, axis=1)

From 689ac3be7eb25fd2462c2cdc6db81d25bcdbe5b9 Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Wed, 19 May 2021 00:30:38 +0000
Subject: [PATCH 04/10] change in comment in sequential test

---
 tests/python/frontend/tensorflow2/test_sequential_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/frontend/tensorflow2/test_sequential_models.py b/tests/python/frontend/tensorflow2/test_sequential_models.py
index e0869ef48bd9..d5ba67be846e 100644
--- a/tests/python/frontend/tensorflow2/test_sequential_models.py
+++ b/tests/python/frontend/tensorflow2/test_sequential_models.py
@@ -16,7 +16,7 @@
 # under the License.
 # pylint: disable=import-self, invalid-name, unused-argument, too-many-lines, len-as-condition, broad-except
 # pylint: disable=import-outside-toplevel, redefined-builtin
-"""TF2 to relay converter test: x+1 as a very basic example"""
+"""TF2 to relay converter test: testing models built with tf.keras.Sequential()"""
 
 import tempfile
 import numpy as np

From 663748123f20f77daeb7bbe3a35ef57412331e2f Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Thu, 20 May 2021 02:33:34 +0000
Subject: [PATCH 05/10] addressed the comments

---
 tests/python/frontend/tensorflow2/common.py   |  38 +-
 .../tensorflow2/test_functional_models.py     | 449 +++++++++---------
 .../tensorflow2/test_sequential_models.py     |  83 ++--
 3 files changed, 279 insertions(+), 291 deletions(-)

diff --git a/tests/python/frontend/tensorflow2/common.py b/tests/python/frontend/tensorflow2/common.py
index 0eda019062c2..feffa3e69f17 100644
--- a/tests/python/frontend/tensorflow2/common.py
+++ b/tests/python/frontend/tensorflow2/common.py
@@ -65,19 +65,14 @@ def run_tf_code(func, input_):
     return a
 
 
-def compile_graph_runtime(
-    mod, params, target="llvm", target_host="llvm", opt_level=3, output_sig=None
-):
+def compile_graph_executor(mod, params, target="llvm", target_host="llvm", opt_level=3):
     with tvm.transform.PassContext(opt_level):
         lib = relay.build(mod, target=target, target_host=target_host, params=params)
     return lib
 
 
-def compile_vm(
-    mod, params, target="llvm", target_host="llvm", opt_level=3, disabled_pass=None, output_sig=None
-):
+def compile_vm(mod, params, target="llvm", target_host="llvm", opt_level=3, disabled_pass=None):
     with tvm.transform.PassContext(opt_level, disabled_pass=disabled_pass):
-        mod = relay.transform.InferType()(mod)
         vm_exec = relay.vm.compile(mod, target, target_host, params=params)
     return vm_exec
 
@@ -88,7 +83,7 @@ def run_vm(vm_exec, input_, ctx=tvm.cpu(0)):
     return vmobj_to_list(_out)
 
 
-def run_graph(lib, input_, ctx=tvm.cpu(0)):
+def run_graph_executor(lib, input_, ctx=tvm.cpu(0)):
     mod = runtime.GraphModule(lib["default"](ctx))
     mod.set_input(0, input_)
     mod.run()
@@ -96,25 +91,30 @@ def run_graph(lib, input_, ctx=tvm.cpu(0)):
     return _out
 
 
-def compare_tf_tvm(gdef, input_, output_, vm=True, output_sig=None):
+def compare_tf_tvm(gdef, input_, output_, runtime="vm", output_tensors=None):
     """compare tf and tvm execution for the same input.
 
     Parameters
     ----------
-    func: tf function. can be from saved model or not. different ways to pass input
-        from saved model: <class 'tensorflow.python.saved_model.load._WrapperFunction'>
-        not from saved model:  <class 'tensorflow.python.eager.def_function.Function'>
-
-    mod: compiled relay module (vm or graph runtime). converted from tf func.
+    gdef: TF2 graph def extracted to be fed into from_tensorflow parser.
+    (https://www.tensorflow.org/code/tensorflow/core/framework/graph.proto)
 
     input_: a single numpy array object
 
+    output_: the expected output from TF to match TVM output with
+
+    runtime: choose TVM runtime; either "vm" for VirtualMachine or "graph" for GraphExecutor
+
+    output_tensors : List of output tensor names (Optional)
+            if not specified then the last node is assumed as graph output.
     """
-    mod, params = from_tensorflow(gdef, outputs=output_sig)
-    if vm:
-        exec_ = compile_vm(mod, params, output_sig=output_sig)
+    mod, params = from_tensorflow(gdef, outputs=output_tensors)
+    if runtime == "vm":
+        exec_ = compile_vm(mod, params)
         tvm_out = run_vm(exec_, input_)
+    elif runtime == "graph":
+        lib = compile_graph_executor(mod, params)
+        tvm_out = run_graph_executor(lib, input_)
     else:
-        lib = compile_graph_runtime(mod, params, output_sig=output_sig)
-        tvm_out = run_graph(lib, input_)
+        raise RuntimeError("Runtime input not supported: %s" % runtime)
     tvm.testing.assert_allclose(output_, tvm_out, atol=1e-5)
diff --git a/tests/python/frontend/tensorflow2/test_functional_models.py b/tests/python/frontend/tensorflow2/test_functional_models.py
index db19f88fc65e..40d42a28025a 100644
--- a/tests/python/frontend/tensorflow2/test_functional_models.py
+++ b/tests/python/frontend/tensorflow2/test_functional_models.py
@@ -26,342 +26,335 @@
 from common import run_tf_code
 
 
-class AddOne(tf.Module):
-    """ simple function to test x=x+1; scalar as input"""
-
-    def get_input(self):
-        return np.array(1.0, dtype="float32")
+def _function_graph(TestClass):
+    f = TestClass().func
+    gdef = f.get_concrete_function().graph.as_graph_def()
+    gdef_ops = list(set([n.op for n in gdef.node]))
+    input_ = TestClass().get_input()
+    output = run_tf_code(f, input_)
+    return gdef, input_, output
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
-    def func(self, x):
-        return x + 1
 
+def _model_graph(TestClass):
+    model = TestClass()
+    with tempfile.TemporaryDirectory() as model_path:
+        tf.saved_model.save(model, model_path)
+        imported = tf.saved_model.load(model_path)
 
-class AddOne2D(AddOne):
-    """2D array as input"""
+    f = imported.signatures["serving_default"]
+    gdef = f.graph.as_graph_def(add_shapes=True)
 
-    def get_input(self):
-        return np.ones((2, 2), dtype="float32")
+    input_ = model.get_input()
+    output = run_tf_code(f, input_)
+    return gdef, input_, output
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
-    def func(self, x):
-        return x + 1
 
+def run_all(TestClass):
+    def run_func_graph(TestClass, runtime="vm"):
+        compare_tf_tvm(*_function_graph(TestClass), runtime=runtime)
 
-class AddOne2DConstant(AddOne):
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+    def run_model_graph(TestClass):
+        compare_tf_tvm(*_model_graph(TestClass), runtime="vm")
 
-    def get_input(self):
-        return np.ones((2, 2), dtype="float32")
+    run_model_graph(TestClass)
+    for runtime_ in ["vm", "graph"]:
+        run_func_graph(TestClass, runtime=runtime_)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
-    def func(self, x):
-        return x + np.ones((2, 2), dtype="float32")
 
+def test_add_one():
+    class AddOne(tf.Module):
+        """ simple function to test x=x+1; scalar as input"""
 
-class SubOne2DConstant(tf.Module):
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+        def get_input(self):
+            return np.array(1.0, dtype="float32")
 
-    def get_input(self):
-        return np.ones((2, 2), dtype="float32")
+        @tf.function(input_signature=[tf.TensorSpec(shape=(), dtype=tf.float32)])
+        def func(self, x):
+            return x + 1
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
-    def func(self, x):
-        return x - np.ones((2, 2), dtype="float32")
+    run_all(AddOne)
 
 
-class MulOne2DConstant(tf.Module):
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
+def test_add_one_2d():
+    class AddOne2D(tf.Module):
+        """2D array as input"""
 
-    def get_input(self):
-        return np.ones((2, 2), dtype="float32")
+        def get_input(self):
+            return np.ones((2, 2), dtype="float32")
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
-    def func(self, x):
-        return x * np.ones((2, 2), dtype="float32")
+        @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+        def func(self, x):
+            return x + 1
 
+    run_all(AddOne2D)
 
-class DivOne2DConstant(tf.Module):
-    """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
-    def get_input(self):
-        return np.ones((2, 2), dtype="float32")
+def test_add_one_2d_constant():
+    class AddOne2DConstant(tf.Module):
+        """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
-    def func(self, x):
-        return x / np.ones((2, 2), dtype="float32")
+        def get_input(self):
+            return np.ones((2, 2), dtype="float32")
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+        def func(self, x):
+            return x + np.ones((2, 2), dtype="float32")
 
-class StridedSlice(tf.Module):
-    def get_input(self):
-        return np.ones((3, 2, 3), dtype=np.float32)
+    run_all(AddOne2DConstant)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
-    def func(self, x):
-        return tf.strided_slice(x, [1, 0, 0], [2, 1, 3], [1, 1, 1])
 
+def test_sub_one_2d_constant():
+    class SubOne2DConstant(tf.Module):
+        """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
-class Shape(tf.Module):
-    def get_input(self):
-        return np.ones((3, 2, 3), dtype=np.float32)
+        def get_input(self):
+            return np.ones((2, 2), dtype="float32")
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
-    def func(self, x):
-        a = tf.ones_like(tf.raw_ops.Shape(input=x), dtype=tf.float32)
-        return a + x
+        @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+        def func(self, x):
+            return x - np.ones((2, 2), dtype="float32")
 
+    run_all(SubOne2DConstant)
 
-class Pack(tf.Module):
-    def get_input(self):
-        return np.ones((2, 3), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3), dtype=tf.float32)])
-    def func(self, x):
-        return tf.raw_ops.Pack(values=[x, x], axis=0)
+def test_mul_one_2d_constant():
+    class MulOne2DConstant(tf.Module):
+        """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
+        def get_input(self):
+            return np.ones((2, 2), dtype="float32")
 
-class Split(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+        def func(self, x):
+            return x * np.ones((2, 2), dtype="float32")
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        a, b, c = tf.split(x, 3, axis=1)
-        return tf.raw_ops.Pack(values=[a, b, c], axis=1)
+    run_all(MulOne2DConstant)
 
 
-class Maximum(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+def test_div_one_2d_constant():
+    class DivOne2DConstant(tf.Module):
+        """2D array as input with 2D constant as well; 2D constant stored in params after convert"""
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        a, b = tf.split(x, 2, axis=1)
-        return tf.math.maximum(a, b, name=None)
+        def get_input(self):
+            return np.ones((2, 2), dtype="float32")
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)])
+        def func(self, x):
+            return x / np.ones((2, 2), dtype="float32")
 
-class Less(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+    run_all(DivOne2DConstant)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        a, b = tf.split(x, 2, axis=1)
-        return tf.math.less(a, b, name=None)
 
+def test_strided_slice():
+    class StridedSlice(tf.Module):
+        def get_input(self):
+            return np.ones((3, 2, 3), dtype=np.float32)
 
-class Equal(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
+        def func(self, x):
+            return tf.strided_slice(x, [1, 0, 0], [2, 1, 3], [1, 1, 1])
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        a, b = tf.split(x, 2, axis=1)
-        return tf.math.equal(a, b, name=None)
+    run_all(StridedSlice)
 
 
-class Cast(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+def test_split():
+    class Split(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.cast(x, tf.int32)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            a, b, c = tf.split(x, 3, axis=1)
+            return tf.raw_ops.Pack(values=[a, b, c], axis=1)
 
+    run_all(Split)
 
-class ExpandDims(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.expand_dims(x, axis=2)
+def test_shape():
+    class Shape(tf.Module):
+        def get_input(self):
+            return np.ones((3, 2, 3), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(3, 2, 3), dtype=tf.float32)])
+        def func(self, x):
+            a = tf.ones_like(tf.raw_ops.Shape(input=x), dtype=tf.float32)
+            return a + x
 
-class Transpose(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+    run_all(Shape)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        x = tf.expand_dims(x, axis=2)
-        return tf.transpose(x, perm=[0, 2, 1])
 
+def test_pack():
+    class Pack(tf.Module):
+        def get_input(self):
+            return np.ones((2, 3), dtype=np.float32)
 
-class Reshape(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(2, 3), dtype=tf.float32)])
+        def func(self, x):
+            return tf.raw_ops.Pack(values=[x, x], axis=0)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.reshape(x, (1, 2, 15))
+    run_all(Pack)
 
 
-class Tanh(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+def test_max():
+    class Maximum(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.math.tanh(x)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            a, b = tf.split(x, 2, axis=1)
+            return tf.math.maximum(a, b, name=None)
 
+    run_all(Maximum)
 
-class Sigmoid(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.math.sigmoid(x)
+def test_less():
+    class Less(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            a, b = tf.split(x, 2, axis=1)
+            return tf.math.less(a, b, name=None)
 
-class Relu(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+    run_all(Less)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.nn.relu(x)
 
+def test_equal():
+    class Equal(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
-class Floor(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            a, b = tf.split(x, 2, axis=1)
+            return tf.math.equal(a, b, name=None)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        return tf.math.floor(x)
+    run_all(Equal)
 
 
-class FloorMod(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
+def test_cast():
+    class Cast(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        a, b = tf.split(x, 2, axis=1)
-        return tf.math.floormod(a, b)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.cast(x, tf.int32)
 
+    run_all(Cast)
 
-class ConcatV2(tf.Module):
-    def get_input(self):
-        return np.ones((1, 30), dtype=np.float32)
 
-    @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
-    def func(self, x):
-        a, b, c = tf.split(x, 3, axis=1)
-        return tf.raw_ops.ConcatV2(values=[a, b, c], axis=1)
+def test_expand_dims():
+    class ExpandDims(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.expand_dims(x, axis=2)
 
-def _function_graph(TestClass):
-    f = TestClass().func
-    gdef = f.get_concrete_function().graph.as_graph_def()
-    gdef_ops = list(set([n.op for n in gdef.node]))
-    input_ = TestClass().get_input()
-    output = run_tf_code(f, input_)
-    return gdef, input_, output
+    run_all(ExpandDims)
 
 
-def _model_graph(TestClass):
-    model = TestClass()
-    with tempfile.TemporaryDirectory() as model_path:
-        tf.saved_model.save(model, model_path)
-        imported = tf.saved_model.load(model_path)
+def test_transpose():
+    class Transpose(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
-    f = imported.signatures["serving_default"]
-    gdef = f.graph.as_graph_def(add_shapes=True)
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            x = tf.expand_dims(x, axis=2)
+            return tf.transpose(x, perm=[0, 2, 1])
 
-    input_ = model.get_input()
-    output = run_tf_code(f, input_)
-    return gdef, input_, output
+    run_all(Transpose)
 
 
-def run_func_graph(TestClass, use_vm=False):
-    compare_tf_tvm(*_function_graph(TestClass), vm=use_vm)
+def test_reshape():
+    class Reshape(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.reshape(x, (1, 2, 15))
 
-def run_model_graph(TestClass, output_sig=None):
-    compare_tf_tvm(*_model_graph(TestClass), vm=True, output_sig=output_sig)
+    run_all(Reshape)
 
 
-def run_all(TestClass):
-    run_model_graph(TestClass)
-    for use_vm in [True, False]:
-        run_func_graph(TestClass, use_vm=use_vm)
+def test_tanh():
+    class Tanh(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.math.tanh(x)
 
-def test_basic_ops():
-    run_all(AddOne)
-    run_all(AddOne2D)
-    run_all(AddOne2DConstant)
-    run_all(SubOne2DConstant)
-    run_all(MulOne2DConstant)
-    run_all(DivOne2DConstant)
+    run_all(Tanh)
 
 
-def test_strided_slice():
-    run_all(StridedSlice)
+def test_sigmoid():
+    class Sigmoid(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.math.sigmoid(x)
 
-def test_shape():
-    run_all(Shape)
+    run_all(Sigmoid)
 
 
-def test_pack():
-    run_all(Pack)
+def test_relu():
+    class Relu(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.nn.relu(x)
 
-def test_split():
-    run_all(Split)
+    run_all(Relu)
 
 
-def test_max():
-    run_all(Maximum)
+def test_floor():
+    class Floor(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            return tf.math.floor(x)
 
-def test_less():
-    run_all(Less)
+    run_all(Floor)
 
 
-def test_equal():
-    run_all(Equal)
+def test_floor_mod():
+    class FloorMod(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            a, b = tf.split(x, 2, axis=1)
+            return tf.math.floormod(a, b)
 
-def test_floor():
-    run_all(Floor)
     run_all(FloorMod)
 
 
 def test_concat_v2():
-    run_all(ConcatV2)
-
-
-def test_cast():
-    run_all(Cast)
-
-
-def test_expand_dims():
-    run_all(ExpandDims)
-
-
-def test_transpose():
-    run_all(Transpose)
+    class ConcatV2(tf.Module):
+        def get_input(self):
+            return np.ones((1, 30), dtype=np.float32)
 
+        @tf.function(input_signature=[tf.TensorSpec(shape=(1, 30), dtype=tf.float32)])
+        def func(self, x):
+            a, b, c = tf.split(x, 3, axis=1)
+            return tf.raw_ops.ConcatV2(values=[a, b, c], axis=1)
 
-def test_reshape():
-    run_all(Reshape)
-
-
-def test_tanh():
-    run_all(Tanh)
-
-
-def test_sigmoid():
-    run_all(Sigmoid)
-
-
-def test_relu():
-    run_all(Relu)
+    run_all(ConcatV2)
 
 
 if __name__ == "__main__":
diff --git a/tests/python/frontend/tensorflow2/test_sequential_models.py b/tests/python/frontend/tensorflow2/test_sequential_models.py
index d5ba67be846e..394a49d0f2e9 100644
--- a/tests/python/frontend/tensorflow2/test_sequential_models.py
+++ b/tests/python/frontend/tensorflow2/test_sequential_models.py
@@ -48,69 +48,64 @@ def model_graph(model, input_shape):
         gdef = f.graph.as_graph_def(add_shapes=True)
         return gdef, _input, _output
 
-    compare_tf_tvm(*model_graph(model_fn, input_shape), vm=True, output_sig=None)
-
-
-def dense_model(input_shape, num_units=128):
-    return tf.keras.Sequential(
-        [tf.keras.layers.Flatten(input_shape=input_shape[1:]), tf.keras.layers.Dense(num_units)]
-    )
-
-
-def mnist_model(input_shape):
-    return tf.keras.Sequential(
-        [
-            tf.keras.layers.Flatten(input_shape=input_shape[1:]),
-            tf.keras.layers.Dense(128, activation="relu"),
-            tf.keras.layers.Dense(10),
-        ]
-    )
-
-
-def conv2d_model(input_shape, kernel=(3, 3), filters=16):
-    model = tf.keras.Sequential(
-        [
-            tf.keras.layers.Input(shape=input_shape[1:], batch_size=1),
-            tf.keras.layers.Conv2D(filters, kernel),
-        ]
-    )
-    return model
-
-
-def maxpool_model(input_shape, pool_size=(2, 2)):
-    model = tf.keras.Sequential(
-        [tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:])]
-    )
-    return model
-
-
-def maxpool_batchnorm_model(input_shape, pool_size=(2, 2)):
-    model = tf.keras.Sequential(
-        [
-            tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:]),
-            tf.keras.layers.BatchNormalization(),
-        ]
-    )
-    return model
+    compare_tf_tvm(*model_graph(model_fn, input_shape), runtime="vm")
 
 
 def test_dense_model():
+    def dense_model(input_shape, num_units=128):
+        return tf.keras.Sequential(
+            [tf.keras.layers.Flatten(input_shape=input_shape[1:]), tf.keras.layers.Dense(num_units)]
+        )
+
     run_sequential_model(dense_model, input_shape=(1, 28, 28))
 
 
 def test_mnist_model():
+    def mnist_model(input_shape):
+        return tf.keras.Sequential(
+            [
+                tf.keras.layers.Flatten(input_shape=input_shape[1:]),
+                tf.keras.layers.Dense(128, activation="relu"),
+                tf.keras.layers.Dense(10),
+            ]
+        )
+
     run_sequential_model(mnist_model, input_shape=(1, 28, 28))
 
 
 def test_conv2d_model():
+    def conv2d_model(input_shape, kernel=(3, 3), filters=16):
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.Input(shape=input_shape[1:], batch_size=1),
+                tf.keras.layers.Conv2D(filters, kernel),
+            ]
+        )
+        return model
+
     run_sequential_model(conv2d_model, input_shape=(1, 32, 32, 3))
 
 
 def test_maxpool_model():
+    def maxpool_model(input_shape, pool_size=(2, 2)):
+        model = tf.keras.Sequential(
+            [tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:])]
+        )
+        return model
+
     run_sequential_model(maxpool_model, input_shape=(1, 32, 32, 3))
 
 
 def test_maxpool_batchnorm_model():
+    def maxpool_batchnorm_model(input_shape, pool_size=(2, 2)):
+        model = tf.keras.Sequential(
+            [
+                tf.keras.layers.MaxPool2D(pool_size=pool_size, input_shape=input_shape[1:]),
+                tf.keras.layers.BatchNormalization(),
+            ]
+        )
+        return model
+
     run_sequential_model(maxpool_batchnorm_model, input_shape=(1, 32, 32, 3))
 
 

From 0b04d40fc06faa7151b63f1c924d17231831c2a4 Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Thu, 20 May 2021 03:04:29 +0000
Subject: [PATCH 06/10] refactored to place vmobj_to_list in a common file

---
 tests/python/frontend/tensorflow2/common.py | 26 +++++----------------
 1 file changed, 6 insertions(+), 20 deletions(-)

diff --git a/tests/python/frontend/tensorflow2/common.py b/tests/python/frontend/tensorflow2/common.py
index feffa3e69f17..4d2fbf0b2a0a 100644
--- a/tests/python/frontend/tensorflow2/common.py
+++ b/tests/python/frontend/tensorflow2/common.py
@@ -24,42 +24,28 @@
 from tvm.runtime.vm import VirtualMachine
 import tvm.contrib.graph_executor as runtime
 from tvm.relay.frontend.tensorflow import from_tensorflow
+
 import tvm.testing
+from tvm.relay.testing.tf import vmobj_to_list as vmobj_to_list
 
 import tensorflow as tf
 from tensorflow.python.eager.def_function import Function
 
 
-def vmobj_to_list(o):
-    if isinstance(o, tvm.nd.NDArray):
-        out = o.asnumpy().tolist()
-    elif isinstance(o, tvm.runtime.container.ADT):
-        result = []
-        for f in o:
-            result.append(vmobj_to_list(f))
-        out = result
-    else:
-        raise RuntimeError("Unknown object type: %s" % type(o))
-    return out
-
-
 def run_tf_code(func, input_):
+    print(type(func))
     if type(func) is Function:
         out = func(input_)
         if isinstance(out, list):
             a = [x.numpy() for x in out]
         else:
-            a = out.numpy()
+            a = [out.numpy()]
     else:
         a = func(tf.constant(input_))
         if type(a) is dict:
             a = [x.numpy() for x in a.values()]
-            if len(a) == 1:
-                a = a[0]
         elif type(a) is list:
             a = [x.numpy() for x in a]
-            if len(a) == 1:
-                a = a[0]
         else:
             a = a.numpy()
     return a
@@ -87,8 +73,7 @@ def run_graph_executor(lib, input_, ctx=tvm.cpu(0)):
     mod = runtime.GraphModule(lib["default"](ctx))
     mod.set_input(0, input_)
     mod.run()
-    _out = mod.get_output(0).asnumpy()
-    return _out
+    return [mod.get_output(0).asnumpy()]
 
 
 def compare_tf_tvm(gdef, input_, output_, runtime="vm", output_tensors=None):
@@ -117,4 +102,5 @@ def compare_tf_tvm(gdef, input_, output_, runtime="vm", output_tensors=None):
         tvm_out = run_graph_executor(lib, input_)
     else:
         raise RuntimeError("Runtime input not supported: %s" % runtime)
+
     tvm.testing.assert_allclose(output_, tvm_out, atol=1e-5)

From 4dd61677abe4cf598cd79af8690f66f46f2defd7 Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Thu, 20 May 2021 20:25:49 +0000
Subject: [PATCH 07/10] Added helper function in python/tvm/relay/testing/tf.py

Co-authored-by: David Huang <davhuan@amazon.com>
Co-authored-by: Rohan Mukherjee <mukrohan@amazon.com>
Co-authored-by: Xiao <weix@amazon.com>
---
 python/tvm/relay/testing/tf.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/python/tvm/relay/testing/tf.py b/python/tvm/relay/testing/tf.py
index d20c0e0ab9dd..45afdb7ff512 100644
--- a/python/tvm/relay/testing/tf.py
+++ b/python/tvm/relay/testing/tf.py
@@ -28,6 +28,8 @@
 # Tensorflow imports
 import tensorflow as tf
 from tensorflow.core.framework import graph_pb2
+
+import tvm
 from tvm.contrib.download import download_testdata
 
 try:
@@ -73,6 +75,32 @@ def convert_to_list(x):
     return x
 
 
+def vmobj_to_list(o):
+    if isinstance(o, tvm.nd.NDArray):
+        return [o.asnumpy()]
+    elif isinstance(o, tvm.runtime.container.ADT):
+        result = []
+        for f in o:
+            result.extend(vmobj_to_list(f))
+        return result
+    elif isinstance(o, tvm.relay.backend.interpreter.ConstructorValue):
+        if o.constructor.name_hint == "Cons":
+            tl = vmobj_to_list(o.fields[1])
+            hd = vmobj_to_list(o.fields[0])
+            hd.extend(tl)
+            return hd
+        elif o.constructor.name_hint == "Nil":
+            return []
+        elif "tensor_nil" in o.constructor.name_hint:
+            return [0]
+        elif "tensor" in o.constructor.name_hint:
+            return [o.fields[0].asnumpy()]
+        else:
+            raise RuntimeError("Unknown object type: %s" % o.constructor.name_hint)
+    else:
+        raise RuntimeError("Unknown object type: %s" % type(o))
+
+
 def AddShapesToGraphDef(session, out_node):
     """Add shapes attribute to nodes of the graph.
         Input graph here is the default graph in context.

From a14c386391271a9cc7de8907a495254416424aad Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Thu, 20 May 2021 20:38:37 +0000
Subject: [PATCH 08/10] Refactor tf according to CI error

Co-authored-by: David Huang <davhuan@amazon.com>
Co-authored-by: Rohan Mukherjee <mukrohan@amazon.com>
Co-authored-by: Xiao <weix@amazon.com>
---
 python/tvm/relay/testing/tf.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/python/tvm/relay/testing/tf.py b/python/tvm/relay/testing/tf.py
index 45afdb7ff512..b5779dcdd8ed 100644
--- a/python/tvm/relay/testing/tf.py
+++ b/python/tvm/relay/testing/tf.py
@@ -77,28 +77,28 @@ def convert_to_list(x):
 
 def vmobj_to_list(o):
     if isinstance(o, tvm.nd.NDArray):
-        return [o.asnumpy()]
+        result = [o.asnumpy()]
     elif isinstance(o, tvm.runtime.container.ADT):
         result = []
         for f in o:
             result.extend(vmobj_to_list(f))
-        return result
     elif isinstance(o, tvm.relay.backend.interpreter.ConstructorValue):
         if o.constructor.name_hint == "Cons":
             tl = vmobj_to_list(o.fields[1])
             hd = vmobj_to_list(o.fields[0])
             hd.extend(tl)
-            return hd
+            result = hd
         elif o.constructor.name_hint == "Nil":
-            return []
+            result = []
         elif "tensor_nil" in o.constructor.name_hint:
-            return [0]
+            result = [0]
         elif "tensor" in o.constructor.name_hint:
-            return [o.fields[0].asnumpy()]
+            result = [o.fields[0].asnumpy()]
         else:
             raise RuntimeError("Unknown object type: %s" % o.constructor.name_hint)
     else:
         raise RuntimeError("Unknown object type: %s" % type(o))
+    return result
 
 
 def AddShapesToGraphDef(session, out_node):

From 512ab4e3e3a867506886075fc131627b274e0836 Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Thu, 20 May 2021 20:54:10 +0000
Subject: [PATCH 09/10] Added docstring

Co-authored-by: David Huang <davhuan@amazon.com>
Co-authored-by: Rohan Mukherjee <mukrohan@amazon.com>
Co-authored-by: Xiao <weix@amazon.com>
---
 python/tvm/relay/testing/tf.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/python/tvm/relay/testing/tf.py b/python/tvm/relay/testing/tf.py
index b5779dcdd8ed..9fb3f1102137 100644
--- a/python/tvm/relay/testing/tf.py
+++ b/python/tvm/relay/testing/tf.py
@@ -76,6 +76,20 @@ def convert_to_list(x):
 
 
 def vmobj_to_list(o):
+    """Converts TVM objects returned by VM execution to Python List.
+
+    Parameters
+    ----------
+    o : Obj
+        VM Object as output from VM runtime executor.
+
+    Returns
+    -------
+    result : list
+        Numpy objects as list with equivalent values to the input object.
+
+    """
+
     if isinstance(o, tvm.nd.NDArray):
         result = [o.asnumpy()]
     elif isinstance(o, tvm.runtime.container.ADT):

From fd8f8e5b379a8a46325336b389c69b9734537c4d Mon Sep 17 00:00:00 2001
From: Rohan <mukrohan@amazon.com>
Date: Mon, 24 May 2021 23:44:53 +0000
Subject: [PATCH 10/10] removing print

---
 tests/python/frontend/tensorflow2/common.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/python/frontend/tensorflow2/common.py b/tests/python/frontend/tensorflow2/common.py
index 4d2fbf0b2a0a..e30ee7b0c993 100644
--- a/tests/python/frontend/tensorflow2/common.py
+++ b/tests/python/frontend/tensorflow2/common.py
@@ -33,7 +33,6 @@
 
 
 def run_tf_code(func, input_):
-    print(type(func))
     if type(func) is Function:
         out = func(input_)
         if isinstance(out, list):