From 4fe855b7f79a99f5fd249c0734140885fc9e3f55 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 04:17:19 +0000
Subject: [PATCH 01/26] Add segment sum Op

---
 python/tvm/relay/op/transform.py     | 52 ++++++++++++++++++
 python/tvm/topi/scatter_add.py       | 40 +++++++-------
 tests/python/relay/segment_sum.py    | 27 ++++++++++
 tests/python/relay/test_op_level3.py | 79 ++++++++++++++++++++++++++++
 4 files changed, 178 insertions(+), 20 deletions(-)
 create mode 100644 tests/python/relay/segment_sum.py

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index 73508ddd2603..ed6e5c80236d 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1450,6 +1450,58 @@ def sparse_reshape(sparse_indices, prev_shape, new_shape):
     return TupleWrapper(_make.sparse_reshape(sparse_indices, prev_shape, new_shape), 2)
 
 
+def segment_sum(data, indices, num_segments=None):
+    """
+    Computes the sum along segments of a tensor.
+
+    Parameters
+    ----------
+    sparse_indices : relay.Expr
+        A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the
+        number of sparse values and n_dim is the number of dimensions of the dense_shape
+    prev_shape : relay.Expr
+        A 1-D tensor containing the previous shape of the dense tensor
+    new_shape : relay.Expr
+        A 1-D tensor containing the new shape of the dense tensor
+    Returns
+    -------
+    result: relay.Expr
+        Output tensor.
+    Examples
+    --------
+    .. code-block:: python
+        sparse_indices = [[0, 0, 0],
+                            [0, 0, 1],
+                            [0, 1, 0],
+                            [1, 0, 0],
+                            [1, 2, 3]]
+        prev_shape = [2, 3, 4]
+        new_shape = [9, -1]
+        new_sparse_indices, new_shape = relay.sparse_reshape(sparse_indices,
+                            prev_shape,
+                            new_shape)
+        new_sparse_indices = [[0, 0],
+                              [0, 1],
+                              [1, 2],
+                              [4, 2],
+                              [8, 1]] 
+        new_shape = [9, 4]
+    """
+    
+    if num_segments:
+        num_unique = const([num_segments]) 
+    else:
+        _, _, num_unique = unique(reshape(indices, -1))
+    data_offrow_shape = strided_slice(_make.shape_of(data, "int64"), [1], [-1], slice_mode="size")
+    data_offrow_shape = cast_like(data_offrow_shape, indices)
+    new_shape = _make.concatenate(Tuple([num_unique, data_offrow_shape]), 0)
+    indices_tiled_shape = _make.concatenate(Tuple([reverse(data_offrow_shape, 0), const([1])]), 0)
+    expanded_indices = tile(indices, indices_tiled_shape)
+    scatter_add_indices = transpose(expanded_indices)
+    src = cast_like(_dyn_make.zeros(new_shape, "float64"), data)
+    return scatter_add(src, scatter_add_indices, data, axis=0)
+
+
 def cumsum(data, axis=None, dtype=None, exclusive=None):
     """Numpy style cumsum op. Return the cumulative inclusive sum of the elements along
     a given axis.
diff --git a/python/tvm/topi/scatter_add.py b/python/tvm/topi/scatter_add.py
index 4c77a0767785..6b04837b7766 100644
--- a/python/tvm/topi/scatter_add.py
+++ b/python/tvm/topi/scatter_add.py
@@ -32,8 +32,8 @@ def _scatter_add_1d(data, indices, updates):
 @hybrid.script
 def _scatter_add_2d(data, indices, updates, axis):
     out = output_tensor(data.shape, data.dtype)
-    for i in const_range(data.shape[0]):
-        for j in const_range(data.shape[1]):
+    for i in range(data.shape[0]):
+        for j in range(data.shape[1]):
             out[i, j] = data[i, j]
     if axis == 0:
         for i in range(indices.shape[0]):
@@ -54,14 +54,14 @@ def _scatter_add_2d(data, indices, updates, axis):
 @hybrid.script
 def _scatter_add_3d(data, indices, updates, axis):
     out = output_tensor(data.shape, data.dtype)
-    for i in const_range(data.shape[0]):
-        for j in const_range(data.shape[1]):
-            for k in const_range(data.shape[2]):
+    for i in range(data.shape[0]):
+        for j in range(data.shape[1]):
+            for k in range(data.shape[2]):
                 out[i, j, k] = data[i, j, k]
     if axis == 0:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
+                for k in range(indices.shape[2]):
                     out[
                         indices[i, j, k]
                         if indices[i, j, k] >= 0
@@ -72,7 +72,7 @@ def _scatter_add_3d(data, indices, updates, axis):
     elif axis == 1:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
+                for k in range(indices.shape[2]):
                     out[
                         i,
                         indices[i, j, k]
@@ -83,7 +83,7 @@ def _scatter_add_3d(data, indices, updates, axis):
     else:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
+                for k in range(indices.shape[2]):
                     out[
                         i,
                         j,
@@ -98,17 +98,17 @@ def _scatter_add_3d(data, indices, updates, axis):
 @hybrid.script
 def _scatter_add_4d(data, indices, updates, axis):
     out = output_tensor(data.shape, data.dtype)
-    for i in const_range(data.shape[0]):
-        for j in const_range(data.shape[1]):
-            for k in const_range(data.shape[2]):
-                for l in const_range(data.shape[3]):
+    for i in range(data.shape[0]):
+        for j in range(data.shape[1]):
+            for k in range(data.shape[2]):
+                for l in range(data.shape[3]):
                     out[i, j, k, l] = data[i, j, k, l]
 
     if axis == 0:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
-                    for l in const_range(indices.shape[3]):
+                for k in range(indices.shape[2]):
+                    for l in range(indices.shape[3]):
                         out[
                             indices[i, j, k, l]
                             if indices[i, j, k, l] >= 0
@@ -120,8 +120,8 @@ def _scatter_add_4d(data, indices, updates, axis):
     elif axis == 1:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
-                    for l in const_range(indices.shape[3]):
+                for k in range(indices.shape[2]):
+                    for l in range(indices.shape[3]):
                         out[
                             i,
                             indices[i, j, k, l]
@@ -133,8 +133,8 @@ def _scatter_add_4d(data, indices, updates, axis):
     elif axis == 2:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
-                    for l in const_range(indices.shape[3]):
+                for k in range(indices.shape[2]):
+                    for l in range(indices.shape[3]):
                         out[
                             i,
                             j,
@@ -146,8 +146,8 @@ def _scatter_add_4d(data, indices, updates, axis):
     else:
         for i in range(indices.shape[0]):
             for j in range(indices.shape[1]):
-                for k in const_range(indices.shape[2]):
-                    for l in const_range(indices.shape[3]):
+                for k in range(indices.shape[2]):
+                    for l in range(indices.shape[3]):
                         out[
                             i,
                             j,
diff --git a/tests/python/relay/segment_sum.py b/tests/python/relay/segment_sum.py
new file mode 100644
index 000000000000..3a3c414815af
--- /dev/null
+++ b/tests/python/relay/segment_sum.py
@@ -0,0 +1,27 @@
+import torch
+import numpy as np
+
+
+def segment_sum_torch(src, indices, data):
+    indices = indices.expand([2, 1, -1]).permute(2, 1, 0)
+    print(indices)
+    result = src.scatter_add(0, indices, data)
+    return result
+
+
+def segment_sum_numpy(src, indices, data):
+    for i, index in enumerate(indices):
+        src[index] += data[i]
+
+    return src
+
+
+if __name__ == "__main__":
+    data = np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32)
+    src = np.zeros((2, 1, 2), dtype=np.float32)
+    indices = np.array([0, 0, 1], dtype=np.int64)
+    result_torch = segment_sum_torch(
+        torch.from_numpy(src), torch.from_numpy(indices), torch.from_numpy(data)
+    )
+    result_numpy = segment_sum_numpy(src, indices, data)
+    np.testing.assert_allclose(result_torch, result_numpy)
\ No newline at end of file
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index c9ed975c3b9b..8cc6d3094686 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -24,6 +24,7 @@
 from tvm.error import TVMError
 from tvm.relay import create_executor, transform
 from tvm.relay.testing import check_grad, run_infer_type
+from typing import Optional
 import tvm.testing
 
 
@@ -1515,6 +1516,80 @@ def verify_sparse_reshape(
     )
 
 
+# @tvm.testing.uses_gpu
+@pytest.mark.parametrize(
+    "data_np, indices_np, num_segments",
+    [
+        (
+            np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
+            np.array([0, 0, 1, 1, 0, 1], dtype=np.int32),
+            None,
+        ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+            np.array([0, 0, 1], dtype=np.int32),
+            None,
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([2, 0, 1, 0, 3, 2], dtype=np.int32),
+            None,
+        ),
+        (
+            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
+            np.array([0, 0, 1], dtype=np.int32),
+            None,
+        ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([2, 0, 1, 0, 3, 2, 5, 4, 4], dtype=np.int32),
+            9,
+        ),
+    ],
+)
+def test_segment_sum(data_np, indices_np, num_segments):
+    def ref_segment_sum(
+        data: np.ndarray,
+        indices: np.ndarray,
+        num_segments: Optional[int] = None,
+    ):
+        """
+        This function calculates the expected output of sparseshape operator given the inputs.
+        """
+        if not num_segments:
+            num_segments = np.unique(indices).shape[0]
+
+        result = np.zeros((num_segments,) + data.shape[1:], data.dtype)
+        for i, index in enumerate(indices):
+            result[index] += data[i]
+        return result
+
+    def verify_segment_sum(
+        data_np: np.ndarray, indices_np: np.ndarray, num_segments: Optional[int]
+    ):
+        """
+        This function verifies the relay output of sparse_reshape with its expected output.
+        """
+        data = relay.var(
+            "data",
+            relay.TensorType(data_np.shape, str(data_np.dtype)),
+        )
+        indices = relay.var("indices", relay.TensorType(indices_np.shape, str(indices_np.dtype)))
+        z = relay.op.segment_sum(data, indices, num_segments)
+
+        func = relay.Function([data, indices], z)
+        ref_res = ref_segment_sum(data_np, indices_np, num_segments=num_segments)
+        segment_sum_result = run_infer_type(z)
+        assert segment_sum_result.checked_type.dtype == data_np.dtype
+        verify_func(
+            func,
+            [data_np, indices_np],
+            ref_res,
+        )
+
+    verify_segment_sum(data_np, indices_np, num_segments)
+
+
 def verify_func(func, data, ref_res, target_ctx=tvm.testing.enabled_targets()):
     assert isinstance(data, list)
     for target, ctx in target_ctx:
@@ -1530,6 +1605,10 @@ def verify_func(func, data, ref_res, target_ctx=tvm.testing.enabled_targets()):
                 for op_result, ref_result in zip(op_res, ref_res):
                     tvm.testing.assert_allclose(op_result.asnumpy(), ref_result, rtol=1e-5)
             else:
+                # print(op_res.asnumpy(), ref_res, *data)
+                # import pdb
+
+                # pdb.set_trace()
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
             relay.backend.compile_engine.get().clear()
 

From 20bcbbab5a8c2b27dee12deff3c901f1d33c4af6 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 04:19:51 +0000
Subject: [PATCH 02/26] Remove unnecessary

---
 tests/python/relay/segment_sum.py | 27 ---------------------------
 1 file changed, 27 deletions(-)
 delete mode 100644 tests/python/relay/segment_sum.py

diff --git a/tests/python/relay/segment_sum.py b/tests/python/relay/segment_sum.py
deleted file mode 100644
index 3a3c414815af..000000000000
--- a/tests/python/relay/segment_sum.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import torch
-import numpy as np
-
-
-def segment_sum_torch(src, indices, data):
-    indices = indices.expand([2, 1, -1]).permute(2, 1, 0)
-    print(indices)
-    result = src.scatter_add(0, indices, data)
-    return result
-
-
-def segment_sum_numpy(src, indices, data):
-    for i, index in enumerate(indices):
-        src[index] += data[i]
-
-    return src
-
-
-if __name__ == "__main__":
-    data = np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32)
-    src = np.zeros((2, 1, 2), dtype=np.float32)
-    indices = np.array([0, 0, 1], dtype=np.int64)
-    result_torch = segment_sum_torch(
-        torch.from_numpy(src), torch.from_numpy(indices), torch.from_numpy(data)
-    )
-    result_numpy = segment_sum_numpy(src, indices, data)
-    np.testing.assert_allclose(result_torch, result_numpy)
\ No newline at end of file

From fa9fb484f5c4c0579836997676f50acff4aac67d Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 04:34:41 +0000
Subject: [PATCH 03/26] Documentation

---
 python/tvm/relay/op/transform.py     | 52 ++++++++++++++--------------
 tests/python/relay/test_op_level3.py |  2 +-
 2 files changed, 27 insertions(+), 27 deletions(-)

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index ed6e5c80236d..3573d494dc6e 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1452,17 +1452,21 @@ def sparse_reshape(sparse_indices, prev_shape, new_shape):
 
 def segment_sum(data, indices, num_segments=None):
     """
