giuseros commented on a change in pull request #6928:
URL: https://github.com/apache/incubator-tvm/pull/6928#discussion_r525237522
##########
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:
Could you add a bit more comments through this file? This would make it
easier to read and also more future proof
##########
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:
Is there a better way to do this? In theory you could pre-compute also
the segment sizes. Or you could calculate the sizes on the fly. At least, I
would put something like: `sys.float_info.min`
----------------------------------------------------------------
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:
us...@infra.apache.org
