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 f792a1d1aa [Relax][Frontend][TFLite] Support StableHLO shape ops
(#20114)
f792a1d1aa is described below
commit f792a1d1aa631ee8498600d0398d219df40d816c
Author: Hongyi Wu <[email protected]>
AuthorDate: Tue Aug 18 12:52:53 2026 +0800
[Relax][Frontend][TFLite] Support StableHLO shape ops (#20114)
## Summary
This PR adds Relax TFLite frontend support for the remaining StableHLO
shape
operators tracked by #19519:
- `STABLEHLO_RESHAPE` -> `R.reshape`
- `STABLEHLO_SLICE` -> `R.strided_slice`
- `STABLEHLO_TRANSPOSE` -> `R.permute_dims`
It carries forward the implementation from #19869 by @Mohxen onto the
current
`main` branch and addresses the outstanding review feedback by using
explicit
`ValueError` checks for the input and output arity of all three new
converters.
## Design
### StableHLO reshape
`STABLEHLO_RESHAPE` has one tensor input and a statically described
result
shape. The converter reads that shape from the TFLite output tensor
metadata
and emits `relax.op.reshape`.
### StableHLO slice
`STABLEHLO_SLICE` stores `start_indices`, `limit_indices`, and `strides`
in
`StablehloSliceOptions`. The converter parses those vectors, applies
them to all
input axes, and emits `relax.op.strided_slice`.
### StableHLO transpose
`STABLEHLO_TRANSPOSE` stores its permutation in
`StablehloTransposeOptions`. The converter parses the permutation and
emits
`relax.op.permute_dims`.
## Operator Support
| Operator | TFLite metadata | Relax lowering | Supported subset |
|---|---|---|---|
| `STABLEHLO_RESHAPE` | output tensor shape | `R.reshape` | static
result shape |
| `STABLEHLO_SLICE` | start, limit, and stride vectors |
`R.strided_slice` | static slice attributes |
| `STABLEHLO_TRANSPOSE` | permutation vector | `R.permute_dims` | static
permutation |
## Tests
The tests manually build minimal TFLite flatbuffers for each StableHLO
operator and compare the imported Relax IR with
`tvm.ir.assert_structural_equal`. The slice fixture exercises non-unit
strides,
and the transpose fixture uses a nontrivial three-dimensional
permutation.
Local validation:
```bash
python -m ruff format --check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m ruff check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m py_compile \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m pytest tests/python/relax/test_frontend_tflite.py \
-k "stablehlo_reshape or stablehlo_slice or stablehlo_transpose" -q
```
Result:
```text
ruff format --check: 2 files already formatted
ruff check: All checks passed
py_compile: passed
targeted StableHLO shape tests: 3 passed, 551 deselected
```
## References
- Completes the remaining StableHLO shape-operator items in #19519.
- Continues and supersedes #19869 by @Mohxen.
---------
Co-authored-by: Mohxen <[email protected]>
---
.../tvm/relax/frontend/tflite/tflite_frontend.py | 70 ++++++
tests/python/relax/test_frontend_tflite.py | 257 +++++++++++++++++++++
2 files changed, 327 insertions(+)
diff --git a/python/tvm/relax/frontend/tflite/tflite_frontend.py
b/python/tvm/relax/frontend/tflite/tflite_frontend.py
index d271221fc7..13d78c4c30 100644
--- a/python/tvm/relax/frontend/tflite/tflite_frontend.py
+++ b/python/tvm/relax/frontend/tflite/tflite_frontend.py
@@ -391,6 +391,7 @@ class OperatorConverter:
"STABLEHLO_REDUCE": self._convert_stablehlo_reduce,
"STABLEHLO_REDUCE_WINDOW": self._convert_stablehlo_reduce_window,
"STABLEHLO_REMAINDER": self._convert_stablehlo_remainder,
+ "STABLEHLO_RESHAPE": self._convert_stablehlo_reshape,
"STABLEHLO_RNG_BIT_GENERATOR":
self._convert_stablehlo_rng_bit_generator,
"STABLEHLO_RSQRT":
functools.partial(self._convert_stablehlo_unary, relax_op=_op.rsqrt),
"STABLEHLO_SCATTER": self._convert_stablehlo_scatter,
@@ -400,11 +401,13 @@ class OperatorConverter:
"STABLEHLO_SHIFT_LEFT": functools.partial(
self._convert_stablehlo_binary, relax_op=_op.left_shift
),
+ "STABLEHLO_SLICE": self._convert_stablehlo_slice,
"STABLEHLO_SORT": self._convert_stablehlo_sort,
"STABLEHLO_SUBTRACT": functools.partial(
self._convert_stablehlo_binary, relax_op=_op.subtract
),
"STABLEHLO_TANH": functools.partial(self._convert_stablehlo_unary,
relax_op=_op.tanh),
+ "STABLEHLO_TRANSPOSE": self._convert_stablehlo_transpose,
"STABLEHLO_WHILE": self._convert_stablehlo_while,
"SQUEEZE": self.convert_squeeze,
"STRIDED_SLICE": self.convert_strided_slice,
@@ -3188,6 +3191,73 @@ class OperatorConverter:
reshaped = self.bb.normalize(relax.op.reshape(in_expr,
intermediate_shape))
return self.bb.normalize(relax.op.broadcast_to(reshaped, output_shape))
+ def _convert_stablehlo_reshape(self, op):
+ """Convert STABLEHLO_RESHAPE to Relax."""
+ input_tensors = self.get_input_tensors(op)
+ if len(input_tensors) != 1:
+ raise ValueError(
+ f"STABLEHLO_RESHAPE expects exactly 1 input tensor, but got
{len(input_tensors)}"
+ )
+ output_tensors = self.get_output_tensors(op)
+ if len(output_tensors) != 1:
+ raise ValueError(
+ f"STABLEHLO_RESHAPE expects exactly 1 output tensor, but got
{len(output_tensors)}"
+ )
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+ output_shape = self._get_relax_tensor_shape(output_tensors[0])
+ return self.bb.normalize(relax.op.reshape(in_expr, output_shape))
+
+ def _convert_stablehlo_slice(self, op):
+ """Convert STABLEHLO_SLICE to Relax."""
+ from tflite.StablehloSliceOptions import StablehloSliceOptions
+
+ input_tensors = self.get_input_tensors(op)
+ if len(input_tensors) != 1:
+ raise ValueError(
+ f"STABLEHLO_SLICE expects exactly 1 input tensor, but got
{len(input_tensors)}"
+ )
+ output_tensors = self.get_output_tensors(op)
+ if len(output_tensors) != 1:
+ raise ValueError(
+ f"STABLEHLO_SLICE expects exactly 1 output tensor, but got
{len(output_tensors)}"
+ )
+
+ opts = self._get_stablehlo_options(op, StablehloSliceOptions)
+ begin = [int(d) for d in opts.StartIndicesAsNumpy()]
+ end = [int(d) for d in opts.LimitIndicesAsNumpy()]
+ strides = [int(d) for d in opts.StridesAsNumpy()]
+ axes = list(range(len(begin)))
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+ return self.bb.normalize(
+ relax.op.strided_slice(in_expr, axes=axes, begin=begin, end=end,
strides=strides)
+ )
+
+ def _convert_stablehlo_transpose(self, op):
+ """Convert STABLEHLO_TRANSPOSE to Relax."""
+ from tflite.StablehloTransposeOptions import StablehloTransposeOptions
+
+ input_tensors = self.get_input_tensors(op)
+ if len(input_tensors) != 1:
+ raise ValueError(
+ f"STABLEHLO_TRANSPOSE expects exactly 1 input tensor, but got
{len(input_tensors)}"
+ )
+ output_tensors = self.get_output_tensors(op)
+ if len(output_tensors) != 1:
+ raise ValueError(
+ "STABLEHLO_TRANSPOSE expects exactly 1 output tensor, "
+ f"but got {len(output_tensors)}"
+ )
+
+ opts = self._get_stablehlo_options(op, StablehloTransposeOptions)
+ permutation = [int(d) for d in opts.PermutationAsNumpy()]
+ if self._is_tflite_complex64_type(input_tensors[0].tensor.Type()):
+ permutation.append(len(permutation))
+
+ in_expr = self.get_tensor_expr(input_tensors[0])
+ return self.bb.normalize(relax.op.permute_dims(in_expr,
axes=permutation))
+
def _convert_stablehlo_iota(self, op):
"""Convert STABLEHLO_IOTA to Relax (arange + broadcast)."""
from tflite.StablehloIotaOptions import StablehloIotaOptions
diff --git a/tests/python/relax/test_frontend_tflite.py
b/tests/python/relax/test_frontend_tflite.py
index 1561826b17..e7ebbeaf92 100644
--- a/tests/python/relax/test_frontend_tflite.py
+++ b/tests/python/relax/test_frontend_tflite.py
@@ -4662,7 +4662,9 @@ _tfl_stablehlo_gather_opts =
_get_tflite_schema_module("StablehloGatherOptions")
_tfl_stablehlo_reduce_opts =
_get_tflite_schema_module("StablehloReduceOptions")
_tfl_stablehlo_reduce_window_opts =
_get_tflite_schema_module("StablehloReduceWindowOptions")
_tfl_stablehlo_scatter_opts =
_get_tflite_schema_module("StablehloScatterOptions")
+_tfl_stablehlo_slice_opts = _get_tflite_schema_module("StablehloSliceOptions")
_tfl_stablehlo_sort_opts = _get_tflite_schema_module("StablehloSortOptions")
+_tfl_stablehlo_transpose_opts =
_get_tflite_schema_module("StablehloTransposeOptions")
_tfl_stablehlo_while_opts = _get_tflite_schema_module("StablehloWhileOptions")
_tfl_stablehlo_rng_opts =
_get_tflite_schema_module("StablehloRngBitGeneratorOptions")
_tfl_call_options = _get_tflite_schema_module("CallOptions")
@@ -9512,6 +9514,261 @@ def test_stablehlo_concatenate(dimension):
tvm.ir.assert_structural_equal(mod, Expected)
+def _build_stablehlo_reshape_model(input_shape, output_shape,
tensor_type=_tfl_tensor_type.FLOAT32):
+ """STABLEHLO_RESHAPE with given input and output shapes."""
+ builder = flatbuffers.Builder(1024)
+
+ builtin_op = _get_stablehlo_builtin_operator("STABLEHLO_RESHAPE")
+ op_code = _build_operator_code(builder, builtin_op)
+
+ tensors = [
+ _build_tensor(builder, 0, input_shape, tensor_type=tensor_type),
+ _build_tensor(builder, 1, output_shape, tensor_type=tensor_type),
+ ]
+ op = _build_operator(builder, 0, [0], [1])
+ subgraph = _build_subgraph(
+ builder,
+ tensors=tensors,
+ operators=[op],
+ inputs=[0],
+ outputs=[1],
+ )
+ buffers = [_build_buffer(builder) for _ in range(2)]
+ return _finish_tflite_model(
+ builder, subgraph=subgraph, operator_codes=[op_code], buffers=buffers
+ )
+
+
+def test_stablehlo_reshape():
+ """TFLite StableHLO RESHAPE lowers to Relax reshape."""
+ mod = _load_model_from_buffer(
+ _build_stablehlo_reshape_model(input_shape=[2, 3], output_shape=[3, 2])
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((3, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((3, 2), dtype="float32") = R.reshape(x, (3, 2))
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_stablehlo_reshape_scalar():
+ """TFLite StableHLO RESHAPE supports a rank-0 output."""
+ mod =
_load_model_from_buffer(_build_stablehlo_reshape_model(input_shape=[1],
output_shape=[]))
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1,), dtype="float32")) -> R.Tensor((),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((), dtype="float32") = R.reshape(x, ())
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_stablehlo_reshape_complex64():
+ """TFLite StableHLO RESHAPE preserves the Relax complex pair axis."""
+ mod = _load_model_from_buffer(
+ _build_stablehlo_reshape_model(
+ input_shape=[2, 3],
+ output_shape=[3, 2],
+ tensor_type=_tfl_tensor_type.COMPLEX64,
+ )
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 2), dtype="float32")) -> R.Tensor((3, 2,
2), dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((3, 2, 2), dtype="float32") = R.reshape(x, (3, 2,
2))
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def _build_stablehlo_slice_model(input_shape, start_indices, limit_indices,
strides, output_shape):
+ """STABLEHLO_SLICE with static start, limit, and stride attributes."""
+ builder = flatbuffers.Builder(1024)
+
+ start_vec = _tflite_int64_vector(
+ builder,
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsStartStartIndicesVector,
+ start_indices,
+ )
+ limit_vec = _tflite_int64_vector(
+ builder,
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsStartLimitIndicesVector,
+ limit_indices,
+ )
+ strides_vec = _tflite_int64_vector(
+ builder,
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsStartStridesVector,
+ strides,
+ )
+
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsStart(builder)
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsAddStartIndices(builder,
start_vec)
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsAddLimitIndices(builder,
limit_vec)
+ _tfl_stablehlo_slice_opts.StablehloSliceOptionsAddStrides(builder,
strides_vec)
+ slice_opts = _tfl_stablehlo_slice_opts.StablehloSliceOptionsEnd(builder)
+
+ builtin_op = _get_stablehlo_builtin_operator("STABLEHLO_SLICE")
+ op_code = _build_operator_code(builder, builtin_op)
+
+ tensors = [
+ _build_tensor(builder, 0, input_shape),
+ _build_tensor(builder, 1, output_shape),
+ ]
+ op = _build_operator(
+ builder,
+ 0,
+ [0],
+ [1],
+ builtin_options2_type=_tfl_builtin_options2.StablehloSliceOptions,
+ builtin_options2=slice_opts,
+ )
+ subgraph = _build_subgraph(
+ builder,
+ tensors=tensors,
+ operators=[op],
+ inputs=[0],
+ outputs=[1],
+ )
+ buffers = [_build_buffer(builder) for _ in range(2)]
+ return _finish_tflite_model(
+ builder, subgraph=subgraph, operator_codes=[op_code], buffers=buffers
+ )
+
+
+def test_stablehlo_slice():
+ """TFLite StableHLO SLICE lowers to Relax strided_slice."""
+ mod = _load_model_from_buffer(
+ _build_stablehlo_slice_model(
+ input_shape=[4, 5],
+ start_indices=[1, 0],
+ limit_indices=[4, 4],
+ strides=[2, 2],
+ output_shape=[2, 2],
+ )
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((4, 5), dtype="float32")) -> R.Tensor((2, 2),
dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((2, 2), dtype="float32") = R.strided_slice(
+ x, axes=[0, 1], begin=[1, 0], end=[4, 4], strides=[2, 2]
+ )
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def _build_stablehlo_transpose_model(
+ input_shape, permutation, output_shape,
tensor_type=_tfl_tensor_type.FLOAT32
+):
+ """STABLEHLO_TRANSPOSE with a static permutation."""
+ builder = flatbuffers.Builder(1024)
+
+ perm_vec = _tflite_int64_vector(
+ builder,
+
_tfl_stablehlo_transpose_opts.StablehloTransposeOptionsStartPermutationVector,
+ permutation,
+ )
+ _tfl_stablehlo_transpose_opts.StablehloTransposeOptionsStart(builder)
+
_tfl_stablehlo_transpose_opts.StablehloTransposeOptionsAddPermutation(builder,
perm_vec)
+ transpose_opts =
_tfl_stablehlo_transpose_opts.StablehloTransposeOptionsEnd(builder)
+
+ builtin_op = _get_stablehlo_builtin_operator("STABLEHLO_TRANSPOSE")
+ op_code = _build_operator_code(builder, builtin_op)
+
+ tensors = [
+ _build_tensor(builder, 0, input_shape, tensor_type=tensor_type),
+ _build_tensor(builder, 1, output_shape, tensor_type=tensor_type),
+ ]
+ op = _build_operator(
+ builder,
+ 0,
+ [0],
+ [1],
+ builtin_options2_type=_tfl_builtin_options2.StablehloTransposeOptions,
+ builtin_options2=transpose_opts,
+ )
+ subgraph = _build_subgraph(
+ builder,
+ tensors=tensors,
+ operators=[op],
+ inputs=[0],
+ outputs=[1],
+ )
+ buffers = [_build_buffer(builder) for _ in range(2)]
+ return _finish_tflite_model(
+ builder, subgraph=subgraph, operator_codes=[op_code], buffers=buffers
+ )
+
+
+def test_stablehlo_transpose():
+ """TFLite StableHLO TRANSPOSE lowers to Relax permute_dims."""
+ mod = _load_model_from_buffer(
+ _build_stablehlo_transpose_model(
+ input_shape=[2, 3, 4], permutation=[1, 2, 0], output_shape=[3, 4,
2]
+ )
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), dtype="float32")) -> R.Tensor((3, 4,
2), dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((3, 4, 2), dtype="float32") = R.permute_dims(x,
axes=[1, 2, 0])
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_stablehlo_transpose_complex64():
+ """TFLite StableHLO TRANSPOSE leaves the Relax complex pair axis
trailing."""
+ mod = _load_model_from_buffer(
+ _build_stablehlo_transpose_model(
+ input_shape=[2, 3, 4],
+ permutation=[1, 2, 0],
+ output_shape=[3, 4, 2],
+ tensor_type=_tfl_tensor_type.COMPLEX64,
+ )
+ )
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 2), dtype="float32")) -> R.Tensor(
+ (3, 4, 2, 2), dtype="float32"
+ ):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((3, 4, 2, 2), dtype="float32") =
R.permute_dims(x, axes=[1, 2, 0, 3])
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
def _build_stablehlo_broadcast_in_dim_model(input_shape, broadcast_dims,
output_shape):
"""STABLEHLO_BROADCAST_IN_DIM with given broadcast dimensions."""
builder = flatbuffers.Builder(1024)