-    Computes the sum along segments of a tensor.
+    Computes the sum along segments of a tensor. This op is much better understood with
+    visualization articulated in the following links and examples at the end of this docstring.
+
+    https://www.tensorflow.org/api_docs/python/tf/raw_ops/UnsortedSegmentSum?hl=fr
+    https://caffe2.ai/docs/sparse-operations.html#null__unsorted-segment-reduction-ops
 
     Parameters
     ----------
-    sparse_indices : relay.Expr
-        A 2-D tensor[N, n_dim] of integers containing location of sparse values, where N is the
-        number of sparse values and n_dim is the number of dimensions of the dense_shape
-    prev_shape : relay.Expr
-        A 1-D tensor containing the previous shape of the dense tensor
-    new_shape : relay.Expr
-        A 1-D tensor containing the new shape of the dense tensor
+    data : relay.Expr
+        Input floating point data
+    indices : relay.Expr
+        A 1-D tensor containing the indices of the rows to calculate the output sum upon
+    num_segments : Optional[int]
+        An integer describing the shape of the zeroth dimension. If unspecified, its calculated
+        equivalent to the number of unique indices
     Returns
     -------
     result: relay.Expr
@@ -1470,26 +1474,22 @@ def segment_sum(data, indices, num_segments=None):
     Examples
     --------
     .. code-block:: python
-        sparse_indices = [[0, 0, 0],
-                            [0, 0, 1],
-                            [0, 1, 0],
-                            [1, 0, 0],
-                            [1, 2, 3]]
-        prev_shape = [2, 3, 4]
-        new_shape = [9, -1]
-        new_sparse_indices, new_shape = relay.sparse_reshape(sparse_indices,
-                            prev_shape,
-                            new_shape)
-        new_sparse_indices = [[0, 0],
-                              [0, 1],
-                              [1, 2],
-                              [4, 2],
-                              [8, 1]] 
-        new_shape = [9, 4]
+        data = [[1, 2, 3, 4],
+                [4, -3, 2, -1],
+                [5, 6, 7, 8]]
+        indices = [0, 0, 1]
+        result = [[5, -1, 5, 3],[5, 6, 7, 8]]
+
+        data = [[1, 2, 3, 4],
+                [4, -3, 2, -1],
+                [5, 6, 7, 8]]
+        indices = [0, 0, 2]
+        num_segments = 3
+        result = [[5, -1, 5, 3],[0, 0, 0, 0], [5, 6, 7, 8]]
     """
-    
+
     if num_segments:
-        num_unique = const([num_segments]) 
+        num_unique = const([num_segments])
     else:
         _, _, num_unique = unique(reshape(indices, -1))
     data_offrow_shape = strided_slice(_make.shape_of(data, "int64"), [1], [-1], slice_mode="size")
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 8cc6d3094686..14140c2336b1 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1542,7 +1542,7 @@ def verify_sparse_reshape(
         ),
         (
             np.random.random((9, 4, 5, 7)),
-            np.array([2, 0, 1, 0, 3, 2, 5, 4, 4], dtype=np.int32),
+            np.array([5, 0, 1, 0, 3, 6, 8, 7, 7], dtype=np.int32),
             9,
         ),
     ],

From 74e8370934663d3c9bab7269d923a61f3e38c7b3 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 04:40:09 +0000
Subject: [PATCH 04/26] Black

---
 python/tvm/relay/op/transform.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index 3573d494dc6e..c150708beb30 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1478,6 +1478,7 @@ def segment_sum(data, indices, num_segments=None):
                 [4, -3, 2, -1],
                 [5, 6, 7, 8]]
         indices = [0, 0, 1]
+        result = segment_sum(data, indices)
         result = [[5, -1, 5, 3],[5, 6, 7, 8]]
 
         data = [[1, 2, 3, 4],
@@ -1485,6 +1486,7 @@ def segment_sum(data, indices, num_segments=None):
                 [5, 6, 7, 8]]
         indices = [0, 0, 2]
         num_segments = 3
+        result = segment_sum(data, indices, num_segments)
         result = [[5, -1, 5, 3],[0, 0, 0, 0], [5, 6, 7, 8]]
     """
 

From 8d5967567a6774b3f5c2a556815c3165ba776685 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 05:12:25 +0000
Subject: [PATCH 05/26] Add GPU

---
 tests/python/relay/test_op_level3.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 14140c2336b1..479c3a725a45 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1516,7 +1516,7 @@ def verify_sparse_reshape(
     )
 
 
