[Relax][Frontend][TFLite] Add CUMSUM operator mapping#19434
Conversation
This commit adds frontend support for the TFLite `CUMSUM` operator by lowering it to `relax.op.cumsum`. Specifically, it handles: - Extracting the `axis` parameter from a constant tensor and converting it to an integer. - Parsing the `exclusive` flag from `CumsumOptions` via FlatBuffers. - Deriving the target `dtype` from the output tensor. - Raising a `NotImplementedError` for the `reverse` flag as it is not yet supported by the Relax op. Tracked in apache#19412.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for the TFLite CUMSUM operator to the Relax frontend, including the conversion logic and a test suite. The review feedback identifies two necessary improvements: implementing a check to prevent dynamic axis values which are currently unsupported, and adding a check for quantized inputs to prevent potential overflows or incorrect computations.
…s in CUMSUM operator
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request implements the CUMSUM operator in the TFLite to Relax frontend, including the conversion logic and a verification test case. The implementation supports basic cumulative sums but excludes support for quantization, dynamic axes, and reverse operations. A review comment suggests refactoring the axis extraction logic into a shared helper method to improve maintainability and reduce code duplication.
| axis = self.get_tensor_value(input_tensors[1]) | ||
| if isinstance(axis, np.ndarray): | ||
| assert axis.size == 1, "only one value is expected." | ||
| axis = int(axis.flat[0]) |
There was a problem hiding this comment.
The logic to extract the axis value from a numpy array is repeated in multiple places in this file. Consider creating a helper method to handle this extraction to improve maintainability.
2 participants
This commit adds frontend support for the TFLite
CUMSUMoperator by lowering it torelax.op.cumsum.Specifically, it handles:
axisparameter from a constant tensor and converting it to an integer.exclusiveflag fromCumsumOptionsvia FlatBuffers.dtypefrom the output tensor.NotImplementedErrorfor thereverseflag as it is not yet supported by the Relax op.Tracked in #19412.