Aharrypotter commented on code in PR #19491:
URL: https://github.com/apache/tvm/pull/19491#discussion_r3174773948
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -2583,6 +2592,100 @@ def convert_scatter_nd(self, op):
data = relax.op.zeros(shape, updates_dtype)
return relax.op.scatter_nd(data, indices, updates, "update")
+ def _get_segment_scatter_base(self, output_shape, output_dtype, reduction):
+ """Create the identity base tensor for scatter-based segment
reductions."""
+ if reduction == "add":
+ return relax.op.zeros(output_shape, output_dtype)
+ if reduction == "mul":
+ return relax.op.full(output_shape, relax.const(1, output_dtype),
output_dtype)
+ if reduction == "min":
+ np_dtype = np.dtype(output_dtype)
+ if np.issubdtype(np_dtype, np.floating):
+ identity = np.finfo(np_dtype).max
+ elif np.issubdtype(np_dtype, np.integer):
+ identity = np.iinfo(np_dtype).max
+ else:
+ raise tvm.error.OpNotImplemented(
+ f"UNSORTED_SEGMENT_MIN does not support output dtype
{output_dtype}."
+ )
+ return relax.op.full(output_shape, relax.const(identity,
output_dtype), output_dtype)
+
+ raise ValueError(f"Unsupported segment reduction mode: {reduction}")
+
+ def _get_segment_num_segments(self, op_name, input_tensors):
+ if op_name == "SEGMENT_SUM":
+ segment_ids_tensor = input_tensors[1]
+ if self.has_expr(segment_ids_tensor.tensor_idx):
+ raise tvm.error.OpNotImplemented(
+ "TFLite SEGMENT_SUM with runtime segment_ids is not
supported, "
+ "because TFLite does not encode a reliable output segment
count."
+ )
Review Comment:
Verified that the original code parses correctly. I’ll include this
indentation cleanup if another commit update is needed.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