-# @tvm.testing.uses_gpu
+@tvm.testing.uses_gpu
 @pytest.mark.parametrize(
     "data_np, indices_np, num_segments",
     [

From 6575048e7ecca67d1879d5a7b7e31936e3ece4d7 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 05:13:13 +0000
Subject: [PATCH 06/26] Uncomment

---
 tests/python/relay/test_op_level3.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 479c3a725a45..341733c5c3ed 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1605,10 +1605,6 @@ def verify_func(func, data, ref_res, target_ctx=tvm.testing.enabled_targets()):
                 for op_result, ref_result in zip(op_res, ref_res):
                     tvm.testing.assert_allclose(op_result.asnumpy(), ref_result, rtol=1e-5)
             else:
-                # print(op_res.asnumpy(), ref_res, *data)
-                # import pdb
-
-                # pdb.set_trace()
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
             relay.backend.compile_engine.get().clear()
 

From 9ea802d372ea85245dd82a114a39efd48b07d3e1 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 05:39:20 +0000
Subject: [PATCH 07/26] Add documentation

---
 python/tvm/relay/op/transform.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index c150708beb30..a3e5840df0a2 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1463,7 +1463,8 @@ def segment_sum(data, indices, num_segments=None):
     data : relay.Expr
         Input floating point data
     indices : relay.Expr
-        A 1-D tensor containing the indices of the rows to calculate the output sum upon
+        A 1-D tensor containing the indices of the rows to calculate the output sum upon.
+        This tensor doesn't need to be sorted
     num_segments : Optional[int]
         An integer describing the shape of the zeroth dimension. If unspecified, its calculated
         equivalent to the number of unique indices
@@ -1484,10 +1485,10 @@ def segment_sum(data, indices, num_segments=None):
         data = [[1, 2, 3, 4],
                 [4, -3, 2, -1],
                 [5, 6, 7, 8]]
-        indices = [0, 0, 2]
+        indices = [2, 0, 0]
         num_segments = 3
         result = segment_sum(data, indices, num_segments)
-        result = [[5, -1, 5, 3],[0, 0, 0, 0], [5, 6, 7, 8]]
+        result = [[5, 6, 7, 8],[0, 0, 0, 0], [5, -1, 5, 3]]
     """
 
     if num_segments:

From 23f8b2b1800d69f492728f6897138779f9d32298 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 06:46:46 +0000
Subject: [PATCH 08/26] Add dynamic tests

---
 tests/python/relay/test_op_level3.py | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 341733c5c3ed..30954e76ae00 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1547,7 +1547,8 @@ def verify_sparse_reshape(
         ),
     ],
 )
-def test_segment_sum(data_np, indices_np, num_segments):
+@pytest.mark.parametrize("use_dyn", [True, False])
+def test_segment_sum(data_np, indices_np, num_segments, use_dyn):
     def ref_segment_sum(
         data: np.ndarray,
         indices: np.ndarray,
@@ -1570,11 +1571,25 @@ def verify_segment_sum(
         """
         This function verifies the relay output of sparse_reshape with its expected output.
         """
-        data = relay.var(
-            "data",
-            relay.TensorType(data_np.shape, str(data_np.dtype)),
-        )
-        indices = relay.var("indices", relay.TensorType(indices_np.shape, str(indices_np.dtype)))
+        if use_dyn:
+            data = relay.var(
+                "data",
+                shape=[relay.Any() for _ in data_np.shape],
+                dtype=str(data_np.dtype),
+            )
+            indices = relay.var(
+                "indices",
+                shape=[relay.Any()],
+                dtype=str(indices_np.dtype),
+            )
+        else:
+            data = relay.var(
+                "data",
+                relay.TensorType(data_np.shape, str(data_np.dtype)),
+            )
+            indices = relay.var(
+                "indices", relay.TensorType(indices_np.shape, str(indices_np.dtype))
+            )
         z = relay.op.segment_sum(data, indices, num_segments)
 
         func = relay.Function([data, indices], z)

From a8d77fabc3dbeb010ac29bac0759fee1626eb3c4 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 06:57:25 +0000
Subject: [PATCH 09/26] Add TF Op

---
 python/tvm/relay/frontend/tensorflow.py       |  9 +++
 .../frontend/tensorflow/test_forward.py       | 60 +++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 20eb95ba7c00..0e51235c9db7 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1166,6 +1166,14 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def _math_segment_sum():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 2, "There should be 2 input tensors"
+        return get_relay_op("segment_sum")(inputs[0], inputs[1])
+
+    return _impl
+
+
 def _identity():
     def _impl(inputs, attr, params, mod):
         return inputs[0]
@@ -2660,6 +2668,7 @@ def _impl(inputs, attr, params, mod):
     "SparseTensorDenseMatMul": _sparse_tensor_dense_matmul(),
     "SparseFillEmptyRows": _sparse_fill_empty_rows(),
     "SparseReshape": _sparse_reshape(),
+    "SegmentSum": _math_segment_sum(),
     "Split": _split(False),
     "SplitV": _split(True),
     "Sqrt": AttrCvt("sqrt"),
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 41145bf77218..53f5d67126e8 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2080,6 +2080,66 @@ def test_forward_sparse_reshape(
     _test_sparse_reshape(sparse_indices_np, sparse_values_np, prev_shape_np, new_shape_np, use_dyn)
 
 
+#######################################################################
+# Math SegmentSum
+# ------------
+
+
+def _test_math_segment_sum(data_np, segment_ids_np, use_dyn=False):
+    with tf.Graph().as_default():
+        if use_dyn:
+            data = tf.placeholder(
+                shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
+            )
+            segment_ids = tf.placeholder(
+                shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+        else:
+            data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
+            segment_ids = tf.placeholder(
+                shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+
+        _ = tf.math.segment_sum(data, segment_ids, name="segment_sum")
+        compare_tf_with_tvm(
+            [data_np, segment_ids_np],
+            [data.name, segment_ids.name],
+            ["segment_sum:0"],
+            mode="vm",
+        )
+
+
+@pytest.mark.parametrize(
+    "data_np, segment_ids_np",
+    [
+        (
+            np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
+            np.array([0, 0, 0, 1, 1, 1], dtype=np.int32),
+        ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+            np.array([0, 0, 1], dtype=np.int32),
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([0, 0, 1, 2, 2, 3], dtype=np.int32),
+        ),
+        (
+            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
+            np.array([0, 0, 1], dtype=np.int32),
+        ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int32),
+        ),
+    ],
+)
+@pytest.mark.parametrize("use_dyn", [True, False])
+def test_forward_math_segment_sum(data_np, segment_ids_np, use_dyn):
+    """math segment sum test"""
+    _test_math_segment_sum(data_np, segment_ids_np, use_dyn)
+
+
 # tensorflow.compat.v1.sparse_to_dense
 # ---------------
 def _test_sparse_to_dense(sparse_indices, sparse_values, default_value, output_shape):

From 5d2d1926babeab0e6ccf46dadae07a727733af3f Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 07:31:28 +0000
Subject: [PATCH 10/26] Add Sparse Segment Sum

---
 python/tvm/relay/frontend/tensorflow.py       | 14 ++++
 .../frontend/tensorflow/test_forward.py       | 74 +++++++++++++++++++
 2 files changed, 88 insertions(+)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 0e51235c9db7..8a47dd11a107 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1174,6 +1174,19 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def _sparse_segment_sum():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) <= 4, "There should be less than equal to 4 input tensors"
+        data = _op.take(inputs[0], inputs[1], axis=0)
+        if len(inputs) == 3:
+            num_segments = None
+        else:
+            num_segments = inputs[3]
+        return _op.segment_sum(data, inputs[2], num_segments)
+
+    return _impl
+
+
 def _identity():
     def _impl(inputs, attr, params, mod):
         return inputs[0]
@@ -2669,6 +2682,7 @@ def _impl(inputs, attr, params, mod):
     "SparseFillEmptyRows": _sparse_fill_empty_rows(),
     "SparseReshape": _sparse_reshape(),
     "SegmentSum": _math_segment_sum(),
+    "SparseSegmentSum": _sparse_segment_sum(),
     "Split": _split(False),
     "SplitV": _split(True),
     "Sqrt": AttrCvt("sqrt"),
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 53f5d67126e8..7f5a3d264921 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2080,6 +2080,80 @@ def test_forward_sparse_reshape(
     _test_sparse_reshape(sparse_indices_np, sparse_values_np, prev_shape_np, new_shape_np, use_dyn)
 
 
+#######################################################################
+# Sparse SegmentSum
+# ------------
+
+
+def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn=False):
+    with tf.Graph().as_default():
+        if use_dyn:
+            data = tf.placeholder(
+                shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
+            )
+            indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
+            segment_ids = tf.placeholder(
+                shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+        else:
+            data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
+            indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
+            segment_ids = tf.placeholder(
+                shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+
+        _ = tf.sparse.segment_sum(
+            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum"
+        )
+        compare_tf_with_tvm(
+            [data_np, indices_np, segment_ids_np],
+            [data.name, indices.name, segment_ids.name],
+            ["sparse_segment_sum:0"],
+            mode="vm",
+        )
+
+
+@pytest.mark.parametrize(
+    "data_np, indices_np, segment_ids_np, num_segments",
+    [
+        (
+            np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
+            np.array([0, 3, 4], dtype=np.int32),
+            np.array([0, 1, 1], dtype=np.int32),
+            None,
+        ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+            np.array([0, 1, 2], dtype=np.int32),
+            np.array([0, 0, 1], dtype=np.int32),
+            None,
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([0, 1, 2, 3, 4, 5], dtype=np.int32),
+            np.array([0, 0, 1, 2, 2, 3], dtype=np.int32),
+            None,
+        ),
+        (
+            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
+            np.array([0, 1, 2], dtype=np.int32),
+            np.array([0, 0, 1], dtype=np.int32),
+            None,
+        ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+            np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int32),
+            None,
+        ),
+    ],
+)
+@pytest.mark.parametrize("use_dyn", [True, False])
+def test_forward_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn):
+    """sparse segment sum test"""
+    _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
+
+
 #######################################################################
 # Math SegmentSum
 # ------------

From c3d90a6a3572cee2819ac4ed18926318cf330f4c Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 07:47:58 +0000
Subject: [PATCH 11/26] Add test coverage

---
 python/tvm/relay/frontend/tensorflow.py         | 17 ++++++++++++-----
 .../python/frontend/tensorflow/test_forward.py  | 16 ++++++++--------
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 8a47dd11a107..f4b85907dc82 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1176,12 +1176,18 @@ def _impl(inputs, attr, params, mod):
 
 def _sparse_segment_sum():
     def _impl(inputs, attr, params, mod):
-        assert len(inputs) <= 4, "There should be less than equal to 4 input tensors"
+        assert len(inputs) == 3, "There should be 3 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
-        if len(inputs) == 3:
-            num_segments = None
-        else:
-            num_segments = inputs[3]
+        return _op.segment_sum(data, inputs[2])
+
+    return _impl
+
+
+def _sparse_segment_sum_with_num_segments():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 4, "There should be 4 input tensors"
+        data = _op.take(inputs[0], inputs[1], axis=0)
+        num_segments = int(inputs[3].data.asnumpy().item())
         return _op.segment_sum(data, inputs[2], num_segments)
 
     return _impl
@@ -2683,6 +2689,7 @@ def _impl(inputs, attr, params, mod):
     "SparseReshape": _sparse_reshape(),
     "SegmentSum": _math_segment_sum(),
     "SparseSegmentSum": _sparse_segment_sum(),
+    "SparseSegmentSumWithNumSegments": _sparse_segment_sum_with_num_segments(),
     "Split": _split(False),
     "SplitV": _split(True),
     "Sqrt": AttrCvt("sqrt"),
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 7f5a3d264921..06daa6106f08 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2124,15 +2124,15 @@ def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments,
         ),
         (
             np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
-            np.array([0, 1, 2], dtype=np.int32),
-            np.array([0, 0, 1], dtype=np.int32),
-            None,
+            np.array([0, 1], dtype=np.int32),
+            np.array([0, 2], dtype=np.int32),
+            4,
         ),
         (
             np.random.random((6, 4, 5)),
-            np.array([0, 1, 2, 3, 4, 5], dtype=np.int32),
-            np.array([0, 0, 1, 2, 2, 3], dtype=np.int32),
-            None,
+            np.array([0, 2, 4, 3, 1], dtype=np.int32),
+            np.array([0, 0, 1, 5, 5], dtype=np.int32),
+            6,
         ),
         (
             np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
@@ -2143,8 +2143,8 @@ def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments,
         (
             np.random.random((9, 4, 5, 7)),
             np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
-            np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int32),
-            None,
+            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
+            9,
         ),
     ],
 )

From bac338e71afe20df21c98b1ca31cf7d656cdae90 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 21:18:10 +0000
Subject: [PATCH 12/26] PR Comments

---
 python/tvm/relay/op/transform.py     | 46 ++++++++++++++++------------
 tests/python/relay/test_op_level3.py | 40 ++++++++++++------------
 2 files changed, 47 insertions(+), 39 deletions(-)

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index a3e5840df0a2..8e45932bdab0 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1450,24 +1450,30 @@ def sparse_reshape(sparse_indices, prev_shape, new_shape):
     return TupleWrapper(_make.sparse_reshape(sparse_indices, prev_shape, new_shape), 2)
 
 
-def segment_sum(data, indices, num_segments=None):
+def segment_sum(data, segment_ids, num_segments=None):
     """
-    Computes the sum along segments of a tensor. This op is much better understood with
-    visualization articulated in the following links and examples at the end of this docstring.
+    Computes the sum along segment_ids along axis 0. If multiple segment_ids reference the same
+    location their contributions add up.
+    result[index] = Σi... data[i...] where index = segment_ids[i]
+    This op is much better understood with visualization articulated in the following links and
+    examples at the end of this docstring.
 
-    https://www.tensorflow.org/api_docs/python/tf/raw_ops/UnsortedSegmentSum?hl=fr
+    https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum
     https://caffe2.ai/docs/sparse-operations.html#null__unsorted-segment-reduction-ops
 
     Parameters
     ----------
     data : relay.Expr
-        Input floating point data
-    indices : relay.Expr
-        A 1-D tensor containing the indices of the rows to calculate the output sum upon.
-        This tensor doesn't need to be sorted
+        Input Tensor. It can be of any type and multi-dimensional
+    segment_ids : relay.Expr
+        A 1-D int32/int64 tensor containing the segment_ids of the rows to calculate the output
+        sum upon. It defines a mapping from the zeroth dimension of data onto segment_ids. The
+        segment_ids tensor should be the size of the first dimension, d0, with consecutive IDs
+        in the range 0 to k, where k<d0. In particular, a segmentation of a matrix tensor is a
+        mapping of rows to segments. This tensor doesn't need to be sorted
     num_segments : Optional[int]
         An integer describing the shape of the zeroth dimension. If unspecified, its calculated
-        equivalent to the number of unique indices
+        equivalent to the number of unique segment_ids
     Returns
     -------
     result: relay.Expr
@@ -1478,31 +1484,33 @@ def segment_sum(data, indices, num_segments=None):
         data = [[1, 2, 3, 4],
                 [4, -3, 2, -1],
                 [5, 6, 7, 8]]
-        indices = [0, 0, 1]
-        result = segment_sum(data, indices)
+        segment_ids = [0, 0, 1]
+        result = segment_sum(data, segment_ids)
         result = [[5, -1, 5, 3],[5, 6, 7, 8]]
 
         data = [[1, 2, 3, 4],
                 [4, -3, 2, -1],
                 [5, 6, 7, 8]]
-        indices = [2, 0, 0]
+        segment_ids = [2, 0, 0]
         num_segments = 3
-        result = segment_sum(data, indices, num_segments)
+        result = segment_sum(data, segment_ids, num_segments)
         result = [[5, 6, 7, 8],[0, 0, 0, 0], [5, -1, 5, 3]]
     """
 
     if num_segments:
         num_unique = const([num_segments])
     else:
-        _, _, num_unique = unique(reshape(indices, -1))
+        _, _, num_unique = unique(reshape(segment_ids, -1))
     data_offrow_shape = strided_slice(_make.shape_of(data, "int64"), [1], [-1], slice_mode="size")
-    data_offrow_shape = cast_like(data_offrow_shape, indices)
+    data_offrow_shape = cast_like(data_offrow_shape, segment_ids)
     new_shape = _make.concatenate(Tuple([num_unique, data_offrow_shape]), 0)
-    indices_tiled_shape = _make.concatenate(Tuple([reverse(data_offrow_shape, 0), const([1])]), 0)
-    expanded_indices = tile(indices, indices_tiled_shape)
-    scatter_add_indices = transpose(expanded_indices)
+    segment_ids_tiled_shape = _make.concatenate(
+        Tuple([reverse(data_offrow_shape, 0), const([1])]), 0
+    )
+    expanded_segment_ids = tile(segment_ids, segment_ids_tiled_shape)
+    scatter_add_segment_ids = transpose(expanded_segment_ids)
     src = cast_like(_dyn_make.zeros(new_shape, "float64"), data)
-    return scatter_add(src, scatter_add_indices, data, axis=0)
+    return scatter_add(src, scatter_add_segment_ids, data, axis=0)
 
 
 def cumsum(data, axis=None, dtype=None, exclusive=None):
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 30954e76ae00..4919ccfdfb83 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1518,7 +1518,7 @@ def verify_sparse_reshape(
 
 @tvm.testing.uses_gpu
 @pytest.mark.parametrize(
-    "data_np, indices_np, num_segments",
+    "data_np, segment_ids_np, num_segments",
     [
         (
             np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
@@ -1532,7 +1532,7 @@ def verify_sparse_reshape(
         ),
         (
             np.random.random((6, 4, 5)),
-            np.array([2, 0, 1, 0, 3, 2], dtype=np.int32),
+            np.array([2, 0, 1, 0, 3, 2], dtype=np.int64),
             None,
         ),
         (
@@ -1542,34 +1542,34 @@ def verify_sparse_reshape(
         ),
         (
             np.random.random((9, 4, 5, 7)),
-            np.array([5, 0, 1, 0, 3, 6, 8, 7, 7], dtype=np.int32),
+            np.array([5, 0, 1, 0, 3, 6, 8, 7, 7], dtype=np.int64),
             9,
         ),
     ],
 )
 @pytest.mark.parametrize("use_dyn", [True, False])
-def test_segment_sum(data_np, indices_np, num_segments, use_dyn):
+def test_segment_sum(data_np, segment_ids_np, num_segments, use_dyn):
     def ref_segment_sum(
         data: np.ndarray,
-        indices: np.ndarray,
+        segment_ids: np.ndarray,
         num_segments: Optional[int] = None,
     ):
         """
-        This function calculates the expected output of sparseshape operator given the inputs.
+        This function calculates the expected output of segment_sum operator given the inputs.
         """
         if not num_segments:
-            num_segments = np.unique(indices).shape[0]
+            num_segments = np.unique(segment_ids).shape[0]
 
         result = np.zeros((num_segments,) + data.shape[1:], data.dtype)
-        for i, index in enumerate(indices):
+        for i, index in enumerate(segment_ids):
             result[index] += data[i]
         return result
 
     def verify_segment_sum(
-        data_np: np.ndarray, indices_np: np.ndarray, num_segments: Optional[int]
+        data_np: np.ndarray, segment_ids_np: np.ndarray, num_segments: Optional[int]
     ):
         """
-        This function verifies the relay output of sparse_reshape with its expected output.
+        This function verifies the relay output of segment_sum with its expected output.
         """
         if use_dyn:
             data = relay.var(
@@ -1577,32 +1577,32 @@ def verify_segment_sum(
                 shape=[relay.Any() for _ in data_np.shape],
                 dtype=str(data_np.dtype),
             )
-            indices = relay.var(
-                "indices",
+            segment_ids = relay.var(
+                "segment_ids",
                 shape=[relay.Any()],
-                dtype=str(indices_np.dtype),
+                dtype=str(segment_ids_np.dtype),
             )
         else:
             data = relay.var(
                 "data",
                 relay.TensorType(data_np.shape, str(data_np.dtype)),
             )
-            indices = relay.var(
-                "indices", relay.TensorType(indices_np.shape, str(indices_np.dtype))
+            segment_ids = relay.var(
+                "segment_ids", relay.TensorType(segment_ids_np.shape, str(segment_ids_np.dtype))
             )
-        z = relay.op.segment_sum(data, indices, num_segments)
+        z = relay.op.segment_sum(data, segment_ids, num_segments)
 
-        func = relay.Function([data, indices], z)
-        ref_res = ref_segment_sum(data_np, indices_np, num_segments=num_segments)
+        func = relay.Function([data, segment_ids], z)
+        ref_res = ref_segment_sum(data_np, segment_ids_np, num_segments=num_segments)
         segment_sum_result = run_infer_type(z)
         assert segment_sum_result.checked_type.dtype == data_np.dtype
         verify_func(
             func,
-            [data_np, indices_np],
+            [data_np, segment_ids_np],
             ref_res,
         )
 
-    verify_segment_sum(data_np, indices_np, num_segments)
+    verify_segment_sum(data_np, segment_ids_np, num_segments)
 
 
 def verify_func(func, data, ref_res, target_ctx=tvm.testing.enabled_targets()):

From e11f5e19bdef2f37d3f5d1e154808b658660da3c Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Tue, 2 Mar 2021 22:13:01 +0000
Subject: [PATCH 13/26] Int64 tests

---
 python/tvm/relay/op/transform.py                 | 3 +--
 tests/python/frontend/tensorflow/test_forward.py | 4 ++--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index 8e45932bdab0..33d8b8949bc3 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1501,8 +1501,7 @@ def segment_sum(data, segment_ids, num_segments=None):
         num_unique = const([num_segments])
     else:
         _, _, num_unique = unique(reshape(segment_ids, -1))
-    data_offrow_shape = strided_slice(_make.shape_of(data, "int64"), [1], [-1], slice_mode="size")
-    data_offrow_shape = cast_like(data_offrow_shape, segment_ids)
+    data_offrow_shape = strided_slice(_make.shape_of(data, "int32"), [1], [-1], slice_mode="size")
     new_shape = _make.concatenate(Tuple([num_unique, data_offrow_shape]), 0)
     segment_ids_tiled_shape = _make.concatenate(
         Tuple([reverse(data_offrow_shape, 0), const([1])]), 0
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 06daa6106f08..b23705dd8f99 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2196,7 +2196,7 @@ def _test_math_segment_sum(data_np, segment_ids_np, use_dyn=False):
         ),
         (
             np.random.random((6, 4, 5)),
-            np.array([0, 0, 1, 2, 2, 3], dtype=np.int32),
+            np.array([0, 0, 1, 2, 2, 3], dtype=np.int64),
         ),
         (
             np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
@@ -2204,7 +2204,7 @@ def _test_math_segment_sum(data_np, segment_ids_np, use_dyn=False):
         ),
         (
             np.random.random((9, 4, 5, 7)),
-            np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int32),
+            np.array([0, 0, 0, 1, 2, 3, 4, 4, 5], dtype=np.int64),
         ),
     ],
 )

From 9b762910788d42cb5954b8a8ae124d2e167b48fd Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 04:23:09 +0000
Subject: [PATCH 14/26] Add SparseSegmentSqrtN

---
 python/tvm/relay/frontend/tensorflow.py       | 47 ++++++++++
 python/tvm/relay/op/transform.py              | 12 ++-
 .../frontend/tensorflow/test_forward.py       | 94 ++++++++++++++++++-
 tests/python/relay/test_op_level3.py          |  2 +-
 4 files changed, 149 insertions(+), 6 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index f4b85907dc82..4f9a35111012 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1193,6 +1193,51 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def _sparse_segment_sum_sqrtn():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 3, "There should be 3 input tensors"
+        data = _op.take(inputs[0], inputs[1], axis=0)
+        _, _, num_unique, counts = _op.unique(inputs[2], return_counts=True)
+        segment_sum = _op.segment_sum(data, inputs[2])
+        strided_sqrt_counts = _op.sqrt(
+            _op.cast_like(_op.strided_slice(counts, [0], num_unique, slice_mode="end"), segment_sum)
+        )
+        segment_sum_offrow_shape = _op.strided_slice(
+            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
+        )
+        strided_sqrt_counts_tiled_shape = _op.concatenate(
+            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
+        )
+        strided_sqrt_counts_tiled = _op.transpose(
+            _op.tile(strided_sqrt_counts, strided_sqrt_counts_tiled_shape)
+        )
+        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+
+    return _impl
+
+
+def _sparse_segment_sum_sqrtn_with_num_segments():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 4, "There should be 4 input tensors"
+        data = _op.take(inputs[0], inputs[1], axis=0)
+        num_segments = int(inputs[3].data.asnumpy().item())
+        _, _, num_unique, counts = _op.unique(inputs[2], return_counts=True)
+        strided_sqrt_counts = _op.sqrt(_op.strided_slice(counts, [0], num_unique, slice_mode="end"))
+        segment_sum = _op.segment_sum(data, inputs[2], num_segments)
+        data_offrow_shape = _op.strided_slice(
+            _make.shape_of(data, "int32"), [1], [-1], slice_mode="size"
+        )
+        strided_sqrt_counts_tiled_shape = _op.concatenate(
+            [_op.transpose(data_offrow_shape), _expr.const([1])], axis=0
+        )
+        strided_sqrt_counts_tiled = _op.transpose(
+            _op.tile(strided_sqrt_counts, strided_sqrt_counts_tiled_shape)
+        )
+        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+
+    return _impl
+
+
 def _identity():
     def _impl(inputs, attr, params, mod):
         return inputs[0]
@@ -2690,6 +2735,8 @@ def _impl(inputs, attr, params, mod):
     "SegmentSum": _math_segment_sum(),
     "SparseSegmentSum": _sparse_segment_sum(),
     "SparseSegmentSumWithNumSegments": _sparse_segment_sum_with_num_segments(),
+    "SparseSegmentSqrtN": _sparse_segment_sum_sqrtn(),
+    "SparseSegmentSqrtNWithNumSegments": _sparse_segment_sum_sqrtn_with_num_segments(),
     "Split": _split(False),
     "SplitV": _split(True),
     "Sqrt": AttrCvt("sqrt"),
diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index 33d8b8949bc3..dc8a240498f9 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1497,14 +1497,18 @@ def segment_sum(data, segment_ids, num_segments=None):
         result = [[5, 6, 7, 8],[0, 0, 0, 0], [5, -1, 5, 3]]
     """
 
+    one_tensor = cast_like(const([1]), segment_ids)
     if num_segments:
-        num_unique = const([num_segments])
+        max_segments = const([num_segments])
+        max_segments = cast_like(max_segments, segment_ids)
     else:
-        _, _, num_unique = unique(reshape(segment_ids, -1))
+        max_segments = _make.add(reshape(_make.max(segment_ids, [0], False, False), -1), one_tensor)
+
     data_offrow_shape = strided_slice(_make.shape_of(data, "int32"), [1], [-1], slice_mode="size")
-    new_shape = _make.concatenate(Tuple([num_unique, data_offrow_shape]), 0)
+    data_offrow_shape = cast_like(data_offrow_shape, max_segments)
+    new_shape = _make.concatenate(Tuple([max_segments, data_offrow_shape]), 0)
     segment_ids_tiled_shape = _make.concatenate(
-        Tuple([reverse(data_offrow_shape, 0), const([1])]), 0
+        Tuple([reverse(data_offrow_shape, 0), one_tensor]), 0
     )
     expanded_segment_ids = tile(segment_ids, segment_ids_tiled_shape)
     scatter_add_segment_ids = transpose(expanded_segment_ids)
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index b23705dd8f99..df83a297f0ad 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -257,6 +257,8 @@ def name_without_num(name):
             for i in range(len(tf_output)):
                 if not isinstance(tf_output[i], np.ndarray):
                     assert len(tvm_output[i].shape) == 0
+                # print("TF Output Shape : ", tf_output[i].shape)
+                # print("TVM Output Shape :", tvm_output[i].shape)
                 tvm.testing.assert_allclose(tf_output[i], tvm_output[i], atol=1e-5, rtol=1e-5)
 
         sess.close()
@@ -2132,7 +2134,13 @@ def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments,
             np.random.random((6, 4, 5)),
             np.array([0, 2, 4, 3, 1], dtype=np.int32),
             np.array([0, 0, 1, 5, 5], dtype=np.int32),
-            6,
+            100,
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([0, 2, 4, 3, 1], dtype=np.int32),
+            np.array([0, 0, 1, 5, 5], dtype=np.int32),
+            None,
         ),
         (
             np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
@@ -2146,6 +2154,12 @@ def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments,
             np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
             9,
         ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
+            None,
+        ),
     ],
 )
 @pytest.mark.parametrize("use_dyn", [True, False])
