[Relax][Frontend][TFLite] Support StableHLO shape ops

python/tvm/relax/frontend/tflite/tflite_frontend.py

@@ -391,6 +391,7 @@ def __init__(self, model, subgraph, exp_tab, ctx, conversion_state=None):
             "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 @@ def __init__(self, model, subgraph, exp_tab, ctx, conversion_state=None):
             "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 @@ def _convert_stablehlo_broadcast_in_dim(self, op):
         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

tests/python/relax/test_frontend_tflite.py

@@ -4662,7 +4662,9 @@ def _get_tflite_schema_enum(enum_name):
 _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 main(
     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)