[GitHub] [incubator-tvm] alter-xp commented on a change in pull request #6928: [TF frontend] add some "Segment" and "UnsortedSegment" ops

[GitBox <gi...@apache.org>]

alter-xp commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r526690254



##########
File path: python/tvm/topi/tensor.py
##########
@@ -73,3 +75,341 @@ def full_like(x, fill_value):
         The result.
     """
     return cpp.full_like(x, fill_value)
+
+
+def segment_max(data, segment_ids, num_out):

Review comment:
       👌

##########
File path: python/tvm/topi/tensor.py
##########
@@ -73,3 +75,341 @@ def full_like(x, fill_value):
         The result.
     """
     return cpp.full_like(x, fill_value)
+
+
+def segment_max(data, segment_ids, num_out):
+    """segment_max operator.
+
+    Parameters
+    ----------
+    data : tvm.te.Tensor
+        input data
+
+    segment_ids : tvm.te.Tensor
+        input segment ids
+
+    num_out : int
+        number of output
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        Tensor with shape determined by the segment ids.
+    """
+
+    def _segment_max(data, segment_ids, out_buf):
+
+        ib = tir.ir_builder.create()
+        input_data = ib.buffer_ptr(data)
+        seg_ids = ib.buffer_ptr(segment_ids)
+        out = ib.buffer_ptr(out_buf)
+
+        shape = get_const_tuple(data.shape)
+        num_segment = get_const_tuple(out_buf.shape)[0]
+        inner_size = 1
+        for s in range(1, len(shape)):
+            inner_size = inner_size * shape[s]
+
+        with ib.for_range(0, num_segment) as n:
+            with ib.for_range(0, inner_size) as j:
+                out_index = n * inner_size + j
+                out[out_index] = -3.4028235e38

Review comment:
       `sys.float_info.min` is a number like 2.2250738585072014e-308, , which is not very suitable here, because it is always greater than 0. I replaced it with float("int")




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org