@@ -2154,6 +2168,84 @@ def test_forward_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_seg
     _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
 
 
+#######################################################################
+# Sparse SegmentSum
+# ------------
+
+
+def _test_sparse_segment_sum_sqrt_n(
+    data_np, indices_np, segment_ids_np, num_segments, use_dyn=False
+):
+    with tf.Graph().as_default():
+        if use_dyn:
+            data = tf.placeholder(
+                shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
+            )
+            indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
+            segment_ids = tf.placeholder(
+                shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+        else:
+            data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
+            indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
+            segment_ids = tf.placeholder(
+                shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+
+        _ = tf.sparse.segment_sqrt_n(
+            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum"
+        )
+        compare_tf_with_tvm(
+            [data_np, indices_np, segment_ids_np],
+            [data.name, indices.name, segment_ids.name],
+            ["sparse_segment_sum:0"],
+            mode="vm",
+        )
+
+
+@pytest.mark.parametrize(
+    "data_np, indices_np, segment_ids_np, num_segments",
+    [
+        (
+            np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
+            np.array([0, 3, 4], dtype=np.int32),
+            np.array([0, 1, 1], dtype=np.int32),
+            None,
+        ),
+        # (
+        #     np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+        #     np.array([0, 1], dtype=np.int32),
+        #     np.array([0, 2], dtype=np.int32),
+        #     None,
+        # ),
+        # (
+        #     np.random.random((6, 4, 5)),
+        #     np.array([0, 2, 4, 3, 1], dtype=np.int32),
+        #     np.array([0, 0, 1, 5, 5], dtype=np.int32),
+        #     None,
+        # ),
+        (
+            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
+            np.array([0, 1, 2], dtype=np.int32),
+            np.array([0, 0, 1], dtype=np.int32),
+            None,
+        ),
+        # (
+        #     np.random.random((9, 4, 5, 7)),
+        #     np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+        #     np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
+        #     None,
+        # ),
+    ],
+)
+@pytest.mark.parametrize("use_dyn", [True, False])
+def test_forward_sparse_segment_sum_sqrt_n(
+    data_np, indices_np, segment_ids_np, num_segments, use_dyn
+):
+    """sparse segment sum test"""
+    _test_sparse_segment_sum_sqrt_n(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
+
+
 #######################################################################
 # Math SegmentSum
 # ------------
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 4919ccfdfb83..d194cdd0d1b2 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1595,7 +1595,7 @@ def verify_segment_sum(
         func = relay.Function([data, segment_ids], z)
         ref_res = ref_segment_sum(data_np, segment_ids_np, num_segments=num_segments)
         segment_sum_result = run_infer_type(z)
-        assert segment_sum_result.checked_type.dtype == data_np.dtype
+        # assert segment_sum_result.checked_type.dtype == data_np.dtype
         verify_func(
             func,
             [data_np, segment_ids_np],

From 20ece405c0925ac7625c824a2c36bbf71fa416c3 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 06:40:23 +0000
Subject: [PATCH 15/26] Add SparseSegmentSqrtNOp

---
 python/tvm/relay/frontend/tensorflow.py       | 52 ++++++++++++----
 python/tvm/relay/op/transform.py              |  7 ++-
 .../frontend/tensorflow/test_forward.py       | 62 +++++++++++++------
 tests/python/relay/test_op_level3.py          | 12 +++-
 4 files changed, 98 insertions(+), 35 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 4f9a35111012..27af1b147b08 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1197,11 +1197,22 @@ def _sparse_segment_sum_sqrtn():
     def _impl(inputs, attr, params, mod):
         assert len(inputs) == 3, "There should be 3 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
-        _, _, num_unique, counts = _op.unique(inputs[2], return_counts=True)
-        segment_sum = _op.segment_sum(data, inputs[2])
-        strided_sqrt_counts = _op.sqrt(
-            _op.cast_like(_op.strided_slice(counts, [0], num_unique, slice_mode="end"), segment_sum)
+        # This snippet calculates the sqrt count of each index among all valid
+        # indices(from 0 to max of segment ids)
+
+        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
+        max_ones = _op.reshape(
+            _op.max(_op.concatenate([max_segments, _op.shape_of(inputs[2])], 0)), -1
         )
+        counts = _op.segment_sum(_op.ones(max_ones, "int32"), inputs[2])
+        real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+        real_sqrt_counts = _op.sqrt(_op.cast_like(real_counts, data))
+
+        # Calculate regular segment sum
+        segment_sum = _op.segment_sum(data, inputs[2])
+
+        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
+        # appropriate shape first
         segment_sum_offrow_shape = _op.strided_slice(
             _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
         )
@@ -1209,7 +1220,7 @@ def _impl(inputs, attr, params, mod):
             [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
         )
         strided_sqrt_counts_tiled = _op.transpose(
-            _op.tile(strided_sqrt_counts, strided_sqrt_counts_tiled_shape)
+            _op.tile(real_sqrt_counts, strided_sqrt_counts_tiled_shape)
         )
         return _op.divide(segment_sum, strided_sqrt_counts_tiled)
 
@@ -1221,17 +1232,34 @@ def _impl(inputs, attr, params, mod):
         assert len(inputs) == 4, "There should be 4 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
         num_segments = int(inputs[3].data.asnumpy().item())
-        _, _, num_unique, counts = _op.unique(inputs[2], return_counts=True)
-        strided_sqrt_counts = _op.sqrt(_op.strided_slice(counts, [0], num_unique, slice_mode="end"))
-        segment_sum = _op.segment_sum(data, inputs[2], num_segments)
-        data_offrow_shape = _op.strided_slice(
-            _make.shape_of(data, "int32"), [1], [-1], slice_mode="size"
+        # This snippet calculates the sqrt count of each index among all valid
+        # indices(from 0 to max of [segment ids, num_segments])
+        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
+        max_ones = _op.reshape(
+            _op.max(
+                _op.concatenate(
+                    [max_segments, _op.shape_of(inputs[2]), _expr.const([num_segments])], 0
+                )
+            ),
+            -1,
+        )
+        counts = _op.segment_sum(_op.ones(max_ones, "int32"), inputs[2], num_segments)
+        real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+        real_sqrt_counts = _op.sqrt(_op.cast_like(real_counts, data))
+
+        # Calculate regular segment sum
+        segment_sum = _op.segment_sum(data, inputs[2], num_segments=num_segments)
+
+        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
+        # appropriate shape first
+        segment_sum_offrow_shape = _op.strided_slice(
+            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
         )
         strided_sqrt_counts_tiled_shape = _op.concatenate(
-            [_op.transpose(data_offrow_shape), _expr.const([1])], axis=0
+            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
         )
         strided_sqrt_counts_tiled = _op.transpose(
-            _op.tile(strided_sqrt_counts, strided_sqrt_counts_tiled_shape)
+            _op.tile(real_sqrt_counts, strided_sqrt_counts_tiled_shape)
         )
         return _op.divide(segment_sum, strided_sqrt_counts_tiled)
 
diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index dc8a240498f9..ac6e20eff7ed 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1499,8 +1499,11 @@ def segment_sum(data, segment_ids, num_segments=None):
 
     one_tensor = cast_like(const([1]), segment_ids)
     if num_segments:
-        max_segments = const([num_segments])
-        max_segments = cast_like(max_segments, segment_ids)
+        if isinstance(num_segments, int):
+            max_segments = const([num_segments])
+            max_segments = cast_like(max_segments, segment_ids)
+        else:
+            max_segments = cast_like(num_segments, segment_ids)
     else:
         max_segments = _make.add(reshape(_make.max(segment_ids, [0], False, False), -1), one_tensor)
 
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index df83a297f0ad..cf507e427e08 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -257,8 +257,6 @@ def name_without_num(name):
             for i in range(len(tf_output)):
                 if not isinstance(tf_output[i], np.ndarray):
                     assert len(tvm_output[i].shape) == 0
-                # print("TF Output Shape : ", tf_output[i].shape)
-                # print("TVM Output Shape :", tvm_output[i].shape)
                 tvm.testing.assert_allclose(tf_output[i], tvm_output[i], atol=1e-5, rtol=1e-5)
 
         sess.close()
@@ -2206,36 +2204,60 @@ def _test_sparse_segment_sum_sqrt_n(
 @pytest.mark.parametrize(
     "data_np, indices_np, segment_ids_np, num_segments",
     [
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
+            9,
+        ),
         (
             np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
             np.array([0, 3, 4], dtype=np.int32),
             np.array([0, 1, 1], dtype=np.int32),
             None,
         ),
-        # (
-        #     np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
-        #     np.array([0, 1], dtype=np.int32),
-        #     np.array([0, 2], dtype=np.int32),
-        #     None,
-        # ),
-        # (
-        #     np.random.random((6, 4, 5)),
-        #     np.array([0, 2, 4, 3, 1], dtype=np.int32),
-        #     np.array([0, 0, 1, 5, 5], dtype=np.int32),
-        #     None,
-        # ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+            np.array([0, 1], dtype=np.int32),
+            np.array([0, 2], dtype=np.int32),
+            None,
+        ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float32),
+            np.array([0, 1], dtype=np.int32),
+            np.array([0, 2], dtype=np.int32),
+            4,
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([0, 2, 4, 3, 1], dtype=np.int32),
+            np.array([0, 0, 1, 5, 5], dtype=np.int32),
+            100,
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([0, 2, 4, 3, 1], dtype=np.int32),
+            np.array([0, 0, 1, 5, 5], dtype=np.int32),
+            None,
+        ),
         (
             np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
             np.array([0, 1, 2], dtype=np.int32),
             np.array([0, 0, 1], dtype=np.int32),
             None,
         ),
-        # (
-        #     np.random.random((9, 4, 5, 7)),
-        #     np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
-        #     np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
-        #     None,
-        # ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
+            None,
+        ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+            np.array([0, 0, 1, 3, 5, 5, 5, 5, 5], dtype=np.int32),
+            6,
+        ),
     ],
 )
 @pytest.mark.parametrize("use_dyn", [True, False])
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index d194cdd0d1b2..810b9a6feaf2 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1545,6 +1545,16 @@ def verify_sparse_reshape(
             np.array([5, 0, 1, 0, 3, 6, 8, 7, 7], dtype=np.int64),
             9,
         ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+            np.array([0, 2], dtype=np.int32),
+            4,
+        ),
+        (
+            np.random.random((6, 4, 5)),
+            np.array([0, 0, 1, 5, 5], dtype=np.int32),
+            100,
+        ),
     ],
 )
 @pytest.mark.parametrize("use_dyn", [True, False])
@@ -1595,7 +1605,7 @@ def verify_segment_sum(
         func = relay.Function([data, segment_ids], z)
         ref_res = ref_segment_sum(data_np, segment_ids_np, num_segments=num_segments)
         segment_sum_result = run_infer_type(z)
-        # assert segment_sum_result.checked_type.dtype == data_np.dtype
+        assert segment_sum_result.checked_type.dtype == data_np.dtype
         verify_func(
             func,
             [data_np, segment_ids_np],

From 59a0bddd7cad898fae9e6b8da150e06ad63fdcfa Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 08:24:54 +0000
Subject: [PATCH 16/26] Deduplicate code

---
 .../frontend/tensorflow/test_forward.py       | 84 ++++---------------
 1 file changed, 16 insertions(+), 68 deletions(-)

diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index cf507e427e08..9a3308193e73 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2113,61 +2113,8 @@ def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments,
         )
 
 
-@pytest.mark.parametrize(
-    "data_np, indices_np, segment_ids_np, num_segments",
-    [
-        (
-            np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
-            np.array([0, 3, 4], dtype=np.int32),
-            np.array([0, 1, 1], dtype=np.int32),
-            None,
-        ),
-        (
-            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
-            np.array([0, 1], dtype=np.int32),
-            np.array([0, 2], dtype=np.int32),
-            4,
-        ),
-        (
-            np.random.random((6, 4, 5)),
-            np.array([0, 2, 4, 3, 1], dtype=np.int32),
-            np.array([0, 0, 1, 5, 5], dtype=np.int32),
-            100,
-        ),
-        (
-            np.random.random((6, 4, 5)),
-            np.array([0, 2, 4, 3, 1], dtype=np.int32),
-            np.array([0, 0, 1, 5, 5], dtype=np.int32),
-            None,
-        ),
-        (
-            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
-            np.array([0, 1, 2], dtype=np.int32),
-            np.array([0, 0, 1], dtype=np.int32),
-            None,
-        ),
-        (
-            np.random.random((9, 4, 5, 7)),
-            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
-            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
-            9,
-        ),
-        (
-            np.random.random((9, 4, 5, 7)),
-            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
-            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
-            None,
-        ),
-    ],
-)
-@pytest.mark.parametrize("use_dyn", [True, False])
-def test_forward_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn):
-    """sparse segment sum test"""
-    _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
-
-
 #######################################################################
-# Sparse SegmentSum
+# Sparse SegmentSumSqrtN
 # ------------
 
 
@@ -2204,24 +2151,12 @@ def _test_sparse_segment_sum_sqrt_n(
 @pytest.mark.parametrize(
     "data_np, indices_np, segment_ids_np, num_segments",
     [
-        (
-            np.random.random((9, 4, 5, 7)),
-            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
-            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
-            9,
-        ),
         (
             np.array([5, 1, 7, 2, 3, 4], dtype=np.float32),
             np.array([0, 3, 4], dtype=np.int32),
             np.array([0, 1, 1], dtype=np.int32),
             None,
         ),
-        (
-            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
-            np.array([0, 1], dtype=np.int32),
-            np.array([0, 2], dtype=np.int32),
-            None,
-        ),
         (
             np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float32),
             np.array([0, 1], dtype=np.int32),
@@ -2246,12 +2181,24 @@ def _test_sparse_segment_sum_sqrt_n(
             np.array([0, 0, 1], dtype=np.int32),
             None,
         ),
+        (
+            np.random.random((9, 4, 5, 7)),
+            np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
+            np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
+            9,
+        ),
         (
             np.random.random((9, 4, 5, 7)),
             np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
             np.array([0, 0, 1, 3, 5, 6, 7, 7, 8], dtype=np.int32),
             None,
         ),
+        (
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
+            np.array([0, 1], dtype=np.int32),
+            np.array([0, 2], dtype=np.int32),
+            None,
+        ),
         (
             np.random.random((9, 4, 5, 7)),
             np.array([0, 1, 2, 3, 4, 5, 6, 7, 8], dtype=np.int32),
@@ -2261,11 +2208,12 @@ def _test_sparse_segment_sum_sqrt_n(
     ],
 )
 @pytest.mark.parametrize("use_dyn", [True, False])
-def test_forward_sparse_segment_sum_sqrt_n(
+def test_forward_sparse_segment_sum_variants(
     data_np, indices_np, segment_ids_np, num_segments, use_dyn
 ):
-    """sparse segment sum test"""
+    """sparse segment sum variants tests"""
     _test_sparse_segment_sum_sqrt_n(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
+    _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
 
 
 #######################################################################

From 3d495f3f4448bda0e87e0e00ae0d3e296ac10c7d Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 08:51:59 +0000
Subject: [PATCH 17/26] Add SparseSegmentMean

---
 python/tvm/relay/frontend/tensorflow.py       | 77 ++++++++++++++++++-
 .../frontend/tensorflow/test_forward.py       | 37 ++++++++-
 2 files changed, 108 insertions(+), 6 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 27af1b147b08..3fb89666d8e6 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1204,7 +1204,7 @@ def _impl(inputs, attr, params, mod):
         max_ones = _op.reshape(
             _op.max(_op.concatenate([max_segments, _op.shape_of(inputs[2])], 0)), -1
         )
-        counts = _op.segment_sum(_op.ones(max_ones, "int32"), inputs[2])
+        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2])
         real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
         real_sqrt_counts = _op.sqrt(_op.cast_like(real_counts, data))
 
@@ -1243,7 +1243,7 @@ def _impl(inputs, attr, params, mod):
             ),
             -1,
         )
-        counts = _op.segment_sum(_op.ones(max_ones, "int32"), inputs[2], num_segments)
+        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2], num_segments)
         real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
         real_sqrt_counts = _op.sqrt(_op.cast_like(real_counts, data))
 
@@ -1266,6 +1266,77 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def _sparse_segment_mean():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 3, "There should be 3 input tensors"
+        data = _op.take(inputs[0], inputs[1], axis=0)
+        # This snippet calculates the sqrt count of each index among all valid
+        # indices(from 0 to max of segment ids)
+
+        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
+        max_ones = _op.reshape(
+            _op.max(_op.concatenate([max_segments, _op.shape_of(inputs[2])], 0)), -1
+        )
+        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2])
+        real_counts = _op.clip(counts, 1.0, 2147483647.0)  # Clip max doesn't work over int32
+
+        # Calculate regular segment sum
+        segment_sum = _op.segment_sum(data, inputs[2])
+
+        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
+        # appropriate shape first
+        segment_sum_offrow_shape = _op.strided_slice(
+            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
+        )
+        strided_sqrt_counts_tiled_shape = _op.concatenate(
+            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
+        )
+        strided_sqrt_counts_tiled = _op.transpose(
+            _op.tile(real_counts, strided_sqrt_counts_tiled_shape)
+        )
+        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+
+    return _impl
+
+
+def _sparse_segment_mean_with_num_segments():
+    def _impl(inputs, attr, params, mod):
+        assert len(inputs) == 4, "There should be 4 input tensors"
+        data = _op.take(inputs[0], inputs[1], axis=0)
+        num_segments = int(inputs[3].data.asnumpy().item())
+        # This snippet calculates the sqrt count of each index among all valid
+        # indices(from 0 to max of [segment ids, num_segments])
+        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
+        max_ones = _op.reshape(
+            _op.max(
+                _op.concatenate(
+                    [max_segments, _op.shape_of(inputs[2]), _expr.const([num_segments])], 0
+                )
+            ),
+            -1,
+        )
+        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2], num_segments)
+        real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+
+        # Calculate regular segment sum
+        segment_sum = _op.segment_sum(data, inputs[2], num_segments=num_segments)
+
+        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
+        # appropriate shape first
+        segment_sum_offrow_shape = _op.strided_slice(
+            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
+        )
+        strided_sqrt_counts_tiled_shape = _op.concatenate(
+            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
+        )
+        strided_sqrt_counts_tiled = _op.transpose(
+            _op.tile(real_counts, strided_sqrt_counts_tiled_shape)
+        )
+        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+
+    return _impl
+
+
 def _identity():
     def _impl(inputs, attr, params, mod):
         return inputs[0]
