This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 0b0afd8dd3 [Relax][Frontend][TFLite] Add CUMSUM operator mapping
(#19434)
0b0afd8dd3 is described below
commit 0b0afd8dd3e3cc1e17512a9a8e11d7db55285b2c
Author: Peruere1828 <[email protected]>
AuthorDate: Sat Apr 25 04:54:18 2026 +0800
[Relax][Frontend][TFLite] Add CUMSUM operator mapping (#19434)
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.
---
.../tvm/relax/frontend/tflite/tflite_frontend.py | 44 ++++++++++++++++++++++
tests/python/relax/test_frontend_tflite.py | 35 +++++++++++++++++
2 files changed, 79 insertions(+)
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py
b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index 334021b903..d773d8d7ce 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -129,6 +129,7 @@ class OperatorConverter:
"CONCATENATION": self.convert_concatenation,
"CONV_2D": functools.partial(self.convert_conv,
conv_type="conv2d"),
"COS": functools.partial(self._convert_unary_elemwise,
relax_op=_op.cos),
+ "CUMSUM": self.convert_cumsum,
"DENSIFY": self.convert_densify,
"DEPTH_TO_SPACE": self.convert_depth_to_space,
"DEPTHWISE_CONV_2D": functools.partial(self.convert_conv,
conv_type="depthwise"),
@@ -1426,6 +1427,49 @@ class OperatorConverter:
rhs_expr = self.get_tensor_expr(rhs_tensor)
lhs_expr = relax.op.add(lhs_expr, rhs_expr)
return lhs_expr
+
+ def convert_cumsum(self, op):
+ """Convert TFLite CUMSUM"""
+ if self.is_quantized(op):
+ raise tvm.error.OpNotImplemented(
+ "The TFLite to Relax converter does not support quantized
CUMSUM operator yet."
+ )
+
+ from tflite.BuiltinOptions import BuiltinOptions
+ from tflite.CumsumOptions import CumsumOptions
+
+ input_tensors = self.get_input_tensors(op)
+ assert len(input_tensors) == 2, "input tensors length should be 2"
+
+ input_expr = self.get_tensor_expr(input_tensors[0])
+
+ if self.has_expr(input_tensors[1].tensor_idx):
+ raise tvm.error.OpNotImplemented(
+ "The TFLite to Relax converter does not support dynamic axis
for CUMSUM yet."
+ )
+ 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])
+
+ assert op.BuiltinOptionsType() == BuiltinOptions.CumsumOptions
+ op_options = op.BuiltinOptions()
+ cumsum_options = CumsumOptions()
+ cumsum_options.Init(op_options.Bytes, op_options.Pos)
+ exclusive = cumsum_options.Exclusive()
+ if cumsum_options.Reverse():
+ raise tvm.error.OpNotImplemented(
+ "The TFLite to Relax converter does not support reverse CUMSUM
operator yet."
+ )
+
+ output_tensors = self.get_output_tensors(op)
+ assert len(output_tensors) == 1, "output tensors length should be 1"
+
+ out_dtype = self.get_tensor_type_str(output_tensors[0].tensor.Type())
+
+ out = relax.op.cumsum(input_expr, axis, out_dtype, exclusive)
+
+ return out
def convert_squared_difference(self, op):
"""Convert TFLite SQUARED DIFFERENCE"""
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index 92080634e2..15ca1cacf1 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -142,6 +142,41 @@ def test_add_n():
verify(AddN, Expected)
+def test_cumsum():
+ class Cumsum(tf.Module):
+ @tf.function(
+ input_signature=[
+ tf.TensorSpec(shape=(3, 4), dtype=tf.float32),
+ tf.TensorSpec(shape=(5, 6), dtype=tf.int32)
+ ]
+ )
+ def func(self, x, y):
+ out1 = tf.math.cumsum(x, axis=0)
+ out2 = tf.math.cumsum(y, axis=1, exclusive=True)
+ return out1, out2
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((3, 4), dtype="float32"),
+ y: R.Tensor((5, 6), dtype="int32"),
+ ) -> R.Tuple(R.Tensor((3, 4), dtype="float32"), R.Tensor((5, 6),
dtype="int32")):
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ gv1: R.Tensor((3, 4), dtype="float32") = R.cumsum(
+ x, axis=0, dtype="float32", exclusive=False
+ )
+ gv2: R.Tensor((5, 6), dtype="int32") = R.cumsum(
+ y, axis=1, dtype="int32", exclusive=True
+ )
+ gv = (gv1, gv2)
+ R.output(gv)
+ return gv
+
+ verify(Cumsum, Expected)
+
+
def test_split():
class Split(tf.Module):
@tf.function(input_signature=[tf.TensorSpec(shape=(1, 30),
dtype=tf.float32)])