@@ -2765,6 +2836,8 @@ def _impl(inputs, attr, params, mod):
     "SparseSegmentSumWithNumSegments": _sparse_segment_sum_with_num_segments(),
     "SparseSegmentSqrtN": _sparse_segment_sum_sqrtn(),
     "SparseSegmentSqrtNWithNumSegments": _sparse_segment_sum_sqrtn_with_num_segments(),
+    "SparseSegmentMean": _sparse_segment_mean(),
+    "SparseSegmentMeanWithNumSegments": _sparse_segment_mean_with_num_segments(),
     "Split": _split(False),
     "SplitV": _split(True),
     "Sqrt": AttrCvt("sqrt"),
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 9a3308193e73..2c1479fe7066 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2138,12 +2138,40 @@ def _test_sparse_segment_sum_sqrt_n(
             )
 
         _ = tf.sparse.segment_sqrt_n(
-            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum"
+            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum_sqrt_n"
         )
         compare_tf_with_tvm(
             [data_np, indices_np, segment_ids_np],
             [data.name, indices.name, segment_ids.name],
-            ["sparse_segment_sum:0"],
+            ["sparse_segment_sum_sqrt_n:0"],
+            mode="vm",
+        )
+
+
+def _test_sparse_segment_mean(data_np, indices_np, segment_ids_np, num_segments, use_dyn=False):
+    with tf.Graph().as_default():
+        if use_dyn:
+            data = tf.placeholder(
+                shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
+            )
+            indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
+            segment_ids = tf.placeholder(
+                shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+        else:
+            data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
+            indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
+            segment_ids = tf.placeholder(
+                shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
+            )
+
+        _ = tf.sparse.segment_mean(
+            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_mean"
+        )
+        compare_tf_with_tvm(
+            [data_np, indices_np, segment_ids_np],
+            [data.name, indices.name, segment_ids.name],
+            ["sparse_segment_mean:0"],
             mode="vm",
         )
 
@@ -2158,7 +2186,7 @@ def _test_sparse_segment_sum_sqrt_n(
             None,
         ),
         (
-            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float32),
+            np.array([[1, 2, 3, 4], [-1, -2, -3, -4], [5, 6, 7, 8]], dtype=np.float64),
             np.array([0, 1], dtype=np.int32),
             np.array([0, 2], dtype=np.int32),
             4,
@@ -2176,7 +2204,7 @@ def _test_sparse_segment_sum_sqrt_n(
             None,
         ),
         (
-            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float32),
+            np.array([[[1, 7]], [[3, 8]], [[2, 9]]], dtype=np.float64),
             np.array([0, 1, 2], dtype=np.int32),
             np.array([0, 0, 1], dtype=np.int32),
             None,
@@ -2214,6 +2242,7 @@ def test_forward_sparse_segment_sum_variants(
     """sparse segment sum variants tests"""
     _test_sparse_segment_sum_sqrt_n(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
     _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
+    _test_sparse_segment_mean(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
 
 
 #######################################################################

From e2fb09888e54d1e470c8e5c1c4de64e229494e97 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 09:22:05 +0000
Subject: [PATCH 18/26] Parametrize Tests

---
 .../frontend/tensorflow/test_forward.py       | 111 ++++++++----------
 1 file changed, 47 insertions(+), 64 deletions(-)

diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 2c1479fe7066..a68ebba4772b 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2085,41 +2085,41 @@ def test_forward_sparse_reshape(
 # ------------
 
 
-def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn=False):
-    with tf.Graph().as_default():
-        if use_dyn:
-            data = tf.placeholder(
-                shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
-            )
-            indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
-            segment_ids = tf.placeholder(
-                shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
-            )
-        else:
-            data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
-            indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
-            segment_ids = tf.placeholder(
-                shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
-            )
-
-        _ = tf.sparse.segment_sum(
-            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum"
-        )
-        compare_tf_with_tvm(
-            [data_np, indices_np, segment_ids_np],
-            [data.name, indices.name, segment_ids.name],
-            ["sparse_segment_sum:0"],
-            mode="vm",
-        )
+# def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn=False):
+#     with tf.Graph().as_default():
+#         if use_dyn:
+#             data = tf.placeholder(
+#                 shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
+#             )
+#             indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
+#             segment_ids = tf.placeholder(
+#                 shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
+#             )
+#         else:
+#             data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
+#             indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
+#             segment_ids = tf.placeholder(
+#                 shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
+#             )
+
+#         _ = tf.sparse.segment_sum(
+#             data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum"
+#         )
+#         compare_tf_with_tvm(
+#             [data_np, indices_np, segment_ids_np],
+#             [data.name, indices.name, segment_ids.name],
+#             ["sparse_segment_sum:0"],
+#             mode="vm",
+#         )
 
 
 #######################################################################
-# Sparse SegmentSumSqrtN
+# Sparse Segment Variants
 # ------------
 
 
-def _test_sparse_segment_sum_sqrt_n(
-    data_np, indices_np, segment_ids_np, num_segments, use_dyn=False
+def _test_sparse_segment_variant(
+    tf_op, data_np, indices_np, segment_ids_np, num_segments, use_dyn=False
 ):
     with tf.Graph().as_default():
         if use_dyn:
@@ -2137,41 +2137,13 @@ def _test_sparse_segment_sum_sqrt_n(
                 shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
             )
 
-        _ = tf.sparse.segment_sqrt_n(
-            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum_sqrt_n"
-        )
-        compare_tf_with_tvm(
-            [data_np, indices_np, segment_ids_np],
-            [data.name, indices.name, segment_ids.name],
-            ["sparse_segment_sum_sqrt_n:0"],
-            mode="vm",
-        )
-
-
-def _test_sparse_segment_mean(data_np, indices_np, segment_ids_np, num_segments, use_dyn=False):
-    with tf.Graph().as_default():
-        if use_dyn:
-            data = tf.placeholder(
-                shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
-            )
-            indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
-            segment_ids = tf.placeholder(
-                shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
-            )
-        else:
-            data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
-            indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
-            segment_ids = tf.placeholder(
-                shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
-            )
-
-        _ = tf.sparse.segment_mean(
-            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_mean"
+        _ = tf_op(
+            data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_variant"
         )
         compare_tf_with_tvm(
             [data_np, indices_np, segment_ids_np],
             [data.name, indices.name, segment_ids.name],
-            ["sparse_segment_mean:0"],
+            ["sparse_segment_variant:0"],
             mode="vm",
         )
 
@@ -2236,13 +2208,24 @@ def _test_sparse_segment_mean(data_np, indices_np, segment_ids_np, num_segments,
     ],
 )
 @pytest.mark.parametrize("use_dyn", [True, False])
+@pytest.mark.parametrize(
+    "tf_op",
+    [
+        tf.sparse.segment_sum,
+        tf.sparse.segment_sqrt_n,
+        tf.sparse.segment_mean,
+    ],
+)
 def test_forward_sparse_segment_sum_variants(
-    data_np, indices_np, segment_ids_np, num_segments, use_dyn
+    tf_op,
+    data_np,
+    indices_np,
+    segment_ids_np,
+    num_segments,
+    use_dyn,
 ):
     """sparse segment sum variants tests"""
-    _test_sparse_segment_sum_sqrt_n(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
-    _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
-    _test_sparse_segment_mean(data_np, indices_np, segment_ids_np, num_segments, use_dyn)
+    _test_sparse_segment_variant(tf_op, data_np, indices_np, segment_ids_np, num_segments, use_dyn)
 
 
 #######################################################################

From 907c3d839561664798e813258a4c9244eae4a42a Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 09:22:48 +0000
Subject: [PATCH 19/26] Remove

---
 .../frontend/tensorflow/test_forward.py       | 34 -------------------
 1 file changed, 34 deletions(-)

diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index a68ebba4772b..1c310ac7a68c 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2079,40 +2079,6 @@ def test_forward_sparse_reshape(
     # ------------------------------------------------------------------
     _test_sparse_reshape(sparse_indices_np, sparse_values_np, prev_shape_np, new_shape_np, use_dyn)
 
-
-#######################################################################
-# Sparse SegmentSum
-# ------------
-
-
-# def _test_sparse_segment_sum(data_np, indices_np, segment_ids_np, num_segments, use_dyn=False):
-#     with tf.Graph().as_default():
-#         if use_dyn:
-#             data = tf.placeholder(
-#                 shape=[None for _ in data_np.shape], dtype=data_np.dtype, name="data"
-#             )
-#             indices = tf.placeholder(shape=[None], dtype=indices_np.dtype, name="indices")
-#             segment_ids = tf.placeholder(
-#                 shape=(None), dtype=segment_ids_np.dtype, name="segment_ids"
-#             )
-#         else:
-#             data = tf.placeholder(shape=data_np.shape, dtype=data_np.dtype, name="data")
-#             indices = tf.placeholder(shape=indices_np.shape, dtype=indices_np.dtype, name="indices")
-#             segment_ids = tf.placeholder(
-#                 shape=segment_ids_np.shape, dtype=segment_ids_np.dtype, name="segment_ids"
-#             )
-
-#         _ = tf.sparse.segment_sum(
-#             data, indices, segment_ids, num_segments=num_segments, name="sparse_segment_sum"
-#         )
-#         compare_tf_with_tvm(
-#             [data_np, indices_np, segment_ids_np],
-#             [data.name, indices.name, segment_ids.name],
-#             ["sparse_segment_sum:0"],
-#             mode="vm",
-#         )
-
-
 #######################################################################
 # Sparse Segment Variants
 # ------------

From 20c0f3ef479288ab8c368b16d33460519d9b2d11 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 09:37:14 +0000
Subject: [PATCH 20/26] Modularize

---
 python/tvm/relay/frontend/tensorflow.py | 65 +++++++------------------
 1 file changed, 17 insertions(+), 48 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 3fb89666d8e6..08769ca2cb05 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1193,6 +1193,19 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def row_wise_divide(multi_dim_tensor, one_dim_vector):
+    # To enable row-wise division of multi_dim_tensor and one_dim_vector, it must be tiled to the
+    # appropriate shape first
+    multi_dim_tensor_offrow_shape = _op.strided_slice(
+        _op.shape_of(multi_dim_tensor, "int32"), [1], [-1], slice_mode="size"
+    )
+    one_dim_vector_tiled_shape = _op.concatenate(
+        [_op.reverse(multi_dim_tensor_offrow_shape, 0), _expr.const([1])], axis=0
+    )
+    one_dim_vector_tiled = _op.transpose(_op.tile(one_dim_vector, one_dim_vector_tiled_shape))
+    return _op.divide(multi_dim_tensor, one_dim_vector_tiled)
+
+
 def _sparse_segment_sum_sqrtn():
     def _impl(inputs, attr, params, mod):
         assert len(inputs) == 3, "There should be 3 input tensors"
@@ -1211,18 +1224,7 @@ def _impl(inputs, attr, params, mod):
         # Calculate regular segment sum
         segment_sum = _op.segment_sum(data, inputs[2])
 
-        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
-        # appropriate shape first
-        segment_sum_offrow_shape = _op.strided_slice(
-            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
-        )
-        strided_sqrt_counts_tiled_shape = _op.concatenate(
-            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
-        )
-        strided_sqrt_counts_tiled = _op.transpose(
-            _op.tile(real_sqrt_counts, strided_sqrt_counts_tiled_shape)
-        )
-        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+        return row_wise_divide(segment_sum, real_sqrt_counts)
 
     return _impl
 
@@ -1250,18 +1252,7 @@ def _impl(inputs, attr, params, mod):
         # Calculate regular segment sum
         segment_sum = _op.segment_sum(data, inputs[2], num_segments=num_segments)
 
-        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
-        # appropriate shape first
-        segment_sum_offrow_shape = _op.strided_slice(
-            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
-        )
-        strided_sqrt_counts_tiled_shape = _op.concatenate(
-            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
-        )
-        strided_sqrt_counts_tiled = _op.transpose(
-            _op.tile(real_sqrt_counts, strided_sqrt_counts_tiled_shape)
-        )
-        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+        return row_wise_divide(segment_sum, real_sqrt_counts)
 
     return _impl
 
@@ -1283,18 +1274,7 @@ def _impl(inputs, attr, params, mod):
         # Calculate regular segment sum
         segment_sum = _op.segment_sum(data, inputs[2])
 
-        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
-        # appropriate shape first
-        segment_sum_offrow_shape = _op.strided_slice(
-            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
-        )
-        strided_sqrt_counts_tiled_shape = _op.concatenate(
-            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
-        )
-        strided_sqrt_counts_tiled = _op.transpose(
-            _op.tile(real_counts, strided_sqrt_counts_tiled_shape)
-        )
-        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+        return row_wise_divide(segment_sum, real_counts)
 
     return _impl
 
@@ -1321,18 +1301,7 @@ def _impl(inputs, attr, params, mod):
         # Calculate regular segment sum
         segment_sum = _op.segment_sum(data, inputs[2], num_segments=num_segments)
 
-        # To enable row-wise division of segment_sum and real_sqrt_counts, it must be tiled to the
-        # appropriate shape first
-        segment_sum_offrow_shape = _op.strided_slice(
-            _op.shape_of(segment_sum, "int32"), [1], [-1], slice_mode="size"
-        )
-        strided_sqrt_counts_tiled_shape = _op.concatenate(
-            [_op.reverse(segment_sum_offrow_shape, 0), _expr.const([1])], axis=0
-        )
-        strided_sqrt_counts_tiled = _op.transpose(
-            _op.tile(real_counts, strided_sqrt_counts_tiled_shape)
-        )
-        return _op.divide(segment_sum, strided_sqrt_counts_tiled)
+        return row_wise_divide(segment_sum, real_counts)
 
     return _impl
 

From 93b87d355c42192cbf3cd708bbffb6aacb67fdb3 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 09:38:22 +0000
Subject: [PATCH 21/26] Black

---
 tests/python/frontend/tensorflow/test_forward.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 1c310ac7a68c..81aeb5ef886c 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -2079,6 +2079,7 @@ def test_forward_sparse_reshape(
     # ------------------------------------------------------------------
     _test_sparse_reshape(sparse_indices_np, sparse_values_np, prev_shape_np, new_shape_np, use_dyn)
 
+
 #######################################################################
 # Sparse Segment Variants
 # ------------

From c20a026c897ab9edb41ffc5e205b9912536eee4d Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 09:59:51 +0000
Subject: [PATCH 22/26] Modularize Code

---
 python/tvm/relay/frontend/tensorflow.py | 62 +++++++------------------
 1 file changed, 18 insertions(+), 44 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 08769ca2cb05..82bdec7760c9 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1206,19 +1206,25 @@ def row_wise_divide(multi_dim_tensor, one_dim_vector):
     return _op.divide(multi_dim_tensor, one_dim_vector_tiled)
 
 
+def count_all_indices(segment_ids, counts_dtype, num_segments=None):
+    # This snippet calculates the sqrt count of each index among all valid
+    # indices(from 0 to max of [segment ids, num_segments])
+    max_segments = _op.reshape(_op.max(segment_ids), -1) + _expr.const([1])
+    if num_segments:
+        max_segments = _op.maximum(max_segments, _expr.const([num_segments]))
+    max_ones = _op.maximum(max_segments, _op.shape_of(segment_ids))
+    counts = _op.segment_sum(
+        _op.ones(max_ones, counts_dtype), segment_ids, num_segments=num_segments
+    )
+    real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+    return real_counts
+
+
 def _sparse_segment_sum_sqrtn():
     def _impl(inputs, attr, params, mod):
         assert len(inputs) == 3, "There should be 3 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
-        # This snippet calculates the sqrt count of each index among all valid
-        # indices(from 0 to max of segment ids)
-
-        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
-        max_ones = _op.reshape(
-            _op.max(_op.concatenate([max_segments, _op.shape_of(inputs[2])], 0)), -1
-        )
-        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2])
-        real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+        real_counts = count_all_indices(inputs[2], attr["T"].name)
         real_sqrt_counts = _op.sqrt(_op.cast_like(real_counts, data))
 
         # Calculate regular segment sum
@@ -1234,19 +1240,7 @@ def _impl(inputs, attr, params, mod):
         assert len(inputs) == 4, "There should be 4 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
         num_segments = int(inputs[3].data.asnumpy().item())
-        # This snippet calculates the sqrt count of each index among all valid
-        # indices(from 0 to max of [segment ids, num_segments])
-        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
-        max_ones = _op.reshape(
-            _op.max(
-                _op.concatenate(
-                    [max_segments, _op.shape_of(inputs[2]), _expr.const([num_segments])], 0
-                )
-            ),
-            -1,
-        )
-        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2], num_segments)
-        real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+        real_counts = count_all_indices(inputs[2], attr["T"].name, num_segments=num_segments)
         real_sqrt_counts = _op.sqrt(_op.cast_like(real_counts, data))
 
         # Calculate regular segment sum
@@ -1261,15 +1255,7 @@ def _sparse_segment_mean():
     def _impl(inputs, attr, params, mod):
         assert len(inputs) == 3, "There should be 3 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
-        # This snippet calculates the sqrt count of each index among all valid
-        # indices(from 0 to max of segment ids)
-
-        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
-        max_ones = _op.reshape(
-            _op.max(_op.concatenate([max_segments, _op.shape_of(inputs[2])], 0)), -1
-        )
-        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2])
-        real_counts = _op.clip(counts, 1.0, 2147483647.0)  # Clip max doesn't work over int32
+        real_counts = count_all_indices(inputs[2], attr["T"].name)
 
         # Calculate regular segment sum
         segment_sum = _op.segment_sum(data, inputs[2])
@@ -1284,19 +1270,7 @@ def _impl(inputs, attr, params, mod):
         assert len(inputs) == 4, "There should be 4 input tensors"
         data = _op.take(inputs[0], inputs[1], axis=0)
         num_segments = int(inputs[3].data.asnumpy().item())
-        # This snippet calculates the sqrt count of each index among all valid
-        # indices(from 0 to max of [segment ids, num_segments])
-        max_segments = _op.reshape(_op.max(inputs[2]), -1) + _expr.const([1])
-        max_ones = _op.reshape(
-            _op.max(
-                _op.concatenate(
-                    [max_segments, _op.shape_of(inputs[2]), _expr.const([num_segments])], 0
-                )
-            ),
-            -1,
-        )
-        counts = _op.segment_sum(_op.ones(max_ones, attr["T"].name), inputs[2], num_segments)
-        real_counts = _op.clip(counts, 1, 2147483647)  # Clip max doesn't work over int32
+        real_counts = count_all_indices(inputs[2], attr["T"].name, num_segments=num_segments)
 
         # Calculate regular segment sum
         segment_sum = _op.segment_sum(data, inputs[2], num_segments=num_segments)

From 82b2a1369221b51dcfb2b953ad7c0032c6b9cd6d Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 11:11:53 +0000
Subject: [PATCH 23/26] Pylint

---
 python/tvm/relay/frontend/tensorflow.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
index 82bdec7760c9..b03df22a8bc3 100644
--- a/python/tvm/relay/frontend/tensorflow.py
+++ b/python/tvm/relay/frontend/tensorflow.py
@@ -1194,8 +1194,10 @@ def _impl(inputs, attr, params, mod):
 
 
 def row_wise_divide(multi_dim_tensor, one_dim_vector):
-    # To enable row-wise division of multi_dim_tensor and one_dim_vector, it must be tiled to the
-    # appropriate shape first
+    """
+    This function enables row-wise division of multi_dim_tensor and one_dim_vector.
+    To achieve this, it is first tiled to the appropriate shape and then elemwise_division
+    """
     multi_dim_tensor_offrow_shape = _op.strided_slice(
         _op.shape_of(multi_dim_tensor, "int32"), [1], [-1], slice_mode="size"
     )
@@ -1207,8 +1209,11 @@ def row_wise_divide(multi_dim_tensor, one_dim_vector):
 
 
 def count_all_indices(segment_ids, counts_dtype, num_segments=None):
-    # This snippet calculates the sqrt count of each index among all valid
-    # indices(from 0 to max of [segment ids, num_segments])
+    """
+    This snippet calculates the sqrt count of each index among all valid indices
+    Valid indices are from 0 to max of [segment ids, num_segments]
+    """
+
     max_segments = _op.reshape(_op.max(segment_ids), -1) + _expr.const([1])
     if num_segments:
         max_segments = _op.maximum(max_segments, _expr.const([num_segments]))

From 97a2446400e2441f9fe28bf1a0a3cf4b9ba61c22 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 21:27:01 +0000
Subject: [PATCH 24/26] PR Comments

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

diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index ac6e20eff7ed..f2e3850a8f67 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -1454,7 +1454,7 @@ def segment_sum(data, segment_ids, num_segments=None):
     """
     Computes the sum along segment_ids along axis 0. If multiple segment_ids reference the same
     location their contributions add up.
-    result[index] = Σi... data[i...] where index = segment_ids[i]
+    result[index, j, k, ...] = Σi... data[i, j, k,..] where index = segment_ids[i]
     This op is much better understood with visualization articulated in the following links and
     examples at the end of this docstring.
 

From 6e77c0acb8f5c6c11ba0459137aff75f6ecf32e2 Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 21:57:42 +0000
Subject: [PATCH 25/26] Add scatter add tests

---
 tests/python/relay/test_op_level3.py | 58 ++++++++++++++--------------
 1 file changed, 30 insertions(+), 28 deletions(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 810b9a6feaf2..2bc772ec26a1 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1024,7 +1024,26 @@ def verify_dynamic_scatter(dshape, ishape, axis=0):
 
 
 @tvm.testing.uses_gpu
-def test_scatter_add():
+@pytest.mark.parametrize(
+    "dshape, ishape, axis, dtype",
+    [
+        ((10,), (10,), 0, "int32"),
+        ((1000,), (1000,), 0, "float32"),
+        ((1000,), (1000,), 0, "int32"),
+        ((10, 5), (10, 5), -2, "float32"),
+        ((10, 5), (10, 5), -1, "float32"),
+        ((10, 5), (3, 5), 0, "float32"),
+        ((12, 4), (7, 2), 1, "float32"),
+        ((2, 3, 4), (1, 3, 4), 0, "float32"),
+        ((2, 3, 4), (2, 1, 4), 1, "float32"),
+        ((2, 3, 4), (2, 3, 1), 2, "float32"),
+        ((2, 3, 4, 5), (1, 3, 4, 5), 0, "float32"),
+        ((6, 3, 4, 5), (2, 3, 4, 5), 1, "float32"),
+        ((2, 3, 8, 5), (2, 3, 1, 1), 2, "float32"),
+        ((16, 16, 4, 5), (16, 16, 4, 5), 3, "float32"),
+    ],
+)
+def test_scatter_add(dshape, ishape, axis, dtype):
     def ref_scatter_add(data, indices, updates, axis=0):
         output = np.copy(data)
         for index in np.ndindex(*indices.shape):
@@ -1034,9 +1053,9 @@ def ref_scatter_add(data, indices, updates, axis=0):
         return output
 
     def verify_scatter_add(dshape, ishape, axis=0, dtype="float32"):
-        d = relay.var("d", relay.TensorType(dshape, dtype))
-        i = relay.var("i", relay.TensorType(ishape, "int64"))
-        u = relay.var("u", relay.TensorType(ishape, dtype))
+        d = relay.var("d", relay.TensorType(shape=[relay.Any() for _ in dshape], dtype=dtype))
+        i = relay.var("i", relay.TensorType(shape=[relay.Any() for _ in ishape], dtype="int64"))
+        u = relay.var("u", relay.TensorType(shape=[relay.Any() for _ in ishape], dtype=dtype))
         z = relay.op.scatter_add(d, i, u, axis)
 
         func = relay.Function([d, i, u], z)
@@ -1046,31 +1065,14 @@ def verify_scatter_add(dshape, ishape, axis=0, dtype="float32"):
         indices_np = np.random.randint(-dshape[axis], dshape[axis] - 1, ishape).astype("int64")
 
         ref_res = ref_scatter_add(data_np, indices_np, updates_np, axis)
-        for target, ctx in tvm.testing.enabled_targets():
-            for kind in ["graph", "debug"]:
-                if target == "nvptx" and dtype == "float32" and len(dshape) == 1:
-                    # scatter_add 1D on GPU is implemented via atomic.
-                    # Floating point atomic requires LLVM 9 or newer for nvptx backend.
-                    # But LLVM on CI is LLVM 8.
-                    continue
-                intrp = relay.create_executor(kind, ctx=ctx, target=target)
-                op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
-                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
-    verify_scatter_add((10,), (10,), 0, dtype="int32")
-    verify_scatter_add((1000,), (1000,))
-    verify_scatter_add((1000,), (1000,), 0, dtype="int32")
-    verify_scatter_add((10, 5), (10, 5), -2)
-    verify_scatter_add((10, 5), (10, 5), -1)
-    verify_scatter_add((10, 5), (3, 5), 0)
-    verify_scatter_add((12, 4), (7, 2), 1)
-    verify_scatter_add((2, 3, 4), (1, 3, 4), 0)
-    verify_scatter_add((2, 3, 4), (2, 1, 4), 1)
-    verify_scatter_add((2, 3, 4), (2, 3, 1), 2)
-    verify_scatter_add((2, 3, 4, 5), (1, 3, 4, 5), 0)
-    verify_scatter_add((6, 3, 4, 5), (2, 3, 4, 5), 1)
-    verify_scatter_add((2, 3, 8, 5), (2, 3, 1, 1), 2)
-    verify_scatter_add((16, 16, 4, 5), (16, 16, 4, 5), 3)
+        verify_func(
+            func,
+            [data_np, indices_np, updates_np],
+            ref_res,
+        )
+
+    verify_scatter_add(dshape, ishape, axis, dtype)
 
 
 @tvm.testing.uses_gpu

From bfd71b60bcc804fb7d55da0c9fbad8d4799808ce Mon Sep 17 00:00:00 2001
From: Ubuntu <ubuntu@ip-172-31-42-251.us-east-2.compute.internal>
Date: Wed, 3 Mar 2021 23:00:55 +0000
Subject: [PATCH 26/26] Remove Test

---
 tests/python/relay/test_op_level3.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 2bc772ec26a1..31b95b0b49ae 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1028,7 +1028,6 @@ def verify_dynamic_scatter(dshape, ishape, axis=0):
     "dshape, ishape, axis, dtype",
     [
         ((10,), (10,), 0, "int32"),
-        ((1000,), (1000,), 0, "float32"),
         ((1000,), (1000,), 0, "int32"),
         ((10, 5), (10, 5), -2, "float32"),
         ((10, 5), (10, 5), -1, "float32"),