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

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name, unused-argument, too-many-lines
# pylint: disable=import-outside-toplevel, use-list-literal
# pylint: disable=no-value-for-parameter, unused-variable
# pylint: disable=unexpected-keyword-arg, unused-import, too-many-function-args
# ruff: noqa: RUF005
"""Tensorflow lite frontend."""

import functools
import itertools
import math

import numpy as np

import tvm
from tvm import relax, tirx
from tvm.relax import op as _op

from .tflite_flexbuffer import FlexBufferDecoder

__all__ = ["from_tflite"]


def to_int_list(np_array):
    """Convert a np array to a python int list.

    Note: This function converts np.int32 to python's int.
    If we don't do this conversion, numpy's automatic upcast will make
    the shape / parameters be converted to int64 IntImm in relax and
    cause problems in relax/TOPI.
    """
    return [int(x) for x in np_array]


class ExprTable:
    """Table storing Relax expressions by names."""

    def __init__(self):
        self.exprs = {}
        self.params = {}
        self.const_ctr = 1
        self.in_padding = False

    def new_const(self, value, shape=None, dtype="float32", source_name=None):
        """Construct a new var expr and add to exprs dictionary"""
        name = f"_param_{self.const_ctr}"
        self.const_ctr += 1
        self.exprs[name] = relax.const(value, dtype)
        self.params[name] = (self.exprs[name], value)

        return self.exprs[name]

    def get_expr(self, name):
        return self.exprs[name]

    def set_expr(self, name, expr, force_override=False):
        # assert isinstance(expr, _expr.Expr)
        # if name exists, we should override the value
        # otherwise, we can not get like x = func(x) work.
        # One example is CoreML preprocess, which will override
        # the same name of input.
        # However, according to git log, Find keras frontend depends
        # on this property, so we add one force_override to control it.
        if name not in self.exprs or force_override:
            self.exprs[name] = expr

    def has_expr(self, name):
        return name in self.exprs


class TensorWrapper:
    """Tensor wrapper for TFLite Tensor"""

    def __init__(self, tensor_idx, tensor, buffer, qnn_params=None):
        self.tensor_idx = tensor_idx
        self.tensor = tensor
        self.buffer = buffer
        self.qnn_params = qnn_params


class OperatorConverter:
    """Operator Converted for converting TFLite ops to Relax ops"""

    _SUPPORTED_QUANTIZED_OPS = frozenset(
        {
            "ABS",
            "ADD",
            "ATAN2",
            "CEIL",
            "CONCATENATION",
            "CONV_2D",
            "COS",
            "DEPTHWISE_CONV_2D",
            "DEQUANTIZE",
            "DETECTION_POSTPROCESS",
            "DIV",
            "EQUAL",
            "EXP",
            "FLOOR",
            "FLOOR_DIV",
            "FLOOR_MOD",
            "FULLY_CONNECTED",
            "GREATER",
            "GREATER_EQUAL",
            "HARD_SWISH",
            "LEAKY_RELU",
            "LESS",
            "LESS_EQUAL",
            "LOG",
            "LOGISTIC",
            "LOG_SOFTMAX",
            "MAXIMUM",
            "MEAN",
            "MINIMUM",
            "MUL",
            "NEG",
            "NOT_EQUAL",
            "POW",
            "QUANTIZE",
            "REDUCE_MAX",
            "REDUCE_MIN",
            "REDUCE_PROD",
            "RELU",
            "RELU6",
            "RELU_N1_TO_1",
            "RESHAPE",
            "RESIZE_BILINEAR",
            "ROUND",
            "RSQRT",
            "SIN",
            "SOFTMAX",
            "SQRT",
            "SQUARED_DIFFERENCE",
            "SUB",
            "SUM",
            "TAN",
            "TANH",
            "TRANSPOSE_CONV",
        }
    )

    def __init__(self, model, subgraph, exp_tab, ctx, conversion_state=None):
        from tflite.ActivationFunctionType import ActivationFunctionType
        from tflite.BuiltinOperator import BuiltinOperator
        from tflite.BuiltinOptions import BuiltinOptions

        self.model = model
        self.subgraph = subgraph
        self.exp_tab = exp_tab
        self.builtin_op_code = build_str_map(BuiltinOperator())
        self.activation_fn_type = build_str_map(ActivationFunctionType())
        self.builtin_options = build_str_map(BuiltinOptions())
        self.prefetched_nodes = {}
        self.allow_custom_ops = False
        self.bb = ctx
        if conversion_state is None:
            conversion_state = {
                "lowered_subgraphs": {},
                "lowered_if_functions": {},
                "lowered_while_functions": {},
                "lowering_stack": [],
                "module_builder": ctx,
                "resource_values": {},
                "hashtable_values": {},
                "in_call_once_init": False,
            }
        else:
            conversion_state.setdefault("module_builder", ctx)
            conversion_state.setdefault("resource_values", {})
            conversion_state.setdefault("hashtable_values", {})
            conversion_state.setdefault("in_call_once_init", False)
        self.conversion_state = conversion_state
        self.resource_handles = {}
        self.hashtable_handles = {}

        # Add more operators
        self.convert_map = {
            "ABS": functools.partial(self._convert_unary_elemwise, relax_op=_op.abs),
            "ADD": functools.partial(self._convert_elemwise, relax_op=_op.add),
            "ADD_N": self.convert_add_n,
            "ARG_MAX": functools.partial(self._convert_arg_min_max, relax_op=_op.argmax),
            "ARG_MIN": functools.partial(self._convert_arg_min_max, relax_op=_op.argmin),
            "ASSIGN_VARIABLE": self.convert_assign_variable,
            "ATAN2": functools.partial(self._convert_elemwise, relax_op=_op.atan2),
            "AVERAGE_POOL_2D": functools.partial(self.convert_pool2d, pool_type="average"),
            "BATCH_TO_SPACE_ND": self.convert_batch_to_space_nd,
            "BATCH_MATMUL": self.convert_batch_matmul,
            "BIDIRECTIONAL_SEQUENCE_LSTM": self.convert_bidirectional_sequence_lstm,
            "BIDIRECTIONAL_SEQUENCE_RNN": self.convert_bidirectional_sequence_rnn,
            "BITCAST": self.convert_bitcast,
            "BITWISE_XOR": functools.partial(self._convert_elemwise, relax_op=_op.bitwise_xor),
            "BROADCAST_ARGS": self.convert_broadcast_args,
            "BROADCAST_TO": self.convert_broadcast_to,
            "BUCKETIZE": self.convert_bucketize,
            "CALL": self.convert_call,
            "CALL_ONCE": self.convert_call_once,
            "COMPLEX_ABS": self.convert_complex_abs,
            "CAST": self.convert_cast,
            "CEIL": functools.partial(self._convert_unary_elemwise, relax_op=_op.ceil),
            "CONCATENATION": self.convert_concatenation,
            "CONV_2D": functools.partial(self.convert_conv, conv_type="conv2d"),
            "CONV_3D": self.convert_conv3d,
            "CONV_3D_TRANSPOSE": self.convert_conv3d_transpose,
            "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"),
            "DELEGATE": functools.partial(self.convert_operator_marker, op_name="DELEGATE"),
            "DEQUANTIZE": self.convert_dequantize,
            "DETECTION_POSTPROCESS": self.convert_detection_postprocess,
            "DILATE": self.convert_dilate,
            "DIV": functools.partial(self._convert_elemwise, relax_op=_op.divide),
            "ELU": self.convert_elu,
            "EMBEDDING_LOOKUP": self.convert_embedding_lookup,
            "EMBEDDING_LOOKUP_SPARSE": self.convert_embedding_lookup_sparse,
            "EQUAL": functools.partial(
                self._convert_elemwise, relax_op=_op.equal, comparison_op=True
            ),
            "EXP": functools.partial(self._convert_unary_elemwise, relax_op=_op.exp),
            "EXPAND_DIMS": self.convert_expand_dims,
            "FAKE_QUANT": self.convert_fake_quant,
            "FILL": self.convert_fill,
            "FLOOR_DIV": functools.partial(self._convert_elemwise, relax_op=_op.floor_divide),
            "FLOOR_MOD": functools.partial(self._convert_elemwise, relax_op=_op.floor_mod),
            "FLOOR": functools.partial(self._convert_unary_elemwise, relax_op=_op.floor),
            "FULLY_CONNECTED": self.convert_fully_connected,
            "GATHER": self.convert_gather,
            "GATHER_ND": self.convert_gather_nd,
            "GREATER_EQUAL": functools.partial(
                self._convert_elemwise, relax_op=_op.greater_equal, comparison_op=True
            ),
            "GREATER": functools.partial(
                self._convert_elemwise, relax_op=_op.greater, comparison_op=True
            ),
            "GELU": self.convert_gelu,
            "HARD_SWISH": self.convert_hard_swish,
            "HASHTABLE": self.convert_hashtable,
            "HASHTABLE_FIND": self.convert_hashtable_find,
            "HASHTABLE_IMPORT": self.convert_hashtable_import,
            "HASHTABLE_LOOKUP": self.convert_hashtable_lookup,
            "HASHTABLE_SIZE": self.convert_hashtable_size,
            "IF": self.convert_if,
            "IMAG": self.convert_imag,
            "L2_NORMALIZATION": self.convert_l2_normalization,
            "L2_POOL_2D": functools.partial(self.convert_pool2d, pool_type="l2"),
            "LEAKY_RELU": self.convert_leaky_relu,
            "LESS_EQUAL": functools.partial(
                self._convert_elemwise, relax_op=_op.less_equal, comparison_op=True
            ),
            "LESS": functools.partial(
                self._convert_elemwise, relax_op=_op.less, comparison_op=True
            ),
            "LOCAL_RESPONSE_NORMALIZATION": self.convert_lrn,
            "LOG": functools.partial(self._convert_unary_elemwise, relax_op=_op.log),
            "LOG_SOFTMAX": self.convert_log_softmax,
            "LOGICAL_AND": functools.partial(
                self._convert_logical_binary, relax_op=_op.logical_and
            ),
            "LOGICAL_NOT": self.convert_logical_not,
            "LOGICAL_OR": functools.partial(self._convert_logical_binary, relax_op=_op.logical_or),
            "LOGISTIC": self.convert_logistic,
            "LSTM": self.convert_lstm,
            "MATRIX_DIAG": self.convert_matrix_diag,
            "MATRIX_SET_DIAG": self.convert_matrix_set_diag,
            "MAX_POOL_2D": functools.partial(self.convert_pool2d, pool_type="max"),
            "MAXIMUM": functools.partial(self._convert_elemwise, relax_op=_op.maximum),
            "MEAN": functools.partial(self._convert_reduce, relax_op=_op.mean),
            "MINIMUM": functools.partial(self._convert_elemwise, relax_op=_op.minimum),
            "MIRROR_PAD": self.convert_mirror_pad,
            "MUL": functools.partial(self._convert_elemwise, relax_op=_op.multiply),
            "MULTINOMIAL": self.convert_multinomial,
            "NEG": functools.partial(self._convert_unary_elemwise, relax_op=_op.negative),
            "NOT_EQUAL": functools.partial(
                self._convert_elemwise, relax_op=_op.not_equal, comparison_op=True
            ),
            "ONE_HOT": self.convert_one_hot,
            "PACK": self.convert_pack,
            "PAD": self.convert_pad,
            "PADV2": self.convert_pad,
            "PLACEHOLDER_FOR_GREATER_OP_CODES": functools.partial(
                self.convert_operator_marker, op_name="PLACEHOLDER_FOR_GREATER_OP_CODES"
            ),
            "POW": functools.partial(self._convert_elemwise, relax_op=_op.power),
            "PRELU": self.convert_prelu,
            "RANGE": self.convert_range,
            "RANK": self.convert_rank,
            "QUANTIZE": self.convert_quantize,
            "RANDOM_STANDARD_NORMAL": self.convert_random_standard_normal,
            "RANDOM_UNIFORM": self.convert_random_uniform,
            "READ_VARIABLE": self.convert_read_variable,
            "REAL": self.convert_real,
            "REDUCE_ALL": functools.partial(self._convert_reduce_bool, relax_op=_op.min),
            "REDUCE_ANY": functools.partial(self._convert_reduce_bool, relax_op=_op.max),
            "REDUCE_MAX": functools.partial(self._convert_reduce, relax_op=_op.max),
            "REDUCE_MIN": functools.partial(self._convert_reduce, relax_op=_op.min),
            "REDUCE_PROD": functools.partial(self._convert_reduce, relax_op=_op.prod),
            "RELU": self.convert_relu,
            "RELU_0_TO_1": self.convert_relu_0_to_1,
            "RELU6": self.convert_relu6,
            "RELU_N1_TO_1": self.convert_relu_n1_to_1,
            "RESHAPE": self.convert_reshape,
            "RESIZE_BILINEAR": self.convert_resize_bilinear,
            "RESIZE_NEAREST_NEIGHBOR": self.convert_resize_nearest_neighbor,
            "RFFT2D": self.convert_rfft2d,
            "RIGHT_SHIFT": functools.partial(self._convert_elemwise, relax_op=_op.right_shift),
            "ROUND": functools.partial(self._convert_unary_elemwise, relax_op=_op.round),
            "RSQRT": functools.partial(self._convert_unary_elemwise, relax_op=_op.rsqrt),
            "REVERSE_SEQUENCE": self.convert_reverse_sequence,
            "REVERSE_V2": self.convert_reverse_v2,
            "SCATTER_ND": self.convert_scatter_nd,
            "SELECT": self.convert_select,
            "SELECT_V2": self.convert_select,
            "SEGMENT_SUM": functools.partial(
                self._convert_segment_op, op_name="SEGMENT_SUM", reduction="add"
            ),
            "SHAPE": self.convert_shape,
            "SIGN": functools.partial(self._convert_unary_elemwise, relax_op=_op.sign),
            "SIN": functools.partial(self._convert_unary_elemwise, relax_op=_op.sin),
            "SLICE": self.convert_slice,
            "SOFTMAX": self.convert_softmax,
            "SPACE_TO_BATCH_ND": self.convert_space_to_batch_nd,
            "SPACE_TO_DEPTH": self.convert_space_to_depth,
            "SPARSE_TO_DENSE": self.convert_sparse_to_dense,
            "SPLIT": self.convert_split,
            "SPLIT_V": self.convert_split_v,
            "SQRT": functools.partial(self._convert_unary_elemwise, relax_op=_op.sqrt),
            "SQUARE": self.convert_square,
            "SQUARED_DIFFERENCE": self.convert_squared_difference,
            "STABLEHLO_ABS": functools.partial(self._convert_stablehlo_unary, relax_op=_op.abs),
            "STABLEHLO_ADD": functools.partial(self._convert_stablehlo_binary, relax_op=_op.add),
            "STABLEHLO_AND": self._convert_stablehlo_and,
            "STABLEHLO_BROADCAST_IN_DIM": self._convert_stablehlo_broadcast_in_dim,
            "STABLEHLO_CBRT": self._convert_stablehlo_cbrt,
            "STABLEHLO_CLAMP": self._convert_stablehlo_clamp,
            "STABLEHLO_COMPARE": self._convert_stablehlo_compare,
            "STABLEHLO_COMPOSITE": self._convert_stablehlo_composite,
            "STABLEHLO_CONCATENATE": self._convert_stablehlo_concatenate,
            "STABLEHLO_CONVOLUTION": self._convert_stablehlo_convolution,
            "STABLEHLO_CONVERT": self._convert_stablehlo_convert,
            "STABLEHLO_COSINE": functools.partial(self._convert_stablehlo_unary, relax_op=_op.cos),
            "STABLEHLO_CUSTOM_CALL": self._convert_stablehlo_custom_call,
            "STABLEHLO_DIVIDE": functools.partial(
                self._convert_stablehlo_binary, relax_op=_op.divide
            ),
            "STABLEHLO_DOT_GENERAL": self._convert_stablehlo_dot_general,
            "STABLEHLO_DYNAMIC_SLICE": self._convert_stablehlo_dynamic_slice,
            "STABLEHLO_DYNAMIC_UPDATE_SLICE": self._convert_stablehlo_dynamic_update_slice,
            "STABLEHLO_EXPONENTIAL": functools.partial(
                self._convert_stablehlo_unary, relax_op=_op.exp
            ),
            "STABLEHLO_FLOOR": functools.partial(self._convert_stablehlo_unary, relax_op=_op.floor),
            "STABLEHLO_GATHER": self._convert_stablehlo_gather,
            "STABLEHLO_IOTA": self._convert_stablehlo_iota,
            "STABLEHLO_LOG": functools.partial(self._convert_stablehlo_unary, relax_op=_op.log),
            "STABLEHLO_LOGISTIC": functools.partial(
                self._convert_stablehlo_unary, relax_op=_op.sigmoid
            ),
            "STABLEHLO_MAXIMUM": functools.partial(
                self._convert_stablehlo_binary, relax_op=_op.maximum
            ),
            "STABLEHLO_MINIMUM": functools.partial(
                self._convert_stablehlo_binary, relax_op=_op.minimum
            ),
            "STABLEHLO_MULTIPLY": functools.partial(
                self._convert_stablehlo_binary, relax_op=_op.multiply
            ),
            "STABLEHLO_NEGATE": functools.partial(
                self._convert_stablehlo_unary, relax_op=_op.negative
            ),
            "STABLEHLO_OR": self._convert_stablehlo_or,
            "STABLEHLO_PAD": self._convert_stablehlo_pad,
            "STABLEHLO_POWER": functools.partial(
                self._convert_stablehlo_binary, relax_op=_op.power
            ),
            "STABLEHLO_REDUCE": self._convert_stablehlo_reduce,
            "STABLEHLO_REDUCE_WINDOW": self._convert_stablehlo_reduce_window,
            "STABLEHLO_REMAINDER": self._convert_stablehlo_remainder,
            "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,
            "STABLEHLO_SELECT": functools.partial(
                self._convert_stablehlo_ternary, relax_op=_op.where
            ),
            "STABLEHLO_SHIFT_LEFT": functools.partial(
                self._convert_stablehlo_binary, relax_op=_op.left_shift
            ),
            "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_WHILE": self._convert_stablehlo_while,
            "SQUEEZE": self.convert_squeeze,
            "STRIDED_SLICE": self.convert_strided_slice,
            "SUB": functools.partial(self._convert_elemwise, relax_op=_op.subtract),
            "SUM": functools.partial(self._convert_reduce, relax_op=_op.sum),
            "SVDF": self.convert_svdf,
            "TAN": functools.partial(self._convert_unary_elemwise, relax_op=_op.tan),
            "TANH": self.convert_tanh,
            "TILE": self.convert_tile,
            "TOPK_V2": self.convert_topk_v2,
            "TRANSPOSE_CONV": self.convert_transpose_conv,
            "TRANSPOSE": self.convert_transpose,
            "UNPACK": self.convert_unpack,
            "UNIDIRECTIONAL_SEQUENCE_RNN": self.convert_unidirectional_sequence_rnn,
            "UNIQUE": self.convert_unique,
            "UNSORTED_SEGMENT_MAX": functools.partial(
                self._convert_segment_op, op_name="UNSORTED_SEGMENT_MAX", reduction="max"
            ),
            "UNSORTED_SEGMENT_MIN": functools.partial(
                self._convert_segment_op, op_name="UNSORTED_SEGMENT_MIN", reduction="min"
            ),
            "UNSORTED_SEGMENT_PROD": functools.partial(
                self._convert_segment_op, op_name="UNSORTED_SEGMENT_PROD", reduction="mul"
            ),
            "UNSORTED_SEGMENT_SUM": functools.partial(
                self._convert_segment_op, op_name="UNSORTED_SEGMENT_SUM", reduction="add"
            ),
            "UNIDIRECTIONAL_SEQUENCE_LSTM": self.convert_unidirectional_sequence_lstm,
            "VAR_HANDLE": self.convert_var_handle,
            "WHERE": self.convert_select,
            "WHILE": self.convert_while,
            "ZEROS_LIKE": self.convert_zeros_like,
            "NON_MAX_SUPPRESSION_V4": self.convert_nms_v4,
            "NON_MAX_SUPPRESSION_V5": self.convert_nms_v5,
        }

    def check_unsupported_ops(self):
        """Check unsupported TFLite ops in our converter."""
        unsupported_ops_set = set()
        dynamic_range_ops_set = set()
        unsupported_quantized_ops_set = set()
        for op_idx in range(self.subgraph.OperatorsLength()):
            op = self.subgraph.Operators(op_idx)
            op_code_str = self.get_op_code_str(op)
            if op_code_str not in self.convert_map:
                unsupported_ops_set.add(op_code_str)
                continue

            # Trying to exclude "dynamic range quantization" optimized ops as not supported in TVM
            input_tensors = self.get_input_tensors(op)
            output_tensors = self.get_output_tensors(op)
            qnn_in_cnt = len([_.qnn_params for _ in input_tensors[0:1] if _.qnn_params is not None])
            qnn_weight_cnt = len(
                [_.qnn_params for _ in input_tensors[1:] if _.qnn_params is not None]
            )
            qnn_out_cnt = len([_.qnn_params for _ in output_tensors if _.qnn_params is not None])

            if qnn_in_cnt == 0 and qnn_out_cnt == 0 and qnn_weight_cnt > 0:
                dynamic_range_ops_set.add(op_code_str)

            if (
                qnn_in_cnt + qnn_weight_cnt + qnn_out_cnt > 0
                and op_code_str not in self._SUPPORTED_QUANTIZED_OPS
            ):
                unsupported_quantized_ops_set.add(op_code_str)

        raise_msg = ""

        if unsupported_ops_set:
            ops = str(list(unsupported_ops_set)).strip("[,]")
            raise_msg += f"The following operators are not supported in frontend TFLite: {ops}\n"

        if dynamic_range_ops_set:
            ops = str(list(dynamic_range_ops_set)).strip("[,]")
            raise_msg += (
                f"The following operators are likely to have dynamic range quantization: {ops}. "
                f"If you are running an optimized graph, please turn off dynamic range "
                f"quantization or use full integer quantization\n"
            )

        if unsupported_quantized_ops_set:
            ops = ", ".join(f"'{op}'" for op in sorted(unsupported_quantized_ops_set))
            raise_msg += (
                f"The following quantized TFLite operators are not supported in frontend "
                f"TFLite yet: {ops}. Quantized operators require explicit QDQ lowering "
                f"to avoid applying Relax ops directly to quantized integer tensors.\n"
            )

        if len(raise_msg) > 0:
            raise tvm.error.OpNotImplemented(raise_msg)

    def convert_operator_marker(self, op, op_name):
        """Reject TFLite marker builtins with an explicit diagnostic."""
        raise tvm.error.OpNotImplemented(
            f"TFLite operator marker {op_name} is not a Relax tensor operator"
        )

    def unbind(self, data, axis=1):
        """
        This is a modified version compared to the one in common.py.
        The onnx version takes a relax.Expr.Call, the tflite
        version a TensorWrapper. Also this version by default splits
        along axis 1 and not axis 0 as the onnx version.

         Parameters
         ----------
         data : tvm.relax.frontend.tflite.TensorWrapper
             Input tensor
         axis : int
             Axis along which tensor is split.
         Returns
         -------
         result : List[relax.Expr]
             The sequence of computed tensors
        """
        shape = to_int_list(self.get_tensor_shape(data))
        if axis >= len(shape):
            msg = "Please check input dim, it shouldn't be greater than or equal to rank."
            raise AttributeError(msg)

        selections = shape[axis]
        shape.pop(axis)
        timestep = 0  # Reshape to make time step as the first dim
        shape.insert(timestep, selections)
        res_split = relax.op.split(
            relax.op.reshape(self.get_expr(data.tensor_idx), tuple(shape)), selections, timestep
        )
        ret = []
        for i in range(selections):
            ret.append(_op.squeeze(res_split[i], axis=[timestep]))
        return relax.Tuple(relax.Tuple(ret), selections)

    def _infer_shape(self, arg):
        return self.bb.normalize(arg).ty.shape

    def convert_op_to_relax(self):
        """Convert TFLite ops to relax ops"""
        for op_idx in range(self.subgraph.OperatorsLength()):
            op = self.subgraph.Operators(op_idx)
            op_code_str = self.get_op_code_str(op)
            output_tensors = self.get_output_tensors(op)

            from tflite.Operator import Operator

            assert isinstance(op, Operator)
            ret = self.convert_map[op_code_str](op=op)

            # In case the Op can be prefetched, the output can be optimized out
            if ret is None:
                continue

            ret = self.bb.normalize(ret)

            if len(output_tensors) == 1:
                tensor_idx = output_tensors[0].tensor_idx
                self.exp_tab.set_expr(get_tensor_name(self.subgraph, tensor_idx), ret)
            else:
                for idx, output_tensor in enumerate(output_tensors):
                    self.exp_tab.set_expr(
                        get_tensor_name(self.subgraph, output_tensor.tensor_idx), ret[idx]
                    )

    @staticmethod
    def _decode_tflite_string(value):
        """Decode a TFLite string field."""
        if value is None:
            return ""
        if isinstance(value, bytes | bytearray):
            return value.decode("utf-8")
        return str(value)

    def _get_var_handle_resource_key(self, op, fallback_tensor=None):
        """Return a stable resource key for a VAR_HANDLE op."""
        container = ""
        shared_name = ""
        if op.BuiltinOptions() is not None:
            try:
                from tflite.VarHandleOptions import VarHandleOptions

                opts = self._get_builtin_options(op, VarHandleOptions)
                if hasattr(opts, "Container"):
                    container = self._decode_tflite_string(opts.Container())
                if hasattr(opts, "SharedName"):
                    shared_name = self._decode_tflite_string(opts.SharedName())
            except (ImportError, ModuleNotFoundError):
                pass

        if container or shared_name:
            return (container, shared_name)
        if fallback_tensor is not None:
            return ("", get_tensor_name(self.subgraph, fallback_tensor.tensor_idx))
        raise tvm.error.OpNotImplemented("VAR_HANDLE requires VarHandleOptions")

    def _get_resource_key_for_handle(self, tensor, op_name):
        tensor_name = get_tensor_name(self.subgraph, tensor.tensor_idx)
        if tensor_name not in self.resource_handles:
            raise tvm.error.OpNotImplemented(
                f"{op_name} requires a VAR_HANDLE in the same TFLite subgraph"
            )
        return self.resource_handles[tensor_name]

    def convert_var_handle(self, op):
        """Convert a TFLite VAR_HANDLE into an importer-local resource handle."""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 0 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented("VAR_HANDLE expects no inputs and one output")

        resource_key = self._get_var_handle_resource_key(op, output_tensors[0])
        resource_tensor_name = get_tensor_name(self.subgraph, output_tensors[0].tensor_idx)
        self.resource_handles[resource_tensor_name] = resource_key
        return None

    def convert_assign_variable(self, op):
        """Convert the CALL_ONCE initialization subset of ASSIGN_VARIABLE."""
        if not self.conversion_state["in_call_once_init"]:
            raise tvm.error.OpNotImplemented(
                "ASSIGN_VARIABLE outside CALL_ONCE initialization is not supported by the "
                "Relax TFLite frontend yet because it requires mutable resource state modeling."
            )

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 2 or len(output_tensors) != 0:
            raise tvm.error.OpNotImplemented(
                "ASSIGN_VARIABLE expects a resource handle and value input with no outputs"
            )

        resource_key = self._get_resource_key_for_handle(input_tensors[0], "ASSIGN_VARIABLE")
        self.conversion_state["resource_values"][resource_key] = self.get_tensor_expr(
            input_tensors[1]
        )
        return None

    def convert_read_variable(self, op):
        """Convert READ_VARIABLE for resources initialized by CALL_ONCE."""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 1 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented("READ_VARIABLE expects one input and one output")

        resource_key = self._get_resource_key_for_handle(input_tensors[0], "READ_VARIABLE")
        resource_values = self.conversion_state["resource_values"]
        if resource_key not in resource_values:
            raise tvm.error.OpNotImplemented(
                "READ_VARIABLE requires a resource initialized by a supported CALL_ONCE subgraph"
            )
        return resource_values[resource_key]

    def _is_tflite_string_type(self, tensor_type):
        from tflite.TensorType import TensorType

        return hasattr(TensorType, "STRING") and tensor_type == TensorType.STRING

    def _is_supported_hashtable_type_pair(self, key_dtype, value_dtype):
        from tflite.TensorType import TensorType

        return (key_dtype == TensorType.INT64 and self._is_tflite_string_type(value_dtype)) or (
            self._is_tflite_string_type(key_dtype) and value_dtype == TensorType.INT64
        )

    def _get_hashtable_key(self, op, fallback_tensor=None):
        """Return a stable key and TFLite dtype pair for a HASHTABLE resource."""
        table_id = None
        key_dtype = None
        value_dtype = None
        if op.BuiltinOptions() is not None:
            try:
                from tflite.HashtableOptions import HashtableOptions

                opts = self._get_builtin_options(op, HashtableOptions)
                table_id = int(opts.TableId())
                key_dtype = int(opts.KeyDtype())
                value_dtype = int(opts.ValueDtype())
            except (ImportError, ModuleNotFoundError):
                pass

        if key_dtype is None or value_dtype is None:
            raise tvm.error.OpNotImplemented("HASHTABLE requires HashtableOptions")
        if not self._is_supported_hashtable_type_pair(key_dtype, value_dtype):
            raise tvm.error.OpNotImplemented(
                "TFLite HASHTABLE only supports int64/string or string/int64 tables"
            )

        if table_id is not None:
            return table_id, key_dtype, value_dtype
        if fallback_tensor is not None:
            return (
                get_tensor_name(self.subgraph, fallback_tensor.tensor_idx),
                key_dtype,
                value_dtype,
            )
        raise tvm.error.OpNotImplemented("HASHTABLE requires HashtableOptions")

    def _get_hashtable_info_for_handle(self, tensor, op_name):
        tensor_name = get_tensor_name(self.subgraph, tensor.tensor_idx)
        if tensor_name not in self.hashtable_handles:
            raise tvm.error.OpNotImplemented(
                f"{op_name} requires a HASHTABLE in the same TFLite subgraph"
            )
        return self.hashtable_handles[tensor_name]

    @staticmethod
    def _get_tensor_shape_tuple(tensor_wrapper):
        if tensor_wrapper.tensor.ShapeLength() == 0:
            return ()
        return tuple(int(dim) for dim in tensor_wrapper.tensor.ShapeAsNumpy())

    @staticmethod
    def _has_tensor_buffer_data(tensor_wrapper):
        return (
            tensor_wrapper.buffer is not None
            and hasattr(tensor_wrapper.buffer, "DataLength")
            and tensor_wrapper.buffer.DataLength() > 0
        )

    def _get_string_tensor_value(self, tensor_wrapper, op_name):
        """Decode a constant TFLite string tensor buffer."""
        if not self._is_tflite_string_type(tensor_wrapper.tensor.Type()):
            raise tvm.error.OpNotImplemented(f"{op_name} requires a TensorType.STRING tensor")
        if not self._has_tensor_buffer_data(tensor_wrapper):
            raise tvm.error.OpNotImplemented(f"{op_name} requires a constant string tensor")

        data = bytes(tensor_wrapper.buffer.DataAsNumpy())
        if len(data) < 4:
            raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor buffer")

        count = int(np.frombuffer(data, dtype="<i4", count=1)[0])
        if count < 0:
            raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor count")

        header_size = 4 * (count + 2)
        if len(data) < header_size:
            raise tvm.error.OpNotImplemented(f"{op_name} has an invalid string tensor offsets")

        offsets = np.frombuffer(data, dtype="<i4", count=count + 1, offset=4).astype(np.int64)
        if np.any(offsets < header_size) or np.any(offsets > len(data)):
            raise tvm.error.OpNotImplemented(f"{op_name} has out-of-bounds string tensor offsets")
        if np.any(offsets[:-1] > offsets[1:]):
            raise tvm.error.OpNotImplemented(f"{op_name} has non-monotonic string tensor offsets")

        try:
            values = [
                data[int(offsets[i]) : int(offsets[i + 1])].decode("utf-8") for i in range(count)
            ]
        except UnicodeDecodeError as e:
            raise tvm.error.OpNotImplemented(f"{op_name} has invalid UTF-8 string data: {e}") from e
        shape = self._get_tensor_shape_tuple(tensor_wrapper)
        expected_count = math.prod(shape) if shape else 1
        if expected_count != count:
            raise tvm.error.OpNotImplemented(
                f"{op_name} string tensor buffer count does not match its shape"
            )
        return np.array(values, dtype=object).reshape(shape)

    def convert_hashtable(self, op):
        """Convert a TFLite HASHTABLE into an importer-local table handle."""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 0 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented("HASHTABLE expects no inputs and one output")

        table_key, key_dtype, value_dtype = self._get_hashtable_key(op, output_tensors[0])
        table_tensor_name = get_tensor_name(self.subgraph, output_tensors[0].tensor_idx)
        self.hashtable_handles[table_tensor_name] = {
            "table_key": table_key,
            "key_dtype": key_dtype,
            "value_dtype": value_dtype,
        }
        return None

    def convert_hashtable_import(self, op):
        """Convert static metadata for the CALL_ONCE HASHTABLE_IMPORT subset."""
        if not self.conversion_state["in_call_once_init"]:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_IMPORT outside CALL_ONCE initialization is not supported by the "
                "Relax TFLite frontend yet because it requires mutable resource state modeling."
            )

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 3 or len(output_tensors) != 0:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_IMPORT expects table, keys, and values inputs with no outputs"
            )

        table_info = self._get_hashtable_info_for_handle(input_tensors[0], "HASHTABLE_IMPORT")
        key_tensor = input_tensors[1]
        value_tensor = input_tensors[2]
        if (
            key_tensor.tensor.Type() != table_info["key_dtype"]
            or value_tensor.tensor.Type() != table_info["value_dtype"]
        ):
            raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT key/value dtypes mismatch")
        key_shape = self._get_tensor_shape_tuple(key_tensor)
        value_shape = self._get_tensor_shape_tuple(value_tensor)
        if key_shape != value_shape:
            raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT requires keys and values same shape")
        if not self._has_tensor_buffer_data(key_tensor) or not self._has_tensor_buffer_data(
            value_tensor
        ):
            raise tvm.error.OpNotImplemented("HASHTABLE_IMPORT requires constant keys and values")

        if self._is_tflite_string_type(table_info["key_dtype"]):
            keys = self._get_string_tensor_value(key_tensor, "HASHTABLE_IMPORT")
        else:
            keys = self.get_tensor_value(key_tensor)
        if self._is_tflite_string_type(table_info["value_dtype"]):
            values = self._get_string_tensor_value(value_tensor, "HASHTABLE_IMPORT")
        else:
            values = self.get_tensor_value(value_tensor)

        if np.unique(keys).size != keys.size:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_IMPORT with duplicate keys is not supported"
            )

        hashtable_values = self.conversion_state["hashtable_values"]
        table_key = table_info["table_key"]
        if table_key not in hashtable_values:
            hashtable_values[table_key] = {
                "size": math.prod(key_shape) if key_shape else 1,
                "key_dtype": table_info["key_dtype"],
                "value_dtype": table_info["value_dtype"],
                "keys": keys,
                "values": values,
            }
        return None

    def convert_hashtable_find(self, op):
        """Convert the constant-foldable string-to-int64 HASHTABLE_FIND subset."""
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 3 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_FIND expects table, query, and default inputs with one output"
            )

        table_tensor, query_tensor, default_tensor = input_tensors
        output_tensor = output_tensors[0]
        table_info = self._get_hashtable_info_for_handle(table_tensor, "HASHTABLE_FIND")
        table_key = table_info["table_key"]
        hashtable_values = self.conversion_state["hashtable_values"]
        if table_key not in hashtable_values:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_FIND requires a table initialized by a supported CALL_ONCE subgraph"
            )
        table_values = hashtable_values[table_key]

        if (
            query_tensor.tensor.Type() != table_values["key_dtype"]
            or default_tensor.tensor.Type() != table_values["value_dtype"]
            or output_tensor.tensor.Type() != table_values["value_dtype"]
        ):
            raise tvm.error.OpNotImplemented("HASHTABLE_FIND key/value dtypes mismatch")

        if not (
            self._is_tflite_string_type(table_values["key_dtype"])
            and table_values["value_dtype"] == TensorType.INT64
        ):
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_FIND only supports constant string -> int64 tables"
            )
        if not self._has_tensor_buffer_data(query_tensor):
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_FIND with runtime string queries is not supported"
            )
        if not self._has_tensor_buffer_data(default_tensor):
            raise tvm.error.OpNotImplemented("HASHTABLE_FIND requires constant default values")

        query_shape = self._get_tensor_shape_tuple(query_tensor)
        output_shape = self._get_tensor_shape_tuple(output_tensor)
        if output_shape != query_shape:
            raise tvm.error.OpNotImplemented("HASHTABLE_FIND output shape must match query shape")

        query_values = self._get_string_tensor_value(query_tensor, "HASHTABLE_FIND")
        default_values = self.get_tensor_value(default_tensor)
        default_shape = self._get_tensor_shape_tuple(default_tensor)
        if default_shape == () or default_values.size == 1:
            result = np.full(output_shape, int(default_values.item()), dtype=np.int64)
        elif default_shape == query_shape:
            result = default_values.astype(np.int64).copy()
        else:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_FIND default value must be scalar or match query shape"
            )

        table_map = dict(
            zip(
                table_values["keys"].reshape(-1).tolist(),
                table_values["values"].reshape(-1).astype(np.int64).tolist(),
            )
        )
        for index, key in np.ndenumerate(query_values):
            if key in table_map:
                result[index] = table_map[key]
        return relax.const(result.astype(np.int64), "int64")

    def convert_hashtable_lookup(self, op):
        """Convert TFLite HASHTABLE_LOOKUP for non-string value tensors."""
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 3 or len(output_tensors) != 2:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP expects lookup, key, and value inputs with two outputs"
            )

        lookup_tensor, key_tensor, value_tensor = input_tensors
        output_tensor, hits_tensor = output_tensors

        if (
            lookup_tensor.tensor.Type() != TensorType.INT32
            or key_tensor.tensor.Type() != TensorType.INT32
        ):
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP requires int32 lookup and key tensors"
            )
        if self._is_tflite_string_type(value_tensor.tensor.Type()):
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP with TensorType.STRING values is not supported"
            )
        if value_tensor.tensor.Type() != output_tensor.tensor.Type():
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP output dtype must match the value tensor dtype"
            )
        if hits_tensor.tensor.Type() != TensorType.UINT8:
            raise tvm.error.OpNotImplemented("HASHTABLE_LOOKUP hits output must be uint8")

        lookup_shape = to_int_list(self.get_tensor_shape(lookup_tensor))
        key_shape = to_int_list(self.get_tensor_shape(key_tensor))
        value_shape = to_int_list(self.get_tensor_shape(value_tensor))
        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
        hits_shape = to_int_list(self.get_tensor_shape(hits_tensor))

        if len(lookup_shape) != 1 or len(key_shape) != 1 or len(value_shape) < 1:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP requires rank-1 lookup/key and rank>=1 value tensors"
            )
        if key_shape[0] != value_shape[0]:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP requires key and value tensors to agree on row count"
            )
        if key_shape[0] == 0:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP requires a non-empty key/value table"
            )
        if output_shape != [lookup_shape[0]] + value_shape[1:]:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP output shape must match lookup count and value tail shape"
            )
        if hits_shape != [lookup_shape[0]]:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_LOOKUP hits output shape must match lookup count"
            )

        lookup = self.get_tensor_expr(lookup_tensor)
        key = self.get_tensor_expr(key_tensor)
        value = self.get_tensor_expr(value_tensor)

        positions = relax.op.bucketize(lookup, key, out_int32=True, right=False)
        candidate_keys = relax.op.take(key, positions, axis=0, mode="clip")
        in_range = relax.op.less(positions, relax.const(key_shape[0], "int32"))
        found = relax.op.logical_and(in_range, relax.op.equal(candidate_keys, lookup))

        gathered_values = relax.op.take(value, positions, axis=0, mode="clip")
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
        zero_values = relax.op.zeros(output_shape, output_dtype)

        if len(value_shape) > 1:
            found_values = relax.op.expand_dims(found, axis=list(range(1, len(value_shape))))
            found_values = relax.op.broadcast_to(found_values, output_shape)
        else:
            found_values = found

        output = relax.op.where(found_values, gathered_values, zero_values)
        hits = relax.op.astype(found, "uint8")
        return relax.Tuple([output, hits])

    def convert_hashtable_size(self, op):
        """Convert HASHTABLE_SIZE for a statically imported TFLite hashtable."""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 1 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented("HASHTABLE_SIZE expects one input and one output")

        from tflite.TensorType import TensorType

        if output_tensors[0].tensor.Type() != TensorType.INT64:
            raise tvm.error.OpNotImplemented("HASHTABLE_SIZE output must be int64")
        table_info = self._get_hashtable_info_for_handle(input_tensors[0], "HASHTABLE_SIZE")
        table_key = table_info["table_key"]
        hashtable_values = self.conversion_state["hashtable_values"]
        if table_key not in hashtable_values:
            raise tvm.error.OpNotImplemented(
                "HASHTABLE_SIZE requires a table initialized by a supported CALL_ONCE subgraph"
            )
        return relax.const(np.array([hashtable_values[table_key]["size"]], dtype=np.int64), "int64")

    def get_op_code_str(self, op):
        """Get TFLite ops string representation"""

        from tflite.BuiltinOperator import BuiltinOperator

        op_code_list_idx = op.OpcodeIndex()

        op_c = self.model.OperatorCodes(op_code_list_idx)
        # In TFlite 2.4.x there was a change where the type of the field that contained
        # the builtin code changed from int8 to int32 in the flat buffer representation.
        # However, to retain support for old flat buffers that were created, they retained
        # the original 8 bit field, but named it "deprecated_builtin_code" in TFLite 2.4.
        # This means that the API function BuiltinCode() which originally returned the value
        # of the 8 bit field would now look for the value in the new int32 field in the
        # schema and DeprecatedBuiltinCode() will look at the old 8 bit field.
        # In TFLite 2.4, if the opcode value is less than 127, it can be in either field
        # (however, if it is only in the "builtin_code" field, the model is not backward
        # compatible), so similarly to TFLite 2.4 reader, we'll pick the higher value of the
        # two fields.
        # Remember however that this value came into existence only after Tensorflow
        # lite 2.4.x and hence encase it in a try -except block.
        # Phew !
        try:
            opc = max(op_c.DeprecatedBuiltinCode(), op_c.BuiltinCode())
        except AttributeError:
            # In versions before 2.4 the int8 field that holds the builtin code is accessed
            # by BuiltinCode() and DeprecatedBuiltinCode() doesn't exist
            opc = op_c.BuiltinCode()

        op_code_id = opc
        try:
            op_code_str = self.builtin_op_code[op_code_id]
        except KeyError:
            raise NotImplementedError(
                "TFLite operator with code "
                + str(op_code_id)
                + " is not supported by this version of the fbs schema."
            )
        if op_code_id == BuiltinOperator.CUSTOM:
            # Custom operator
            custom_op_code_str = self.model.OperatorCodes(op_code_list_idx).CustomCode()

            if self.allow_custom_ops:
                return "CUSTOM"

            if custom_op_code_str == b"TFLite_Detection_PostProcess":
                return "DETECTION_POSTPROCESS"

            raise NotImplementedError("Custom operators are currently not supported")
        return op_code_str

    def get_input_tensors(self, op):
        operator_inputs = op.InputsAsNumpy()
        return self.get_tensors(operator_inputs)

    def get_output_tensors(self, op):
        operator_outputs = op.OutputsAsNumpy()
        return self.get_tensors(operator_outputs)

    def get_tensors(self, tensors_idx_list):
        """Get tensor wrapper list from given TFLite tensor index list"""
        return_list = list()
        for tensor_idx in self._indices_or_empty(tensors_idx_list):
            if tensor_idx < 0:
                return_list.append(TensorWrapper(tensor_idx, 0, 0))
                continue

            tensor = self.subgraph.Tensors(tensor_idx)
            buffer_idx = tensor.Buffer()
            buffer = self.model.Buffers(buffer_idx)

            # Check if the tensors are quantized. Parse if yes.
            qnn_params = None
            tflite_qnn_params = tensor.Quantization()
            if tflite_qnn_params is not None:
                # TFLite supports both per-tensor and per-axis (aka channel) quantization.  For
                # per-tensor quantization, scale and zero points are scalar values.  For per-axis
                # quantization, scale and zero points for the weights are tensors (activations are
                # per-tensor quantized). However, the TFLite quantization spec puts restrictions on
                # zero points for per-axis quantization.  Specifically, the zero point is a tensor
                # but all values are 0. More information can be found here -
                # https://www.tensorflow.org/lite/performance/quantization_spec

                tflite_scale = tflite_qnn_params.ScaleAsNumpy()
                tflite_zero_point = tflite_qnn_params.ZeroPointAsNumpy()
                is_qnn_params_valid = True

                # Handle Per-axis and per-tensor cases
                if isinstance(tflite_scale, np.ndarray):
                    assert isinstance(tflite_zero_point, np.ndarray)

                    # Tensor - Per-axis quantization
                    if tflite_scale.size != 1 and tflite_zero_point.size != 1:
                        scale = tflite_scale
                        # Ensure that all zero points are zeros
                        zero_point = tflite_zero_point
                        if not np.all(zero_point == 0):
                            raise tvm.error.OpAttributeInvalid(
                                "TFLite per-axis quantization restricts all zero points to be"
                                + " 0, but a non-zero value is observed"
                            )
                        zero_point = int(zero_point[0])

                    # Scalar - Per-tensor quantization
                    elif tflite_scale.size == 1 and tflite_zero_point.size == 1:
                        scale = float(tflite_scale[0])
                        zero_point = int(tflite_zero_point[0])

                    else:
                        raise NotImplementedError(
                            f"Quantized type {type(tflite_scale)} (scale) and  "
                            f"{type(tflite_zero_point)} (zero point) not supported"
                        )
                elif tflite_scale == 0 and tflite_zero_point == 0:
                    # Handle corner case for ops like quantized reshape whose second operand (shape)
                    # has zero scale and zero zero point. This is not used.
                    is_qnn_params_valid = False
                else:
                    raise NotImplementedError(f"Quantized type {type(tflite_scale)} not supported")

                # Check that the scale and zero points are valid.
                if is_qnn_params_valid:
                    qnn_params = dict()
                    qnn_params["scale"] = relax.const(scale, "float32")
                    qnn_params["zero_point"] = relax.const(zero_point, "int32")
                    qnn_params["axis"] = int(tflite_qnn_params.QuantizedDimension())
            return_list.append(TensorWrapper(tensor_idx, tensor, buffer, qnn_params))
        return return_list

    def get_tensor_type_as_numpy(self, tensor_wrapper):
        """Returns np.dtype out of TensorType"""
        assert isinstance(tensor_wrapper, TensorWrapper)

        from tflite.TensorType import TensorType

        return {
            TensorType.UINT8: np.uint8,
            TensorType.INT8: np.int8,
            TensorType.INT16: np.int16,
            TensorType.FLOAT16: np.float16,
            TensorType.FLOAT32: np.float32,
            TensorType.INT32: np.int32,
            TensorType.INT64: np.int64,
            TensorType.UINT32: np.uint32,
            TensorType.UINT64: np.uint64,
            TensorType.BOOL: np.bool_,
            TensorType.COMPLEX64: np.complex64,
        }[tensor_wrapper.tensor.Type()]

    # pylint: disable=no-else-return
    def get_tensor_value(self, tensor_wrapper, is_sparse=False):
        """Get tensor buffer value from given tensor wrapper.

        Returns the raw TFLite buffer reinterpreted via ``get_tensor_type_as_numpy``: for
        COMPLEX64 this is a ``np.complex64`` ndarray with the TFLite shape (no pair axis).
        This is distinct from ``get_tensor_expr``, which returns the *Relax* representation
        (float32 real/imag pair with a trailing (2,) axis for COMPLEX64). Use
        ``get_tensor_expr`` for Relax IR construction; use ``get_tensor_value`` only when
        you need the literal TFLite buffer (e.g. fft_length constant parsing).
        """
        assert isinstance(tensor_wrapper, TensorWrapper)

        dtype = self.get_tensor_type_as_numpy(tensor_wrapper)
        data = tensor_wrapper.buffer.DataAsNumpy()

        if tensor_wrapper.tensor.ShapeLength() != 0:
            shape = to_int_list(self.get_tensor_shape(tensor_wrapper))
        else:
            shape = []

        if is_sparse:
            return np.frombuffer(data, dtype=dtype)
        else:
            return np.frombuffer(data, dtype=dtype).reshape(shape)

    def get_tensor_type_str(self, tensor_type):
        """Get tensor type string representation when given TFLite tensor type"""

        from tflite.TensorType import TensorType

        if tensor_type == TensorType.INT8:
            return "int8"
        if tensor_type == TensorType.INT16:
            return "int16"
        if tensor_type == TensorType.UINT8:
            return "uint8"
        if tensor_type == TensorType.FLOAT16:
            return "float16"
        if tensor_type == TensorType.FLOAT32:
            return "float32"
        if tensor_type == TensorType.INT32:
            return "int32"
        if tensor_type == TensorType.INT64:
            return "int64"
        if tensor_type == TensorType.UINT32:
            return "uint32"
        if tensor_type == TensorType.UINT64:
            return "uint64"
        if tensor_type == TensorType.BOOL:
            return "bool"
        if tensor_type == TensorType.COMPLEX64:
            return "complex64"
        raise NotImplementedError(f"Tensor type {tensor_type!s} is currently not supported")

    def _is_tflite_complex64_type(self, tensor_type):
        """Return whether a TFLite tensor type is COMPLEX64."""
        from tflite.TensorType import TensorType

        return tensor_type == TensorType.COMPLEX64

    def _unwrap_tflite_tensor(self, tensor):
        """Return the underlying tflite.Tensor, unwrapping TensorWrapper if needed."""
        if isinstance(tensor, TensorWrapper):
            return tensor.tensor
        return tensor

    def _get_relax_tensor_dtype(self, tensor):
        """Return the Relax dtype used to represent a TFLite tensor."""
        tensor = self._unwrap_tflite_tensor(tensor)
        tensor_type = tensor.Type() if hasattr(tensor, "Type") else tensor
        if self._is_tflite_complex64_type(tensor_type):
            return "float32"
        return self.get_tensor_type_str(tensor_type)

    def _get_relax_tensor_shape(self, tensor):
        """Return the Relax shape used to represent a TFLite tensor.

        For COMPLEX64 tensors, the trailing (2,) axis encodes the real/imag pair.
        Returns an empty tuple () for rank-0 tensors. Shape elements are Python ints
        (not numpy scalars) so the result is safe to feed into TIRX ``T.Buffer(shape, ...)``.
        """
        tensor = self._unwrap_tflite_tensor(tensor)
        shape = to_int_list(tensor.ShapeAsNumpy()) if tensor.ShapeLength() > 0 else ()
        if self._is_tflite_complex64_type(tensor.Type()):
            return tuple(shape) + (2,)
        return tuple(shape)

    def _get_shape_expr_from_tensor(self, shape_tensor, prefix):
        """Convert a TFLite shape tensor to a Relax shape expression."""
        if self.has_expr(shape_tensor.tensor_idx):
            dims_expr = self.get_expr(shape_tensor.tensor_idx)
            dims_ndim = int(self.get_tensor_shape(shape_tensor)[0])
            dims_dtype = self.get_tensor_type_str(shape_tensor.tensor.Type())
            dims_expr = self.bb.match_cast(dims_expr, relax.TensorType([dims_ndim], dims_dtype))
            dims_expr = self.bb.normalize(relax.op.astype(dims_expr, "int64"))
            shape_dataflow_var = self.bb.emit(relax.op.tensor_to_shape(dims_expr))
            shape_vars = [tirx.Var(f"{prefix}_{i}", "int64") for i in range(dims_ndim)]
            self.bb.match_cast(shape_dataflow_var, relax.ShapeType(shape_vars))
            return relax.ShapeExpr(shape_vars), shape_vars

        dims = to_int_list(self.get_tensor_value(shape_tensor))
        return dims, dims

    def flatten_to_nd(self, x, nd=3):
        """Flatten input tensor to nd rank"""
        shape = x.ty.shape
        ndims = len(shape)
        if ndims == nd:
            return x
        new_shape = [-1] + [int(shape[i]) for i in range(ndims - nd + 1, ndims)]
        return relax.op.reshape(x, new_shape)

    def has_same_qnn_params(self, lhs_tensor, rhs_tensor):
        lhs_scale = lhs_tensor.qnn_params["scale"]
        rhs_scale = rhs_tensor.qnn_params["scale"]
        lhs_zero_point = lhs_tensor.qnn_params["zero_point"]
        rhs_zero_point = rhs_tensor.qnn_params["zero_point"]
        # 0.1 + 0.2 != 0.3
        return np.allclose(
            lhs_scale.data.numpy(), rhs_scale.data.numpy(), rtol=1e-5, atol=1e-5
        ) and np.allclose(
            lhs_zero_point.data.numpy(), rhs_zero_point.data.numpy(), rtol=1e-5, atol=1e-5
        )

    def quantize(self, expr, tensor_to_quantize):
        """Helper function to quantize a tensor with Relax"""
        tensor_type = tensor_to_quantize.tensor.Type()
        tensor_type_str = self.get_tensor_type_str(tensor_type)
        quantized = relax.op.quantize(
            data=expr,
            scale=tensor_to_quantize.qnn_params["scale"],
            zero_point=tensor_to_quantize.qnn_params["zero_point"],
            axis=tensor_to_quantize.qnn_params["axis"],
            out_dtype=tensor_type_str,
        )
        return quantized

    def dequantize(self, expr, tensor):
        """Helper function to dequantize a tensor with Relax"""
        dequantized = relax.op.dequantize(
            data=expr,
            scale=tensor.qnn_params["scale"],
            zero_point=tensor.qnn_params["zero_point"],
            axis=tensor.qnn_params["axis"],
        )
        return dequantized

    def is_quantized(self, op):
        """Check if an input tensor is quantized."""
        input_tensors = self.get_input_tensors(op)
        first_tensor = input_tensors[0]
        return first_tensor.qnn_params is not None

    def convert_qnn_fused_activation_function(
        self, expr, fused_activation_fn, scale, zero_point, dtype
    ):
        """Convert TFLite fused activation function. The expr is an input quantized tensor with
        scale and zero point"""

        from tflite.ActivationFunctionType import ActivationFunctionType

        # Quantize a float value to an quantized integer value
        def quantize(x):
            return float(round(x / scale) + zero_point)

        # Get min/max of the output dtype. This will be used to ensure that clip a_min/a_max are not
        # beyond the dtype range.
        qmin = float(tvm.tirx.min_value(dtype).value)
        qmax = float(tvm.tirx.max_value(dtype).value)

        # The input expr is a quantized tensor with its scale and zero point. We calculate the
        # suitable clip off points based on these scale and zero point.
        if fused_activation_fn == ActivationFunctionType.NONE:
            return expr
        if fused_activation_fn == ActivationFunctionType.RELU6:
            return relax.op.clip(expr, min=max(qmin, quantize(0)), max=min(qmax, quantize(6.0)))
        if fused_activation_fn == ActivationFunctionType.RELU_N1_TO_1:
            return relax.op.clip(expr, min=max(qmin, quantize(-1.0)), max=min(qmax, quantize(1.0)))
        if fused_activation_fn == ActivationFunctionType.RELU:
            return relax.op.clip(expr, min=max(qmin, quantize(0.0)), max=qmax)

        fused_activation_fn_str = self.activation_fn_type[fused_activation_fn]
        raise tvm.error.OpNotImplemented(
            f"Quantized activation {fused_activation_fn_str} is not supported yet."
        )

    def convert_reshape(self, op):
        """Convert TFLite reshape"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ReshapeOptions import ReshapeOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) in (1, 2), "input tensors should not be empty"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "There should be only 1 output tensor"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx

        if len(input_tensors) == 2:
            shape_tensor = input_tensors[1]
            if self.has_expr(shape_tensor.tensor_idx):
                target_expr = self.get_expr(shape_tensor.tensor_idx)
                """
                target_value, success = try_infer_value(
                    target_expr,
                    parameters={
                        k: tvm.runtime.tensor(np.array(v)) for k, v in self.exp_tab.params.items()
                    },
                )
                """
                target_value, success = target_expr, False
                if success:
                    # convert to flattened list
                    from itertools import chain

                    try:
                        target_shape = list(chain(*target_value))
                    except TypeError:
                        target_shape = list(chain(target_value))
                else:
                    target_shape = target_expr
            else:
                target_shape = self.get_tensor_value(shape_tensor)
                # convert to flattened list
                from itertools import chain

                try:
                    target_shape = list(chain(*target_shape))
                except TypeError:
                    target_shape = list(chain(target_shape))

        else:
            assert op.BuiltinOptionsType() == BuiltinOptions.ReshapeOptions
            op_options = op.BuiltinOptions()
            reshape_options = ReshapeOptions()
            reshape_options.Init(op_options.Bytes, op_options.Pos)
            target_shape = to_int_list(reshape_options.NewShapeAsNumpy())

        in_expr = self.get_expr(input_tensor_idx)

        # If the tensors are quantized, ensure that input/output qnn params are same.

        input_tensor_type_str = self.get_tensor_type_str(input_tensor.tensor.Type())
        if input_tensor.qnn_params and input_tensor_type_str == "int8":
            # TFLite 2.x quantization spec requires qnn params to be same and dtype to be int8.
            # For TFLite 1.x, dtype can be uint8 and qnn params can be different
            output_tensor = output_tensors[0]
            assert self.has_same_qnn_params(input_tensor, output_tensor), (
                "TFLite reshape requires input and output scale and zero points to be equal"
            )

        if input_tensor.qnn_params and input_tensor_type_str == "uint8":
            output_tensor = output_tensors[0]
            if not self.has_same_qnn_params(input_tensor, output_tensor):
                in_f32 = self.dequantize(in_expr, input_tensor)
                out = relax.op.reshape(in_f32, shape=relax.ShapeExpr(target_shape))
                out = self.quantize(out, output_tensor)
                return out

        out = relax.op.reshape(in_expr, shape=relax.ShapeExpr(target_shape))
        return out

    def _convert_resize(self, method, op):
        """Generic method to Convert TFLite RESIZE operators"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ResizeBilinearOptions import ResizeBilinearOptions

        # ResizeNearestNeighborOptions was added in tflite v1.13
        tflite_ver = 1120
        if hasattr(BuiltinOptions, "ResizeNearestNeighborOptions"):
            tflite_ver = 1130

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        # images, 4-D Tensor with shape NHWC.
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        # size - 1-D int32 Tensor of 2 elements: new_height, new_width
        target_size = tuple(self.get_tensor_value(input_tensors[1]))

        # Options - align_corners (bool)
        resize_options = None
        align_corners = False
        bilinear_method = method == "linear"
        if bilinear_method:
            assert op.BuiltinOptionsType() == BuiltinOptions.ResizeBilinearOptions
            resize_options = ResizeBilinearOptions()
        elif tflite_ver >= 1130:
            assert op.BuiltinOptionsType() == BuiltinOptions.ResizeNearestNeighborOptions
            from tflite.ResizeNearestNeighborOptions import ResizeNearestNeighborOptions

            resize_options = ResizeNearestNeighborOptions()

        if resize_options is not None:
            op_options = op.BuiltinOptions()
            resize_options.Init(op_options.Bytes, op_options.Pos)
            align_corners = resize_options.AlignCorners()
            half_pixel_centers = resize_options.HalfPixelCenters()

        # Use layout NHWC
        coord_trans = "align_corners" if align_corners else "asymmetric"
        coord_trans = "half_pixel" if half_pixel_centers else coord_trans

        rounding_method = ""
        if method == "nearest_neighbor":
            if not align_corners and half_pixel_centers:
                rounding_method = "round_prefer_ceil"

        if bilinear_method and input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)
        out = relax.op.image.resize2d(
            in_expr, target_size, None, "NHWC", method, coord_trans, rounding_method
        )
        if bilinear_method and output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)
        return out

    def convert_resize_bilinear(self, op):
        """Convert TFLite RESIZE_BILINEAR"""
        return self._convert_resize("linear", op)

    def convert_resize_nearest_neighbor(self, op):
        """Convert TFLite RESIZE_NEAREST_NEIGHBOR"""
        return self._convert_resize("nearest_neighbor", op)

    def convert_l2_normalization(self, op):
        """Convert TFLite L2_NORMALIZATION"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.L2NormOptions import L2NormOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        assert op.BuiltinOptionsType() == BuiltinOptions.L2NormOptions
        op_options = op.BuiltinOptions()
        l2_norm_options = L2NormOptions()
        l2_norm_options.Init(op_options.Bytes, op_options.Pos)
        fused_activation_fn = l2_norm_options.FusedActivationFunction()

        # TFLite supports normalization only over the last dim
        input_tensor_rank = len(input_tensor.tensor.ShapeAsNumpy())

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "TFLite quantized L2_NORMALIZATION operator is not supported yet."
            )

        # TFL uses only the default epsilon value
        # Implement L2 normalization: output = input / sqrt(sum(input^2) + eps)
        # L2 normalization is applied along the last axis
        squared = relax.op.square(in_expr)
        sum_squared = relax.op.sum(squared, axis=input_tensor_rank - 1, keepdims=True)
        denom = relax.op.sqrt(relax.op.add(sum_squared, relax.const(1e-12, "float32")))
        out = relax.op.divide(in_expr, denom)

        # if we have fused activation fn
        if output_tensor.qnn_params:
            raise tvm.error.OpNotImplemented(
                "TFLite quantized L2_NORMALIZATION operator is not supported yet."
            )
        out = self.convert_fused_activation_function(out, fused_activation_fn)

        return out

    def convert_lrn(self, op):
        """Convert TFLite LOCAL_RESPONSE_NORMALIZATION"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.LocalResponseNormalizationOptions import LocalResponseNormalizationOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented("TFlite quantized LRN operator is not supported yet.")

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.LocalResponseNormalizationOptions
        op_options = op.BuiltinOptions()
        lrn_options = LocalResponseNormalizationOptions()
        lrn_options.Init(op_options.Bytes, op_options.Pos)
        radius = lrn_options.Radius()
        bias = lrn_options.Bias()
        alpha = lrn_options.Alpha()
        beta = lrn_options.Beta()
        size = (radius * 2) + 1
        alpha = alpha * size
        axis = 3  # NHWC format
        data_shape = to_int_list(self.get_tensor_shape(input_tensor))
        in_type = self.get_tensor_type_str(input_tensor.tensor.Type())

        # Relax currently does not expose a dedicated LRN op. Implement NHWC channel LRN
        # by pooling squared values over the channel axis.
        squared = self.bb.normalize(relax.op.square(in_expr))
        squared_2d = _op.reshape(squared, [-1, data_shape[axis], 1, 1])
        pooled = self.bb.normalize(
            relax.op.nn.avg_pool2d(
                squared_2d,
                pool_size=[size, 1],
                strides=[1, 1],
                padding=[radius, 0, radius, 0],
                layout="NHWC",
                count_include_pad=True,
            )
        )
        pooled = self.bb.normalize(_op.reshape(pooled, data_shape))
        denom = relax.op.power(
            relax.op.add(
                relax.const(bias, in_type), relax.op.multiply(relax.const(alpha, in_type), pooled)
            ),
            relax.const(beta, in_type),
        )
        out = relax.op.divide(in_expr, denom)

        return out

    def convert_logistic(self, op):
        """Convert TFLite LOGISTIC"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)
        out = relax.op.sigmoid(in_expr)
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

        return out

    def convert_softmax(self, op):
        """Convert TFLite softmax"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        params = {"axis": -1}  # -1 is channel
        in_expr = self.get_expr(input_tensor_idx)

        # TODO - Naive softmax int8 implementation leads to bad accuracy. Currently, we can
        # dequantize to FP32 and perform softmax on FP32. We can investigate an integer only softmax
        # implementation in future.
        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)

        out = relax.op.nn.softmax(in_expr, **params)

        # Go back to integer dataype if the original operator was quantized.
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

        return out

    def convert_tanh(self, op):
        """Convert TFLite TANH"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)
        out = relax.op.tanh(in_expr)
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)
        return out

    def convert_range(self, op):
        """Convert TFLite Range.

        Constant bounds lower directly to ``relax.op.arange``. Runtime (dynamic)
        scalar bounds are handled by computing the element count in-graph,
        lifting it to a symbolic output dimension, and rebuilding the values as
        ``arange(0, count) * delta + start`` (see ``_convert_dynamic_range``).
        """

        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"

        start, limit, delta = input_tensors[0], input_tensors[1], input_tensors[2]

        # out type inference
        if delta.tensor.Type() == TensorType.FLOAT32:
            out_type = self.get_tensor_type_str(delta.tensor.Type())
        else:
            out_type = self.get_tensor_type_str(start.tensor.Type())

        def is_dynamic(tensor):
            return self.has_expr(tensor.tensor_idx) and not isinstance(
                self.get_expr(tensor.tensor_idx), relax.Constant
            )

        def static_scalar(tensor):
            if self.has_expr(tensor.tensor_idx):
                value = self.get_expr(tensor.tensor_idx).data.numpy()
            else:
                value = self.get_tensor_value(tensor)
            # TFLite RANGE operands are scalar tensors in the flatbuffer.
            assert value.size == 1, "RANGE scalar input must have exactly one element"
            return value.item()

        if not (is_dynamic(start) or is_dynamic(limit) or is_dynamic(delta)):
            return relax.op.arange(
                static_scalar(start), static_scalar(limit), static_scalar(delta), out_type
            )

        return self._convert_dynamic_range(start, limit, delta, out_type)

    def _scalar_tensor_to_dim(self, expr, name):
        """Lift a runtime scalar Relax expr to a symbolic ``tirx.Var`` dimension.

        Mirrors the ``tensor_to_shape`` + ``match_cast`` bridge used by
        ``_get_shape_expr_from_tensor`` so a data-dependent scalar can be used as
        a ``Expr`` (e.g. an output length). The scalar is cast to int64 first.
        """
        expr = self.bb.normalize(relax.op.astype(expr, "int64"))
        expr = self.bb.normalize(relax.op.reshape(expr, (1,)))
        expr = self.bb.match_cast(expr, relax.TensorType([1], "int64"))
        shape_var = self.bb.emit(relax.op.tensor_to_shape(expr))
        dim = tirx.Var(name, "int64")
        self.bb.match_cast(shape_var, relax.ShapeType([dim]))
        return dim

    def _convert_dynamic_range(self, start, limit, delta, out_type):
        """RANGE with dynamic (runtime) scalar bounds, for int and float dtypes.

        ``relax.op.arange`` only accepts compile-time ``Expr`` bounds, and its
        struct-info length formula lacks a negative-step branch, so feeding
        symbolic bounds directly would mis-declare descending ranges. Instead the
        element count ``max(0, ceil((limit - start) / delta))`` is computed
        in-graph and lifted to one symbolic dimension ``L`` (so the declared and
        runtime lengths match by construction); values are rebuilt as
        ``arange(0, L) * delta + start``.
        """
        # int ranges work in int64 for an exact, sign-agnostic count; float
        # ranges work in the output float dtype.
        work_type = out_type if out_type.startswith("float") else "int64"

        def scalar_expr(tensor):
            return self.bb.normalize(relax.op.astype(self.get_tensor_expr(tensor), work_type))

        start_e = scalar_expr(start)
        limit_e = scalar_expr(limit)
        delta_e = scalar_expr(delta)

        if work_type.startswith("float"):
            count = relax.op.ceil(relax.op.divide(relax.op.subtract(limit_e, start_e), delta_e))
        else:
            # ceil((limit - start) / delta) == -floordiv(start - limit, delta),
            # which stays exact and handles negative delta without a float cast.
            count = relax.op.negative(
                relax.op.floor_divide(relax.op.subtract(start_e, limit_e), delta_e)
            )
        count = relax.op.maximum(count, relax.const(0, work_type))
        dim = self._scalar_tensor_to_dim(count, "range_len")

        positions = self.bb.normalize(
            relax.op.astype(relax.op.arange(0, dim, 1, "int64"), work_type)
        )
        out = relax.op.add(relax.op.multiply(positions, delta_e), start_e)
        return out if work_type == out_type else relax.op.astype(out, out_type)

    def convert_rank(self, op):
        """Convert TFLite RANK."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_dtype = self.get_tensor_type_str(output_tensors[0].tensor.Type())

        rank = len(self.get_tensor_shape(input_tensors[0]))
        return relax.const(rank, dtype=output_dtype)

    def convert_shape(self, op):
        """Convert TFLite Shape"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ShapeOptions import ShapeOptions
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.ShapeOptions
        op_options = op.BuiltinOptions()
        shape_options = ShapeOptions()
        shape_options.Init(op_options.Bytes, op_options.Pos)

        # SHAPE must materialize as a tensor output in Relax, not just symbolic shape info.
        out = relax.op.shape_to_tensor(relax.op.shape_of(self.get_tensor_expr(input_tensors[0])))
        if shape_options.OutType() == TensorType.INT32:
            out = relax.op.astype(out, "int32")

        return out

    def convert_bucketize(self, op):
        """Convert TFLite BUCKETIZE."""
        from tflite.BucketizeOptions import BucketizeOptions
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.BucketizeOptions
        op_options = op.BuiltinOptions()
        bucketize_options = BucketizeOptions()
        bucketize_options.Init(op_options.Bytes, op_options.Pos)

        boundaries = self.bb.normalize(
            relax.const(bucketize_options.BoundariesAsNumpy(), dtype="float32")
        )
        out_tensor = self.get_output_tensors(op)[0]
        out_int32 = out_tensor.tensor.Type() == TensorType.INT32
        return relax.op.bucketize(
            self.get_tensor_expr(input_tensors[0]), boundaries, out_int32=out_int32, right=False
        )

    def convert_relu(self, op):
        """Convert TFLite ReLU"""

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_f32 = self.dequantize(in_expr, input_tensor)
            out = relax.op.nn.relu(in_f32)
            out = self.quantize(out, output_tensor)
        else:
            out = relax.op.nn.relu(in_expr)

        return out

    def convert_relu_0_to_1(self, op):
        """Convert TFLite RELU_0_TO_1."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_f32 = self.dequantize(in_expr, input_tensor)
            out = relax.op.clip(in_f32, 0, 1)
            out = self.quantize(out, output_tensor)
        else:
            out = relax.op.clip(in_expr, 0, 1)

        return out

    def convert_hard_swish(self, op):
        """Convert TFLite Hard swish"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        def _relu6(data):
            return relax.op.clip(data, min=0.0, max=6.0)

        def _hard_swish(data):
            return data * _relu6(data + relax.const(3.0)) / relax.const(6.0)

        # Dequantize if the input is quantized.
        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)

        # Perform hardswish
        out = _hard_swish(in_expr)

        # Go back to integer dataype if the original operator was quantized.
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

        return out

    def convert_relu6(self, op):
        """Convert TFLite ReLU6"""

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_f32 = self.dequantize(in_expr, input_tensor)
            out = relax.op.clip(in_f32, min=0, max=6)
            out = self.quantize(out, output_tensor)
        else:
            out = relax.op.clip(in_expr, min=0, max=6)

        return out

    def convert_leaky_relu(self, op):
        """Convert TFLite LEAKY_RELU"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.LeakyReluOptions import LeakyReluOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        assert op.BuiltinOptionsType() == BuiltinOptions.LeakyReluOptions
        op_options = op.BuiltinOptions()
        leaky_relu_options = LeakyReluOptions()
        leaky_relu_options.Init(op_options.Bytes, op_options.Pos)
        alpha_tensor = leaky_relu_options.Alpha()

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)
        out = relax.op.nn.leakyrelu(in_expr, alpha_tensor)
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

        return out

    def convert_relu_n1_to_1(self, op):
        """Convert TFLite RELU_N1_TO_1"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_f32 = self.dequantize(in_expr, input_tensor)
            out = relax.op.clip(in_f32, min=-1, max=1)
            out = self.quantize(out, output_tensor)
        else:
            out = relax.op.clip(in_expr, min=-1, max=1)

        return out

    def convert_log_softmax(self, op):
        """Convert TFLite LOG_SOFTMAX"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)
        out = relax.op.nn.log_softmax(in_expr)
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

        return out

    def convert_concatenation(self, op):
        """Convert TFLite concatenation"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ConcatenationOptions import ConcatenationOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) >= 1, "input tensors should greater than 1"
        in_exprs = [self.get_tensor_expr(_) for _ in input_tensors]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        assert op.BuiltinOptionsType() == BuiltinOptions.ConcatenationOptions
        op_options = op.BuiltinOptions()
        concatenation_options = ConcatenationOptions()
        concatenation_options.Init(op_options.Bytes, op_options.Pos)
        concatenation_axis = concatenation_options.Axis()
        fused_activation_fn = concatenation_options.FusedActivationFunction()

        if not input_tensors[0].qnn_params:
            out = relax.op.concat(in_exprs, axis=concatenation_axis)
        else:
            in_f32s = [
                self.dequantize(expr, tensor) for expr, tensor in zip(in_exprs, input_tensors)
            ]
            out = relax.op.concat(in_f32s, axis=concatenation_axis)
            out = self.quantize(out, output_tensor)

        # Handle fused activations
        if output_tensor.qnn_params:
            scale_val = get_scalar_from_constant(output_tensor.qnn_params["scale"])
            zero_point_val = get_scalar_from_constant(output_tensor.qnn_params["zero_point"])
            output_tensor_type_str = self.get_tensor_type_str(output_tensor.tensor.Type())
            out = self.convert_qnn_fused_activation_function(
                expr=out,
                fused_activation_fn=fused_activation_fn,
                scale=scale_val,
                zero_point=zero_point_val,
                dtype=output_tensor_type_str,
            )
        else:
            out = self.convert_fused_activation_function(out, fused_activation_fn)

        return out

    def _convert_unary_elemwise(self, op, relax_op):
        """Generic method to convert TFLite unary elemwise functions"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)
        out = relax_op(in_expr)
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)
        return out

    def _convert_stablehlo_unary(self, op, relax_op):
        """Convert a unary StableHLO TFLite builtin operator.

        StableHLO builtins do not have TFLite fused activation attributes. Keep
        this path independent from the regular TFLite elemwise/QNN helpers so
        StableHLO semantics are mapped directly to Relax operators.
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        assert len(self.get_output_tensors(op)) == 1, "output tensors length should be 1"

        in_expr = self.get_tensor_expr(input_tensors[0])
        return relax_op(in_expr)

    def _convert_stablehlo_binary(self, op, relax_op):
        """Convert a binary StableHLO TFLite builtin operator.

        StableHLO builtins do not have TFLite fused activation attributes. Keep
        this path independent from the regular TFLite elemwise/QNN helpers so
        StableHLO semantics are mapped directly to Relax operators.
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        assert len(self.get_output_tensors(op)) == 1, "output tensors length should be 1"

        lhs_expr = self.get_tensor_expr(input_tensors[0])
        rhs_expr = self.get_tensor_expr(input_tensors[1])
        return relax_op(lhs_expr, rhs_expr)

    def _convert_stablehlo_and(self, op):
        """Convert StableHLO AND for bool and integer tensors."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        assert len(self.get_output_tensors(op)) == 1, "output tensors length should be 1"

        lhs = self.get_tensor_expr(input_tensors[0])
        rhs = self.get_tensor_expr(input_tensors[1])
        dtype = lhs.ty.dtype.dtype
        if dtype == "bool":
            op_fn = _op.logical_and
        elif dtype.startswith(("int", "uint")):
            op_fn = _op.bitwise_and
        else:
            raise tvm.error.OpNotImplemented(f"STABLEHLO_AND with dtype {dtype} is not supported")
        return self.bb.normalize(op_fn(lhs, rhs))

    def _convert_stablehlo_or(self, op):
        """Convert StableHLO OR for bool and integer tensors."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        assert len(self.get_output_tensors(op)) == 1, "output tensors length should be 1"

        lhs = self.get_tensor_expr(input_tensors[0])
        rhs = self.get_tensor_expr(input_tensors[1])
        dtype = lhs.ty.dtype.dtype
        if dtype == "bool":
            op_fn = _op.logical_or
        elif dtype.startswith(("int", "uint")):
            op_fn = _op.bitwise_or
        else:
            raise tvm.error.OpNotImplemented(f"STABLEHLO_OR with dtype {dtype} is not supported")
        return self.bb.normalize(op_fn(lhs, rhs))

    def _convert_stablehlo_ternary(self, op, relax_op):
        """Convert a ternary StableHLO TFLite builtin operator.

        StableHLO builtins do not have TFLite fused activation attributes. Keep
        this path independent from the regular TFLite elemwise/QNN helpers so
        StableHLO semantics are mapped directly to Relax operators.
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"

        assert len(self.get_output_tensors(op)) == 1, "output tensors length should be 1"

        arg0 = self.get_tensor_expr(input_tensors[0])
        arg1 = self.get_tensor_expr(input_tensors[1])
        arg2 = self.get_tensor_expr(input_tensors[2])
        return relax_op(arg0, arg1, arg2)

    def _get_stablehlo_options(self, op, options_cls):
        """Parse BuiltinOptions2 for a StableHLO TFLite builtin operator.

        Returns an initialized options object of the given class.
        """
        from tflite.BuiltinOptions2 import BuiltinOptions2

        op_options = op.BuiltinOptions2()
        if op_options is None:
            # A malformed flatbuffer may declare a BuiltinOptions2 type without
            # carrying the actual options table. Fail cleanly instead of raising
            # an opaque AttributeError when accessing the missing payload.
            raise tvm.error.OpNotImplemented(
                f"{options_cls.__name__} is required but missing from the operator"
            )
        # Look up the expected BuiltinOptions2 enum value by matching the class
        # name to an enum member (e.g. StablehloConcatenateOptions → 1).
        options_type = getattr(BuiltinOptions2, options_cls.__name__, None)
        if options_type is not None:
            assert op.BuiltinOptions2Type() == options_type, (
                f"Unexpected BuiltinOptions2 type: expected "
                f"{options_cls.__name__}, got {op.BuiltinOptions2Type()}"
            )
        result = options_cls()
        result.Init(op_options.Bytes, op_options.Pos)
        return result

    def _get_static_tensor_shape(self, tensor, op_name):
        """Return a statically-known TFLite tensor shape as Python ints."""
        try:
            return [int(dim) for dim in self.get_tensor_shape(tensor)]
        except (TypeError, ValueError) as err:
            raise tvm.error.OpNotImplemented(
                f"{op_name} requires statically-known tensor shapes"
            ) from err

    def _get_stablehlo_i64_vector(self, vector, default):
        """Convert an optional StableHLO int64 vector field to a Python int list."""
        if vector is None or isinstance(vector, int):
            return list(default)
        return [int(v) for v in vector]

    def _ensure_stablehlo_float_dtype(self, expr, op_name):
        """Return expr dtype if the StableHLO subset supports it."""
        dtype = expr.ty.dtype.dtype
        if not dtype.startswith("float"):
            raise tvm.error.OpNotImplemented(f"{op_name} with dtype {dtype} is not supported")
        return dtype

    def _convert_stablehlo_cbrt(self, op):
        """Convert STABLEHLO_CBRT to a sign-preserving Relax expression."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        assert len(self.get_output_tensors(op)) == 1

        data = self.get_tensor_expr(input_tensors[0])
        dtype = self._ensure_stablehlo_float_dtype(data, "STABLEHLO_CBRT")
        zero = relax.const(0, dtype)
        exponent = relax.const(1.0 / 3.0, dtype)

        is_negative = self.bb.normalize(relax.op.less(data, zero))
        negative_base = self.bb.normalize(relax.op.negative(data))
        negative_root = self.bb.normalize(relax.op.power(negative_base, exponent))
        negative_result = self.bb.normalize(relax.op.negative(negative_root))
        positive_result = self.bb.normalize(relax.op.power(data, exponent))
        return self.bb.normalize(relax.op.where(is_negative, negative_result, positive_result))

    def _convert_stablehlo_remainder(self, op):
        """Convert STABLEHLO_REMAINDER to truncating remainder for float tensors."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(self.get_output_tensors(op)) == 1

        lhs = self.get_tensor_expr(input_tensors[0])
        rhs = self.get_tensor_expr(input_tensors[1])
        self._ensure_stablehlo_float_dtype(lhs, "STABLEHLO_REMAINDER")
        self._ensure_stablehlo_float_dtype(rhs, "STABLEHLO_REMAINDER")

        quotient = self.bb.normalize(relax.op.divide(lhs, rhs))
        truncated = self.bb.normalize(relax.op.trunc(quotient))
        product = self.bb.normalize(relax.op.multiply(rhs, truncated))
        return self.bb.normalize(relax.op.subtract(lhs, product))

    def _get_stablehlo_simple_body_op(self, body_subgraph_index, parent_op_name, input_count):
        """Return the single operator from a simple StableHLO body subgraph."""
        if body_subgraph_index <= 0 or body_subgraph_index >= self.model.SubgraphsLength():
            raise tvm.error.OpNotImplemented(
                f"{parent_op_name} requires a valid non-main body subgraph"
            )

        body_subgraph = self.model.Subgraphs(body_subgraph_index)
        if (
            body_subgraph.InputsLength() != input_count
            or body_subgraph.OutputsLength() != 1
            or body_subgraph.OperatorsLength() != 1
        ):
            raise tvm.error.OpNotImplemented(
                f"{parent_op_name} only supports single-op body subgraphs"
            )

        return body_subgraph.Operators(0)

    def _check_stablehlo_reduce_init(
        self, init_tensor, reducer_name, parent_op_name="STABLEHLO_REDUCE"
    ):
        """Validate that the StableHLO reduce init value matches the Relax identity."""
        if self.has_expr(init_tensor.tensor_idx):
            raise tvm.error.OpNotImplemented(
                f"{parent_op_name} with dynamic init values is not supported"
            )

        init_value = np.asarray(self.get_tensor_value(init_tensor))
        if init_value.shape not in [(), (1,)]:
            raise tvm.error.OpNotImplemented(f"{parent_op_name} requires scalar init values")

        dtype = init_value.dtype
        scalar = init_value.item()
        if reducer_name == "STABLEHLO_ADD":
            is_identity = bool(np.isclose(scalar, 0))
        elif reducer_name == "STABLEHLO_MULTIPLY":
            is_identity = bool(np.isclose(scalar, 1))
        elif reducer_name == "STABLEHLO_MAXIMUM":
            if np.issubdtype(dtype, np.floating):
                is_identity = bool(np.isneginf(scalar))
            elif np.issubdtype(dtype, np.integer):
                is_identity = scalar == np.iinfo(dtype).min
            else:
                is_identity = False
        elif reducer_name == "STABLEHLO_MINIMUM":
            if np.issubdtype(dtype, np.floating):
                is_identity = bool(np.isposinf(scalar))
            elif np.issubdtype(dtype, np.integer):
                is_identity = scalar == np.iinfo(dtype).max
            else:
                is_identity = False
        else:
            raise tvm.error.OpNotImplemented(
                f"{parent_op_name} reducer {reducer_name} is not supported"
            )

        if not is_identity:
            raise tvm.error.OpNotImplemented(
                f"{parent_op_name} init value must match the reducer identity"
            )

    def _convert_stablehlo_reduce(self, op):
        """Convert the single-input STABLEHLO_REDUCE subset to Relax reductions."""
        from tflite.StablehloReduceOptions import StablehloReduceOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloReduceOptions)
        dimensions = self._get_stablehlo_i64_vector(opts.DimensionsAsNumpy(), [])
        body_op = self._get_stablehlo_simple_body_op(
            int(opts.BodySubgraphIndex()), "STABLEHLO_REDUCE", 2
        )
        reducer_name = self.get_op_code_str(body_op)

        reducers = {
            "STABLEHLO_ADD": relax.op.sum,
            "STABLEHLO_MAXIMUM": relax.op.max,
            "STABLEHLO_MINIMUM": relax.op.min,
            "STABLEHLO_MULTIPLY": relax.op.prod,
        }
        if reducer_name not in reducers:
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_REDUCE reducer {reducer_name} is not supported"
            )

        self._check_stablehlo_reduce_init(input_tensors[1], reducer_name)
        data = self.get_tensor_expr(input_tensors[0])
        return self.bb.normalize(reducers[reducer_name](data, axis=dimensions, keepdims=False))

    def _convert_stablehlo_reduce_window(self, op):
        """Convert the NHWC 2D max-pool STABLEHLO_REDUCE_WINDOW subset."""
        from tflite.StablehloReduceWindowOptions import StablehloReduceWindowOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloReduceWindowOptions)
        body_op = self._get_stablehlo_simple_body_op(
            int(opts.BodySubgraphIndex()), "STABLEHLO_REDUCE_WINDOW", 2
        )
        reducer_name = self.get_op_code_str(body_op)
        if reducer_name != "STABLEHLO_MAXIMUM":
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_REDUCE_WINDOW only supports MAXIMUM reducer windows"
            )
        self._check_stablehlo_reduce_init(input_tensors[1], reducer_name, "STABLEHLO_REDUCE_WINDOW")

        data_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_REDUCE_WINDOW")
        if len(data_shape) != 4:
            raise tvm.error.OpNotImplemented("STABLEHLO_REDUCE_WINDOW only supports 4D input")

        window_dimensions = self._get_stablehlo_i64_vector(opts.WindowDimensionsAsNumpy(), [])
        window_strides = self._get_stablehlo_i64_vector(
            opts.WindowStridesAsNumpy(), [1] * len(window_dimensions)
        )
        base_dilations = self._get_stablehlo_i64_vector(
            opts.BaseDilationsAsNumpy(), [1] * len(window_dimensions)
        )
        window_dilations = self._get_stablehlo_i64_vector(
            opts.WindowDilationsAsNumpy(), [1] * len(window_dimensions)
        )
        padding = self._get_stablehlo_i64_vector(
            opts.PaddingAsNumpy(), [0] * (2 * len(window_dimensions))
        )

        if (
            len(window_dimensions) != 4
            or len(window_strides) != 4
            or len(base_dilations) != 4
            or len(window_dilations) != 4
            or len(padding) != 8
        ):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_REDUCE_WINDOW only supports rank-4 window attributes"
            )
        if window_dimensions[0] != 1 or window_dimensions[3] != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_REDUCE_WINDOW only supports pooling over spatial dimensions"
            )
        if window_strides[0] != 1 or window_strides[3] != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_REDUCE_WINDOW only supports unit batch/channel strides"
            )
        if base_dilations != [1, 1, 1, 1]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_REDUCE_WINDOW with base dilation is not supported"
            )
        if padding[0] != 0 or padding[1] != 0 or padding[6] != 0 or padding[7] != 0:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_REDUCE_WINDOW only supports spatial padding"
            )

        data = self.get_tensor_expr(input_tensors[0])
        return self.bb.normalize(
            relax.op.nn.max_pool2d(
                data,
                pool_size=[window_dimensions[1], window_dimensions[2]],
                strides=[window_strides[1], window_strides[2]],
                padding=[padding[2], padding[4], padding[3], padding[5]],
                dilation=[window_dilations[1], window_dilations[2]],
                layout="NHWC",
                out_layout="NHWC",
            )
        )

    def _convert_stablehlo_scatter(self, op):
        """Convert the canonical point-update STABLEHLO_SCATTER subset."""
        from tflite.StablehloScatterOptions import StablehloScatterOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloScatterOptions)
        operand_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_SCATTER")
        indices_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_SCATTER")
        updates_shape = self._get_static_tensor_shape(input_tensors[2], "STABLEHLO_SCATTER")
        operand_rank = len(operand_shape)
        indices_rank = len(indices_shape)

        update_window_dims = self._get_stablehlo_i64_vector(opts.UpdateWindowDimsAsNumpy(), [])
        inserted_window_dims = self._get_stablehlo_i64_vector(opts.InsertedWindowDimsAsNumpy(), [])
        scatter_dims_to_operand_dims = self._get_stablehlo_i64_vector(
            opts.ScatterDimsToOperandDimsAsNumpy(), []
        )
        index_vector_dim = int(opts.IndexVectorDim())

        if indices_rank == 0 or index_vector_dim != indices_rank - 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SCATTER only supports trailing index-vector dimensions"
            )
        if update_window_dims:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SCATTER only supports point updates without update windows"
            )
        if inserted_window_dims != list(range(operand_rank)):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SCATTER only supports point updates for every operand dimension"
            )
        if scatter_dims_to_operand_dims != list(range(operand_rank)):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SCATTER only supports canonical scatter-to-operand dimensions"
            )
        if indices_shape[-1] != operand_rank or updates_shape != indices_shape[:-1]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SCATTER requires point update shapes to match scatter indices"
            )

        body_op = self._get_stablehlo_simple_body_op(
            int(opts.UpdateComputationSubgraphIndex()), "STABLEHLO_SCATTER", 2
        )
        reducer_name = self.get_op_code_str(body_op)
        reductions = {
            "STABLEHLO_ADD": "add",
            "STABLEHLO_MAXIMUM": "max",
            "STABLEHLO_MINIMUM": "min",
            "STABLEHLO_MULTIPLY": "mul",
        }
        if reducer_name not in reductions:
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_SCATTER reducer {reducer_name} is not supported"
            )

        operand = self.get_tensor_expr(input_tensors[0])
        indices = self.get_tensor_expr(input_tensors[1])
        updates = self.get_tensor_expr(input_tensors[2])
        return self.bb.normalize(
            relax.op.scatter_nd(operand, indices, updates, reductions[reducer_name])
        )

    def _convert_stablehlo_composite(self, op):
        """Convert STABLEHLO_COMPOSITE by inlining a simple decomposition subgraph."""
        from tflite.StableHLOCompositeOptions import StableHLOCompositeOptions

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_COMPOSITE only supports single-output decompositions"
            )

        opts = self._get_stablehlo_options(op, StableHLOCompositeOptions)
        composite_name = opts.Name()
        composite_name = (
            composite_name.decode("utf-8") if composite_name is not None else "<unnamed>"
        )
        if opts.CompositeAttributesLength() != 0:
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_COMPOSITE {composite_name} with composite attributes is not supported"
            )

        decomposition_subgraph_index = int(opts.DecompositionSubgraphIndex())
        if (
            decomposition_subgraph_index <= 0
            or decomposition_subgraph_index >= self.model.SubgraphsLength()
        ):
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_COMPOSITE {composite_name} requires a valid decomposition subgraph"
            )
        decomposition_subgraph = self.model.Subgraphs(decomposition_subgraph_index)
        if decomposition_subgraph.InputsLength() != len(input_tensors):
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_COMPOSITE {composite_name} decomposition input count mismatch"
            )
        if decomposition_subgraph.OutputsLength() != 1:
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_COMPOSITE {composite_name} only supports single-output decompositions"
            )

        decomposition_exp_tab = ExprTable()
        decomposition_converter = OperatorConverter(
            self.model, decomposition_subgraph, decomposition_exp_tab, self.bb
        )
        for decomposition_input_idx, composite_input in zip(
            decomposition_subgraph.InputsAsNumpy(), input_tensors
        ):
            decomposition_input_name = get_tensor_name(
                decomposition_subgraph, int(decomposition_input_idx)
            )
            decomposition_exp_tab.set_expr(
                decomposition_input_name,
                self.get_tensor_expr(composite_input),
                force_override=True,
            )

        decomposition_converter.check_unsupported_ops()
        decomposition_converter.convert_op_to_relax()
        decomposition_output_idx = int(decomposition_subgraph.Outputs(0))
        decomposition_output_tensor = decomposition_converter.get_tensors(
            [decomposition_output_idx]
        )[0]
        for const_expr, value in decomposition_exp_tab.params.values():
            param_name = f"_param_{self.exp_tab.const_ctr}"
            self.exp_tab.const_ctr += 1
            self.exp_tab.params[param_name] = (const_expr, value)
        return decomposition_converter.get_tensor_expr(decomposition_output_tensor)

    def _convert_stablehlo_sort(self, op):
        """Convert the single-input STABLEHLO_SORT subset to Relax sort."""
        from tflite.StablehloCompareOptions import StablehloCompareOptions
        from tflite.StablehloComparisonDirection import StablehloComparisonDirection
        from tflite.StablehloComparisonType import StablehloComparisonType
        from tflite.StablehloSortOptions import StablehloSortOptions

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 1 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SORT only supports single-input single-output sort"
            )

        opts = self._get_stablehlo_options(op, StablehloSortOptions)
        if opts.IsStable():
            raise tvm.error.OpNotImplemented("STABLEHLO_SORT stable sort is not supported")

        body_op = self._get_stablehlo_simple_body_op(
            int(opts.ComparatorSubgraphIndex()), "STABLEHLO_SORT", 2
        )
        comparator_name = self.get_op_code_str(body_op)
        if comparator_name != "STABLEHLO_COMPARE":
            raise tvm.error.OpNotImplemented(
                f"STABLEHLO_SORT comparator {comparator_name} is not supported"
            )

        compare_opts = self._get_stablehlo_options(body_op, StablehloCompareOptions)
        if (
            compare_opts.CompareType()
            == StablehloComparisonType.STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER
        ):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_SORT with TOTALORDER comparator is not supported"
            )

        direction = compare_opts.ComparisonDirection()
        _DIR = StablehloComparisonDirection
        if direction == _DIR.STABLEHLO_COMPARISON_DIRECTION_LT:
            descending = False
        elif direction == _DIR.STABLEHLO_COMPARISON_DIRECTION_GT:
            descending = True
        else:
            raise tvm.error.OpNotImplemented("STABLEHLO_SORT only supports LT or GT comparators")

        data = self.get_tensor_expr(input_tensors[0])
        return self.bb.normalize(
            relax.op.sort(data, axis=int(opts.Dimension()), descending=descending)
        )

    def _convert_stablehlo_custom_call(self, op):
        """Convert supported annotation-only STABLEHLO_CUSTOM_CALL targets."""
        from tflite.StablehloCustomCallOptions import StablehloCustomCallOptions

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        opts = self._get_stablehlo_options(op, StablehloCustomCallOptions)
        call_target_name = self._decode_tflite_string(opts.CallTargetName())

        if call_target_name == "Sharding":
            # TensorFlow treats Sharding custom calls as metadata annotations
            # and may erase them by replacing the op with its input. Mirror
            # that identity semantics for the safe single-input/single-output
            # subset. The sharding spec in backend_config is intentionally
            # dropped for single-device import. TFLite has no runtime kernel
            # for general STABLEHLO_CUSTOM_CALL targets.
            if opts.HasSideEffect():
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_CUSTOM_CALL Sharding with side effects is not supported"
                )
            if opts.CalledComputationsLength() != 0:
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_CUSTOM_CALL Sharding with called computations is not supported"
                )
            if len(input_tensors) != 1 or len(output_tensors) != 1:
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_CUSTOM_CALL Sharding requires one input and one output"
                )
            self._check_tensor_metadata_match(
                input_tensors[0], output_tensors[0], "STABLEHLO_CUSTOM_CALL", "Sharding"
            )
            return self.get_tensor_expr(input_tensors[0])

        target = call_target_name or "<empty>"
        raise tvm.error.OpNotImplemented(f"STABLEHLO_CUSTOM_CALL target {target} is not supported")

    def _convert_stablehlo_rng_bit_generator(self, op):
        """Convert STABLEHLO_RNG_BIT_GENERATOR to a bit-exact call_tir kernel."""
        from tflite.RngAlgorithm import RngAlgorithm
        from tflite.StablehloRngBitGeneratorOptions import StablehloRngBitGeneratorOptions

        op_name = "STABLEHLO_RNG_BIT_GENERATOR"
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 1 or len(output_tensors) != 2:
            raise tvm.error.OpNotImplemented(f"{op_name} expects one input and two outputs")

        opts = self._get_stablehlo_options(op, StablehloRngBitGeneratorOptions)
        algorithm_enum = opts.Algorithm()
        # DEFAULT resolves to PHILOX in the TFLite runtime kernel.
        if algorithm_enum == RngAlgorithm.THREEFRY:
            algorithm = "threefry"
        elif algorithm_enum in (RngAlgorithm.PHILOX, RngAlgorithm.DEFAULT):
            algorithm = "philox"
        else:
            raise tvm.error.OpNotImplemented(
                f"{op_name} algorithm {algorithm_enum} is not supported"
            )

        state_tensor = input_tensors[0]
        if self.get_tensor_type_str(state_tensor.tensor.Type()) != "uint64":
            raise tvm.error.OpNotImplemented(f"{op_name} requires a uint64 initial state")
        state_shape = self._get_static_tensor_shape(state_tensor, op_name)
        if len(state_shape) != 1:
            raise tvm.error.OpNotImplemented(f"{op_name} requires a 1-D initial state")
        state_len = int(state_shape[0])
        # State-length constraints mirror the TFLite runtime kernel.
        if algorithm == "threefry" and state_len != 2:
            raise tvm.error.OpNotImplemented(f"{op_name} THREEFRY requires a u64[2] state")
        if algorithm == "philox" and state_len not in (2, 3):
            raise tvm.error.OpNotImplemented(f"{op_name} PHILOX requires a u64[2] or u64[3] state")

        out_state_tensor, out_tensor = output_tensors
        if self.get_tensor_type_str(out_state_tensor.tensor.Type()) != "uint64":
            raise tvm.error.OpNotImplemented(f"{op_name} output state must be uint64")
        out_state_shape = self._get_static_tensor_shape(out_state_tensor, op_name)
        if list(out_state_shape) != list(state_shape):
            raise tvm.error.OpNotImplemented(
                f"{op_name} output state shape must match the initial state"
            )
        out_dtype = self.get_tensor_type_str(out_tensor.tensor.Type())
        if out_dtype not in ("int32", "int64", "uint32", "uint64"):
            raise tvm.error.OpNotImplemented(f"{op_name} output dtype {out_dtype} is not supported")
        out_shape = tuple(self._get_static_tensor_shape(out_tensor, op_name))

        prim_func = _build_stablehlo_rng_bit_generator_primfunc(
            algorithm, state_len, out_dtype, out_shape
        )
        module_builder = self.conversion_state["module_builder"]
        func_name = f"tflite_stablehlo_rng_{algorithm}_{out_state_tensor.tensor_idx}"
        gv = module_builder.add_func(prim_func, func_name)
        state_expr = self.get_tensor_expr(state_tensor)
        call = relax.call_tir(
            gv,
            [state_expr],
            [
                relax.TensorType(tuple(state_shape), "uint64"),
                relax.TensorType(out_shape, out_dtype),
            ],
        )
        return self.bb.normalize(call)

    def _convert_stablehlo_while(self, op):
        """Convert STABLEHLO_WHILE to a recursive Relax private function."""
        from tflite.StablehloWhileOptions import StablehloWhileOptions

        opts = self._get_stablehlo_options(op, StablehloWhileOptions)
        return self._convert_while_like(
            op,
            "STABLEHLO_WHILE",
            int(opts.CondSubgraphIndex()),
            int(opts.BodySubgraphIndex()),
            "tflite_stablehlo_while",
        )

    def _get_builtin_options(self, op, options_cls):
        """Parse BuiltinOptions for a TFLite builtin operator."""
        from tflite.BuiltinOptions import BuiltinOptions

        op_options = op.BuiltinOptions()
        if op_options is None:
            raise tvm.error.OpNotImplemented(f"{options_cls.__name__} is required")

        options_type = getattr(BuiltinOptions, options_cls.__name__, None)
        if options_type is not None and op.BuiltinOptionsType() != options_type:
            raise tvm.error.OpNotImplemented(
                f"Unexpected BuiltinOptions type: expected "
                f"{options_cls.__name__}, got {op.BuiltinOptionsType()}"
            )
        result = options_cls()
        result.Init(op_options.Bytes, op_options.Pos)
        return result

    def _get_subgraph(self, subgraph_index, op_name, allow_main=False):
        """Return a validated TFLite subgraph by index."""
        if subgraph_index < 0 or subgraph_index >= self.model.SubgraphsLength():
            raise tvm.error.OpNotImplemented(f"{op_name} requires a valid subgraph index")
        if not allow_main and subgraph_index == 0:
            raise tvm.error.OpNotImplemented(f"{op_name} cannot target the main subgraph")
        return self.model.Subgraphs(subgraph_index)

    def _make_tuple_or_single(self, exprs):
        """Return a single expression or Relax tuple for a list of expressions."""
        if len(exprs) == 1:
            return exprs[0]
        return relax.Tuple(exprs)

    def _indices_or_empty(self, indices):
        """Return a TFLite index vector, using an empty list for absent vectors."""
        return indices if indices is not None else []

    def _check_subgraph_io(self, subgraph_index, op_name, input_count=None, output_count=None):
        """Validate a referenced subgraph's input and output counts."""
        subgraph = self._get_subgraph(subgraph_index, op_name)
        if input_count is not None and subgraph.InputsLength() != input_count:
            raise tvm.error.OpNotImplemented(f"{op_name} subgraph input count mismatch")
        if output_count is not None and subgraph.OutputsLength() != output_count:
            raise tvm.error.OpNotImplemented(f"{op_name} subgraph output count mismatch")
        return subgraph

    def _check_subgraph_interface(
        self,
        subgraph_index,
        op_name,
        input_tensors=None,
        output_tensors=None,
        input_count=None,
        output_count=None,
    ):
        """Validate a referenced subgraph's arity and tensor metadata."""
        if input_tensors is not None:
            input_count = len(input_tensors)
        if output_tensors is not None:
            output_count = len(output_tensors)

        subgraph = self._check_subgraph_io(
            subgraph_index, op_name, input_count=input_count, output_count=output_count
        )
        if input_tensors is not None:
            self._check_subgraph_tensor_metadata(
                subgraph,
                op_name,
                "subgraph input",
                subgraph.InputsAsNumpy(),
                input_tensors,
            )
        if output_tensors is not None:
            self._check_subgraph_tensor_metadata(
                subgraph,
                op_name,
                "subgraph output",
                subgraph.OutputsAsNumpy(),
                output_tensors,
            )
        return subgraph

    def _get_tensor_metadata(self, tensor):
        """Return static shape and dtype metadata for a TFLite tensor."""
        if isinstance(tensor, TensorWrapper):
            tensor = tensor.tensor
        shape = self._get_relax_tensor_shape(tensor)
        dtype = self._get_relax_tensor_dtype(tensor)
        return shape, dtype

    def _check_tensor_metadata_match(self, actual, expected, op_name, tensor_role):
        """Validate that two TFLite tensors have matching static metadata."""
        if self._get_tensor_metadata(actual) != self._get_tensor_metadata(expected):
            raise tvm.error.OpNotImplemented(f"{op_name} {tensor_role} tensor metadata mismatch")

    def _check_subgraph_tensor_metadata(
        self, subgraph, op_name, tensor_role, subgraph_indices, expected_tensors
    ):
        """Validate referenced subgraph tensor metadata against caller tensors."""
        for subgraph_index, expected_tensor in zip(
            self._indices_or_empty(subgraph_indices), expected_tensors
        ):
            self._check_tensor_metadata_match(
                subgraph.Tensors(int(subgraph_index)),
                expected_tensor,
                op_name,
                tensor_role,
            )

    def _require_scalar_bool_tensor(self, tensor, op_name):
        """Validate that a TFLite tensor is a scalar bool tensor."""
        if isinstance(tensor, TensorWrapper):
            tensor = tensor.tensor
        dtype = self.get_tensor_type_str(tensor.Type())
        if dtype != "bool" or tensor.ShapeLength() != 0:
            raise tvm.error.OpNotImplemented(f"{op_name} requires a scalar bool condition")

    def _get_subgraph_params(self, subgraph):
        """Create Relax parameters for a TFLite subgraph."""
        params = []
        exp_tab = ExprTable()
        for input_index in self._indices_or_empty(subgraph.InputsAsNumpy()):
            tensor = subgraph.Tensors(int(input_index))
            input_name = get_tensor_name(subgraph, int(input_index))
            shape = self._get_relax_tensor_shape(tensor)
            dtype = self._get_relax_tensor_dtype(tensor)
            param = relax.Var(input_name, relax.TensorType(shape=shape, dtype=dtype))
            exp_tab.set_expr(input_name, param)
            params.append(param)
        return params, exp_tab

    def _get_tensor_param(self, tensor_wrapper):
        """Create a Relax parameter from TFLite tensor metadata."""
        name = get_tensor_name(self.subgraph, tensor_wrapper.tensor_idx)
        shape = self._get_relax_tensor_shape(tensor_wrapper)
        dtype = self._get_relax_tensor_dtype(tensor_wrapper)
        return relax.Var(name, relax.TensorType(shape=shape, dtype=dtype))

    def _lower_subgraph_to_function(self, subgraph_index, function_name_hint, op_name="CALL"):
        """Lower a TFLite subgraph into a private Relax function."""
        lowered_subgraphs = self.conversion_state["lowered_subgraphs"]
        if subgraph_index in lowered_subgraphs:
            return lowered_subgraphs[subgraph_index]

        lowering_stack = self.conversion_state["lowering_stack"]
        if subgraph_index in lowering_stack:
            raise tvm.error.OpNotImplemented(
                f"Recursive TFLite {op_name} subgraphs are not supported"
            )

        subgraph = self._get_subgraph(subgraph_index, op_name)
        lowering_stack.append(subgraph_index)
        try:
            params, subgraph_exp_tab = self._get_subgraph_params(subgraph)
            subgraph_bb = relax.BlockBuilder()
            with subgraph_bb.function(function_name_hint, params=params, private=True):
                with subgraph_bb.dataflow():
                    subgraph_converter = type(self)(
                        self.model,
                        subgraph,
                        subgraph_exp_tab,
                        subgraph_bb,
                        self.conversion_state,
                    )
                    subgraph_converter.check_unsupported_ops()
                    subgraph_converter.convert_op_to_relax()
                    output_tensors = subgraph_converter.get_tensors(subgraph.OutputsAsNumpy())
                    outputs = [
                        subgraph_converter.get_tensor_expr(tensor) for tensor in output_tensors
                    ]
                    output = subgraph_bb.emit_output(self._make_tuple_or_single(outputs))
                subgraph_bb.emit_func_output(output)

            subgraph_mod = subgraph_bb.get()
            module_builder = self.conversion_state["module_builder"]
            gv = module_builder.add_func(subgraph_mod[function_name_hint], function_name_hint)
            lowered_subgraphs[subgraph_index] = gv
            return gv
        finally:
            lowering_stack.pop()

    def _bind_call_outputs(self, call, output_count):
        """Return per-output expressions from a single or tuple-valued call."""
        if output_count == 1:
            return [call]
        return [call[index] for index in range(output_count)]

    def _lower_if_to_function(
        self,
        then_subgraph_index,
        else_subgraph_index,
        input_tensors,
        branch_input_count,
        output_count,
    ):
        """Lower a TFLite IF op into a private Relax function."""
        cache_key = (then_subgraph_index, else_subgraph_index, branch_input_count, output_count)
        lowered_if_functions = self.conversion_state["lowered_if_functions"]
        if cache_key in lowered_if_functions:
            return lowered_if_functions[cache_key]

        then_func = self._lower_subgraph_to_function(
            then_subgraph_index,
            f"tflite_if_then_subgraph_{then_subgraph_index}",
            op_name="IF",
        )
        else_func = self._lower_subgraph_to_function(
            else_subgraph_index,
            f"tflite_if_else_subgraph_{else_subgraph_index}",
            op_name="IF",
        )
        if_name = f"tflite_if_subgraph_{then_subgraph_index}_{else_subgraph_index}"
        params = [self._get_tensor_param(tensor) for tensor in input_tensors]
        cond = params[0]
        branch_args = params[1:]

        if_bb = relax.BlockBuilder()
        with if_bb.function(if_name, params=params, private=True):
            result = relax.If(
                cond,
                relax.Call(then_func, branch_args),
                relax.Call(else_func, branch_args),
            )
            if_bb.emit_func_output(result)
        if_func = if_bb.get()[if_name]
        module_builder = self.conversion_state["module_builder"]
        gv = module_builder.add_func(if_func, if_name)
        lowered_if_functions[cache_key] = gv
        return gv

    def _lower_while_to_function(
        self,
        cond_subgraph_index,
        body_subgraph_index,
        loop_var_count,
        cond_func,
        body_func,
        body_subgraph,
        function_prefix="tflite_while",
    ):
        """Lower a TFLite WHILE op into a recursive private Relax function."""
        cache_key = (function_prefix, cond_subgraph_index, body_subgraph_index, loop_var_count)
        lowered_while_functions = self.conversion_state["lowered_while_functions"]
        if cache_key in lowered_while_functions:
            return lowered_while_functions[cache_key]

        loop_name = f"{function_prefix}_subgraph_{cond_subgraph_index}_{body_subgraph_index}"
        params, _ = self._get_subgraph_params(body_subgraph)
        dummy_body = self._make_tuple_or_single(params)
        module_builder = self.conversion_state["module_builder"]
        loop_gv = module_builder.add_func(relax.Function(params, dummy_body), loop_name)
        lowered_while_functions[cache_key] = loop_gv

        loop_bb = relax.BlockBuilder()
        with loop_bb.function(loop_name, params=params, private=True):
            cond = loop_bb.emit(relax.Call(cond_func, params), "while_cond")
            next_state = relax.Call(body_func, params)
            next_args = self._bind_call_outputs(next_state, loop_var_count)
            true_branch = relax.Call(loop_gv, next_args)
            false_branch = self._make_tuple_or_single(params)
            result = relax.If(cond, true_branch, false_branch)
            loop_bb.emit_func_output(result)
        loop_func = loop_bb.get()[loop_name]
        module_builder.update_func(loop_gv, loop_func)
        return loop_gv

    def convert_call(self, op):
        """Convert TFLite CALL to a Relax private function call."""
        from tflite.CallOptions import CallOptions

        opts = self._get_builtin_options(op, CallOptions)
        subgraph_index = int(opts.Subgraph())
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        self._check_subgraph_interface(
            subgraph_index,
            "CALL",
            input_tensors=input_tensors,
            output_tensors=output_tensors,
        )

        callee = self._lower_subgraph_to_function(
            subgraph_index, f"tflite_call_subgraph_{subgraph_index}", op_name="CALL"
        )
        args = [self.get_tensor_expr(tensor) for tensor in input_tensors]
        return relax.Call(callee, args)

    def convert_if(self, op):
        """Convert TFLite IF to Relax If with private branch functions."""
        from tflite.IfOptions import IfOptions

        opts = self._get_builtin_options(op, IfOptions)
        then_subgraph_index = int(opts.ThenSubgraphIndex())
        else_subgraph_index = int(opts.ElseSubgraphIndex())
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) < 1:
            raise tvm.error.OpNotImplemented("IF requires a condition input")

        self._require_scalar_bool_tensor(input_tensors[0], "IF")
        branch_input_count = len(input_tensors) - 1
        output_count = len(output_tensors)
        branch_input_tensors = input_tensors[1:]
        self._check_subgraph_interface(
            then_subgraph_index,
            "IF",
            input_tensors=branch_input_tensors,
            output_tensors=output_tensors,
        )
        self._check_subgraph_interface(
            else_subgraph_index,
            "IF",
            input_tensors=branch_input_tensors,
            output_tensors=output_tensors,
        )

        if_func = self._lower_if_to_function(
            then_subgraph_index,
            else_subgraph_index,
            input_tensors,
            branch_input_count,
            output_count,
        )
        args = [self.get_tensor_expr(tensor) for tensor in input_tensors]
        return relax.Call(if_func, args)

    def _convert_while_like(
        self, op, op_name, cond_subgraph_index, body_subgraph_index, function_prefix
    ):
        """Convert a TFLite while-like operator with referenced cond/body subgraphs."""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        loop_var_count = len(input_tensors)
        if loop_var_count == 0:
            raise tvm.error.OpNotImplemented(f"{op_name} requires loop-carried inputs")
        if len(output_tensors) != loop_var_count:
            raise tvm.error.OpNotImplemented(f"{op_name} output count must match input count")

        cond_subgraph = self._check_subgraph_interface(
            cond_subgraph_index,
            op_name,
            input_tensors=input_tensors,
            output_count=1,
        )
        body_subgraph = self._check_subgraph_interface(
            body_subgraph_index,
            op_name,
            input_tensors=input_tensors,
            output_tensors=input_tensors,
        )
        for input_tensor, output_tensor in zip(input_tensors, output_tensors):
            self._check_tensor_metadata_match(input_tensor, output_tensor, op_name, "loop state")
        cond_output = cond_subgraph.Tensors(int(cond_subgraph.Outputs(0)))
        self._require_scalar_bool_tensor(cond_output, op_name)

        cond_func = self._lower_subgraph_to_function(
            cond_subgraph_index,
            f"{function_prefix}_cond_subgraph_{cond_subgraph_index}",
            op_name=op_name,
        )
        body_func = self._lower_subgraph_to_function(
            body_subgraph_index,
            f"{function_prefix}_body_subgraph_{body_subgraph_index}",
            op_name=op_name,
        )

        loop_gv = self._lower_while_to_function(
            cond_subgraph_index,
            body_subgraph_index,
            loop_var_count,
            cond_func,
            body_func,
            body_subgraph,
            function_prefix=function_prefix,
        )

        args = [self.get_tensor_expr(tensor) for tensor in input_tensors]
        return relax.Call(loop_gv, args)

    def convert_while(self, op):
        """Convert TFLite WHILE to a recursive Relax private function."""
        from tflite.WhileOptions import WhileOptions

        opts = self._get_builtin_options(op, WhileOptions)
        return self._convert_while_like(
            op,
            "WHILE",
            int(opts.CondSubgraphIndex()),
            int(opts.BodySubgraphIndex()),
            "tflite_while",
        )

    def convert_call_once(self, op):
        """Convert TFLite CALL_ONCE for no-op and resource-variable initialization subsets."""
        from tflite.CallOnceOptions import CallOnceOptions

        opts = self._get_builtin_options(op, CallOnceOptions)
        init_subgraph_index = int(opts.InitSubgraphIndex())
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 0 or len(output_tensors) != 0:
            raise tvm.error.OpNotImplemented("CALL_ONCE with inputs or outputs is not supported")

        init_subgraph = self._get_subgraph(init_subgraph_index, "CALL_ONCE")
        if init_subgraph.InputsLength() != 0 or init_subgraph.OutputsLength() != 0:
            raise tvm.error.OpNotImplemented(
                "CALL_ONCE with non-empty init subgraph I/O is not supported"
            )
        if init_subgraph.OperatorsLength() != 0:
            self._convert_call_once_init_subgraph(init_subgraph)
        return None

    def _convert_call_once_init_subgraph(self, init_subgraph):
        """Convert the resource-variable initialization subset of a CALL_ONCE subgraph."""
        supported_init_ops = {"VAR_HANDLE", "ASSIGN_VARIABLE", "HASHTABLE", "HASHTABLE_IMPORT"}
        for op_idx in range(init_subgraph.OperatorsLength()):
            op_name = self.get_op_code_str(init_subgraph.Operators(op_idx))
            if op_name not in supported_init_ops:
                raise tvm.error.OpNotImplemented(
                    f"CALL_ONCE init subgraph operator {op_name} is not supported"
                )

        old_in_call_once_init = self.conversion_state["in_call_once_init"]
        self.conversion_state["in_call_once_init"] = True
        try:
            # The supported init ops below only update importer state and return None.
            # If future CALL_ONCE ops emit Relax bindings, revisit sharing the parent builder.
            subgraph_converter = type(self)(
                self.model,
                init_subgraph,
                ExprTable(),
                self.bb,
                self.conversion_state,
            )
            subgraph_converter.check_unsupported_ops()
            subgraph_converter.convert_op_to_relax()
        finally:
            self.conversion_state["in_call_once_init"] = old_in_call_once_init

    def _convert_stablehlo_convert(self, op):
        """Convert STABLEHLO_CONVERT to Relax (astype).

        Reads the output tensor dtype from the TFLite schema and applies
        relax.op.astype.  This path is intentionally separate from the
        generic _convert_stablehlo_unary helper because the output dtype
        is operator-level metadata, not a Relax op parameter.
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"

        in_expr = self.get_tensor_expr(input_tensors[0])
        output_dtype = self.get_tensor_type_str(output_tensors[0].tensor.Type())
        return self.bb.normalize(relax.op.astype(in_expr, output_dtype))

    def _convert_stablehlo_clamp(self, op):
        """Convert STABLEHLO_CLAMP to Relax.

        StableHLO clamp(min, operand, max) → R.minimum(R.maximum(operand, min), max).
        """
        # NOTE: R.clip is not used here because it only accepts scalar Expr
        # min/max, not tensor inputs.
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"

        assert len(self.get_output_tensors(op)) == 1

        min_expr = self.get_tensor_expr(input_tensors[0])
        operand_expr = self.get_tensor_expr(input_tensors[1])
        max_expr = self.get_tensor_expr(input_tensors[2])

        clamped = self.bb.normalize(relax.op.maximum(operand_expr, min_expr))
        return self.bb.normalize(relax.op.minimum(clamped, max_expr))

    def _convert_stablehlo_concatenate(self, op):
        """Convert STABLEHLO_CONCATENATE to Relax."""
        from tflite.StablehloConcatenateOptions import StablehloConcatenateOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) >= 1, "input tensors length should be >= 1"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloConcatenateOptions)
        dim = opts.Dimension()

        in_exprs = [self.get_tensor_expr(t) for t in input_tensors]
        return self.bb.normalize(relax.op.concat(in_exprs, axis=dim))

    def _convert_stablehlo_broadcast_in_dim(self, op):
        """Convert STABLEHLO_BROADCAST_IN_DIM to Relax."""
        from tflite.StablehloBroadcastInDimOptions import StablehloBroadcastInDimOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1
        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1

        opts = self._get_stablehlo_options(op, StablehloBroadcastInDimOptions)
        broadcast_dims = [int(d) for d in opts.BroadcastDimensionsAsNumpy()]

        in_expr = self.get_tensor_expr(input_tensors[0])
        input_shape = [int(d) for d in self.get_tensor_shape(input_tensors[0])]
        output_shape = [int(d) for d in self.get_tensor_shape(output_tensors[0])]

        # Map input dims to output dims via broadcast_dims, filling
        # unmapped positions with 1 so broadcast_to covers them.
        intermediate_shape = [1] * len(output_shape)
        for i, d in enumerate(broadcast_dims):
            intermediate_shape[d] = input_shape[i]

        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_iota(self, op):
        """Convert STABLEHLO_IOTA to Relax (arange + broadcast)."""
        from tflite.StablehloIotaOptions import StablehloIotaOptions

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1

        opts = self._get_stablehlo_options(op, StablehloIotaOptions)
        iota_dim = opts.IotaDimension()

        output_tensor = output_tensors[0]
        output_shape = [int(d) for d in self.get_tensor_shape(output_tensor)]
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        # arange along the iota dimension
        size = output_shape[iota_dim]
        arange_1d = self.bb.normalize(relax.op.arange(0, size, 1, output_dtype))

        # reshape to [1, ..., size, ..., 1]
        broadcast_shape = [1] * len(output_shape)
        broadcast_shape[iota_dim] = size
        arange_reshaped = self.bb.normalize(relax.op.reshape(arange_1d, broadcast_shape))

        # broadcast to full output shape
        return self.bb.normalize(relax.op.broadcast_to(arange_reshaped, output_shape))

    def _convert_stablehlo_compare(self, op):
        """Convert STABLEHLO_COMPARE to Relax binary comparison ops."""
        from tflite.StablehloCompareOptions import StablehloCompareOptions
        from tflite.StablehloComparisonDirection import StablehloComparisonDirection

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2
        assert len(self.get_output_tensors(op)) == 1

        from tflite.StablehloComparisonType import StablehloComparisonType

        opts = self._get_stablehlo_options(op, StablehloCompareOptions)
        direction = opts.ComparisonDirection()
        compare_type = opts.CompareType()

        # TOTALORDER compare is not expressible via Relax comparison ops.
        if compare_type == StablehloComparisonType.STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_COMPARE with TOTALORDER comparison type is not supported"
            )

        _DIR = StablehloComparisonDirection
        direction_map = {
            _DIR.STABLEHLO_COMPARISON_DIRECTION_EQ: relax.op.equal,
            _DIR.STABLEHLO_COMPARISON_DIRECTION_NE: relax.op.not_equal,
            _DIR.STABLEHLO_COMPARISON_DIRECTION_GE: relax.op.greater_equal,
            _DIR.STABLEHLO_COMPARISON_DIRECTION_GT: relax.op.greater,
            _DIR.STABLEHLO_COMPARISON_DIRECTION_LE: relax.op.less_equal,
            _DIR.STABLEHLO_COMPARISON_DIRECTION_LT: relax.op.less,
        }
        relax_fn = direction_map.get(direction)
        if relax_fn is None:
            raise tvm.error.OpNotImplemented(
                f"Unsupported StableHLO comparison direction: {direction}"
            )

        lhs = self.get_tensor_expr(input_tensors[0])
        rhs = self.get_tensor_expr(input_tensors[1])
        return self.bb.normalize(relax_fn(lhs, rhs))

    def _convert_stablehlo_pad(self, op):
        """Convert STABLEHLO_PAD to Relax (nn.pad).

        Maps edge padding to R.nn.pad with constant mode.  Interior padding
        (dilation) is not supported in the first version.
        """
        from tflite.StablehloPadOptions import StablehloPadOptions

        input_tensors = self.get_input_tensors(op)
        # operand + padding_value
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloPadOptions)
        edge_low = [int(d) for d in opts.EdgePaddingLowAsNumpy()]
        edge_high = [int(d) for d in opts.EdgePaddingHighAsNumpy()]
        interior = [int(d) for d in opts.InteriorPaddingAsNumpy()]

        if any(d != 0 for d in interior):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_PAD with interior (dilation) padding is not supported"
            )
        if any(d < 0 for d in edge_low) or any(d < 0 for d in edge_high):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_PAD with negative edge padding (crop) is not supported"
            )

        operand = self.get_tensor_expr(input_tensors[0])

        # R.nn.pad only supports a static Python float pad_value.
        pad_value_tensor = input_tensors[1]
        if not self.has_expr(pad_value_tensor.tensor_idx):
            pad_val = float(self.get_tensor_value(pad_value_tensor))
        else:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_PAD with dynamic padding value is not supported"
            )

        # R.nn.pad with flat pad_width: [lo0, hi0, lo1, hi1, ...]
        pad_width = []
        for lo, hi in zip(edge_low, edge_high):
            pad_width.extend([lo, hi])

        return self.bb.normalize(relax.op.nn.pad(operand, pad_width=pad_width, pad_value=pad_val))

    def _convert_stablehlo_dynamic_slice(self, op):
        """Convert STABLEHLO_DYNAMIC_SLICE to Relax (dynamic_strided_slice).

        Start indices are assumed to be constant (non-dynamic) values stored
        in the flatbuffer.  Truly dynamic (runtime) start indices require
        Relax arithmetic to compute begin/end from scalar inputs and are not
        yet supported.
        """
        from tflite.StablehloDynamicSliceOptions import StablehloDynamicSliceOptions

        input_tensors = self.get_input_tensors(op)
        # operand + N start-index scalars
        assert len(input_tensors) >= 2
        ndim = len(input_tensors) - 1
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloDynamicSliceOptions)
        slice_sizes = [int(d) for d in opts.SliceSizesAsNumpy()]
        assert len(slice_sizes) == ndim

        operand = self.get_tensor_expr(input_tensors[0])

        # Build constant 1D tensors for begin, end, strides
        # (assumes start values are constant in the flatbuffer)
        # TODO: support dynamic start indices via Relax arithmetic
        if any(self.has_expr(t.tensor_idx) for t in input_tensors[1:]):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_DYNAMIC_SLICE with dynamic start indices is not supported"
            )
        start_vals = [int(self.get_tensor_value(t)) for t in input_tensors[1:]]
        operand_shape = [int(d) for d in self.get_tensor_shape(input_tensors[0])]
        for start, size, dim in zip(start_vals, slice_sizes, operand_shape):
            if start < 0 or start + size > dim:
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_DYNAMIC_SLICE with out-of-bounds start indices is not supported"
                )
        end_vals = [s + sz for s, sz in zip(start_vals, slice_sizes)]
        stride_vals = [1] * ndim

        def _const_1d(values, dtype="int64"):
            arr = np.array(values, dtype=dtype)
            return self.bb.normalize(relax.const(arr, dtype=dtype))

        begin = _const_1d(start_vals)
        end = _const_1d(end_vals)
        strides = _const_1d(stride_vals)

        return self.bb.normalize(relax.op.dynamic_strided_slice(operand, begin, end, strides))

    def _convert_stablehlo_dynamic_update_slice(self, op):
        """Convert STABLEHLO_DYNAMIC_UPDATE_SLICE to Relax.

        Lowers to ``relax.op.scatter_nd``. Constant start indices build the index
        grid at compile time; runtime (dynamic) start indices build it in-graph
        with ``arange`` + broadcast, clamping each start to
        ``[0, operand_dim - update_dim]`` per StableHLO semantics.
        """
        input_tensors = self.get_input_tensors(op)
        # operand + update + N start-index scalars
        assert len(input_tensors) >= 3, "input tensors length should be >= 3"
        assert len(self.get_output_tensors(op)) == 1

        operand_tensor = input_tensors[0]
        update_tensor = input_tensors[1]
        start_tensors = input_tensors[2:]

        op_name = "STABLEHLO_DYNAMIC_UPDATE_SLICE"
        operand_shape = self._get_static_tensor_shape(operand_tensor, op_name)
        update_shape = self._get_static_tensor_shape(update_tensor, op_name)
        rank = len(operand_shape)
        if len(update_shape) != rank or len(start_tensors) != rank:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_DYNAMIC_UPDATE_SLICE requires operand, update, "
                "and start-index ranks to match"
            )
        for dim, size in zip(operand_shape, update_shape):
            if size > dim:
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_DYNAMIC_UPDATE_SLICE update shape must be smaller than "
                    "or equal to operand shape for all dimensions"
                )

        operand = self.get_tensor_expr(operand_tensor)
        update = self.get_tensor_expr(update_tensor)

        if any(self.has_expr(t.tensor_idx) for t in start_tensors):
            indices = self._build_dynamic_update_slice_indices(
                start_tensors, operand_shape, update_shape, rank
            )
            return self.bb.normalize(relax.op.scatter_nd(operand, indices, update, "update"))

        start_vals = [int(np.asarray(self.get_tensor_value(t)).item()) for t in start_tensors]
        for start, size, dim in zip(start_vals, update_shape, operand_shape):
            if start < 0 or start + size > dim:
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_DYNAMIC_UPDATE_SLICE with out-of-bounds update "
                    "indices is not supported"
                )

        update_indices = np.indices(update_shape, dtype=np.int64)
        for axis, start in enumerate(start_vals):
            update_indices[axis] += start
        update_indices = np.moveaxis(update_indices, 0, -1)

        indices = self.bb.normalize(relax.const(update_indices, dtype="int64"))
        return self.bb.normalize(relax.op.scatter_nd(operand, indices, update, "update"))

    def _build_dynamic_update_slice_indices(self, start_tensors, operand_shape, update_shape, rank):
        """Build the scatter_nd index grid for runtime DYNAMIC_UPDATE_SLICE starts.

        Returns an int64 tensor of shape ``(*update_shape, rank)`` where axis ``a``
        holds ``arange(update_shape[a]) + clamp(start[a], 0, operand_dim - update_dim)``,
        broadcast over the other axes (StableHLO clamps out-of-range starts).
        """
        axis_indices = []
        for axis in range(rank):
            start_expr = self.bb.normalize(
                relax.op.astype(self.get_tensor_expr(start_tensors[axis]), "int64")
            )
            max_start = operand_shape[axis] - update_shape[axis]
            start_expr = relax.op.maximum(start_expr, relax.const(0, "int64"))
            start_expr = relax.op.minimum(start_expr, relax.const(max_start, "int64"))

            base = relax.op.arange(0, update_shape[axis], 1, "int64")
            idx = relax.op.add(base, start_expr)

            broadcast_shape = [1] * rank
            broadcast_shape[axis] = update_shape[axis]
            idx = self.bb.normalize(relax.op.reshape(idx, broadcast_shape))
            idx = self.bb.normalize(relax.op.broadcast_to(idx, update_shape))
            idx = self.bb.normalize(relax.op.expand_dims(idx, axis=-1))
            axis_indices.append(idx)

        return self.bb.normalize(relax.op.concat(axis_indices, axis=-1))

    def _convert_stablehlo_dot_general(self, op):
        """Convert the canonical 2D STABLEHLO_DOT_GENERAL subset to Relax matmul."""
        from tflite.StablehloDotGeneralOptions import StablehloDotGeneralOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloDotGeneralOptions)
        lhs_batch_dims = self._get_stablehlo_i64_vector(opts.LhsBatchingDimensionsAsNumpy(), [])
        rhs_batch_dims = self._get_stablehlo_i64_vector(opts.RhsBatchingDimensionsAsNumpy(), [])
        lhs_contract_dims = self._get_stablehlo_i64_vector(
            opts.LhsContractingDimensionsAsNumpy(), []
        )
        rhs_contract_dims = self._get_stablehlo_i64_vector(
            opts.RhsContractingDimensionsAsNumpy(), []
        )

        lhs_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_DOT_GENERAL")
        rhs_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_DOT_GENERAL")
        if len(lhs_shape) != 2 or len(rhs_shape) != 2:
            raise tvm.error.OpNotImplemented("STABLEHLO_DOT_GENERAL only supports 2D matmul")
        if lhs_batch_dims or rhs_batch_dims:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_DOT_GENERAL with batching dimensions is not supported"
            )
        if lhs_contract_dims != [1] or rhs_contract_dims != [0]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_DOT_GENERAL only supports canonical contracting dimensions"
            )

        lhs = self.get_tensor_expr(input_tensors[0])
        rhs = self.get_tensor_expr(input_tensors[1])
        return self.bb.normalize(relax.op.matmul(lhs, rhs))

    def _convert_stablehlo_convolution(self, op):
        """Convert the canonical 2D NHWC/HWIO STABLEHLO_CONVOLUTION subset."""
        from tflite.StablehloConvolutionOptions import StablehloConvolutionOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(self.get_output_tensors(op)) == 1

        opts = self._get_stablehlo_options(op, StablehloConvolutionOptions)
        input_spatial_dims = self._get_stablehlo_i64_vector(
            opts.InputSpatialDimensionsAsNumpy(), []
        )
        kernel_spatial_dims = self._get_stablehlo_i64_vector(
            opts.KernelSpatialDimensionsAsNumpy(), []
        )
        output_spatial_dims = self._get_stablehlo_i64_vector(
            opts.OutputSpatialDimensionsAsNumpy(), []
        )
        if input_spatial_dims != [1, 2]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION only supports NHWC input layout"
            )
        if kernel_spatial_dims != [0, 1]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION only supports HWIO kernel layout"
            )
        if output_spatial_dims != [1, 2]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION only supports NHWC output layout"
            )

        if (
            int(opts.InputBatchDimension()) != 0
            or int(opts.InputFeatureDimension()) != 3
            or int(opts.KernelInputFeatureDimension()) != 2
            or int(opts.KernelOutputFeatureDimension()) != 3
            or int(opts.OutputBatchDimension()) != 0
            or int(opts.OutputFeatureDimension()) != 3
        ):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION only supports canonical NHWC/HWIO dimension numbers"
            )
        if int(opts.BatchGroupCount()) != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION with batch_group_count > 1 is not supported"
            )
        if int(opts.FeatureGroupCount()) != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION with feature_group_count > 1 is not supported"
            )

        data_shape = self._get_static_tensor_shape(input_tensors[0], "STABLEHLO_CONVOLUTION")
        kernel_shape = self._get_static_tensor_shape(input_tensors[1], "STABLEHLO_CONVOLUTION")
        if len(data_shape) != 4 or len(kernel_shape) != 4:
            raise tvm.error.OpNotImplemented("STABLEHLO_CONVOLUTION only supports 2D convolution")
        if data_shape[3] != kernel_shape[2]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION input channels must match kernel input channels"
            )

        window_strides = self._get_stablehlo_i64_vector(opts.WindowStridesAsNumpy(), [1, 1])
        padding = self._get_stablehlo_i64_vector(opts.PaddingAsNumpy(), [0, 0, 0, 0])
        lhs_dilation = self._get_stablehlo_i64_vector(opts.LhsDilationAsNumpy(), [1, 1])
        rhs_dilation = self._get_stablehlo_i64_vector(opts.RhsDilationAsNumpy(), [1, 1])
        window_reversal = opts.WindowReversalAsNumpy()
        window_reversal = (
            [False, False] if window_reversal is None else [bool(v) for v in window_reversal]
        )

        if len(window_strides) != 2 or len(rhs_dilation) != 2:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION only supports two spatial dimensions"
            )
        if lhs_dilation != [1, 1]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION with lhs dilation is not supported"
            )
        if any(window_reversal):
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION with window reversal is not supported"
            )
        if len(padding) != 4:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_CONVOLUTION only supports 2D low/high padding"
            )

        # StableHLO stores padding as [low_h, high_h, low_w, high_w].
        relax_padding = [padding[0], padding[2], padding[1], padding[3]]
        data = self.get_tensor_expr(input_tensors[0])
        kernel = self.get_tensor_expr(input_tensors[1])
        self._ensure_stablehlo_float_dtype(data, "STABLEHLO_CONVOLUTION")
        self._ensure_stablehlo_float_dtype(kernel, "STABLEHLO_CONVOLUTION")
        return self.bb.normalize(
            relax.op.nn.conv2d(
                data,
                kernel,
                strides=window_strides,
                padding=relax_padding,
                dilation=rhs_dilation,
                data_layout="NHWC",
                kernel_layout="HWIO",
            )
        )

    def _convert_stablehlo_gather(self, op):
        """Convert STABLEHLO_GATHER to Relax (take-equivalent subset only).

        Only handles gather patterns equivalent to R.take along a single axis.
        Multi-dimensional gathers, index_vector_dim != rank(indices)-1, and
        non-trivial slice_sizes raise OpNotImplemented.
        """
        from tflite.StablehloGatherOptions import StablehloGatherOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1

        opts = self._get_stablehlo_options(op, StablehloGatherOptions)
        offset_dims = [int(d) for d in opts.OffsetDimsAsNumpy()]
        collapsed_slice_dims = [int(d) for d in opts.CollapsedSliceDimsAsNumpy()]
        start_index_map = [int(d) for d in opts.StartIndexMapAsNumpy()]
        slice_sizes = [int(d) for d in opts.SliceSizesAsNumpy()]
        index_vector_dim = int(opts.IndexVectorDim())

        data_tensor, indices_tensor = input_tensors
        data_shape = [int(d) for d in self.get_tensor_shape(data_tensor)]
        indices_shape = [int(d) for d in self.get_tensor_shape(indices_tensor)]
        output_shape = [int(d) for d in self.get_tensor_shape(output_tensors[0])]

        if len(start_index_map) != 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_GATHER only supports one start_index_map entry"
            )
        axis = start_index_map[0]
        if axis < 0 or axis >= len(data_shape):
            raise tvm.error.OpNotImplemented(f"Unsupported STABLEHLO_GATHER axis: {axis}")
        if collapsed_slice_dims != [axis]:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_GATHER only supports collapsed_slice_dims matching the gather axis"
            )
        if len(slice_sizes) != len(data_shape):
            raise tvm.error.OpNotImplemented("STABLEHLO_GATHER slice_sizes must match operand rank")
        for i, (size, dim) in enumerate(zip(slice_sizes, data_shape)):
            expected = 1 if i == axis else dim
            if size != expected:
                raise tvm.error.OpNotImplemented(
                    "STABLEHLO_GATHER only supports take-equivalent slice_sizes"
                )
        if index_vector_dim != len(indices_shape) - 1:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_GATHER only supports trailing index_vector_dim"
            )
        if not indices_shape or indices_shape[index_vector_dim] != 1:
            raise tvm.error.OpNotImplemented("STABLEHLO_GATHER only supports index vector size 1")

        indices_batch_shape = indices_shape[:index_vector_dim]
        expected_offset_dims = list(range(axis)) + list(
            range(axis + len(indices_batch_shape), len(data_shape) + len(indices_batch_shape) - 1)
        )
        if offset_dims != expected_offset_dims:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_GATHER offset_dims do not match Relax take output layout"
            )

        expected_output_shape = data_shape[:axis] + indices_batch_shape + data_shape[axis + 1 :]
        if output_shape != expected_output_shape:
            raise tvm.error.OpNotImplemented(
                "STABLEHLO_GATHER output shape does not match Relax take semantics"
            )

        data = self.get_tensor_expr(data_tensor)
        indices = self.get_tensor_expr(indices_tensor)
        indices = self.bb.normalize(relax.op.reshape(indices, indices_batch_shape))
        return self.bb.normalize(relax.op.take(data, indices, axis=axis, mode="fast"))

    def convert_elu(self, op):
        """Convert TFLite ELU"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)
        exp_type = self.get_tensor_type_str(input_tensor.tensor.Type())
        out = relax.const(-1.0, exp_type) * relax.op.nn.relu(
            relax.const(1.0, exp_type) - relax.op.exp(in_expr)
        ) + relax.op.nn.relu(in_expr)

        return out

    def convert_gelu(self, op):
        """Convert TFLite GELU"""
        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "The TFLite to Relax converter does not support quantized GELU operator yet."
            )

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)
        in_type = self.get_tensor_type_str(input_tensor.tensor.Type())

        return in_expr * (
            relax.const(0.5, dtype=in_type)
            + relax.op.erf(in_expr * relax.const(0.5**0.5, dtype=in_type))
            * relax.const(0.5, dtype=in_type)
        )

    def convert_square(self, op):
        """Convert TFLite SQUARE"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        exp_type = self.get_tensor_type_str(output_tensor.tensor.Type())
        out = relax.op.power(in_expr, relax.const(2, exp_type))

        return out

    def _convert_elemwise(self, op, relax_op, comparison_op=False):
        """Generic method to Convert TFLite elemwise"""

        from tflite.AddOptions import AddOptions
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.DivOptions import DivOptions
        from tflite.MulOptions import MulOptions
        from tflite.SubOptions import SubOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        lhs_tensor = input_tensors[0]
        rhs_tensor = input_tensors[1]
        lhs_expr = self.get_tensor_expr(lhs_tensor)
        rhs_expr = self.get_tensor_expr(rhs_tensor)
        input_is_quantized = lhs_tensor.qnn_params is not None or rhs_tensor.qnn_params is not None

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        if input_is_quantized:
            if lhs_tensor.qnn_params:
                lhs_expr = self.dequantize(lhs_expr, lhs_tensor)
            if rhs_tensor.qnn_params:
                rhs_expr = self.dequantize(rhs_expr, rhs_tensor)

        out = relax_op(lhs_expr, rhs_expr)

        # Options (fused_activation_function)
        options = None
        if op.BuiltinOptionsType() == BuiltinOptions.AddOptions:
            options = AddOptions()
        elif op.BuiltinOptionsType() == BuiltinOptions.SubOptions:
            options = SubOptions()
        elif op.BuiltinOptionsType() == BuiltinOptions.MulOptions:
            options = MulOptions()
        elif op.BuiltinOptionsType() == BuiltinOptions.DivOptions:
            options = DivOptions()

        if options is not None:
            op_options = op.BuiltinOptions()
            options.Init(op_options.Bytes, op_options.Pos)
            fused_activation_fn = options.FusedActivationFunction()

            out = self.convert_fused_activation_function(out, fused_activation_fn)

        if input_is_quantized and not comparison_op:
            if not output_tensor.qnn_params:
                raise tvm.error.OpAttributeInvalid(
                    "Quantized TFLite elemwise operator output must have quantization parameters"
                )
            out = self.quantize(out, output_tensor)
        return out

    def convert_add_n(self, op):
        """Convert TFLite ADD_N"""
        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"

        input_tensors = self.get_input_tensors(op)
        assert not input_tensors[0].qnn_params, "TFLite does not support quantized ADD_N."
        lhs_expr = self.get_tensor_expr(input_tensors[0])
        for rhs_tensor in input_tensors[1:]:
            assert not rhs_tensor.qnn_params, "TFLite does not support quantized ADD_N"
            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"""
        # Check if the input tensor is quantized, call QNN op
        # (https://github.com/tensorflow/tflite-micro/blob/bc35c3ed9c7ab93b3a13b46fce936f854bcfce2c
        # /tensorflow/lite/micro/kernels/squared_difference.cc#L157)
        if self.is_quantized(op):
            input_tensors = self.get_input_tensors(op)
            output_tensors = self.get_output_tensors(op)
            lhs_expr = self.get_tensor_expr(input_tensors[0])
            rhs_expr = self.get_tensor_expr(input_tensors[1])
            assert len(input_tensors) == 2, "input tensors length should be 2"
            assert len(output_tensors) == 1, "output tensors length should be 1"
            lhs_expr_f32 = self.dequantize(lhs_expr, input_tensors[0])
            rhs_expr_f32 = self.dequantize(rhs_expr, input_tensors[1])
            out_f32 = relax.op.subtract(lhs_expr_f32, rhs_expr_f32)
            return self.quantize(out_f32 * out_f32, output_tensors[0])

        difference = self._convert_elemwise(op, _op.subtract)
        # _convert_elemwise has guaranteed only have one output tensor
        exp_type = self.get_tensor_type_str(self.get_output_tensors(op)[0].tensor.Type())
        out = relax.op.power(difference, relax.const(2, exp_type))
        return out

    def _convert_logical_binary(self, relax_op, op):
        """Generic method to convert logical binary ops"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        lhs_tensor = input_tensors[0]
        lhs_expr = self.get_tensor_expr(lhs_tensor)
        rhs_tensor = input_tensors[1]
        rhs_expr = self.get_tensor_expr(rhs_tensor)
        out = relax_op(lhs_expr, rhs_expr)

        return out

    def convert_logical_not(self, op):
        """Convert tflite LOGICAL_NOT"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        data = self.get_expr(input_tensors[0].tensor_idx)
        out = relax.op.logical_not(data)

        return out

    def convert_gather(self, op):
        """Method to Convert TFLite GATHER operator"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.GatherOptions import GatherOptions
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        data = self.get_tensor_expr(input_tensors[0])
        indices = input_tensors[1]
        indices_type = indices.tensor.Type()
        assert indices_type in (TensorType.INT32, TensorType.INT64)

        assert op.BuiltinOptionsType() == BuiltinOptions.GatherOptions
        op_options = op.BuiltinOptions()
        gather_options = GatherOptions()
        gather_options.Init(op_options.Bytes, op_options.Pos)
        axis = gather_options.Axis()

        # Check the indices are with in bounds.
        data_shape = to_int_list(self.get_tensor_shape(input_tensors[0]))
        data_dim = len(data_shape)

        axis = data_dim + axis if axis < 0 else axis
        assert axis >= 0, "Axis out of bounds"
        assert axis < data_dim, "Axis out of bounds"

        if self.has_expr(indices.tensor_idx):
            indices_expr = relax.op.astype(self.get_expr(indices.tensor_idx), "int32")
        else:
            indices_val = self.get_tensor_value(indices)
            indices_expr = self.exp_tab.new_const(
                indices_val,
                dtype=self.get_tensor_type_str(indices_type),
                source_name=indices.tensor.Name(),
            )
            indices_shape = list(indices_val.shape)
            indices_len = len(indices_shape)

            out_shape = data_shape[:axis] + indices_shape[:] + data_shape[axis + 1 :]

            loopover = [range(s) for s in out_shape]
            for idx in list(itertools.product(*loopover)):
                real_indices = (
                    list(idx[:axis])
                    + [indices_val[idx[axis : axis + indices_len]]]
                    + list(idx[axis + indices_len :])
                )
                if np.any(np.subtract(data_shape, real_indices) < 0):
                    raise ValueError("TFLite out of bound indices are not supported.")

        # Use mode 'fast' since indices are already checked within bounds.
        out = relax.op.take(data, indices_expr, axis=axis, mode="fast")
        return out

    def convert_gather_nd(self, op):
        """Method to Convert TFLite GATHER_ND operator"""

        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        for t in input_tensors:
            assert not t.qnn_params, "Quantized input is not expected."

        data = self.get_tensor_expr(input_tensors[0])
        indices = self.get_tensor_expr(input_tensors[1])

        indices_type = input_tensors[1].tensor.Type()
        assert indices_type in (TensorType.INT32, TensorType.INT64)

        indices_dims = len(self._infer_shape(indices))
        indices_t = relax.op.permute_dims(indices, axes=[-1] + list(range(indices_dims - 1)))
        if indices_type == TensorType.INT32:
            # Relax gather_nd requires int64 indices.
            indices_t = relax.op.astype(indices_t, "int64")

        out = relax.op.gather_nd(data, indices_t)
        return out

    def convert_strided_slice(self, op):
        """Method to Convert TFLite STRIDED_SLICE operator.
        NOTE: Eventhough tensorflow supports begin_mask, end_mask, ellipsis_mask, new_axis_mask
        and shrink_axis_mask, tflite doesn't support these and expect these values to be zero.
        But in future, they may open up the mask implementation, so kept the implementation
        same as tensorflow.

        This op extracts a slice of size (end - begin) / stride from the given input tensor.
        Starting at the location specified by begin the slice continues by adding stride to the
        index until all dimensions are not less than end. Note that a stride can be negative,
        which causes a reverse slice.

        For slice input[val0, val1, ..., valn], begin/end/strides will be vectors of length n.

        In each mask field(begin_mask, end_mask, ellipsis_mask, new_axis_mask, shrink_axis_mask)
        the ith bit will correspond to the ith val.

        If the ith bit of begin_mask is set, begin[i] is ignored and the fullest possible range
        in that dimension is used instead.

        If the ith bit of ellipsis_mask is set, as many unspecified dimensions as needed will be
        inserted between other dimensions. Only one non-zero bit is allowed in ellipsis_mask.

        If the ith bit of new_axis_mask is set, then begin, end, and stride are ignored and a
        new length 1 dimension is added at this point in the output tensor.

        If the ith bit of shrink_axis_mask is set, it implies that the ith specification shrinks
        the dimensionality by 1, taking on the value at index begin[i]. end[i] and strides[i]
        are ignored in this case.
        begin and end are zero-indexed. strides entries must be non-zero.

        TVM Relax implementation of doesn't support mask, so the mask values are processed in
        this function and begin/end/strides are updated accordingly. If any mask is present, and
        since tvm doesn't support mask computation directly, the output need a final reshape.
        """

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.StridedSliceOptions import StridedSliceOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 4, "input tensors length should be 4"

        data_expr = self.get_expr(input_tensors[0].tensor_idx)

        begin = list(self.get_tensor_value(input_tensors[1]))
        end = list(self.get_tensor_value(input_tensors[2]))
        stride = list(self.get_tensor_value(input_tensors[3]))

        assert op.BuiltinOptionsType() == BuiltinOptions.StridedSliceOptions
        op_options = op.BuiltinOptions()
        options = StridedSliceOptions()
        options.Init(op_options.Bytes, op_options.Pos)
        begin_mask = options.BeginMask()
        end_mask = options.EndMask()
        ellipsis_mask = options.EllipsisMask()
        new_axis_mask = options.NewAxisMask()
        shrink_axis_mask = options.ShrinkAxisMask()

        data_shape = to_int_list(self.get_tensor_shape(input_tensors[0]))
        data_dim = len(data_shape)
        stride_dim = len(stride)

        def _transform_mask(stride_dim, ellipsis_mask):
            """Handle mask inputs to create new begin, end, stride and output shape"""
            m_begin = [0] * data_dim
            m_end = [0] * data_dim
            m_stride = [0] * data_dim
            fshape_indices = []
            # Count new axis after ellipsis_mask, consider while applying ellipsis_mask.
            ellipsis_seen = False
            new_axes_after_ellipsis = 0
            for i in range(stride_dim):
                mask = 1 << i
                if ellipsis_seen and (mask & new_axis_mask) != 0:
                    new_axes_after_ellipsis += 1
                if (mask & ellipsis_mask) != 0:
                    ellipsis_seen = True
            if not ellipsis_seen:
                # Used later for extending the stride attributes in the below loop.
                ellipsis_mask |= 1 << stride_dim
                stride_dim += 1
            final_index = 0
            for index in range(stride_dim):
                mask = 1 << index
                if mask & ellipsis_mask:
                    # Identify the end index for applying ellipsis_mask
                    to_index = min(
                        ((data_dim - (stride_dim - index)) + 1 + new_axes_after_ellipsis), data_dim
                    )
                    for i in range(final_index, to_index):
                        m_begin[final_index] = 0
                        m_end[final_index] = data_shape[final_index]
                        m_stride[final_index] = 1
                        fshape_indices.append(final_index)
                        final_index += 1
                elif mask & new_axis_mask:
                    fshape_indices.append(-1)
                elif not mask & new_axis_mask:
                    if final_index == len(m_begin):
                        break
                    if mask & begin_mask:
                        m_begin[final_index] = data_shape[final_index] if stride[index] < 0 else 0
                    elif begin[index]:
                        m_begin[final_index] = begin[index]
                    if mask & end_mask:
                        if stride[index] < 0:
                            # Relax negative-step slicing excludes the end index, so an
                            # unspecified lower bound needs one extra step past index 0.
                            m_end[final_index] = -data_shape[final_index] - 1
                        else:
                            m_end[final_index] = data_shape[final_index]
                    elif end[index]:
                        m_end[final_index] = end[index]
                    m_stride[final_index] = stride[index]
                    if mask & shrink_axis_mask:
                        # Tensorflow make axis with shrink_axis_mask as dimension 1
                        m_begin[final_index] = (
                            data_shape[final_index] + begin[index]
                            if begin[index] < 0
                            else begin[index]
                        )
                        m_end[final_index] = m_begin[final_index] + 1
                        m_stride[final_index] = 1
                        fshape_indices.append(-2)
                    else:
                        fshape_indices.append(final_index)

                    final_index += 1
            return m_begin, m_end, m_stride, fshape_indices

        fshape_indices = None
        if begin_mask or end_mask or ellipsis_mask or new_axis_mask or shrink_axis_mask:
            begin, end, stride, fshape_indices = _transform_mask(stride_dim, ellipsis_mask)

        begin = [int(i) for i in begin]
        end = [int(i) for i in end]
        stride = [int(i) for i in stride]
        axes = list(range(len(begin)))
        out = relax.op.strided_slice(data_expr, axes=axes, begin=begin, end=end, strides=stride)
        out_shape = self.bb.normalize(out).ty.shape
        if not fshape_indices:
            fshape_indices = range(len(out_shape))

        # Create final output shape.
        final_output = []
        final_len = len(fshape_indices)
        for gather_index in fshape_indices:
            if gather_index == -1:
                final_output.append(1)
                final_len += 1
            elif gather_index == -2:
                final_len -= 1
            else:
                final_output.append(out_shape[gather_index])

        if final_len == 0:
            return relax.op.squeeze(out, axis=tuple(range(len(fshape_indices))))

        if not final_output:
            return out
        return relax.op.reshape(out, shape=tuple(final_output))

    def convert_zeros_like(self, op):
        """Convert TFLite ZEROS LIKE"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)
        out = relax.op.zeros_like(in_expr)

        return out

    def convert_fill(self, op):
        """Convert TFLite FILL"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        dims_tensor = input_tensors[0]
        in_value_expr = self.get_expr(input_tensors[1].tensor_idx)

        out_shape, _ = self._get_shape_expr_from_tensor(dims_tensor, "fill_dim")
        out = relax.op.full(out_shape, in_value_expr)

        return out

    def _get_random_options(self, op):
        """Return the seed pair for random TFLite operators.

        The runtime imports seeded TFLite random ops with stateless semantics, so identical
        non-zero seed pairs produce identical results on every invocation. The seed pair
        (0, 0) is forwarded as the TF non-deterministic case.
        """
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.RandomOptions import RandomOptions

        if op.BuiltinOptionsType():
            assert op.BuiltinOptionsType() == BuiltinOptions.RandomOptions
            random_options = RandomOptions()
            op_options = op.BuiltinOptions()
            random_options.Init(op_options.Bytes, op_options.Pos)
            return int(random_options.Seed()), int(random_options.Seed2())
        return 0, 0

    def _check_random_output_dtype(self, op_name, output_dtype, supported_dtypes):
        if output_dtype not in supported_dtypes:
            supported = ", ".join(supported_dtypes)
            raise tvm.error.OpNotImplemented(
                f"The TFLite {op_name} converter currently supports output dtype(s) "
                f"{supported} only, but got {output_dtype}."
            )

    def convert_random_uniform(self, op):
        """Convert TFLite RANDOM_UNIFORM using stateless seeded RNG semantics."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
        self._check_random_output_dtype("RANDOM_UNIFORM", output_dtype, ["float32"])

        out_shape, _ = self._get_shape_expr_from_tensor(input_tensors[0], "random_uniform_dim")
        seed, seed2 = self._get_random_options(op)
        return relax.op.call_dps_packed(
            "tvm.contrib.random.uniform",
            (seed, seed2, 0.0, 1.0),
            out_ty=relax.TensorType(out_shape, output_dtype),
        )

    def convert_random_standard_normal(self, op):
        """Convert TFLite RANDOM_STANDARD_NORMAL using stateless seeded RNG semantics."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
        self._check_random_output_dtype("RANDOM_STANDARD_NORMAL", output_dtype, ["float32"])

        out_shape, _ = self._get_shape_expr_from_tensor(
            input_tensors[0], "random_standard_normal_dim"
        )
        seed, seed2 = self._get_random_options(op)
        return relax.op.call_dps_packed(
            "tvm.contrib.random.normal",
            (seed, seed2, 0.0, 1.0),
            out_ty=relax.TensorType(out_shape, output_dtype),
        )

    def convert_multinomial(self, op):
        """Convert TFLite MULTINOMIAL using stateless seeded RNG semantics."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        logits_tensor, num_samples_tensor = input_tensors
        logits_expr = self.get_tensor_expr(logits_tensor)
        batch_size = self.get_tensor_shape(logits_tensor)[0]
        if self.has_expr(num_samples_tensor.tensor_idx):
            scalar_expr = self.get_expr(num_samples_tensor.tensor_idx)
            scalar_dtype = self.get_tensor_type_str(num_samples_tensor.tensor.Type())
            scalar_expr = self.bb.match_cast(scalar_expr, relax.TensorType([], scalar_dtype))
            scalar_expr = self.bb.normalize(relax.op.astype(scalar_expr, "int64"))
            scalar_expr = self.bb.normalize(relax.op.reshape(scalar_expr, [1]))
            shape_dataflow_var = self.bb.emit(relax.op.tensor_to_shape(scalar_expr))
            num_samples = tirx.Var("multinomial_num_samples", "int64")
            self.bb.match_cast(shape_dataflow_var, relax.ShapeType([num_samples]))
        else:
            value = self.get_tensor_value(num_samples_tensor)
            assert value.size == 1, (
                "TFLite MULTINOMIAL num_samples must be a scalar tensor, "
                f"but got {value.size} values"
            )
            num_samples = int(value.item())
        output_batch = batch_size * num_samples

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())
        self._check_random_output_dtype("MULTINOMIAL", output_dtype, ["int32", "int64"])

        seed, seed2 = self._get_random_options(op)
        uniform_sample = relax.op.call_dps_packed(
            "tvm.contrib.random.uniform",
            (seed, seed2, 0.0, 1.0),
            out_ty=relax.TensorType([output_batch, 1], "float32"),
        )
        sample_indices = relax.op.reshape(
            relax.op.broadcast_to(
                relax.op.expand_dims(relax.op.arange(batch_size, dtype="int64"), axis=[1]),
                relax.ShapeExpr([batch_size, num_samples]),
            ),
            relax.ShapeExpr([output_batch, 1]),
        )
        sampled = relax.op.multinomial_from_uniform(
            relax.op.nn.softmax(logits_expr, axis=-1),
            uniform_sample,
            sample_indices,
            dtype=output_dtype,
        )
        return relax.op.reshape(sampled, relax.ShapeExpr([batch_size, num_samples]))

    def _convert_reduce(self, relax_op, op):
        """Generic method to Convert TFLite REDUCE operators"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ReducerOptions import ReducerOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        # input_tensor
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        # axis
        axis_value = self.get_tensor_value(input_tensors[1])
        axis = tuple(axis_value) if len(axis_value.shape) > 0 else tuple((axis_value.item(),))

        # Options - keep_dims (bool)
        # In case Options are not present, set keep_dims to False(default)
        if op.BuiltinOptionsType():
            assert op.BuiltinOptionsType() == BuiltinOptions.ReducerOptions
            reduce_options = ReducerOptions()
            op_options = op.BuiltinOptions()
            reduce_options.Init(op_options.Bytes, op_options.Pos)
            keep_dims = reduce_options.KeepDims()
        else:
            keep_dims = False

        if input_tensor.qnn_params:
            in_expr = self.dequantize(in_expr, input_tensor)

        out = relax_op(in_expr, axis, keep_dims)

        # Finally if the reduce is quantized. Quantize the output.
        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

        return out

    def _convert_reduce_bool(self, relax_op, op):
        """Convert TFLite REDUCE_ANY / REDUCE_ALL (bool-only ops).

        Relax max/min are undefined on bool, so cast through int8.
        """
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ReducerOptions import ReducerOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        axis_value = self.get_tensor_value(input_tensors[1])
        axis = tuple(axis_value) if len(axis_value.shape) > 0 else tuple((axis_value.item(),))

        if op.BuiltinOptionsType():
            assert op.BuiltinOptionsType() == BuiltinOptions.ReducerOptions
            reduce_options = ReducerOptions()
            op_options = op.BuiltinOptions()
            reduce_options.Init(op_options.Bytes, op_options.Pos)
            keep_dims = reduce_options.KeepDims()
        else:
            keep_dims = False

        in_expr = relax.op.astype(in_expr, "int8")
        out = relax_op(in_expr, axis, keep_dims)
        return relax.op.astype(out, "bool")

    def _convert_arg_min_max(self, op, relax_op):
        """Generic method converting TFLite arg_min_max"""

        from tflite.ArgMaxOptions import ArgMaxOptions
        from tflite.ArgMinOptions import ArgMinOptions
        from tflite.BuiltinOptions import BuiltinOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "two input tensor arguments expected"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "one output tensor expected"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)
        axis_tensor = input_tensors[1]
        # In Tensorflow, `axis` argument is a Tensor, not attribute. We
        # support the case where it inputs from a scalar constant.
        axis_value = self.get_tensor_value(axis_tensor)
        assert axis_value.size == 1
        axis_value = axis_value.item()

        if op.BuiltinOptionsType() == BuiltinOptions.ArgMinOptions:
            arg_min_max_options = ArgMinOptions()
        elif op.BuiltinOptionsType() == BuiltinOptions.ArgMaxOptions:
            arg_min_max_options = ArgMaxOptions()
        op_options = op.BuiltinOptions()
        arg_min_max_options.Init(op_options.Bytes, op_options.Pos)

        # set keepdims to True since tflite 1.13 removes all dims of size 1
        # WARNING: all other versions of tflite > 1.13 need keepdims=False
        out = relax_op(in_expr, axis=axis_value, keepdims=False)

        return out

    def convert_fully_connected(self, op):
        """Convert TFLite fully connected"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.FullyConnectedOptions import FullyConnectedOptions
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) in (2, 3), "input tensors length should be two or three"

        input_tensor = input_tensors[0]
        weight_tensor = input_tensors[1]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_tensor_type = output_tensor.tensor.Type()
        output_tensor_type_str = self.get_tensor_type_str(output_tensor_type)

        weight_tensor_shape = to_int_list(self.get_tensor_shape(weight_tensor))

        # Weight should have only 2 dimensions(TFLite convention)
        assert len(weight_tensor_shape) == 2, "Weight should be only 2-dim"

        # Input shape: [i_batch_size, ..., n_inputs]
        # Filter shape: [n_inputs, n_units]
        #
        # As we will transform Fully_Connected Input to MatMul
        # Weight require a transpose
        in_expr = self.get_tensor_expr(input_tensor)

        # TODO: Change the output shape calculation based on keep_dim option
        assert op.BuiltinOptionsType() == BuiltinOptions.FullyConnectedOptions
        op_options = op.BuiltinOptions()
        fully_connected_options = FullyConnectedOptions()
        fully_connected_options.Init(op_options.Bytes, op_options.Pos)
        fused_activation_fn = fully_connected_options.FusedActivationFunction()
        keep_num_dims = fully_connected_options.KeepNumDims()

        # weight tensor type should be INT8/UINT8 (quantization) or FLOAT32
        weight_tensor_type = weight_tensor.tensor.Type()
        assert weight_tensor_type in (
            TensorType.INT8,
            TensorType.UINT8,
            TensorType.FLOAT32,
        )

        weight_expr = self.get_tensor_expr(weight_tensor)
        weight_expr = relax.op.permute_dims(weight_expr, [1, 0])

        if input_tensor.qnn_params:
            # Dequantize input and weight (OC remapped from axis 0 to 1)
            in_f32 = self.dequantize(in_expr, input_tensor)
            weight_axis = weight_tensor.qnn_params["axis"]
            if weight_axis != 0:
                raise tvm.error.OpAttributeInvalid(
                    f"FC weight QuantizedDimension() must be 0 (output-channel "
                    f"axis in [OC,IC] layout), got {weight_axis}"
                )
            w_f32 = relax.op.dequantize(
                weight_expr,
                scale=weight_tensor.qnn_params["scale"],
                zero_point=weight_tensor.qnn_params["zero_point"],
                axis=1,
            )
            out = relax.op.matmul(in_f32, w_f32)
        else:
            out = relax.op.matmul(in_expr, weight_expr)

        # if we have bias
        if len(input_tensors) == 3:
            bias_tensor = input_tensors[2]
            if bias_tensor.tensor_idx != -1:
                bias_tensor_type = bias_tensor.tensor.Type()
                # bias tensor type should be INT32 (quantization) or FLOAT32
                assert bias_tensor_type in (
                    TensorType.INT32,
                    TensorType.INT64,
                    TensorType.FLOAT32,
                )
                bias_tensor_type_str = self.get_tensor_type_str(bias_tensor_type)
                if self.has_expr(bias_tensor.tensor_idx):
                    bias_expr = self.get_expr(bias_tensor.tensor_idx)
                else:
                    bias_expr = self.exp_tab.new_const(
                        self.get_tensor_value(bias_tensor),
                        dtype=bias_tensor_type_str,
                        source_name=bias_tensor.tensor.Name(),
                    )
                if bias_tensor.qnn_params:
                    bias_expr = self.dequantize(bias_expr, bias_tensor)
                elif input_tensor.qnn_params and bias_tensor_type in (
                    TensorType.INT32,
                    TensorType.INT64,
                ):
                    bias_scale = relax.op.multiply(
                        input_tensor.qnn_params["scale"],
                        weight_tensor.qnn_params["scale"],
                    )
                    bias_expr = relax.op.dequantize(
                        bias_expr,
                        scale=bias_scale,
                        zero_point=relax.const(0, "int32"),
                        axis=0,
                    )
                out = relax.op.add(out, bias_expr)

        # Finally if the dense is quantized. Quantize the output.
        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)

            # Call activation function
            output_scale_val = get_scalar_from_constant(output_tensor.qnn_params["scale"])
            output_zero_point_val = get_scalar_from_constant(output_tensor.qnn_params["zero_point"])
            out = self.convert_qnn_fused_activation_function(
                expr=out,
                fused_activation_fn=fused_activation_fn,
                scale=output_scale_val,
                zero_point=output_zero_point_val,
                dtype=output_tensor_type_str,
            )

        else:
            out = self.convert_fused_activation_function(out, fused_activation_fn)

        # Change the output shape calculation based on keep_dim option
        if keep_num_dims:
            input_shape = self._infer_shape(self.get_tensor_expr(input_tensor))
            output_shape = to_int_list(input_shape)[:-1] + [weight_tensor_shape[0]]
            out = relax.op.reshape(out, output_shape)

        return out

    def convert_squeeze(self, op):
        """Convert TFLite squeeze"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.SqueezeOptions import SqueezeOptions

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        assert len(output_tensors) == 1, "output tensors length should be 1"
        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx

        assert op.BuiltinOptionsType() == BuiltinOptions.SqueezeOptions
        op_options = op.BuiltinOptions()
        squeeze_options = SqueezeOptions()
        squeeze_options.Init(op_options.Bytes, op_options.Pos)
        squeeze_axis = squeeze_options.SqueezeDimsAsNumpy()

        in_expr = self.get_expr(input_tensor_idx)
        out = relax.op.squeeze(in_expr, axis=tuple(squeeze_axis))

        return out

    def convert_fused_activation_function(self, in_expr, fused_activation_fn):
        """Convert TFLite fused activation function"""

        from tflite.ActivationFunctionType import ActivationFunctionType

        if fused_activation_fn == ActivationFunctionType.NONE:
            return in_expr
        if fused_activation_fn == ActivationFunctionType.RELU6:
            return relax.op.clip(in_expr, min=0, max=6)
        if fused_activation_fn == ActivationFunctionType.RELU:
            return relax.op.nn.relu(in_expr)
        if fused_activation_fn == ActivationFunctionType.RELU_N1_TO_1:
            return relax.op.clip(in_expr, min=-1, max=1)
        if fused_activation_fn == ActivationFunctionType.TANH:
            return relax.op.tanh(in_expr)
        fused_activation_fn_str = self.activation_fn_type[fused_activation_fn]
        raise tvm.error.OpNotImplemented(
            f"Fused activation {fused_activation_fn_str} is not supported yet."
        )

    def convert_conv(self, op, conv_type):
        """convolution implementation."""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.Conv2DOptions import Conv2DOptions
        from tflite.DepthwiseConv2DOptions import DepthwiseConv2DOptions
        from tflite.Padding import Padding
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) >= 2, "input tensors length should be >= 2"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        weight_tensor = input_tensors[1]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_tensor_type = output_tensor.tensor.Type()
        output_tensor_type_str = self.get_tensor_type_str(output_tensor_type)

        is_depthwise_conv = False
        if conv_type == "conv2d":
            assert op.BuiltinOptionsType() == BuiltinOptions.Conv2DOptions
            op_options = op.BuiltinOptions()
            conv_options = Conv2DOptions()
            conv_options.Init(op_options.Bytes, op_options.Pos)
        elif conv_type == "depthwise":
            is_depthwise_conv = True
            assert op.BuiltinOptionsType() == BuiltinOptions.DepthwiseConv2DOptions
            op_options = op.BuiltinOptions()
            conv_options = DepthwiseConv2DOptions()
            conv_options.Init(op_options.Bytes, op_options.Pos)
            depth_multiplier = conv_options.DepthMultiplier()
        else:
            raise tvm.error.OpNotImplemented(
                f"Operator {conv_type} is not supported for frontend TFLite."
            )

        stride_h = conv_options.StrideH()
        stride_w = conv_options.StrideW()
        dilation_h = conv_options.DilationHFactor()
        dilation_w = conv_options.DilationWFactor()
        padding = conv_options.Padding()
        fused_activation_fn = conv_options.FusedActivationFunction()

        _, input_h, input_w, input_c = to_int_list(self.get_tensor_shape(input_tensor))

        if is_depthwise_conv:
            # TFLite depthwise convolution kernel layout is:
            # 1 KH KW C(input_c * depth_multiplier)
            _, kernel_h, kernel_w, in_channels = to_int_list(self.get_tensor_shape(weight_tensor))
            assert in_channels == input_c * depth_multiplier
        else:
            output_channels, kernel_h, kernel_w, in_channels = to_int_list(
                self.get_tensor_shape(weight_tensor)
            )

        dilated_kernel_h = dilation_h * (kernel_h - 1) + 1
        dilated_kernel_w = dilation_w * (kernel_w - 1) + 1

        params = {
            # "kernel_size": [kernel_h, kernel_w],
            "strides": [stride_h, stride_w],
            "dilation": [dilation_h, dilation_w],
            "padding": [0, 0],
            "data_layout": "NHWC",
        }

        if is_depthwise_conv:
            params["groups"] = int(input_c)
            # If number of input channels is 1, treat as normal
            # convolution.
            params["kernel_layout"] = "HWIO" if input_c == 1 else "HWOI"
        else:
            # params["channels"] = int(output_channels)
            params["kernel_layout"] = "HWIO"
            if input_c != in_channels:
                assert input_c % in_channels == 0, (
                    "Input channels is not divisible of kernel in_channels."
                )
                params["groups"] = int(input_c / in_channels)

        # weight tensor type should be INT8/UINT8 (quantization) or FLOAT32
        weight_tensor_type = weight_tensor.tensor.Type()
        assert weight_tensor_type in (
            TensorType.INT8,
            TensorType.UINT8,
            TensorType.FLOAT32,
        )
        weight_tensor_type_str = self.get_tensor_type_str(weight_tensor_type)

        in_expr = self.get_expr(input_tensor_idx)

        # TFLite converts float32 models to float16 models by introducing
        # a Dequantize op in every op that contains a float32 values.
        # (weights, biases, and constants etc. )
        # So conv op may have weight and bias as tensors instead of values.
        if self.has_expr(weight_tensor.tensor_idx):
            weight_expr = self.get_expr(weight_tensor.tensor_idx)
            if is_depthwise_conv:
                weight_expr = relax.op.reshape(
                    weight_expr, (kernel_h, kernel_w, input_c, depth_multiplier)
                )
            else:
                weight_expr = relax.op.permute_dims(weight_expr, axes=(1, 2, 3, 0))
        else:
            if self.is_prefetched(weight_tensor.tensor_idx):
                weight_value = self.get_prefetched_node(weight_tensor.tensor_idx)
            else:
                weight_value = self.get_tensor_value(weight_tensor)

            # TFLite kernel layout:
            # convolution:
            # OC KH KW IC, we require KH KW IC OC (HWIO)
            # depthwise convolution:
            # 1 KH KW C(input_c * depth_multiplier), we require
            # KH KW IC M (depth_multiplier) (HWOI)
            if is_depthwise_conv:
                weight_value = weight_value.reshape(kernel_h, kernel_w, input_c, depth_multiplier)
            else:
                weight_value = weight_value.transpose((1, 2, 3, 0))

            weight_expr = self.exp_tab.new_const(
                weight_value, dtype=weight_tensor_type_str, source_name=weight_tensor.tensor.Name()
            )

        if padding == Padding.VALID:
            pass
        elif padding == Padding.SAME:
            pad_top, pad_bottom = get_pad_value(input_h, dilated_kernel_h, stride_h)

            pad_left, pad_right = get_pad_value(input_w, dilated_kernel_w, stride_w)
            do_pad = not (pad_top == 0 and pad_bottom == 0 and pad_left == 0 and pad_right == 0)
            if do_pad:
                params["padding"] = [pad_top, pad_left, pad_bottom, pad_right]

        else:
            raise tvm.error.OpAttributeUnImplemented(
                f"Padding format {padding} is not supported for operator Conv."
            )

        if input_tensor.qnn_params:
            # Dequantize input activation
            in_f32 = self.dequantize(in_expr, input_tensor)
            # Dequantize weight with per-channel axis remap.
            # TFLite weight original layout: [OC, KH, KW, IC]
            # After transpose to HWIO: [KH, KW, IC, OC]
            # QuantizedDimension() == 0 (OC in original) → axis 3 in HWIO.
            weight_axis = weight_tensor.qnn_params["axis"]
            if is_depthwise_conv:
                if weight_axis != 0:
                    raise tvm.error.OpNotImplemented(
                        "Per-channel quantized depthwise convolution is not supported "
                        "because the channel axis changes semantics after the "
                        "[1,KH,KW,C*M] → [KH,KW,C,M] reshape."
                    )
            else:
                if weight_axis != 0:
                    raise tvm.error.OpAttributeInvalid(
                        f"Conv2D weight QuantizedDimension() must be 0 (output-channel "
                        f"axis in [OC,KH,KW,IC] layout), got {weight_axis}"
                    )
                weight_axis = 3
            w_f32 = relax.op.dequantize(
                weight_expr,
                scale=weight_tensor.qnn_params["scale"],
                zero_point=weight_tensor.qnn_params["zero_point"],
                axis=weight_axis,
            )
            # Float convolution
            out = relax.op.nn.conv2d(in_f32, w_f32, **params)
        else:
            out = relax.op.nn.conv2d(in_expr, weight_expr, **params)

        # if we have bias
        if len(input_tensors) == 3:
            bias_tensor = input_tensors[2]
            bias_tensor_type = bias_tensor.tensor.Type()
            # bias tensor type should be INT32 (int8 qnn) or INT64 (int16 qnn) or FLOAT32
            assert bias_tensor_type in (
                TensorType.INT32,
                TensorType.INT64,
                TensorType.FLOAT32,
            )
            bias_tensor_type_str = self.get_tensor_type_str(bias_tensor_type)
            if self.has_expr(bias_tensor.tensor_idx):
                bias_expr = self.get_expr(bias_tensor.tensor_idx)
            else:
                bias_expr = self.exp_tab.new_const(
                    self.get_tensor_value(bias_tensor),
                    dtype=bias_tensor_type_str,
                    source_name=bias_tensor.tensor.Name(),
                )
            # For quantized conv, INT32/INT64 bias must be dequantized
            # to float32 before adding to the float conv output.
            if bias_tensor.qnn_params:
                bias_expr = self.dequantize(bias_expr, bias_tensor)
            elif input_tensor.qnn_params and bias_tensor_type in (
                TensorType.INT32,
                TensorType.INT64,
            ):
                bias_expr = relax.op.dequantize(
                    bias_expr,
                    scale=relax.op.multiply(
                        input_tensor.qnn_params["scale"],
                        weight_tensor.qnn_params["scale"],
                    ),
                    zero_point=relax.const(0, "int32"),
                    axis=0,
                )
            out = relax.op.add(out, bias_expr)

        # Handle fused activation.
        if output_tensor.qnn_params:
            # Quantize the float output using the output tensor's qnn params.
            out = self.quantize(out, output_tensor)

            # Call quantized activation function
            output_scale_val = get_scalar_from_constant(output_tensor.qnn_params["scale"])
            output_zero_point_val = get_scalar_from_constant(output_tensor.qnn_params["zero_point"])
            out = self.convert_qnn_fused_activation_function(
                expr=out,
                fused_activation_fn=fused_activation_fn,
                scale=output_scale_val,
                zero_point=output_zero_point_val,
                dtype=output_tensor_type_str,
            )
        else:
            out = self.convert_fused_activation_function(out, fused_activation_fn)
        return out

    def convert_conv3d(self, op):
        """3D convolution implementation."""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.Conv3DOptions import Conv3DOptions
        from tflite.Padding import Padding
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) >= 2, "input tensors length should be >= 2"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        weight_tensor = input_tensors[1]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        assert op.BuiltinOptionsType() == BuiltinOptions.Conv3DOptions
        op_options = op.BuiltinOptions()
        conv3d_options = Conv3DOptions()
        conv3d_options.Init(op_options.Bytes, op_options.Pos)

        stride_d = conv3d_options.StrideD()
        stride_h = conv3d_options.StrideH()
        stride_w = conv3d_options.StrideW()
        dilation_d = conv3d_options.DilationDFactor()
        dilation_h = conv3d_options.DilationHFactor()
        dilation_w = conv3d_options.DilationWFactor()
        padding = conv3d_options.Padding()
        fused_activation_fn = conv3d_options.FusedActivationFunction()

        _, input_d, input_h, input_w, input_c = to_int_list(self.get_tensor_shape(input_tensor))
        # TFLite Conv3D kernel layout is already DHWIO:
        # KD KH KW IC OC
        kernel_d, kernel_h, kernel_w, in_channels, output_channels = to_int_list(
            self.get_tensor_shape(weight_tensor)
        )

        dilated_kernel_d = dilation_d * (kernel_d - 1) + 1
        dilated_kernel_h = dilation_h * (kernel_h - 1) + 1
        dilated_kernel_w = dilation_w * (kernel_w - 1) + 1

        params = {
            "strides": [stride_d, stride_h, stride_w],
            "dilation": [dilation_d, dilation_h, dilation_w],
            "padding": [0, 0, 0, 0, 0, 0],
            "data_layout": "NDHWC",
        }

        params["kernel_layout"] = "DHWIO"
        if input_c != in_channels:
            assert input_c % in_channels == 0, (
                "Input channels is not divisible by kernel in_channels."
            )
            params["groups"] = int(input_c / in_channels)

        # weight tensor type should be INT8/UINT8 (quantization) or FLOAT32
        weight_tensor_type = weight_tensor.tensor.Type()
        assert weight_tensor_type in (
            TensorType.INT8,
            TensorType.UINT8,
            TensorType.FLOAT32,
        )
        weight_tensor_type_str = self.get_tensor_type_str(weight_tensor_type)

        in_expr = self.get_expr(input_tensor_idx)

        # TFLite Conv3D kernel is already in DHWIO layout, no transpose needed.
        if self.has_expr(weight_tensor.tensor_idx):
            weight_expr = self.get_expr(weight_tensor.tensor_idx)
        else:
            if self.is_prefetched(weight_tensor.tensor_idx):
                weight_value = self.get_prefetched_node(weight_tensor.tensor_idx)
            else:
                weight_value = self.get_tensor_value(weight_tensor)

            weight_expr = self.exp_tab.new_const(
                weight_value, dtype=weight_tensor_type_str, source_name=weight_tensor.tensor.Name()
            )

        if padding == Padding.VALID:
            pass
        elif padding == Padding.SAME:
            pad_front, pad_back = get_pad_value(input_d, dilated_kernel_d, stride_d)
            pad_top, pad_bottom = get_pad_value(input_h, dilated_kernel_h, stride_h)
            pad_left, pad_right = get_pad_value(input_w, dilated_kernel_w, stride_w)

            do_pad = not (
                pad_front == 0
                and pad_back == 0
                and pad_top == 0
                and pad_bottom == 0
                and pad_left == 0
                and pad_right == 0
            )
            if do_pad:
                params["padding"] = [pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right]
        else:
            raise tvm.error.OpAttributeUnImplemented(
                f"Padding format {padding} is not supported for operator Conv3D."
            )

        if input_tensor.qnn_params:
            raise tvm.error.OpNotImplemented(
                "Quantized Conv3D is not yet supported in the Relax frontend."
            )

        out = relax.op.nn.conv3d(in_expr, weight_expr, **params)

        # if we have bias
        if len(input_tensors) == 3:
            bias_tensor = input_tensors[2]
            if bias_tensor.tensor_idx != -1:
                bias_tensor_type = bias_tensor.tensor.Type()
                # bias tensor type should be INT32 (int8 qnn) or INT64 (int16 qnn) or FLOAT32
                assert bias_tensor_type in (TensorType.INT32, TensorType.INT64, TensorType.FLOAT32)
                bias_tensor_type_str = self.get_tensor_type_str(bias_tensor_type)
                if self.has_expr(bias_tensor.tensor_idx):
                    bias_expr = self.get_expr(bias_tensor.tensor_idx)
                else:
                    bias_expr = self.exp_tab.new_const(
                        self.get_tensor_value(bias_tensor),
                        dtype=bias_tensor_type_str,
                        source_name=bias_tensor.tensor.Name(),
                    )
                out = relax.op.add(out, bias_expr)

        # Handle fused activation.
        if output_tensor.qnn_params:
            raise tvm.error.OpNotImplemented(
                "Quantized Conv3D is not yet supported in the Relax frontend."
            )

        out = self.convert_fused_activation_function(out, fused_activation_fn)
        return out

    def convert_conv3d_transpose(self, op):
        """3D transposed convolution implementation."""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.Conv3DOptions import Conv3DOptions
        from tflite.Padding import Padding
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) >= 3, "input tensors length should be >= 3"

        # TFLite CONV_3D_TRANSPOSE input order:
        # [0] output_shape, [1] weight, [2] data, [3] bias (optional)
        weight_tensor = input_tensors[1]
        input_tensor = input_tensors[2]
        input_tensor_idx = input_tensor.tensor_idx

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        assert op.BuiltinOptionsType() == BuiltinOptions.Conv3DOptions
        op_options = op.BuiltinOptions()
        conv3d_options = Conv3DOptions()
        conv3d_options.Init(op_options.Bytes, op_options.Pos)

        stride_d = conv3d_options.StrideD()
        stride_h = conv3d_options.StrideH()
        stride_w = conv3d_options.StrideW()
        dilation_d = conv3d_options.DilationDFactor()
        dilation_h = conv3d_options.DilationHFactor()
        dilation_w = conv3d_options.DilationWFactor()
        padding = conv3d_options.Padding()
        fused_activation_fn = conv3d_options.FusedActivationFunction()

        _, input_d, input_h, input_w, input_c = to_int_list(self.get_tensor_shape(input_tensor))

        # TFLite Conv3DTranspose kernel layout is DHWOI:
        # KD KH KW OC IC
        kernel_d, kernel_h, kernel_w, output_channels, in_channels = to_int_list(
            self.get_tensor_shape(weight_tensor)
        )

        dilated_kernel_d = dilation_d * (kernel_d - 1) + 1
        dilated_kernel_h = dilation_h * (kernel_h - 1) + 1
        dilated_kernel_w = dilation_w * (kernel_w - 1) + 1

        params = {
            "strides": [stride_d, stride_h, stride_w],
            "dilation": [dilation_d, dilation_h, dilation_w],
            "padding": [0, 0, 0, 0, 0, 0],
            "output_padding": [0, 0, 0],
            "data_layout": "NDHWC",
            "kernel_layout": "DHWOI",
        }

        if input_c != in_channels:
            assert input_c % in_channels == 0, (
                "Input channels is not divisible by kernel in_channels."
            )
            params["groups"] = int(input_c / in_channels)

        # weight tensor type should be INT8/UINT8 (quantization) or FLOAT32
        weight_tensor_type = weight_tensor.tensor.Type()
        assert weight_tensor_type in (
            TensorType.INT8,
            TensorType.UINT8,
            TensorType.FLOAT32,
        )
        weight_tensor_type_str = self.get_tensor_type_str(weight_tensor_type)

        in_expr = self.get_expr(input_tensor_idx)

        # TFLite Conv3DTranspose kernel is already in DHWOI layout, no transpose needed.
        if self.has_expr(weight_tensor.tensor_idx):
            weight_expr = self.get_expr(weight_tensor.tensor_idx)
        else:
            if self.is_prefetched(weight_tensor.tensor_idx):
                weight_value = self.get_prefetched_node(weight_tensor.tensor_idx)
            else:
                weight_value = self.get_tensor_value(weight_tensor)

            weight_expr = self.exp_tab.new_const(
                weight_value, dtype=weight_tensor_type_str, source_name=weight_tensor.tensor.Name()
            )

        if padding == Padding.VALID:
            pass
        elif padding == Padding.SAME:
            # For transposed convolution with SAME padding:
            # target output_size = input_size * stride
            # total_pad = max(0, dilated_kernel - stride)
            for dim_kernel, dim_stride, label in [
                (dilated_kernel_d, stride_d, "D"),
                (dilated_kernel_h, stride_h, "H"),
                (dilated_kernel_w, stride_w, "W"),
            ]:
                total_pad = max(0, dim_kernel - dim_stride)
                pad_before = total_pad // 2
                pad_after = total_pad - pad_before
                idx = {"D": 0, "H": 1, "W": 2}[label]
                params["padding"][idx] = pad_before
                params["padding"][idx + 3] = pad_after

                # output_padding handles the case when stride > dilated_kernel
                output_pad = max(0, dim_stride - dim_kernel)
                params["output_padding"][idx] = output_pad
        else:
            raise tvm.error.OpAttributeUnImplemented(
                f"Padding format {padding} is not supported for operator Conv3DTranspose."
            )

        if input_tensor.qnn_params:
            raise tvm.error.OpNotImplemented(
                "Quantized Conv3DTranspose is not yet supported in the Relax frontend."
            )

        out = relax.op.nn.conv3d_transpose(in_expr, weight_expr, **params)

        # if we have bias (input_tensors[3])
        if len(input_tensors) >= 4:
            bias_tensor = input_tensors[3]
            if bias_tensor.tensor_idx != -1:
                bias_tensor_type = bias_tensor.tensor.Type()
                # bias tensor type should be INT32 (int8 qnn) or INT64 (int16 qnn) or FLOAT32
                assert bias_tensor_type in (TensorType.INT32, TensorType.INT64, TensorType.FLOAT32)
                bias_tensor_type_str = self.get_tensor_type_str(bias_tensor_type)
                if self.has_expr(bias_tensor.tensor_idx):
                    bias_expr = self.get_expr(bias_tensor.tensor_idx)
                else:
                    bias_expr = self.exp_tab.new_const(
                        self.get_tensor_value(bias_tensor),
                        dtype=bias_tensor_type_str,
                        source_name=bias_tensor.tensor.Name(),
                    )
                out = relax.op.add(out, bias_expr)

        # Handle fused activation.
        if output_tensor.qnn_params:
            raise tvm.error.OpNotImplemented(
                "Quantized Conv3DTranspose is not yet supported in the Relax frontend."
            )

        out = self.convert_fused_activation_function(out, fused_activation_fn)
        return out

    def convert_split(self, op):
        """split implementation."""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.SplitOptions import SplitOptions

        input_tensors = self.get_input_tensors(op)

        assert len(input_tensors) == 2, "input tensors length should be == 2"

        axis_tensor = input_tensors[0]
        split_axis = self.get_tensor_value(axis_tensor)
        input_tensor = input_tensors[1]
        input_tensor_idx = input_tensor.tensor_idx

        assert op.BuiltinOptionsType() == BuiltinOptions.SplitOptions
        op_options = op.BuiltinOptions()
        split_options = SplitOptions()
        split_options.Init(op_options.Bytes, op_options.Pos)
        num_splits = split_options.NumSplits()

        in_expr = self.get_expr(input_tensor_idx)
        out = relax.op.split(in_expr, num_splits, axis=int(split_axis))
        # Relay does not like a TupleWrapper of 1 element, further this
        # only shows up with tf1.13 if we use a split with num_splits==1.
        # In tf 1.14 this doesn't appear as it is automatically a reshape
        # operation.
        if isinstance(out, relax.Tuple):
            if out.size == 1:
                out = out[0]

        return out

    def convert_split_v(self, op):
        """SPLIT_V implementation."""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)

        assert len(input_tensors) == 3, "input tensors length should be 3"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        in_expr = self.get_expr(input_tensor_idx)

        axis_tensor = input_tensors[2]
        split_axis = int(self.get_tensor_value(axis_tensor))

        size_splits_tensor = input_tensors[1]

        if self.has_expr(size_splits_tensor.tensor_idx):
            # Dynamic size_splits case: decompose into dynamic strided slices.
            size_splits_expr = self.get_expr(size_splits_tensor.tensor_idx)
            cumsum = relax.op.cumsum(size_splits_expr, axis=0, dtype="int64")
            # Pad a leading zero so that cumsum[i-1] can be read uniformly
            # via strided_slice even for i == 0.
            zero = relax.const(np.array([0], dtype="int64"), "int64")
            padded_cumsum = relax.op.concat([zero, cumsum], axis=0)
            # TFLite fixes the tuple arity in the graph, even when the split
            # sizes themselves are supplied at runtime.
            num_splits = len(output_tensors)
            rank = len(in_expr.ty.shape)

            # end_base is the full input shape; only split_axis changes per slice.
            end_base = relax.op.shape_to_tensor(relax.op.shape_of(in_expr))
            begin_base = relax.const(np.zeros((rank,), dtype="int64"), "int64")
            strides = relax.const(np.ones((rank,), dtype="int64"), "int64")
            scatter_idx = relax.const([split_axis], "int64")

            outputs = []
            for i in range(num_splits):
                start_val = relax.op.strided_slice(padded_cumsum, axes=[0], begin=[i], end=[i + 1])
                end_val = relax.op.strided_slice(
                    padded_cumsum, axes=[0], begin=[i + 1], end=[i + 2]
                )

                begin = relax.op.scatter_elements(begin_base, scatter_idx, start_val)
                end = relax.op.scatter_elements(end_base, scatter_idx, end_val)
                slice_i = relax.op.dynamic_strided_slice(in_expr, begin, end, strides)
                outputs.append(slice_i)

            out = relax.Tuple(outputs)
        else:
            # Static size_splits case
            size_splits = list(self.get_tensor_value(size_splits_tensor))
            size_splits = tuple(np.cumsum(size_splits)[:-1])
            out = relax.op.split(in_expr, size_splits, axis=split_axis)

        # Relay does not like a TupleWrapper of 1 element, further this
        # only shows up with tf1.13 if we use a split with num_splits==1.
        # In tf 1.14 this doesn't appear as it is automatically a reshape
        # operation.
        if isinstance(out, relax.Tuple) and len(out.fields) == 1:
            out = out[0]

        return out

    def convert_slice(self, op):
        """Convert TFLite SLICE"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be == 3"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        begin = list(self.get_tensor_value(input_tensors[1]))
        size = list(self.get_tensor_value(input_tensors[2]))
        # strided_slice(Relax) needs the slice's end indices, not the size
        end = size
        input_tensor_shape = to_int_list(self.get_tensor_shape(input_tensor))
        input_tensor_rank = len(input_tensor_shape)
        for i in range(input_tensor_rank):
            if size[i] == -1:
                end[i] = input_tensor_shape[i]
            else:
                end[i] += begin[i]

        # Create axes list for all dimensions being sliced
        axes = list(range(input_tensor_rank))
        begin = [int(v) for v in begin]
        end = [int(v) for v in end]
        out = relax.op.strided_slice(in_expr, axes=axes, begin=begin, end=end)
        return out

    def convert_scatter_nd(self, op):
        """Convert TFLite SCATTER_ND"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "SCATTER_ND should have 3 input tensors"
        indices = self.get_tensor_expr(input_tensors[0])
        updates = self.get_tensor_expr(input_tensors[1])
        shape_tensor = input_tensors[2]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "SCATTER_ND should have 1 output tensor"
        updates_dtype = self.get_tensor_type_str(output_tensors[0].tensor.Type())

        if self.has_expr(shape_tensor.tensor_idx):
            shape_expr = self.get_expr(shape_tensor.tensor_idx)
            shape_expr = self.bb.normalize(relax.op.astype(shape_expr, "int64"))
            shape = self.bb.emit(relax.op.tensor_to_shape(shape_expr))
        else:
            shape = to_int_list(self.get_tensor_value(shape_tensor))

        indices_dims = len(self._infer_shape(indices))
        indices = relax.op.permute_dims(indices, axes=[-1] + list(range(indices_dims - 1)))

        data = relax.op.zeros(shape, updates_dtype)
        return relax.op.scatter_nd(data, indices, updates, "update")

    def convert_unique(self, op):
        """Convert TFLite UNIQUE."""
        from tflite.TensorType import TensorType
        from tflite.UniqueOptions import UniqueOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "UNIQUE should have 1 input tensor"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 2, "UNIQUE should have 2 output tensors"

        unique_options = UniqueOptions()
        op_options = op.BuiltinOptions()
        unique_options.Init(op_options.Bytes, op_options.Pos)

        unique = self.bb.normalize(
            relax.op.unique(
                self.get_tensor_expr(input_tensors[0]),
                sorted=False,
                return_index=False,
                return_inverse=True,
                return_counts=False,
                axis=None,
            )
        )
        values = self.bb.emit(relax.TupleGetItem(unique, 0))
        inverse_indices = self.bb.emit(relax.TupleGetItem(unique, 1))

        idx_out_type = unique_options.IdxOutType()
        if idx_out_type == TensorType.INT32:
            inverse_indices = self.bb.emit(relax.op.astype(inverse_indices, "int32"))

        return relax.Tuple([values, inverse_indices])

    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 == "max":
            np_dtype = np.dtype(output_dtype)
            if np.issubdtype(np_dtype, np.floating):
                identity = np.finfo(np_dtype).min
            elif np.issubdtype(np_dtype, np.integer):
                identity = np.iinfo(np_dtype).min
            else:
                raise tvm.error.OpNotImplemented(
                    f"UNSORTED_SEGMENT_MAX does not support output dtype {output_dtype}."
                )
            return relax.op.full(output_shape, relax.const(identity, 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."
                )
            segment_ids = self.get_tensor_value(segment_ids_tensor)
            if np.any(segment_ids < 0):
                raise tvm.error.OpNotImplemented(
                    "TFLite SEGMENT_SUM with negative segment ids is not supported."
                )
            return int(np.max(segment_ids)) + 1 if segment_ids.size else 0

        num_segments_tensor = input_tensors[2]
        if self.has_expr(num_segments_tensor.tensor_idx):
            raise tvm.error.OpNotImplemented(
                f"TFLite {op_name} with runtime num_segments is not supported."
            )
        num_segments_value = self.get_tensor_value(num_segments_tensor)
        assert num_segments_value.size == 1, f"{op_name} num_segments should be a scalar tensor"
        num_segments = int(num_segments_value.item())
        assert num_segments >= 0, f"{op_name} num_segments should be non-negative"
        return num_segments

    def _convert_segment_op(self, op, op_name, reduction):
        """Convert TFLite segment ops through relax.op.scatter_nd."""
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        expected_inputs = 2 if op_name == "SEGMENT_SUM" else 3
        assert len(input_tensors) == expected_inputs, (
            f"{op_name} should have {expected_inputs} input tensors"
        )

        data_tensor = input_tensors[0]
        segment_ids_tensor = input_tensors[1]
        for t in input_tensors:
            assert not t.qnn_params, "Quantized input is not expected."

        segment_ids_type = segment_ids_tensor.tensor.Type()
        assert segment_ids_type in (TensorType.INT32, TensorType.INT64)
        if op_name != "SEGMENT_SUM":
            num_segments_type = input_tensors[2].tensor.Type()
            assert num_segments_type in (TensorType.INT32, TensorType.INT64)
        if not self.has_expr(segment_ids_tensor.tensor_idx):
            segment_ids_value = self.get_tensor_value(segment_ids_tensor)
            if np.any(segment_ids_value < 0):
                raise tvm.error.OpNotImplemented(
                    f"TFLite {op_name} with negative segment ids is not supported."
                )

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, f"{op_name} should have 1 output tensor"
        output_tensor = output_tensors[0]
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        data_shape = to_int_list(self.get_tensor_shape(data_tensor))
        segment_ids_shape = to_int_list(self.get_tensor_shape(segment_ids_tensor))
        segment_ids_rank = len(segment_ids_shape)
        assert data_shape[:segment_ids_rank] == segment_ids_shape, (
            f"{op_name} requires segment_ids shape to match a prefix of data shape"
        )
        num_segments = self._get_segment_num_segments(op_name, input_tensors)
        output_shape = [num_segments] + data_shape[segment_ids_rank:]

        data = self.get_tensor_expr(data_tensor)
        segment_ids = self.get_tensor_expr(segment_ids_tensor)
        indices = relax.op.expand_dims(segment_ids, axis=[segment_ids_rank])

        base = self._get_segment_scatter_base(output_shape, output_dtype, reduction)
        return relax.op.scatter_nd(base, indices, data, reduction)

    def convert_select(self, op):
        """Convert TFLite SELECT"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be == 3"
        cond = self.get_tensor_expr(input_tensors[0])
        x = self.get_tensor_expr(input_tensors[1])
        y = self.get_tensor_expr(input_tensors[2])

        out = relax.op.where(cond, x, y)

        return out

    def convert_transpose(self, op):
        """transpose implementation."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx

        in_expr = self.get_expr(input_tensor_idx)

        # axis
        in_axis = tuple(self.get_tensor_value(input_tensors[1]))

        if not in_axis:
            out = relax.op.permute_dims(in_expr)
        else:
            out = relax.op.permute_dims(in_expr, in_axis)

        return out

    def convert_reverse_sequence(self, op):
        """Convert TFLite REVERSE_SEQUENCE"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.ReverseSequenceOptions import ReverseSequenceOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "TFLite does not support quantized REVERSE_SEQUENCE operator yet."
            )

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        in_expr = self.get_tensor_expr(input_tensors[0])
        length_expr = self.get_tensor_expr(input_tensors[1])

        assert op.BuiltinOptionsType() == BuiltinOptions.ReverseSequenceOptions
        op_options = op.BuiltinOptions()
        options = ReverseSequenceOptions()
        options.Init(op_options.Bytes, op_options.Pos)
        batch_axis = options.BatchDim()
        seq_axis = options.SeqDim()

        return relax.op.reverse_sequence(in_expr, length_expr, seq_axis, batch_axis)

    def convert_bitcast(self, op):
        """Convert TFLite BITCAST"""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        assert len(output_tensors) == 1, "output tensors length should be 1"

        in_expr = self.get_tensor_expr(input_tensors[0])
        input_dtype = self.get_tensor_type_str(input_tensors[0].tensor.Type())
        output_dtype = self.get_tensor_type_str(output_tensors[0].tensor.Type())
        input_shape = to_int_list(self.get_tensor_shape(input_tensors[0]))
        output_shape = to_int_list(self.get_tensor_shape(output_tensors[0]))

        input_nbytes = int(np.prod(input_shape)) * np.dtype(input_dtype).itemsize
        output_nbytes = int(np.prod(output_shape)) * np.dtype(output_dtype).itemsize
        assert input_nbytes == output_nbytes, (
            "TFLite BITCAST requires input.nbytes == output.nbytes, "
            f"but got input={input_nbytes} bytes, output={output_nbytes} bytes"
        )

        return relax.op.memory.view(in_expr, shape=output_shape, dtype=output_dtype)

    def convert_rfft2d(self, op):
        """Convert TFLite RFFT2D for static no-padding/no-truncation shapes."""
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        if len(input_tensors) != 2 or len(output_tensors) != 1:
            raise tvm.error.OpNotImplemented("RFFT2D expects two inputs and one output")

        data_tensor, fft_length_tensor = input_tensors
        output_tensor = output_tensors[0]
        if data_tensor.tensor.Type() != TensorType.FLOAT32:
            raise tvm.error.OpNotImplemented("RFFT2D input must be float32")
        if not self._is_tflite_complex64_type(output_tensor.tensor.Type()):
            raise tvm.error.OpNotImplemented("RFFT2D output must be COMPLEX64")
        if (
            data_tensor.tensor.Sparsity() is not None
            or fft_length_tensor.tensor.Sparsity() is not None
        ):
            raise tvm.error.OpNotImplemented("RFFT2D does not support sparse inputs")

        input_shape = tuple(to_int_list(self.get_tensor_shape(data_tensor)))
        tflite_output_shape = tuple(to_int_list(self.get_tensor_shape(output_tensor)))
        if len(input_shape) < 2:
            raise tvm.error.OpNotImplemented("RFFT2D input rank must be at least 2")

        try:
            fft_length_value = self.get_tensor_value_or_prefetched(fft_length_tensor)
        except (ValueError, TypeError):
            raise tvm.error.OpNotImplemented("RFFT2D requires a constant fft_length") from None
        if fft_length_value is None:
            raise tvm.error.OpNotImplemented("RFFT2D requires a constant fft_length")
        # Reject non-integer fft_length tensors before astype("int64") can silently
        # truncate (e.g. float32 2.7 -> int64 2).
        if not np.issubdtype(fft_length_value.dtype, np.integer):
            raise tvm.error.OpNotImplemented(
                f"RFFT2D fft_length must be an integer tensor, got dtype {fft_length_value.dtype!r}"
            )
        fft_length = fft_length_value.astype("int64")
        if tuple(fft_length.shape) != (2,):
            raise tvm.error.OpNotImplemented("RFFT2D fft_length must be a length-2 tensor")

        height = int(fft_length[0])
        width = int(fft_length[1])
        if height <= 0 or width <= 0:
            raise tvm.error.OpNotImplemented(
                f"RFFT2D fft_length must be positive, got ({height}, {width})"
            )
        if height != int(input_shape[-2]) or width != int(input_shape[-1]):
            raise tvm.error.OpNotImplemented(
                "RFFT2D currently supports fft_length matching the input spatial shape"
            )
        expected_tflite_output_shape = input_shape[:-2] + (height, width // 2 + 1)
        if tflite_output_shape != expected_tflite_output_shape:
            raise tvm.error.OpNotImplemented("RFFT2D output shape does not match fft_length")

        relax_output_shape = self._get_relax_tensor_shape(output_tensor)
        # Dispatch: power-of-2 sizes use the O(N^2 log N) Cooley-Tukey FFT kernel;
        # the remaining (odd / non-power-of-2) shapes fall back to the O(N^4) DFT
        # reference. Both kernels share the same call_tir contract, so the
        # downstream code is kernel-agnostic.
        if _is_power_of_2(height) and _is_power_of_2(width):
            prim_func = _build_tflite_rfft2d_fft_primfunc(input_shape, relax_output_shape)
        else:
            prim_func = _build_tflite_rfft2d_primfunc(input_shape, relax_output_shape)
        module_builder = self.conversion_state["module_builder"]
        func_name = f"tflite_rfft2d_{output_tensor.tensor_idx}"
        gv = module_builder.add_func(prim_func, func_name)
        data_expr = self.get_tensor_expr(data_tensor)
        call = relax.call_tir(
            gv,
            [data_expr],
            relax.TensorType(relax_output_shape, "float32"),
        )
        return self.bb.normalize(call)

    def convert_broadcast_args(self, op):
        """Convert TFLite BROADCAST_ARGS"""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        assert len(output_tensors) == 1, "output tensors length should be 1"

        s0 = self.get_tensor_expr(input_tensors[0])
        s1 = self.get_tensor_expr(input_tensors[1])
        s0_len = to_int_list(self.get_tensor_shape(input_tensors[0]))[0]
        s1_len = to_int_list(self.get_tensor_shape(input_tensors[1]))[0]
        out_dtype = self.get_tensor_type_str(input_tensors[0].tensor.Type())

        # Left-pad the shorter input with 1s to length target_len.
        target_len = tirx.max(s0_len, s1_len)
        one = relax.const(1, dtype=out_dtype)
        s0 = relax.op.concat(
            [relax.op.full([target_len - s0_len], one, dtype=out_dtype), s0], axis=0
        )
        s1 = relax.op.concat(
            [relax.op.full([target_len - s1_len], one, dtype=out_dtype), s1], axis=0
        )
        # Per-dim broadcast. If either side is 1 take the other, else elementwise max.
        s0_is_one = relax.op.equal(s0, one)
        s1_is_one = relax.op.equal(s1, one)
        return relax.op.where(
            s0_is_one,
            s1,
            relax.op.where(s1_is_one, s0, relax.op.maximum(s0, s1)),
        )

    def convert_cast(self, op):
        """Convert TFLite CAST"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.CastOptions import CastOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        # MLIR-based converter outputs no BuiltinOptions for Cast operator. In this
        # case the output type can be derived from the Cast operator output tensor.
        # When TOCO converter is used there will be "normal" BuiltinOptions.CastOptions
        # with output type.
        if op.BuiltinOptions() is not None:
            assert op.BuiltinOptionsType() == BuiltinOptions.CastOptions
            op_options = op.BuiltinOptions()
            cast_options = CastOptions()
            cast_options.Init(op_options.Bytes, op_options.Pos)
            cast_dtype = cast_options.OutDataType()
        else:
            cast_dtype = self.get_output_tensors(op)[0].tensor.Type()

        out = relax.op.astype(in_expr, self.get_tensor_type_str(cast_dtype))

        return out

    def convert_tile(self, op):
        """tile implementation."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx

        in_expr = self.get_expr(input_tensor_idx)

        reps = tuple(self.get_tensor_value(input_tensors[1]))

        out = relax.op.tile(in_expr, reps)

        return out

    def convert_topk_v2(self, op):
        """Convert TFLite TOPK_v2"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        in_expr = self.get_expr(input_tensor_idx)
        k = self.get_tensor_value(input_tensors[1])
        out = relax.op.topk(in_expr, int(k))

        return out

    def convert_pool2d(self, op, pool_type):
        """pool2d implementation."""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.Padding import Padding
        from tflite.Pool2DOptions import Pool2DOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors should be 1"
        output_tensor = output_tensors[0]
        output_tensor_type = output_tensor.tensor.Type()
        output_tensor_type_str = self.get_tensor_type_str(output_tensor_type)

        assert op.BuiltinOptionsType() == BuiltinOptions.Pool2DOptions
        op_options = op.BuiltinOptions()
        pool2d_options = Pool2DOptions()
        pool2d_options.Init(op_options.Bytes, op_options.Pos)
        stride_h = pool2d_options.StrideH()
        stride_w = pool2d_options.StrideW()
        padding = pool2d_options.Padding()
        filter_h = pool2d_options.FilterHeight()
        filter_w = pool2d_options.FilterWidth()
        fused_activation_fn = pool2d_options.FusedActivationFunction()

        params = {
            "pool_size": (filter_h, filter_w),
            "strides": (stride_h, stride_w),
            "padding": [0, 0],
            "layout": "NHWC",
        }

        in_expr = self.get_expr(input_tensor_idx)

        _, input_h, input_w, _ = to_int_list(self.get_tensor_shape(input_tensor))

        if padding == Padding.VALID:
            pass
        elif padding == Padding.SAME:
            pad_top, pad_bottom = get_pad_value(input_h, filter_h, stride_h)
            pad_left, pad_right = get_pad_value(input_w, filter_w, stride_w)
            params["padding"] = [pad_top, pad_left, pad_bottom, pad_right]
        else:
            raise tvm.error.OpAttributeUnImplemented(
                f"Padding format {padding} for operator Pool2D is not supported."
            )

        if pool_type == "average":
            if input_tensor.qnn_params:
                assert self.has_same_qnn_params(input_tensor, output_tensor), (
                    "TFLite avg_pool2dreshape requires input and output scale"
                    "and zero points to be equal"
                )
                out = relax.op.astype(in_expr, "int32")
                out = relax.op.nn.avg_pool2d(out, **params)
                out = relax.op.astype(out, output_tensor_type_str)
            else:
                out = relax.op.nn.avg_pool2d(in_expr, **params)
        elif pool_type == "max":
            if input_tensor.qnn_params:
                assert self.has_same_qnn_params(input_tensor, output_tensor), (
                    "qnn.op.max_pool2d requires input and output qnn params to be same"
                )
            out = relax.op.nn.max_pool2d(in_expr, **params)
        elif pool_type == "l2":
            # L2_POOL_2D is equivalent to square_root(avg_pool(square(in_data)))
            # TFLite does not have support for quantised L2_POOL_2D op.
            assert not input_tensor.qnn_params, (
                "As TFLite does not have support for quantized L2_POOL_2D, \
                Quantized input is not expected."
            )
            exp_type = self.get_tensor_type_str(output_tensor.tensor.Type())
            square_exp = relax.op.power(in_expr, relax.const(2, exp_type))
            avg_pool_exp = relax.op.nn.avg_pool2d(square_exp, **params)
            out = relax.op.sqrt(avg_pool_exp)
        else:
            raise tvm.error.OpNotImplemented(
                f"Operator {pool_type} pool is not supported for frontend TFLite."
            )

        # Handle fused activations
        if output_tensor.qnn_params:
            scale_val = get_scalar_from_constant(output_tensor.qnn_params["scale"])
            zero_point_val = get_scalar_from_constant(output_tensor.qnn_params["zero_point"])
            out = self.convert_qnn_fused_activation_function(
                expr=out,
                fused_activation_fn=fused_activation_fn,
                scale=scale_val,
                zero_point=zero_point_val,
                dtype=output_tensor_type_str,
            )
        else:
            out = self.convert_fused_activation_function(out, fused_activation_fn)

        return out

    def convert_pad(self, op):
        """Convert TFLite PAD/PADV2 \
           TFLite treats PAD and PADV2 operators identically"""

        input_tensors = self.get_input_tensors(op)

        # TFLite PAD/PADV2 only supports CONSTANT mode
        assert len(input_tensors) == 2 or len(input_tensors) == 3, (
            "input tensor's length should be 2 for PAD and 3 for PADV2"
        )

        if len(input_tensors) == 3:
            assert input_tensors[0].tensor.Type() == input_tensors[2].tensor.Type(), (
                "constant_values tensor must be of same type as input tensor"
            )

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        # paddings
        pad_list = self.get_tensor_value(input_tensors[1])

        # convert list of lists to tuple of tuples
        paddings = []
        for val in pad_list:
            paddings += val.tolist()

        # Set the pad value, by default 0, unless constant_values parameter is provided
        pad_value = 0

        if input_tensor.qnn_params:
            # Check that input and output tensor have same qnn params.
            output_tensors = self.get_output_tensors(op)
            output_tensor = output_tensors[0]
            assert self.has_same_qnn_params(input_tensor, output_tensor), (
                "TFLite PADV2 requires input and output scale and zero points to be equal"
            )

            # The pad value for quantized pad is the input zero point by default.
            pad_value = float(input_tensor.qnn_params["zero_point"].data.numpy())

        if len(input_tensors) == 3:
            pad_value = self.get_tensor_value(input_tensors[2])
            if isinstance(pad_value, np.ndarray):
                pad_value = pad_value.tolist()
            if isinstance(pad_value, list):
                assert len(pad_value) == 1, "Only one constant value is expected."
                pad_value = pad_value[0]
            if input_tensor.qnn_params:
                # Check that input tensor and constant_values have same qnn params.
                assert self.has_same_qnn_params(input_tensor, input_tensors[2]), (
                    "TFLite PADV2 requires input and constant_values tensors' \
                        scale and zero points to be equal"
                )

        out = relax.op.nn.pad(in_expr, pad_width=paddings, pad_value=pad_value)
        return out

    def convert_mirror_pad(self, op):
        """Convert TFLite MIRROR_PAD"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.MirrorPadOptions import MirrorPadOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        # tensor
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        # paddings
        pad_list = self.get_tensor_value(input_tensors[1])
        # convert list of lists to tuple of tuples
        paddings = tuple(tuple(val.astype(np.int32)) for val in pad_list)

        assert op.BuiltinOptionsType() == BuiltinOptions.MirrorPadOptions
        op_options = op.BuiltinOptions()
        mirror_pad_options = MirrorPadOptions()
        mirror_pad_options.Init(op_options.Bytes, op_options.Pos)
        mode_byte = mirror_pad_options.Mode()

        mode = "REFLECT" if mode_byte == 0 else "SYMMETRIC"
        if mode == "SYMMETRIC":
            raise tvm.error.OpAttributeUnImplemented(
                "MIRROR_PAD with SYMMETRIC mode is not yet supported."
            )
        # Flatten tuple-of-tuples to a list for relax.op.nn.pad
        flat_pads = [int(v) for pair in paddings for v in pair]
        out = relax.op.nn.pad(in_expr, flat_pads, pad_mode="reflect")

        return out

    def convert_pack(self, op):
        """Convert TFLite pack"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.PackOptions import PackOptions

        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"

        if input_tensors[0].qnn_params:
            output_tensor = output_tensors[0]
            assert self.has_same_qnn_params(input_tensors[0], output_tensor), (
                "TFLite pack requires input and output scale and zero points to be equal"
            )

            for input_tensor in input_tensors:
                assert self.has_same_qnn_params(input_tensors[0], input_tensor), (
                    "TFLite pack requires all input tensors to have same scale and zero point"
                )

        assert op.BuiltinOptionsType() == BuiltinOptions.PackOptions
        op_options = op.BuiltinOptions()
        pack_options = PackOptions()
        pack_options.Init(op_options.Bytes, op_options.Pos)
        pack_axis = pack_options.Axis()
        pack_values_count = pack_options.ValuesCount()
        assert len(input_tensors) == pack_values_count, "Discordance in input values count"

        in_exprs = [self.get_tensor_expr(_) for _ in input_tensors]
        in_exprs_reshaped = [_op.expand_dims(_, axis=pack_axis) for _ in in_exprs]
        out = relax.op.concat(in_exprs_reshaped, pack_axis)
        return out

    def convert_unpack(self, op):
        """Convert TFLite unpack"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.UnpackOptions import UnpackOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)
        assert op.BuiltinOptionsType() == BuiltinOptions.UnpackOptions
        op_options = op.BuiltinOptions()
        unpack_options = UnpackOptions()
        unpack_options.Init(op_options.Bytes, op_options.Pos)
        num_unpacks = unpack_options.Num()
        unpack_axis = unpack_options.Axis()

        # Relax doesn't support 'unpack' operator so we use 'split' & 'squeeze' instead.
        # We have to do 'squeeze' along the split axis.
        # Relax expects squeeze_axis to be List.
        squeeze_axis = [unpack_axis]

        # Relax doesn't like TupleWrapper of 1 element so we isolate the case of unpacking
        # a tensor by an axis with len(axis) == 1. For reference see convert_split().
        # Such unpacking will result in the same tensor so we omit 'split' and only squeeze
        # along the axis of dim == 1.
        if num_unpacks == 1:
            squeezed = relax.op.squeeze(in_expr, axis=squeeze_axis)
            if isinstance(squeezed, relax.Tuple):
                squeezed = squeezed[0]
        else:
            splitted = relax.op.split(in_expr, indices_or_sections=num_unpacks, axis=unpack_axis)
            squeezed = relax.Tuple(
                [
                    _op.squeeze(relax.TupleGetItem(splitted, i), axis=squeeze_axis)
                    for i in range(num_unpacks)
                ]
            )

        return squeezed

    def convert_lstm(self, op):
        """Convert TFLite LSTM (single-step).

        Standard LSTM cell with FULL kernel and coupled input-forget gate.
        Peephole, projection, and layer norm are not supported.

        Inputs (24 tensors, many optional):
          [0]  input                      [batch, input_size]
          [1]  input_to_input_weights     (optional, -1 => coupled)
          [2]  input_to_forget_weights    [num_units, input_size]
          [3]  input_to_cell_weights      [num_units, input_size]
          [4]  input_to_output_weights    [num_units, input_size]
          [5]  recurrent_to_input_weights (optional)
          [6]  recurrent_to_forget_weights [num_units, num_units]
          [7]  recurrent_to_cell_weights  [num_units, num_units]
          [8]  recurrent_to_output_weights [num_units, num_units]
          [9-11] cell_to_*_weights        (optional, not supported)
          [12] input_gate_bias            (optional)
          [13] forget_gate_bias           [num_units]
          [14] cell_bias                  [num_units]
          [15] output_gate_bias           [num_units]
          [16-17] projection_weights/bias (optional, not supported)
          [18] output_state               [batch, num_units]
          [19] cell_state                 [batch, num_units]
          [20-23] layer_norm              (optional, not supported)

        Output:
          [0] output  [batch, num_units]

        Cell (coupled input-forget):
          f = sigmoid(x @ W_f.T + h @ R_f.T + b_f)
          i = 1 - f
          g = tanh(x @ W_c.T + h @ R_c.T + b_c)
          o = sigmoid(x @ W_o.T + h @ R_o.T + b_o)
          c_new = f * c_prev + i * g
          h_new = fused_activation(o * tanh(c_new))
        """
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.LSTMOptions import LSTMOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented("TFLite quantized LSTM is not supported yet.")

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 24, (
            f"input tensors length should be 24, got {len(input_tensors)}"
        )

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) >= 1, "output tensors length should be at least 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.LSTMOptions
        op_options = op.BuiltinOptions()
        lstm_opts = LSTMOptions()
        lstm_opts.Init(op_options.Bytes, op_options.Pos)

        fused_activation_fn = lstm_opts.FusedActivationFunction()
        cell_clip = lstm_opts.CellClip()
        proj_clip = lstm_opts.ProjClip()

        in_expr = self.get_tensor_expr(input_tensors[0])

        # Only coupled input-forget gate is supported.
        if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx != -1:
            raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")

        # Peephole, projection, and layer norm are not modeled yet.
        if (
            any(t.tensor_idx != -1 for t in input_tensors[9:12])
            or any(t.tensor_idx != -1 for t in input_tensors[16:18])
            or any(t.tensor_idx != -1 for t in input_tensors[20:24])
        ):
            raise tvm.error.OpNotImplemented(
                "Peephole, projection, and layer norm LSTM are not supported yet."
            )

        # Weights.
        w_f = self.get_tensor_expr(input_tensors[2])
        w_c = self.get_tensor_expr(input_tensors[3])
        w_o = self.get_tensor_expr(input_tensors[4])

        r_f = self.get_tensor_expr(input_tensors[6])
        r_c = self.get_tensor_expr(input_tensors[7])
        r_o = self.get_tensor_expr(input_tensors[8])

        # Biases.
        b_f = self.get_tensor_expr(input_tensors[13])
        b_c = self.get_tensor_expr(input_tensors[14])
        b_o = self.get_tensor_expr(input_tensors[15])

        # State inputs.
        h_prev = self.get_tensor_expr(input_tensors[18])
        c_prev = self.get_tensor_expr(input_tensors[19])

        # Coupled input-forget gate.
        f = relax.op.sigmoid(
            relax.op.add(
                relax.op.add(
                    relax.op.matmul(in_expr, relax.op.permute_dims(w_f)),
                    relax.op.matmul(h_prev, relax.op.permute_dims(r_f)),
                ),
                b_f,
            )
        )
        i = relax.op.subtract(
            relax.const(1.0, "float32"),
            f,
        )

        # Cell candidate.
        g = relax.op.tanh(
            relax.op.add(
                relax.op.add(
                    relax.op.matmul(in_expr, relax.op.permute_dims(w_c)),
                    relax.op.matmul(h_prev, relax.op.permute_dims(r_c)),
                ),
                b_c,
            )
        )

        # Output gate.
        o = relax.op.sigmoid(
            relax.op.add(
                relax.op.add(
                    relax.op.matmul(in_expr, relax.op.permute_dims(w_o)),
                    relax.op.matmul(h_prev, relax.op.permute_dims(r_o)),
                ),
                b_o,
            )
        )

        # Cell state update with optional clipping.
        c_new = relax.op.add(
            relax.op.multiply(f, c_prev),
            relax.op.multiply(i, g),
        )
        if cell_clip > 0:
            c_new = relax.op.clip(c_new, -cell_clip, cell_clip)

        # Hidden state.
        # TFLite applies the fused activation to the cell state before the
        # output gate multiply.
        h_new = relax.op.multiply(
            o, self.convert_fused_activation_function(c_new, fused_activation_fn)
        )
        if proj_clip > 0:
            h_new = relax.op.clip(h_new, -proj_clip, proj_clip)

        # Update state tensors in the expression table for subsequent ops.
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[18].tensor_idx),
            h_new,
            force_override=True,
        )
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[19].tensor_idx),
            c_new,
            force_override=True,
        )

        return h_new

    def convert_svdf(self, op):
        """Convert TFLite SVDF (single-step).

        Structured-Vectorized Bidirectional Filter for keyword spotting.

        Inputs (5 tensors):
          [0] input           [batch, input_size]
          [1] feature_weights [num_filters, input_size]
          [2] time_weights    [num_filters, memory_size]
          [3] bias            [num_filters]           (optional)
          [4] state           [batch, num_filters * memory_size]  (variable)

        Output:
          [0] output  [batch, num_units]

        Computation:
          feat = x @ W_feat.T                              # feature projection
          state_r = reshape(state, [B, F, memory_size])    # ring buffer
          time = sum(state_r * time_weights, axis=-1)      # time filtering
          out = activation(sum(reshape(time, [B, U, rank]), axis=-1) + bias)
        """
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.SVDFOptions import SVDFOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented("TFLite quantized SVDF is not supported yet.")

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 5, (
            f"input tensors length should be 5, got {len(input_tensors)}"
        )

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) >= 1, "output tensors length should be at least 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.SVDFOptions
        op_options = op.BuiltinOptions()
        svdf_opts = SVDFOptions()
        svdf_opts.Init(op_options.Bytes, op_options.Pos)

        rank = svdf_opts.Rank()
        fused_activation_fn = svdf_opts.FusedActivationFunction()

        in_expr = self.get_tensor_expr(input_tensors[0])
        feat_weights = self.get_tensor_expr(input_tensors[1])
        time_weights = self.get_tensor_expr(input_tensors[2])

        batch_size = self.get_tensor_shape(input_tensors[0])[0]
        if isinstance(batch_size, np.integer | int):
            batch_size = int(batch_size)
        num_filters = to_int_list(self.get_tensor_shape(input_tensors[1]))[0]
        if num_filters % rank != 0:
            raise tvm.error.OpNotImplemented("SVDF num_filters must be divisible by rank.")
        num_units = num_filters // rank
        memory_size = to_int_list(self.get_tensor_shape(input_tensors[2]))[1]

        # Feature projection: [batch, input_size] @ [input_size, num_filters]
        feat = relax.op.matmul(in_expr, relax.op.permute_dims(feat_weights))

        # Time filtering: reshape state -> weight -> reduce.
        state_expr = self.get_tensor_expr(input_tensors[4])
        state_3d = relax.op.reshape(state_expr, (batch_size, num_filters, memory_size))

        # time_weights: [num_filters, memory_size], broadcast to [1, num_filters, memory_size]
        tw_3d = relax.op.reshape(time_weights, (1, num_filters, memory_size))
        time_weighted = relax.op.multiply(state_3d, tw_3d)
        time_output = relax.op.sum(time_weighted, axis=-1, keepdims=False)
        reduced = relax.op.reshape(time_output, (batch_size, num_units, rank))
        result = relax.op.sum(reduced, axis=-1, keepdims=False)

        # Add bias if present
        if input_tensors[3].tensor_idx != -1:
            bias_expr = self.get_tensor_expr(input_tensors[3])
            result = relax.op.add(result, bias_expr)

        result = self.convert_fused_activation_function(result, fused_activation_fn)

        # Update state tensor in the expression table for subsequent steps.
        # SVDF state is a FIFO ring-buffer: shift left by 1, append new feat.
        feat_3d = relax.op.expand_dims(feat, axis=-1)
        if memory_size > 1:
            shifted_state = relax.op.strided_slice(
                state_3d, axes=[2], begin=[1], end=[int(memory_size)]
            )
            new_state_3d = relax.op.concat([shifted_state, feat_3d], axis=2)
        else:
            new_state_3d = feat_3d
        new_state = relax.op.reshape(new_state_3d, (batch_size, num_filters * memory_size))
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[4].tensor_idx),
            new_state,
            force_override=True,
        )

        return result

    def convert_unidirectional_sequence_rnn(self, op):
        """Convert TFLite UNIDIRECTIONAL_SEQUENCE_RNN.

        Inputs (5 tensors):
          [0] input          [batch, time, input_size]  (or [time, batch, input_size] if time_major)
          [1] input_weights  [num_units, input_size]
          [2] recurrent_weights [num_units, num_units]
          [3] bias           [num_units]
          [4] hidden_state   [batch, num_units]  (variable, zero-initialised)

        Output:
          [0] output  [batch, time, num_units]

        Cell equation:
          h_t = fused_activation(x_t @ W.T + h_{t-1} @ Wr.T + b)
        """
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.SequenceRNNOptions import SequenceRNNOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "TFLite quantized UNIDIRECTIONAL_SEQUENCE_RNN is not supported yet."
            )

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 5, "input tensors length should be 5"

        input_tensor = input_tensors[0]
        weights_tensor = input_tensors[1]
        recurrent_tensor = input_tensors[2]
        bias_tensor = input_tensors[3]
        hidden_state_tensor = input_tensors[4]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) >= 1, "output tensors length should be at least 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.SequenceRNNOptions
        op_options = op.BuiltinOptions()
        seq_rnn_options = SequenceRNNOptions()
        seq_rnn_options.Init(op_options.Bytes, op_options.Pos)
        time_major = seq_rnn_options.TimeMajor()
        fused_activation_fn = seq_rnn_options.FusedActivationFunction()

        # Constant weight/bias expressions.
        weights_expr = self.get_tensor_expr(weights_tensor)  # [num_units, input_size]
        recurrent_expr = self.get_tensor_expr(recurrent_tensor)  # [num_units, num_units]

        # bias is optional (tensor_idx == -1 when absent); default to zeros.
        if bias_tensor.tensor_idx != -1:
            bias_expr = self.get_tensor_expr(bias_tensor)  # [num_units]
        else:
            num_units = int(self.get_tensor_shape(weights_tensor)[0])
            bias_dtype = self.get_tensor_type_str(weights_tensor.tensor.Type())
            bias_expr = relax.op.zeros((num_units,), dtype=bias_dtype)

        # Transpose to [input_size, num_units] and [num_units, num_units] for x @ W.T.
        w_t = relax.op.permute_dims(weights_expr)
        wr_t = relax.op.permute_dims(recurrent_expr)

        # Resolve the input expression; normalise to batch-major [batch, time, input_size].
        # Only the time dimension must be static (needed for unrolling); batch may be dynamic.
        in_expr = self.get_tensor_expr(input_tensor)
        in_shape = self.get_tensor_shape(input_tensor)
        if time_major:
            in_expr = relax.op.permute_dims(in_expr, [1, 0, 2])
            num_steps = int(in_shape[0])
        else:
            num_steps = int(in_shape[1])

        # Initial hidden state: use the model's tensor value when available (non-zero init or
        # graph input), otherwise fall back to zeros for the common variable-tensor case.
        h_dtype = self.get_tensor_type_str(hidden_state_tensor.tensor.Type())
        if self.has_expr(hidden_state_tensor.tensor_idx) or (
            hidden_state_tensor.buffer is not None and hidden_state_tensor.buffer.DataLength() > 0
        ):
            h = self.get_tensor_expr(hidden_state_tensor)
        else:
            h_shape = tuple(to_int_list(self.get_tensor_shape(hidden_state_tensor)))
            h = relax.op.zeros(h_shape, dtype=h_dtype)

        # Unroll over the time axis.
        # relax.op.split with 1 section returns the tensor directly; handle uniformly.
        if num_steps == 1:
            steps = [relax.op.squeeze(in_expr, axis=[1])]
        else:
            splits = relax.op.split(in_expr, num_steps, axis=1)
            steps = [relax.op.squeeze(splits[i], axis=[1]) for i in range(num_steps)]

        outputs = []
        for x_t in steps:  # x_t: [batch, input_size]
            gates = relax.op.add(
                relax.op.add(relax.op.matmul(x_t, w_t), relax.op.matmul(h, wr_t)),
                bias_expr,
            )
            h = self.convert_fused_activation_function(gates, fused_activation_fn)
            outputs.append(h)

        # Stack timestep outputs: [batch, time, num_units].
        return relax.op.stack(outputs, axis=1)

    def convert_unidirectional_sequence_lstm(self, op):
        """Convert TFLite UNIDIRECTIONAL_SEQUENCE_LSTM.

        Inputs (24 tensors, same layout as single-step LSTM):
          [0]  input                       [batch, time, input_size]
          [1]  input_to_input_weights      [num_units, input_size]   (optional)
          [2]  input_to_forget_weights     [num_units, input_size]
          [3]  input_to_cell_weights       [num_units, input_size]
          [4]  input_to_output_weights     [num_units, input_size]
          [5]  recurrent_to_input_weights  [num_units, num_units]   (optional)
          [6]  recurrent_to_forget_weights [num_units, num_units]
          [7]  recurrent_to_cell_weights   [num_units, num_units]
          [8]  recurrent_to_output_weights [num_units, num_units]
          [9]  cell_to_input_weights       [num_units]              (optional)
          [10] cell_to_forget_weights      [num_units]              (optional)
          [11] cell_to_output_weights      [num_units]              (optional)
          [12] input_gate_bias             [num_units]              (optional)
          [13] forget_gate_bias            [num_units]
          [14] cell_gate_bias              [num_units]
          [15] output_gate_bias            [num_units]
          [16] projection_weights          [num_units, num_units]   (optional)
          [17] projection_bias             [num_units]              (optional)
          [18] output_state                [batch, num_units]       (variable)
          [19] cell_state                  [batch, num_units]       (variable)
          [20-23] optional layer norm weights

        Output:
          [0] output  [batch, time, num_units]

        Uses coupled input-forget gate (i = 1 - f) for the FULL kernel.
        """
        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.UnidirectionalSequenceLSTMOptions import UnidirectionalSequenceLSTMOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "TFLite quantized UNIDIRECTIONAL_SEQUENCE_LSTM is not supported yet."
            )

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 24, (
            f"input tensors length should be 24, got {len(input_tensors)}"
        )

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) >= 1, "output tensors length should be at least 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.UnidirectionalSequenceLSTMOptions
        op_options = op.BuiltinOptions()
        lstm_opts = UnidirectionalSequenceLSTMOptions()
        lstm_opts.Init(op_options.Bytes, op_options.Pos)
        time_major = lstm_opts.TimeMajor()
        fused_activation_fn = lstm_opts.FusedActivationFunction()
        cell_clip = lstm_opts.CellClip()
        proj_clip = lstm_opts.ProjClip()

        # Only coupled input-forget gate is supported.
        if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx != -1:
            raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [9, 10, 11]):
            raise tvm.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [16, 17]):
            raise tvm.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [20, 21, 22, 23]):
            raise tvm.error.OpNotImplemented("TFLite layer-norm LSTM is not supported yet.")

        # Weights (transposed once outside the loop).
        w_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[2]))
        w_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[3]))
        w_o_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[4]))
        r_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[6]))
        r_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[7]))
        r_o_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[8]))

        # Biases.
        b_f = self.get_tensor_expr(input_tensors[13])
        b_c = self.get_tensor_expr(input_tensors[14])
        b_o = self.get_tensor_expr(input_tensors[15])

        # Initial states.
        h = self.get_tensor_expr(input_tensors[18])
        c = self.get_tensor_expr(input_tensors[19])

        # Resolve the input expression; normalise to batch-major [batch, time, input_size].
        in_expr = self.get_tensor_expr(input_tensors[0])
        in_shape = self.get_tensor_shape(input_tensors[0])
        if time_major:
            in_expr = relax.op.permute_dims(in_expr, [1, 0, 2])
            num_steps = int(in_shape[0])
        else:
            num_steps = int(in_shape[1])

        # Unroll over the time axis.
        if num_steps == 1:
            steps = [relax.op.squeeze(in_expr, axis=[1])]
        else:
            splits = relax.op.split(in_expr, num_steps, axis=1)
            steps = [relax.op.squeeze(splits[i], axis=[1]) for i in range(num_steps)]

        one = relax.const(1.0, "float32")
        outputs = []
        for x_t in steps:
            f = relax.op.sigmoid(
                relax.op.add(
                    relax.op.add(
                        relax.op.matmul(x_t, w_f_t),
                        relax.op.matmul(h, r_f_t),
                    ),
                    b_f,
                )
            )
            i = relax.op.subtract(one, f)
            g = self.convert_fused_activation_function(
                relax.op.add(
                    relax.op.add(relax.op.matmul(x_t, w_c_t), relax.op.matmul(h, r_c_t)),
                    b_c,
                ),
                fused_activation_fn,
            )
            o = relax.op.sigmoid(
                relax.op.add(
                    relax.op.add(
                        relax.op.matmul(x_t, w_o_t),
                        relax.op.matmul(h, r_o_t),
                    ),
                    b_o,
                )
            )

            c_new = relax.op.add(relax.op.multiply(f, c), relax.op.multiply(i, g))
            if cell_clip > 0.0:
                c_new = relax.op.clip(c_new, -cell_clip, cell_clip)

            h_new = relax.op.multiply(
                o, self.convert_fused_activation_function(c_new, fused_activation_fn)
            )
            if proj_clip > 0.0:
                h_new = relax.op.clip(h_new, -proj_clip, proj_clip)
            outputs.append(h_new)
            h, c = h_new, c_new

        h_out = relax.op.stack(outputs, axis=1)
        if time_major:
            h_out = relax.op.permute_dims(h_out, [1, 0, 2])

        # Update state tensors in the expression table for subsequent ops.
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[18].tensor_idx),
            h,
            force_override=True,
        )
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[19].tensor_idx),
            c,
            force_override=True,
        )

        return h_out

    def convert_bidirectional_sequence_rnn(self, op):
        """Convert TFLite BIDIRECTIONAL_SEQUENCE_RNN.

        Inputs (9 tensors, aux_input not supported):
          [0] input                [batch, time, input_size]
          [1] fw_weights           [num_units, input_size]
          [2] fw_recurrent_weights [num_units, num_units]
          [3] fw_bias              [num_units]
          [4] fw_hidden_state      [batch, num_units]         (variable)
          [5] bw_weights           [num_units, input_size]
          [6] bw_recurrent_weights [num_units, num_units]
          [7] bw_bias              [num_units]
          [8] bw_hidden_state      [batch, num_units]         (variable)

        Output (merge_outputs=True):
          [0] output  [batch, time, 2 * num_units]  (fw and bw concatenated)

        Output (merge_outputs=False):
          [0] fw_output  [batch, time, num_units]
          [1] bw_output  [batch, time, num_units]
        """
        from tflite.BidirectionalSequenceRNNOptions import BidirectionalSequenceRNNOptions
        from tflite.BuiltinOptions import BuiltinOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "TFLite quantized BIDIRECTIONAL_SEQUENCE_RNN is not supported yet."
            )

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 12, (
            f"input tensors length should be 12, got {len(input_tensors)}"
        )

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) >= 1, "output tensors length should be at least 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.BidirectionalSequenceRNNOptions
        op_options = op.BuiltinOptions()
        rnn_opts = BidirectionalSequenceRNNOptions()
        rnn_opts.Init(op_options.Bytes, op_options.Pos)
        time_major = rnn_opts.TimeMajor()
        fused_activation_fn = rnn_opts.FusedActivationFunction()
        merge_outputs = rnn_opts.MergeOutputs()
        if any(input_tensors[idx].tensor_idx != -1 for idx in [9, 10, 11]):
            raise tvm.error.OpNotImplemented(
                "TFLite BIDIRECTIONAL_SEQUENCE_RNN aux input is not supported yet."
            )

        # Forward weights and biases.
        fw_weights_expr = self.get_tensor_expr(input_tensors[1])
        fw_recurrent_expr = self.get_tensor_expr(input_tensors[2])
        fw_bias_expr = self.get_tensor_expr(input_tensors[3])
        fw_w_t = relax.op.permute_dims(fw_weights_expr)
        fw_wr_t = relax.op.permute_dims(fw_recurrent_expr)

        # Backward weights and biases.
        bw_weights_expr = self.get_tensor_expr(input_tensors[5])
        bw_recurrent_expr = self.get_tensor_expr(input_tensors[6])
        bw_bias_expr = self.get_tensor_expr(input_tensors[7])
        bw_w_t = relax.op.permute_dims(bw_weights_expr)
        bw_wr_t = relax.op.permute_dims(bw_recurrent_expr)

        # Resolve the input expression; normalise to batch-major [batch, time, input_size].
        in_expr = self.get_tensor_expr(input_tensors[0])
        in_shape = self.get_tensor_shape(input_tensors[0])
        if time_major:
            in_expr = relax.op.permute_dims(in_expr, [1, 0, 2])
            num_steps = int(in_shape[0])
        else:
            num_steps = int(in_shape[1])

        # Initial hidden states.
        def _get_hidden_state(tensor):
            if self.has_expr(tensor.tensor_idx) or (
                tensor.buffer is not None and tensor.buffer.DataLength() > 0
            ):
                return self.get_tensor_expr(tensor)
            dtype = self.get_tensor_type_str(tensor.tensor.Type())
            h_shape = tuple(to_int_list(self.get_tensor_shape(tensor)))
            return relax.op.zeros(h_shape, dtype=dtype)

        fw_h = _get_hidden_state(input_tensors[4])
        bw_h = _get_hidden_state(input_tensors[8])

        # Unroll over the time axis.
        if num_steps == 1:
            steps = [relax.op.squeeze(in_expr, axis=[1])]
        else:
            splits = relax.op.split(in_expr, num_steps, axis=1)
            steps = [relax.op.squeeze(splits[i], axis=[1]) for i in range(num_steps)]

        # Forward pass.
        fw_outputs = []
        for x_t in steps:
            gates = relax.op.add(
                relax.op.add(relax.op.matmul(x_t, fw_w_t), relax.op.matmul(fw_h, fw_wr_t)),
                fw_bias_expr,
            )
            fw_h = self.convert_fused_activation_function(gates, fused_activation_fn)
            fw_outputs.append(fw_h)

        # Backward pass (process steps in reverse).
        bw_outputs = []
        for x_t in reversed(steps):
            gates = relax.op.add(
                relax.op.add(relax.op.matmul(x_t, bw_w_t), relax.op.matmul(bw_h, bw_wr_t)),
                bw_bias_expr,
            )
            bw_h = self.convert_fused_activation_function(gates, fused_activation_fn)
            bw_outputs.append(bw_h)
        bw_outputs.reverse()

        fw_stacked = relax.op.stack(fw_outputs, axis=1)  # [batch, time, num_units]
        bw_stacked = relax.op.stack(bw_outputs, axis=1)  # [batch, time, num_units]
        if time_major:
            fw_stacked = relax.op.permute_dims(fw_stacked, [1, 0, 2])
            bw_stacked = relax.op.permute_dims(bw_stacked, [1, 0, 2])

        # Update state tensors in the expression table for subsequent ops.
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[4].tensor_idx),
            fw_h,
            force_override=True,
        )
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[8].tensor_idx),
            bw_h,
            force_override=True,
        )

        if merge_outputs:
            return relax.op.concat([fw_stacked, bw_stacked], axis=-1)
        else:
            return relax.Tuple([fw_stacked, bw_stacked])

    def convert_bidirectional_sequence_lstm(self, op):
        """Convert TFLite BIDIRECTIONAL_SEQUENCE_LSTM.

        Inputs (48 tensors, indices 0-17 forward LSTM, 18-34 backward LSTM, 35-38 states,
        39-47 optional aux inputs, which are not supported):

        Forward LSTM cell (indices 0-17, same layout as single-step LSTM):
          [0]  input (shared)              [batch, time, input_size]
          [1]  fw_input_to_input_weights   (optional)
          [2]  fw_input_to_forget_weights
          [3]  fw_input_to_cell_weights
          [4]  fw_input_to_output_weights
          [5]  fw_recurrent_to_input_wts   (optional)
          [6]  fw_recurrent_to_forget_wts
          [7]  fw_recurrent_to_cell_wts
          [8]  fw_recurrent_to_output_wts
          [9-11] fw cell_to_*_weights      (optional, not supported)
          [12] fw_input_gate_bias          (optional)
          [13] fw_forget_gate_bias
          [14] fw_cell_gate_bias
          [15] fw_output_gate_bias
          [16] fw_projection_weights       (optional, not supported)
          [17] fw_projection_bias          (optional, not supported)

        Backward LSTM cell (indices 18-34, same layout as fw):
          [19] bw_input_to_forget_weights
          [20] bw_input_to_cell_weights
          [21] bw_input_to_output_weights
          [23] bw_recurrent_to_forget_wts
          [24] bw_recurrent_to_cell_wts
          [25] bw_recurrent_to_output_wts
          [30] bw_forget_gate_bias
          [31] bw_cell_gate_bias
          [32] bw_output_gate_bias

        State tensors:
          [35] fw_activation_state  [batch, num_units]
          [36] fw_cell_state        [batch, num_units]
          [37] bw_activation_state  [batch, num_units]
          [38] bw_cell_state        [batch, num_units]

        Output (merge_outputs=True):
          [0] output  [batch, time, 2 * num_units]

        Output (merge_outputs=False):
          [0] fw_output  [batch, time, num_units]
          [1] bw_output  [batch, time, num_units]
        """
        from tflite.BidirectionalSequenceLSTMOptions import BidirectionalSequenceLSTMOptions
        from tflite.BuiltinOptions import BuiltinOptions

        if self.is_quantized(op):
            raise tvm.error.OpNotImplemented(
                "TFLite quantized BIDIRECTIONAL_SEQUENCE_LSTM is not supported yet."
            )

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 48, (
            f"input tensors length should be 48, got {len(input_tensors)}"
        )

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) >= 1, "output tensors length should be at least 1"

        assert op.BuiltinOptionsType() == BuiltinOptions.BidirectionalSequenceLSTMOptions
        op_options = op.BuiltinOptions()
        lstm_opts = BidirectionalSequenceLSTMOptions()
        lstm_opts.Init(op_options.Bytes, op_options.Pos)
        time_major = lstm_opts.TimeMajor()
        fused_activation_fn = lstm_opts.FusedActivationFunction()
        merge_outputs = lstm_opts.MergeOutputs()
        cell_clip = lstm_opts.CellClip()
        proj_clip = lstm_opts.ProjClip()

        # ── Forward LSTM weights (transposed once outside the loop) ──
        if input_tensors[1].tensor_idx != -1 or input_tensors[5].tensor_idx != -1:
            raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [9, 10, 11]):
            raise tvm.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [16, 17]):
            raise tvm.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")

        fw_w_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[2]))
        fw_w_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[3]))
        fw_w_o_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[4]))
        fw_r_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[6]))
        fw_r_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[7]))
        fw_r_o_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[8]))
        fw_b_f = self.get_tensor_expr(input_tensors[13])
        fw_b_c = self.get_tensor_expr(input_tensors[14])
        fw_b_o = self.get_tensor_expr(input_tensors[15])

        # ── Backward LSTM weights (transposed once outside the loop) ──
        if input_tensors[18].tensor_idx != -1 or input_tensors[22].tensor_idx != -1:
            raise tvm.error.OpNotImplemented("Only coupled input-forget LSTM is supported.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [26, 27, 28]):
            raise tvm.error.OpNotImplemented("TFLite peephole LSTM is not supported yet.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in [33, 34]):
            raise tvm.error.OpNotImplemented("TFLite projection LSTM is not supported yet.")
        if any(input_tensors[idx].tensor_idx != -1 for idx in range(39, 48)):
            raise tvm.error.OpNotImplemented(
                "TFLite BIDIRECTIONAL_SEQUENCE_LSTM aux input is not supported yet."
            )

        bw_w_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[19]))
        bw_w_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[20]))
        bw_w_o_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[21]))
        bw_r_f_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[23]))
        bw_r_c_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[24]))
        bw_r_o_t = relax.op.permute_dims(self.get_tensor_expr(input_tensors[25]))
        bw_b_f = self.get_tensor_expr(input_tensors[30])
        bw_b_c = self.get_tensor_expr(input_tensors[31])
        bw_b_o = self.get_tensor_expr(input_tensors[32])

        # ── Initial states ──
        fw_h = self.get_tensor_expr(input_tensors[35])
        fw_c = self.get_tensor_expr(input_tensors[36])
        bw_h = self.get_tensor_expr(input_tensors[37])
        bw_c = self.get_tensor_expr(input_tensors[38])

        # ── Unroll input ──
        in_expr = self.get_tensor_expr(input_tensors[0])
        in_shape = self.get_tensor_shape(input_tensors[0])
        if time_major:
            in_expr = relax.op.permute_dims(in_expr, [1, 0, 2])
            num_steps = int(in_shape[0])
        else:
            num_steps = int(in_shape[1])

        if num_steps == 1:
            steps = [relax.op.squeeze(in_expr, axis=[1])]
        else:
            splits = relax.op.split(in_expr, num_steps, axis=1)
            steps = [relax.op.squeeze(splits[i], axis=[1]) for i in range(num_steps)]

        one = relax.const(1.0, "float32")

        def _lstm_step(x_t, h, c, w_f_t, w_c_t, w_o_t, r_f_t, r_c_t, r_o_t, b_f, b_c, b_o):
            """Single LSTM step with coupled input-forget gate."""
            f = relax.op.sigmoid(
                relax.op.add(
                    relax.op.add(
                        relax.op.matmul(x_t, w_f_t),
                        relax.op.matmul(h, r_f_t),
                    ),
                    b_f,
                )
            )
            i = relax.op.subtract(one, f)
            g = self.convert_fused_activation_function(
                relax.op.add(
                    relax.op.add(relax.op.matmul(x_t, w_c_t), relax.op.matmul(h, r_c_t)),
                    b_c,
                ),
                fused_activation_fn,
            )
            o = relax.op.sigmoid(
                relax.op.add(
                    relax.op.add(
                        relax.op.matmul(x_t, w_o_t),
                        relax.op.matmul(h, r_o_t),
                    ),
                    b_o,
                )
            )
            c_new = relax.op.add(relax.op.multiply(f, c), relax.op.multiply(i, g))
            if cell_clip > 0.0:
                c_new = relax.op.clip(c_new, -cell_clip, cell_clip)
            h_new = relax.op.multiply(
                o, self.convert_fused_activation_function(c_new, fused_activation_fn)
            )
            if proj_clip > 0.0:
                h_new = relax.op.clip(h_new, -proj_clip, proj_clip)
            return h_new, c_new

        # ── Forward pass ──
        fw_outputs = []
        for x_t in steps:
            fw_h, fw_c = _lstm_step(
                x_t,
                fw_h,
                fw_c,
                fw_w_f_t,
                fw_w_c_t,
                fw_w_o_t,
                fw_r_f_t,
                fw_r_c_t,
                fw_r_o_t,
                fw_b_f,
                fw_b_c,
                fw_b_o,
            )
            fw_outputs.append(fw_h)

        # ── Backward pass ──
        bw_outputs = []
        for x_t in reversed(steps):
            bw_h, bw_c = _lstm_step(
                x_t,
                bw_h,
                bw_c,
                bw_w_f_t,
                bw_w_c_t,
                bw_w_o_t,
                bw_r_f_t,
                bw_r_c_t,
                bw_r_o_t,
                bw_b_f,
                bw_b_c,
                bw_b_o,
            )
            bw_outputs.append(bw_h)
        bw_outputs.reverse()

        fw_stacked = relax.op.stack(fw_outputs, axis=1)
        bw_stacked = relax.op.stack(bw_outputs, axis=1)
        if time_major:
            fw_stacked = relax.op.permute_dims(fw_stacked, [1, 0, 2])
            bw_stacked = relax.op.permute_dims(bw_stacked, [1, 0, 2])

        # Update state tensors in the expression table for subsequent ops.
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[35].tensor_idx),
            fw_h,
            force_override=True,
        )
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[36].tensor_idx),
            fw_c,
            force_override=True,
        )
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[37].tensor_idx),
            bw_h,
            force_override=True,
        )
        self.exp_tab.set_expr(
            get_tensor_name(self.subgraph, input_tensors[38].tensor_idx),
            bw_c,
            force_override=True,
        )

        if merge_outputs:
            return relax.op.concat([fw_stacked, bw_stacked], axis=-1)
        else:
            return relax.Tuple([fw_stacked, bw_stacked])

    def convert_batch_to_space_nd(self, op):
        """batch_to_space_nd implementation."""

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        in_expr = self.get_expr(input_tensor_idx)

        block_shape = to_int_list(self.get_tensor_value(input_tensors[1]))
        crops = self.get_tensor_value(input_tensors[2])
        crop_begin = to_int_list(crops[:, 0])
        crop_end = to_int_list(crops[:, 1])

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        out = relax.op.call_dps_packed(
            "topi.nn.batch_to_space_nd",
            (
                in_expr,
                relax.ShapeExpr(block_shape),
                relax.ShapeExpr(crop_begin),
                relax.ShapeExpr(crop_end),
            ),
            out_ty=relax.TensorType(output_shape, output_dtype),
        )

        return out

    def convert_broadcast_to(self, op):
        """Convert TFLite BROADCAST_TO"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        data = self.get_tensor_expr(input_tensors[0])
        shape_tensor = input_tensors[1]
        if self.has_expr(shape_tensor.tensor_idx):
            shape_expr = self.get_expr(shape_tensor.tensor_idx)
            shape = self.bb.emit(relax.op.tensor_to_shape(shape_expr))
        else:
            shape = to_int_list(self.get_tensor_value(shape_tensor))
        return relax.op.broadcast_to(data, shape)

    def convert_embedding_lookup(self, op):
        """Convert TFLite EMBEDDING_LOOKUP"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"
        params = self.get_tensor_expr(input_tensors[0])
        indices_tensor = input_tensors[1]
        if self.has_expr(indices_tensor.tensor_idx):
            indices = relax.op.astype(self.get_expr(indices_tensor.tensor_idx), "int32")
        else:
            indices = self.get_tensor_expr(indices_tensor)
        return relax.op.take(params, indices, axis=0)

    def convert_embedding_lookup_sparse(self, op):
        """Convert TFLite EMBEDDING_LOOKUP_SPARSE."""
        from tflite.CombinerType import CombinerType
        from tflite.EmbeddingLookupSparseOptions import EmbeddingLookupSparseOptions
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 5, "EMBEDDING_LOOKUP_SPARSE should have 5 input tensors"
        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "EMBEDDING_LOOKUP_SPARSE should have 1 output tensor"

        ids_tensor, indices_tensor, dense_shape_tensor, weights_tensor, params_tensor = (
            input_tensors
        )
        output_tensor = output_tensors[0]

        for tensor in input_tensors:
            assert not tensor.qnn_params, "Quantized input is not expected."

        assert ids_tensor.tensor.Type() == TensorType.INT32
        assert indices_tensor.tensor.Type() == TensorType.INT32
        assert dense_shape_tensor.tensor.Type() == TensorType.INT32
        assert weights_tensor.tensor.Type() == TensorType.FLOAT32
        assert params_tensor.tensor.Type() == TensorType.FLOAT32
        assert output_tensor.tensor.Type() == TensorType.FLOAT32

        ids_shape = to_int_list(self.get_tensor_shape(ids_tensor))
        indices_shape = to_int_list(self.get_tensor_shape(indices_tensor))
        dense_shape_shape = to_int_list(self.get_tensor_shape(dense_shape_tensor))
        weights_shape = to_int_list(self.get_tensor_shape(weights_tensor))
        params_shape = to_int_list(self.get_tensor_shape(params_tensor))

        assert len(ids_shape) == 1, "EMBEDDING_LOOKUP_SPARSE ids must be rank 1"
        assert len(indices_shape) == 2, "EMBEDDING_LOOKUP_SPARSE indices must be rank 2"
        assert len(dense_shape_shape) == 1, "EMBEDDING_LOOKUP_SPARSE dense_shape must be rank 1"
        assert len(weights_shape) == 1, "EMBEDDING_LOOKUP_SPARSE weights must be rank 1"
        assert len(params_shape) >= 2, "EMBEDDING_LOOKUP_SPARSE params must be rank >= 2"
        assert indices_shape[0] == ids_shape[0], (
            "EMBEDDING_LOOKUP_SPARSE ids and indices must agree on lookup count"
        )
        assert weights_shape[0] == ids_shape[0], (
            "EMBEDDING_LOOKUP_SPARSE ids and weights must agree on lookup count"
        )

        if self.has_expr(dense_shape_tensor.tensor_idx):
            raise tvm.error.OpNotImplemented(
                "TFLite EMBEDDING_LOOKUP_SPARSE with runtime dense_shape is not supported."
            )

        dense_shape = to_int_list(self.get_tensor_value(dense_shape_tensor))
        lookup_rank = indices_shape[1]
        assert len(dense_shape) == lookup_rank, (
            "EMBEDDING_LOOKUP_SPARSE dense_shape length must match indices width"
        )
        assert lookup_rank >= 1, "EMBEDDING_LOOKUP_SPARSE indices width must be positive"
        if not self.has_expr(ids_tensor.tensor_idx):
            ids_value = self.get_tensor_value(ids_tensor)
            if np.any(ids_value < 0):
                raise tvm.error.OpNotImplemented(
                    "TFLite EMBEDDING_LOOKUP_SPARSE with negative ids is not supported."
                )

        params = self.get_tensor_expr(params_tensor)
        ids = self.get_tensor_expr(ids_tensor)
        weights = self.get_tensor_expr(weights_tensor)
        indices = self.get_tensor_expr(indices_tensor)

        ids = relax.op.astype(ids, "int32")
        lookup = relax.op.take(params, ids, axis=0)

        embedding_tail_shape = params_shape[1:]
        output_prefix_shape = dense_shape[:-1]
        output_shape = output_prefix_shape + embedding_tail_shape

        # Aggregation buckets are defined by every sparse index dimension except the last one.
        bucket_indices = relax.op.strided_slice(indices, axes=[1], begin=[0], end=[lookup_rank - 1])

        weight_expand_shape = [ids_shape[0]] + [1] * len(embedding_tail_shape)
        weighted_lookup = relax.op.multiply(lookup, relax.op.reshape(weights, weight_expand_shape))

        value_base = relax.const(np.zeros(output_shape, dtype=np.float32), "float32")
        summed_lookup = relax.op.scatter_nd(value_base, bucket_indices, weighted_lookup, "add")

        op_options = op.BuiltinOptions()
        sparse_options = EmbeddingLookupSparseOptions()
        sparse_options.Init(op_options.Bytes, op_options.Pos)
        combiner = sparse_options.Combiner()
        if combiner == CombinerType.SUM:
            return summed_lookup

        count_shape = output_prefix_shape
        count_base = relax.const(np.zeros(count_shape, dtype=np.float32), "float32")
        bucket_count_updates = relax.const(np.ones(ids_shape, dtype=np.float32), "float32")
        bucket_counts = relax.op.scatter_nd(count_base, bucket_indices, bucket_count_updates, "add")
        if combiner == CombinerType.MEAN:
            denominator_updates = weights
        elif combiner == CombinerType.SQRTN:
            denominator_updates = relax.op.multiply(weights, weights)
        else:
            raise tvm.error.OpNotImplemented(
                f"Unsupported TFLite EMBEDDING_LOOKUP_SPARSE combiner value {combiner}"
            )

        denominator = relax.op.scatter_nd(count_base, bucket_indices, denominator_updates, "add")
        if combiner == CombinerType.SQRTN:
            denominator = relax.op.sqrt(denominator)

        broadcast_shape = count_shape + [1] * len(embedding_tail_shape)
        denominator = relax.op.reshape(denominator, broadcast_shape)
        denominator = relax.op.broadcast_to(denominator, output_shape)
        normalized = relax.op.divide(summed_lookup, denominator)
        bucket_counts = relax.op.reshape(bucket_counts, broadcast_shape)
        bucket_counts = relax.op.broadcast_to(bucket_counts, output_shape)
        return relax.op.where(
            relax.op.greater(bucket_counts, relax.const(0.0, "float32")), normalized, value_base
        )

    def convert_batch_matmul(self, op):
        """batch_matmul implementation."""

        from tflite.BatchMatMulOptions import BatchMatMulOptions

        input_tensors = self.get_input_tensors(op)

        assert len(input_tensors) == 2, "two input tensor arguments expected"

        if self.is_quantized(op):
            raise NotImplementedError(
                "Quantized BATCH_MATMUL is not yet supported in the Relax frontend"
            )

        batch_matmul_options = BatchMatMulOptions()
        op_options = op.BuiltinOptions()
        batch_matmul_options.Init(op_options.Bytes, op_options.Pos)

        input_a = self.get_expr(input_tensors[0].tensor_idx)
        input_b = self.get_expr(input_tensors[1].tensor_idx)

        shape_a = list(input_a.ty.shape)
        shape_b = list(input_b.ty.shape)
        rank_a = len(shape_a)
        rank_b = len(shape_b)

        if rank_a > 2 or rank_b > 2:
            # Broadcast batch dimensions
            new_a_shape = [1] * max(0, rank_b - rank_a) + [int(s) for s in shape_a]
            new_b_shape = [1] * max(0, rank_a - rank_b) + [int(s) for s in shape_b]
            max_rank = max(rank_a, rank_b)

            batch_shape = [max(new_a_shape[i], new_b_shape[i]) for i in range(max_rank - 2)]

            a_broadcast = batch_shape + [int(shape_a[-2]), int(shape_a[-1])]
            b_broadcast = batch_shape + [int(shape_b[-2]), int(shape_b[-1])]

            if [int(s) for s in shape_a] != a_broadcast:
                input_a = relax.op.broadcast_to(input_a, a_broadcast)
            if [int(s) for s in shape_b] != b_broadcast:
                input_b = relax.op.broadcast_to(input_b, b_broadcast)

            input_a = self.flatten_to_nd(input_a, 3)
            input_b = self.flatten_to_nd(input_b, 3)

            adj_x = batch_matmul_options.AdjX()
            adj_y = batch_matmul_options.AdjY()

            if adj_x:
                input_a = relax.op.permute_dims(input_a, [0, 2, 1])
            if adj_y:
                input_b = relax.op.permute_dims(input_b, [0, 2, 1])

            output = relax.op.matmul(input_a, input_b)

            # Compute output matmul dims from original shapes
            m_dim = int(shape_a[-1]) if adj_x else int(shape_a[-2])
            n_dim = int(shape_b[-2]) if adj_y else int(shape_b[-1])
            final_shape = [int(s) for s in shape_a[: rank_a - 2]] + [m_dim, n_dim]
            return relax.op.reshape(output, final_shape)

        # rank <= 2: use matmul directly
        if batch_matmul_options.AdjX():
            input_a = relax.op.permute_dims(input_a)
        if batch_matmul_options.AdjY():
            input_b = relax.op.permute_dims(input_b)
        return relax.op.matmul(input_a, input_b)

    def convert_space_to_batch_nd(self, op):
        """space_to_batch_nd implementation."""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"

        input_tensor = input_tensors[0]
        input_tensor_idx = input_tensor.tensor_idx
        in_expr = self.get_expr(input_tensor_idx)

        block_shape = to_int_list(self.get_tensor_value(input_tensors[1]))
        paddings = self.get_tensor_value(input_tensors[2])
        pad_before = to_int_list(paddings[:, 0])
        pad_after = to_int_list(paddings[:, 1])

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        out = relax.op.call_dps_packed(
            "topi.nn.space_to_batch_nd",
            (
                in_expr,
                relax.ShapeExpr(block_shape),
                relax.ShapeExpr(pad_before),
                relax.ShapeExpr(pad_after),
                0.0,
            ),
            out_ty=relax.TensorType(output_shape, output_dtype),
        )

        return out

    def convert_depth_to_space(self, op):
        """Convert TFLite DEPTH_TO_SPACE"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.DepthToSpaceOptions import DepthToSpaceOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        assert op.BuiltinOptionsType() == BuiltinOptions.DepthToSpaceOptions
        op_options = op.BuiltinOptions()
        depth_to_space_options = DepthToSpaceOptions()
        depth_to_space_options.Init(op_options.Bytes, op_options.Pos)
        block_size = depth_to_space_options.BlockSize()

        # TFLite uses NHWC layout: (N, H, W, C) -> (N, H*bs, W*bs, C/(bs*bs))
        input_shape = self.get_tensor_shape(input_tensor)
        n, h, w, c = input_shape
        out_c = c // (block_size**2)
        out = relax.op.reshape(in_expr, (n, h, w, block_size, block_size, out_c))
        out = relax.op.permute_dims(out, [0, 1, 3, 2, 4, 5])
        out = relax.op.reshape(out, (n, h * block_size, w * block_size, out_c))

        return out

    def convert_space_to_depth(self, op):
        """Convert TFLite SPACE_TO_DEPTH"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.SpaceToDepthOptions import SpaceToDepthOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        assert op.BuiltinOptionsType() == BuiltinOptions.SpaceToDepthOptions
        op_options = op.BuiltinOptions()
        space_to_depth_options = SpaceToDepthOptions()
        space_to_depth_options.Init(op_options.Bytes, op_options.Pos)
        block_size = space_to_depth_options.BlockSize()

        # TFLite uses NHWC layout: (N, H, W, C) -> (N, H/bs, W/bs, C*bs*bs)
        input_shape = self.get_tensor_shape(input_tensor)
        n, h, w, c = input_shape
        out = relax.op.reshape(
            in_expr, (n, h // block_size, block_size, w // block_size, block_size, c)
        )
        out = relax.op.permute_dims(out, [0, 1, 3, 2, 4, 5])
        out = relax.op.reshape(
            out, (n, h // block_size, w // block_size, c * block_size * block_size)
        )

        return out

    def convert_sparse_to_dense(self, op):
        """Convert TFLite SPARSE_TO_DENSE"""
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 4, "input tensors length should be 4"

        indices, values = input_tensors[0], input_tensors[2]
        default_value = input_tensors[3]
        output_shape = input_tensors[1]

        for t in input_tensors:
            assert not t.qnn_params, "Quantized input is not expected."

        for t in [indices, output_shape]:
            t_type = t.tensor.Type()
            assert t_type in (TensorType.INT32, TensorType.INT64)

        output_tensors = self.get_output_tensors(op)
        output_tensor = output_tensors[0]
        output_shape_val = to_int_list(self.get_tensor_shape(output_tensor))
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        indices_expr = self.get_tensor_expr(indices)
        values_expr = self.get_tensor_expr(values)
        default_value_expr = self.get_tensor_expr(default_value)
        output_shape_expr = relax.const(list(self.get_tensor_value(output_shape)), "int32")

        out = relax.op.call_dps_packed(
            "topi.sparse_to_dense",
            (indices_expr, output_shape_expr, values_expr, default_value_expr),
            out_ty=relax.TensorType(output_shape_val, output_dtype),
        )

        return out

    def convert_prelu(self, op):
        """Convert TFLite PReLU"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        input_tensor = input_tensors[0]
        alpha_tensor = input_tensors[1]
        data_shape = to_int_list(self.get_tensor_shape(input_tensor))
        in_expr = self.get_tensor_expr(input_tensor)
        alpha_expr = self.get_tensor_expr(alpha_tensor)
        alpha_expr = self.bb.normalize(relax.op.broadcast_to(alpha_expr, data_shape))
        alpha_expr = self.bb.normalize(relax.op.reshape(alpha_expr, [-1]))
        out = relax.op.nn.prelu(_op.reshape(in_expr, [-1]), alpha_expr, axis=0)
        out = relax.op.reshape(out, data_shape)
        return out

    def convert_transpose_conv(self, op):
        """Convert TFLite TRANSPOSE_CONV"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.Padding import Padding
        from tflite.TensorType import TensorType
        from tflite.TransposeConvOptions import TransposeConvOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) >= 3, "input tensors length should be >= 3"

        # Input (data) Tensor. NHWC layout
        input_tensor = input_tensors[2]
        _, _, _, input_c = to_int_list(self.get_tensor_shape(input_tensor))
        # Weights tensor. TFLite uses OHWI layout
        weights_tensor = input_tensors[1]
        out_channels, kernel_h, kernel_w, in_channels = to_int_list(
            self.get_tensor_shape(weights_tensor)
        )

        assert input_c == in_channels, (
            "Input channel in the filter should match to channel in the input"
        )
        # output_shape Tensor. NHWC layout
        output_shape_tensor = input_tensors[0]

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]
        output_tensor_type = output_tensor.tensor.Type()
        output_tensor_type_str = self.get_tensor_type_str(output_tensor_type)

        assert op.BuiltinOptionsType() == BuiltinOptions.TransposeConvOptions
        op_options = op.BuiltinOptions()
        deconv_options = TransposeConvOptions()
        deconv_options.Init(op_options.Bytes, op_options.Pos)

        padding = deconv_options.Padding()
        stride_h = deconv_options.StrideH()
        stride_w = deconv_options.StrideW()
        assert padding in (
            Padding.VALID,
            Padding.SAME,
        ), f"Padding format {padding} is not supported for operator TRANSPOSE_CONV"

        # Data
        in_expr = self.get_expr(input_tensor.tensor_idx)

        # Weights
        weights_tensor_type = weights_tensor.tensor.Type()
        # weights tensor type should be UINT8 (quantization) or FLOAT32
        assert weights_tensor_type in (
            TensorType.INT8,
            TensorType.UINT8,
            TensorType.FLOAT32,
        )
        weight_tensor_type_str = self.get_tensor_type_str(weights_tensor_type)

        if self.has_expr(weights_tensor.tensor_idx):
            weight_expr_iohw = self.get_expr(weights_tensor.tensor_idx)
            weight_expr_iohw = relax.op.permute_dims(weight_expr_iohw, axes=(3, 0, 1, 2))
        else:
            weight_value_ohwi = self.get_tensor_value_or_prefetched(weights_tensor)
            # Relax kernel_layout should be OIHW
            # Relax weights layout should be different from kernel_layout - it should be IOHW
            weight_value_iohw = np.transpose(weight_value_ohwi, (3, 0, 1, 2))
            weight_expr_iohw = self.exp_tab.new_const(
                weight_value_iohw,
                dtype=weight_tensor_type_str,
                source_name=weights_tensor.tensor.Name(),
            )

        # Output shape value
        output_shape_value = self.get_tensor_value(output_shape_tensor)
        # Relax expects filter output channel to match to output tensor channel.
        assert out_channels == output_shape_value[3], (
            "Output channel in the filter should match to channel in the output_shape"
        )

        if padding == Padding.SAME:
            output_h, output_w = output_shape_value[1], output_shape_value[2]
            pad_top, pad_bottom = get_pad_value(output_h, kernel_h, stride_h)
            pad_left, pad_right = get_pad_value(output_w, kernel_w, stride_w)
            padding = (pad_top, pad_left, pad_bottom, pad_right)
        else:
            padding = (0, 0, 0, 0)

        if input_tensor.qnn_params:
            in_f32 = self.dequantize(in_expr, input_tensor)
            weight_axis = weights_tensor.qnn_params["axis"]
            if weight_axis != 0:
                raise tvm.error.OpAttributeInvalid(
                    f"TransposeConv weight QuantizedDimension() must be 0 "
                    f"(output-channel axis in OHWI layout), got {weight_axis}"
                )
            w_f32 = relax.op.dequantize(
                weight_expr_iohw,
                scale=weights_tensor.qnn_params["scale"],
                zero_point=weights_tensor.qnn_params["zero_point"],
                axis=1,
            )
            out = relax.op.nn.conv2d_transpose(
                in_f32,
                w_f32,
                strides=(stride_h, stride_w),
                padding=padding,
                data_layout="NHWC",
                kernel_layout="IOHW",
                out_dtype="float32",
            )
        else:
            out = relax.op.nn.conv2d_transpose(
                in_expr,
                weight_expr_iohw,
                strides=(stride_h, stride_w),
                padding=padding,
                data_layout="NHWC",
                kernel_layout="IOHW",
                out_dtype=output_tensor_type_str,
            )

        # Checking if there is a fused bias
        if len(input_tensors) == 4:
            bias_tensor = input_tensors[3]
            bias_tensor_type = bias_tensor.tensor.Type()
            # bias tensor type should be INT32 (quantization) or FLOAT32
            assert bias_tensor_type in (
                TensorType.INT32,
                TensorType.INT64,
                TensorType.FLOAT32,
            )
            bias_tensor_type_str = self.get_tensor_type_str(bias_tensor_type)
            if self.has_expr(bias_tensor.tensor_idx):
                bias_expr = self.get_expr(bias_tensor.tensor_idx)
            else:
                bias_expr = self.exp_tab.new_const(
                    self.get_tensor_value(bias_tensor),
                    dtype=bias_tensor_type_str,
                    source_name=bias_tensor.tensor.Name(),
                )
            if bias_tensor.qnn_params:
                bias_expr = self.dequantize(bias_expr, bias_tensor)
            elif input_tensor.qnn_params and bias_tensor_type in (
                TensorType.INT32,
                TensorType.INT64,
            ):
                bias_scale = relax.op.multiply(
                    input_tensor.qnn_params["scale"],
                    weights_tensor.qnn_params["scale"],
                )
                bias_expr = relax.op.dequantize(
                    bias_expr,
                    scale=bias_scale,
                    zero_point=relax.const(0, "int32"),
                    axis=0,
                )
            out = relax.op.add(out, bias_expr)

        if output_tensor.qnn_params:
            out = self.quantize(out, output_tensor)
        return out

    def convert_quantize(self, op):
        """Convert TFLite Quantize"""

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]
        input_tensor_type_str = self.get_tensor_type_str(input_tensor.tensor.Type())
        in_expr = self.get_tensor_expr(input_tensor)

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        # The output must be quantized
        assert output_tensor.qnn_params

        # TFLite Quantize op can also act as Requantize op
        if input_tensor_type_str == "float32":
            out = self.quantize(in_expr, output_tensor)
        else:
            in_f32 = self.dequantize(in_expr, input_tensor)
            out = self.quantize(in_f32, output_tensor)
        return out

    def convert_dequantize(self, op):
        """Convert TFLite Dequantize"""
        from tflite.TensorType import TensorType

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = input_tensors[0]

        if input_tensor.tensor.Type() == TensorType.FLOAT16:
            dtype = self.get_tensor_type_str(input_tensor.tensor.Type())
            input_value = self.get_tensor_value(input_tensor)
            in_expr = self.exp_tab.new_const(
                input_value, dtype=dtype, source_name=input_tensor.tensor.Name()
            )
            out = relax.op.astype(in_expr, "float32")
            return out

        in_expr = self.get_expr(input_tensor.tensor_idx)

        # The input must be quantized
        assert input_tensor.qnn_params
        # Dequantize the input.
        out = self.dequantize(in_expr, input_tensor)

        return out

    def convert_dilate(self, op):
        """Convert TFLite DILATE"""
        input_tensors = self.get_input_tensors(op)
        output_tensors = self.get_output_tensors(op)
        assert len(input_tensors) == 3, "input tensors length should be 3"
        assert len(output_tensors) == 1, "output tensors length should be 1"

        in_expr = self.get_tensor_expr(input_tensors[0])
        in_shape = to_int_list(self.get_tensor_shape(input_tensors[0]))
        in_dtype = self.get_tensor_type_str(input_tensors[0].tensor.Type())
        n_dims = len(in_shape)

        dilations_tensor = input_tensors[1]
        padding_expr = self.get_tensor_expr(input_tensors[2])

        # Runtime dilations bind tensor values to TIR Vars for symbolic
        # per-axis math.
        if self.has_expr(dilations_tensor.tensor_idx):
            dilations_expr = self.get_expr(dilations_tensor.tensor_idx)
            dilations_expr = self.bb.match_cast(dilations_expr, relax.TensorType([n_dims], "int32"))
            dilations_int64 = self.bb.normalize(relax.op.astype(dilations_expr, "int64"))
            shape_var = self.bb.emit(relax.op.tensor_to_shape(dilations_int64))
            stride_vars = [tirx.Var(f"dilate_stride_{i}", "int64") for i in range(n_dims)]
            self.bb.match_cast(shape_var, relax.ShapeType(stride_vars))
            strides = stride_vars
        else:
            strides = to_int_list(self.get_tensor_value(dilations_tensor))

        # Per axis: reshape to add a size-1 stride-axis, concat (s-1) padding
        # values along it, reshape to merge axes (length d*s), trim trailing
        # pad to TFLite's output dim formula (d-1)*s + 1.
        result = in_expr
        current_shape = list(in_shape)
        axes = list(range(n_dims))
        ones = [1] * n_dims
        for axis in range(n_dims):
            d = current_shape[axis]
            s = strides[axis]
            expanded_shape = current_shape[: axis + 1] + [1] + current_shape[axis + 1 :]
            expanded = relax.op.reshape(result, expanded_shape)
            pad_shape = list(expanded_shape)
            pad_shape[axis + 1] = s - 1
            pad = relax.op.full(pad_shape, padding_expr, dtype=in_dtype)
            concatted = relax.op.concat([expanded, pad], axis=axis + 1)
            merged_shape = list(current_shape)
            merged_shape[axis] = d * s
            merged = relax.op.reshape(concatted, merged_shape)
            # (d - 1) * s + 1 is the output dim along this axis.
            final_dim = (d - 1) * s + 1
            end = list(merged_shape)
            end[axis] = final_dim
            result = relax.op.strided_slice(
                merged, axes=axes, begin=[0] * n_dims, end=end, strides=ones
            )
            current_shape = list(merged_shape)
            current_shape[axis] = final_dim

        return result

    def convert_detection_postprocess(self, op):
        """Convert TFLite_Detection_PostProcess"""
        flexbuffer = op.CustomOptionsAsNumpy().tobytes()
        custom_options = FlexBufferDecoder(flexbuffer).decode()

        use_regular_nms = bool(custom_options.get("use_regular_nms", False))

        required_attrs = [
            "num_classes",
            "max_detections",
            "detections_per_class",
            "nms_iou_threshold",
            "nms_score_threshold",
            "x_scale",
            "y_scale",
            "w_scale",
            "h_scale",
        ]
        missing_attrs = [key for key in required_attrs if key not in custom_options]
        if missing_attrs:
            raise ValueError(
                "DETECTION_POSTPROCESS custom options miss required attributes: "
                + ", ".join(missing_attrs)
            )

        num_classes = int(custom_options["num_classes"])
        max_detections = int(custom_options["max_detections"])
        detections_per_class = int(custom_options["detections_per_class"])
        iou_threshold = float(custom_options["nms_iou_threshold"])
        score_threshold = float(custom_options["nms_score_threshold"])
        x_scale = float(custom_options["x_scale"])
        y_scale = float(custom_options["y_scale"])
        w_scale = float(custom_options["w_scale"])
        h_scale = float(custom_options["h_scale"])

        if num_classes <= 0:
            raise ValueError("DETECTION_POSTPROCESS requires num_classes > 0.")
        if max_detections <= 0:
            raise ValueError("DETECTION_POSTPROCESS requires max_detections > 0.")
        if detections_per_class <= 0:
            raise ValueError("DETECTION_POSTPROCESS requires detections_per_class > 0.")
        if not 0.0 <= iou_threshold <= 1.0:
            raise ValueError("DETECTION_POSTPROCESS requires nms_iou_threshold in [0, 1].")
        if x_scale <= 0.0 or y_scale <= 0.0 or w_scale <= 0.0 or h_scale <= 0.0:
            raise ValueError("DETECTION_POSTPROCESS requires x/y/w/h_scale to be > 0.")

        inputs = self.get_input_tensors(op)
        assert len(inputs) == 3, "inputs length should be 3"
        cls_pred = self.get_expr(inputs[1].tensor_idx)
        loc_prob = self.get_expr(inputs[0].tensor_idx)
        batch_size = inputs[1].tensor.Shape(0)
        anchor_values = self.get_tensor_value(inputs[2])
        anchor_boxes = len(anchor_values)
        anchor_type = self.get_tensor_type_str(inputs[2].tensor.Type())
        anchor_expr = self.exp_tab.new_const(
            anchor_values, dtype=anchor_type, source_name=inputs[2].tensor.Name()
        )

        if inputs[0].qnn_params:
            loc_prob = self.dequantize(loc_prob, inputs[0])
        if inputs[1].qnn_params:
            cls_pred = self.dequantize(cls_pred, inputs[1])
        if inputs[2].qnn_params:
            anchor_expr = self.dequantize(anchor_expr, inputs[2])

        # loc_prob coords are in yxhw format
        # need to convert to xywh
        loc_coords = relax.op.split(loc_prob, 4, axis=2)
        loc_prob = relax.op.concat(
            [loc_coords[1], loc_coords[0], loc_coords[3], loc_coords[2]], axis=2
        )
        # reshape loc_prob tensor so is can be consumed by
        # multibox_transform_loc
        loc_prob = relax.op.reshape(loc_prob, [batch_size, anchor_boxes * 4])

        # anchor coords are in yxhw format
        # need to convert to ltrb
        anchor_coords = relax.op.split(anchor_expr, 4, axis=1)
        anchor_y = anchor_coords[0]
        anchor_x = anchor_coords[1]
        anchor_h = anchor_coords[2]
        anchor_w = anchor_coords[3]
        plus_half = relax.const(0.5, dtype="float32")
        minus_half = relax.const(-0.5, dtype="float32")
        anchor_l = relax.op.add(anchor_x, relax.op.multiply(anchor_w, minus_half))
        anchor_r = relax.op.add(anchor_x, relax.op.multiply(anchor_w, plus_half))
        anchor_t = relax.op.add(anchor_y, relax.op.multiply(anchor_h, minus_half))
        anchor_b = relax.op.add(anchor_y, relax.op.multiply(anchor_h, plus_half))
        anchor_expr = relax.op.concat([anchor_l, anchor_t, anchor_r, anchor_b], axis=1)
        anchor_expr = relax.op.expand_dims(anchor_expr, 0)

        # attributes for multibox_transform_loc
        multibox_transform_loc_attrs = {}
        multibox_transform_loc_attrs["clip"] = False
        multibox_transform_loc_attrs["threshold"] = 0.0 if use_regular_nms else score_threshold
        multibox_transform_loc_attrs["variances"] = (
            1 / x_scale,
            1 / y_scale,
            1 / w_scale,
            1 / h_scale,
        )
        multibox_transform_loc_attrs["keep_background"] = use_regular_nms

        multibox_res = self.bb.emit(
            relax.op.vision.multibox_transform_loc(
                # reshape cls_pred so it can be consumed by
                # multibox_transform_loc
                relax.op.permute_dims(cls_pred, [0, 2, 1]),
                loc_prob,
                anchor_expr,
                **multibox_transform_loc_attrs,
            )
        )
        transformed_boxes = self.bb.emit(relax.TupleGetItem(multibox_res, 0))
        transformed_scores = self.bb.emit(relax.TupleGetItem(multibox_res, 1))

        if use_regular_nms:
            nms_out = self.bb.emit(
                relax.op.vision.all_class_non_max_suppression(
                    transformed_boxes,
                    transformed_scores,
                    relax.const(detections_per_class, "int64"),
                    relax.const(iou_threshold, "float32"),
                    relax.const(score_threshold, "float32"),
                    output_format="tensorflow",
                )
            )
            selected_indices = self.bb.emit(relax.TupleGetItem(nms_out, 0))
            selected_scores = self.bb.emit(relax.TupleGetItem(nms_out, 1))
            num_detections = self.bb.emit(relax.TupleGetItem(nms_out, 2))
            class_id_from_score = None
        else:
            topk_res = self.bb.emit(
                relax.op.topk(transformed_scores, k=1, axis=1, ret_type="both", largest=True)
            )
            max_scores = self.bb.emit(relax.TupleGetItem(topk_res, 0))
            class_id_from_score = self.bb.emit(relax.TupleGetItem(topk_res, 1))
            nms_out = self.bb.emit(
                relax.op.vision.all_class_non_max_suppression(
                    transformed_boxes,
                    max_scores,
                    relax.const(max_detections, "int64"),
                    relax.const(iou_threshold, "float32"),
                    relax.const(score_threshold, "float32"),
                    output_format="tensorflow",
                )
            )
            selected_indices = self.bb.emit(relax.TupleGetItem(nms_out, 0))
            selected_scores = self.bb.emit(relax.TupleGetItem(nms_out, 1))
            num_detections = self.bb.emit(relax.TupleGetItem(nms_out, 2))
            class_id_from_score = relax.op.squeeze(class_id_from_score, axis=[1])

        selected_score_slots = selected_scores.ty.shape.values[1]
        selected_detection_positions = relax.op.expand_dims(
            relax.op.arange(selected_score_slots, dtype="int64"), axis=0
        )
        selected_valid_detection_mask = relax.op.less(
            selected_detection_positions, relax.op.expand_dims(num_detections, axis=1)
        )
        masked_selected_scores = relax.op.where(
            selected_valid_detection_mask,
            selected_scores,
            relax.const(-1.0, "float32"),
        )
        topk_scores_res = self.bb.emit(
            relax.op.topk(
                masked_selected_scores, k=max_detections, axis=1, ret_type="both", largest=True
            )
        )
        detection_scores = self.bb.emit(relax.TupleGetItem(topk_scores_res, 0))
        top_positions = self.bb.emit(relax.TupleGetItem(topk_scores_res, 1))
        num_detections = relax.op.minimum(
            num_detections, relax.const([max_detections], dtype="int64")
        )
        detection_positions = relax.op.expand_dims(
            relax.op.arange(max_detections, dtype="int64"), axis=0
        )
        valid_detection_mask = relax.op.less(
            detection_positions, relax.op.expand_dims(num_detections, axis=1)
        )
        top_positions_expanded = relax.op.expand_dims(top_positions, axis=2)
        top_positions_for_pairs = relax.op.repeat(top_positions_expanded, 2, axis=2)
        top_index_pairs = relax.op.gather_elements(
            selected_indices, top_positions_for_pairs, axis=1
        )
        top_box_ids = relax.op.squeeze(
            relax.op.strided_slice(top_index_pairs, axes=[2], begin=[1], end=[2]),
            axis=[2],
        )
        top_box_ids_for_gather = relax.op.expand_dims(relax.op.astype(top_box_ids, "int64"), axis=2)
        detection_boxes = relax.op.gather_nd(
            transformed_boxes, top_box_ids_for_gather, batch_dims=1
        )

        if use_regular_nms:
            detection_classes = relax.op.squeeze(
                relax.op.strided_slice(top_index_pairs, axes=[2], begin=[0], end=[1]),
                axis=[2],
            )
            detection_classes = relax.op.astype(detection_classes, "int32")
        else:
            top_box_ids_for_class = relax.op.expand_dims(
                relax.op.astype(top_box_ids, "int64"), axis=2
            )
            detection_classes = relax.op.gather_nd(
                class_id_from_score, top_box_ids_for_class, batch_dims=1
            )

        detection_mask = relax.op.expand_dims(valid_detection_mask, axis=2)
        detection_boxes = relax.op.where(
            detection_mask,
            detection_boxes,
            relax.op.zeros((batch_size, max_detections, 4), dtype="float32"),
        )
        detection_classes = relax.op.where(
            valid_detection_mask,
            detection_classes,
            relax.op.zeros((batch_size, max_detections), dtype="int32"),
        )
        detection_scores = relax.op.where(
            valid_detection_mask,
            detection_scores,
            relax.op.zeros((batch_size, max_detections), dtype="float32"),
        )
        detection_classes = relax.op.astype(detection_classes, "float32")
        num_detections = relax.op.astype(num_detections, "float32")
        return relax.Tuple([detection_boxes, detection_classes, detection_scores, num_detections])

    def convert_nms_v4(self, op):
        """Convert TFLite NonMaxSuppressionV4"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 5, "input tensor length should be 5"

        boxes = self.get_tensor_expr(input_tensors[0])
        scores = self.get_tensor_expr(input_tensors[1])

        max_output_size = self.get_tensor_value(input_tensors[2])
        iou_threshold = self.get_tensor_value(input_tensors[3])
        score_threshold = self.get_tensor_value(input_tensors[4])

        if isinstance(max_output_size, np.ndarray):
            assert max_output_size.size == 1, "only one value is expected."
            max_output_size = int(max_output_size)

        if isinstance(iou_threshold, np.ndarray):
            assert iou_threshold.size == 1, "only one value is expected."
            iou_threshold = float(iou_threshold)

        if isinstance(score_threshold, np.ndarray):
            assert score_threshold.size == 1, "only one value is expected."
            score_threshold = float(score_threshold)

        scores_expand = relax.op.expand_dims(scores, axis=-1)
        data = relax.op.concat([scores_expand, boxes], axis=-1)
        data = relax.op.expand_dims(data, axis=0)

        valid_counts_ret = relax.op.vision.get_valid_counts(
            data, score_threshold=score_threshold, id_index=-1, score_index=0
        )
        count = valid_counts_ret[0]
        data = valid_counts_ret[1]
        indices = valid_counts_ret[2]

        nms_ret = relax.op.vision.non_max_suppression(
            data=data,
            valid_count=count,
            indices=indices,
            max_output_size=max_output_size,
            iou_threshold=iou_threshold,
            force_suppress=True,
            top_k=-1,
            coord_start=1,
            score_index=0,
            id_index=-1,
            return_indices=True,
            invalid_to_bottom=False,
        )

        selected_indices = relax.op.squeeze(nms_ret[0], axis=[0])
        selected_indices = relax.op.strided_slice(
            selected_indices, axes=[0], begin=[0], end=[max_output_size]
        )
        num_valid = relax.op.reshape(nms_ret[1], [])

        return relax.Tuple([selected_indices, num_valid])

    def convert_nms_v5(self, op):
        """Convert TFLite NonMaxSuppressionV5"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 6, "input tensor length should be 6"

        boxes = self.get_tensor_expr(input_tensors[0])
        scores = self.get_tensor_expr(input_tensors[1])

        max_output_size = self.get_tensor_value(input_tensors[2])
        iou_threshold = self.get_tensor_value(input_tensors[3])
        score_threshold = self.get_tensor_value(input_tensors[4])
        soft_nms_sigma = self.get_tensor_value(input_tensors[5])

        if isinstance(max_output_size, np.ndarray):
            assert max_output_size.size == 1, "only one value is expected."
            max_output_size = int(max_output_size)

        if isinstance(iou_threshold, np.ndarray):
            assert iou_threshold.size == 1, "only one value is expected."
            iou_threshold = float(iou_threshold)

        if isinstance(score_threshold, np.ndarray):
            assert score_threshold.size == 1, "only one value is expected."
            score_threshold = float(score_threshold)

        if isinstance(soft_nms_sigma, np.ndarray):
            assert soft_nms_sigma.size == 1, "only one value is expected."
            soft_nms_sigma = float(soft_nms_sigma)

        scores_expand = relax.op.expand_dims(scores, axis=-1)
        data = relax.op.concat([scores_expand, boxes], axis=-1)
        data = relax.op.expand_dims(data, axis=0)

        valid_counts_ret = relax.op.vision.get_valid_counts(
            data, score_threshold=score_threshold, id_index=-1, score_index=0
        )
        count = valid_counts_ret[0]
        data = valid_counts_ret[1]
        indices = valid_counts_ret[2]

        nms_ret = relax.op.vision.non_max_suppression(
            data=data,
            valid_count=count,
            indices=indices,
            max_output_size=max_output_size,
            iou_threshold=iou_threshold,
            force_suppress=True,
            top_k=-1,
            coord_start=1,
            score_index=0,
            id_index=-1,
            return_indices=True,
            invalid_to_bottom=False,
            soft_nms_sigma=soft_nms_sigma,
            score_threshold=score_threshold,
        )

        if soft_nms_sigma > 0.0:
            processed_data = relax.op.squeeze(nms_ret[0], axis=[0])
            indices_from_nms = nms_ret[1]
            num_valid_from_nms = nms_ret[2]
        else:
            indices_from_nms = nms_ret[0]
            num_valid_from_nms = nms_ret[1]

        selected_indices = relax.op.squeeze(indices_from_nms, axis=[0])
        selected_indices = relax.op.strided_slice(
            selected_indices, axes=[0], begin=[0], end=[max_output_size]
        )
        num_valid = relax.op.reshape(num_valid_from_nms, [])

        if soft_nms_sigma > 0.0:
            # Extract decayed scores from the processed data (score_index=0)
            selected_scores = relax.op.strided_slice(processed_data, axes=[1], begin=[0], end=[1])
            selected_scores = relax.op.squeeze(selected_scores, axis=[1])
            selected_scores = relax.op.strided_slice(
                selected_scores, axes=[0], begin=[0], end=[max_output_size]
            )
            selected_scores = relax.op.clip(
                selected_scores, min=0.0, max=float(np.finfo("float32").max)
            )
        else:
            # Clamp out-of-bound padded indices to prevent take() crash.
            num_boxes = int(self.get_tensor_shape(input_tensors[0])[0])
            safe_indices = relax.op.clip(selected_indices, min=0, max=num_boxes - 1)
            selected_scores = relax.op.take(scores, safe_indices, axis=0)

        out = relax.Tuple([selected_indices, selected_scores, num_valid])
        return out

    def convert_expand_dims(self, op):
        """Convert TFLite EXPAND_DIMS"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensors length should be 2"

        if input_tensors[0].qnn_params:
            # Check that input and output tensor have same qnn params.
            output_tensors = self.get_output_tensors(op)
            assert self.has_same_qnn_params(input_tensors[0], output_tensors[0]), (
                "TFLite EXPAND_DIMS requires input and output tensors' \
                    scale and zero points to be equal"
            )

        input_expr = self.get_tensor_expr(input_tensors[0])
        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])

        ndims = len(input_tensors[0].tensor.ShapeAsNumpy())
        assert -1 - ndims <= axis <= ndims, "axis out of range"

        out = relax.op.expand_dims(input_expr, axis, 1)

        return out

    def convert_one_hot(self, op):
        """Convert TFLite ONE_HOT"""

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.OneHotOptions import OneHotOptions

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 4, "Input tensor's length should be 4"

        # Ensuring input isn't quantized
        assert all(not i.qnn_params for i in input_tensors), "Quantized input is not expected."

        # TFlite ONE_HOT requires both on_value
        # and off_value, making dtype redundant.
        indices = input_tensors[0]
        depth = input_tensors[1]
        on_value = input_tensors[2]
        off_value = input_tensors[3]

        assert on_value.tensor.Type() == off_value.tensor.Type(), (
            "on_value and off_value should be the same type"
        )

        # Getting relax expr for indices
        indices_expr = self.get_expr(indices.tensor_idx)

        # Getting depth value
        depth = self.get_tensor_value(depth)
        if isinstance(depth, np.ndarray):
            depth = int(depth)

        # Getting Axis from Option (Attributes)
        assert op.BuiltinOptionsType() == BuiltinOptions.OneHotOptions
        op_options = op.BuiltinOptions()
        one_hot_options = OneHotOptions()
        one_hot_options.Init(op_options.Bytes, op_options.Pos)
        axis = one_hot_options.Axis()

        # Extract scalar values for on_value and off_value as Expr
        dtype = self.get_tensor_type_str(on_value.tensor.Type())
        on_val = self.get_tensor_value(on_value).item()
        off_val = self.get_tensor_value(off_value).item()
        if "float" in dtype:
            on_prim = relax.prim_value(tvm.tirx.FloatImm(dtype, float(on_val)))
            off_prim = relax.prim_value(tvm.tirx.FloatImm(dtype, float(off_val)))
        else:
            on_prim = relax.prim_value(tvm.tirx.IntImm(dtype, int(on_val)))
            off_prim = relax.prim_value(tvm.tirx.IntImm(dtype, int(off_val)))

        out = relax.op.one_hot(indices_expr, on_prim, off_prim, depth, axis)

        return out

    def convert_reverse_v2(self, op):
        """Convert TFLite REVERSE_V2"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensor's length should be 2"

        input_expr = self.get_expr(input_tensors[0].tensor_idx)

        # Getting axis value
        axis = self.get_tensor_value(input_tensors[1])
        if isinstance(axis, np.ndarray):
            assert len(axis) == 1, "TFLite does not support multi-axis yet"
            axis = int(axis.flat[0])

        out = relax.op.flip(input_expr, axis)
        return out

    def convert_matrix_set_diag(self, op):
        """Convert TFLite MATRIX_SET_DIAG"""

        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 2, "input tensor's length should be 2"

        assert input_tensors[0].tensor.Type() == input_tensors[1].tensor.Type(), (
            "input and diagonal should be the same type of tensors"
        )

        if input_tensors[0].qnn_params:
            # Check that input and output tensor have same qnn params.
            output_tensors = self.get_output_tensors(op)
            assert self.has_same_qnn_params(input_tensors[0], output_tensors[0]), (
                "TFLite MATRIX_SET_DIAG requires input and output tensors' \
                    scale and zero points to be equal"
            )

            # Check that input and diagonal tensor have same qnn params.
            assert self.has_same_qnn_params(input_tensors[0], input_tensors[1]), (
                "TFLite MATRIX_SET_DIAG requires input and diagonal tensors' \
                    scale and zero points to be equal"
            )

        input_expr = self.get_tensor_expr(input_tensors[0])
        diagonal_expr = self.get_tensor_expr(input_tensors[1])

        output_tensors = self.get_output_tensors(op)
        output_tensor = output_tensors[0]
        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        # topi.matrix_set_diag(
        #     input, diagonal, k1, k2, super_diag_right_align, sub_diag_right_align
        # )
        # TFLite MATRIX_SET_DIAG only sets the main diagonal, so k1=0, k2=0
        out = relax.op.call_dps_packed(
            "topi.matrix_set_diag",
            (
                input_expr,
                diagonal_expr,
                relax.const(0),
                relax.const(0),
                relax.const(False),
                relax.const(False),
            ),
            out_ty=relax.TensorType(output_shape, output_dtype),
        )
        return out

    def convert_matrix_diag(self, op):
        """Convert TFLite MATRIX_DIAG"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensor's length should be 1"

        diagonal = input_tensors[0]

        if diagonal.qnn_params:
            # Check that diagonal and output tensor have same qnn params.
            output_tensors = self.get_output_tensors(op)
            assert self.has_same_qnn_params(diagonal, output_tensors[0]), (
                "TFLite MATRIX_DIAG requires diagonal and output tensors' \
                    scale and zero points to be equal"
            )

        output_tensors = self.get_output_tensors(op)
        output_tensor = output_tensors[0]
        output_shape = to_int_list(self.get_tensor_shape(output_tensor))
        output_dtype = self.get_tensor_type_str(output_tensor.tensor.Type())

        diagonal_expr = self.get_tensor_expr(diagonal)
        zeros_expr = relax.op.zeros(output_shape, output_dtype)

        # topi.matrix_set_diag(
        #     input, diagonal, k1, k2, super_diag_right_align, sub_diag_right_align
        # )
        # TFLite MATRIX_DIAG only sets the main diagonal, so k1=0, k2=0
        out = relax.op.call_dps_packed(
            "topi.matrix_set_diag",
            (
                zeros_expr,
                diagonal_expr,
                relax.const(0),
                relax.const(0),
                relax.const(False),
                relax.const(False),
            ),
            out_ty=relax.TensorType(output_shape, output_dtype),
        )
        return out

    def convert_densify(self, op):
        """Convert TFLite DENSIFY"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        output_tensors = self.get_output_tensors(op)
        assert len(output_tensors) == 1, "output tensors length should be 1"
        output_tensor = output_tensors[0]

        sparse_weight_tensor = input_tensors[0]
        sparse_weight_tensor_type_str = self.get_tensor_type_str(sparse_weight_tensor.tensor.Type())

        # NOTE: With current implementation in TFLite, Densify Op does not need to be present
        # in runtime.
        # TODO(ANSHUMAN87): we need to use the sparse_indices output
        # from below function and use that in sparse_to_dense Op.
        # Once the stack corruption issue is resolved in sparse_to_dense Op.
        _, dense_weight = prepare_dense_matrix_from_sparse(
            sparse_weight_tensor.tensor,
            self.get_tensor_value(sparse_weight_tensor, is_sparse=True),
            sparse_weight_tensor_type_str,
        )

        self.set_prefetched_node(output_tensor.tensor_idx, dense_weight)

    def convert_fake_quant(self, op):
        """Convert TFLite FAKE_QUANT"""
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"

        input_tensor = input_tensors[0]
        in_expr = self.get_expr(input_tensor.tensor_idx)

        from tflite.BuiltinOptions import BuiltinOptions
        from tflite.FakeQuantOptions import FakeQuantOptions

        assert op.BuiltinOptionsType() == BuiltinOptions.FakeQuantOptions

        op_options = op.BuiltinOptions()
        fake_quant_options = FakeQuantOptions()
        fake_quant_options.Init(op_options.Bytes, op_options.Pos)

        opt_min = fake_quant_options.Min()
        opt_max = fake_quant_options.Max()
        narrow_range = fake_quant_options.NarrowRange()
        num_bits = fake_quant_options.NumBits()

        assert 2 <= num_bits <= 16

        quant_min = 1 if narrow_range else 0
        quant_max = (1 << num_bits) - 1
        scale = (opt_max - opt_min) / (quant_max - quant_min)

        zero_point_from_min = quant_min - opt_min / scale
        if zero_point_from_min <= quant_min:
            nudged_zero_point = quant_min
        elif zero_point_from_min >= quant_max:
            nudged_zero_point = quant_max
        else:
            nudged_zero_point = round(zero_point_from_min)

        nudged_min = (quant_min - nudged_zero_point) * scale
        nudged_max = (quant_max - nudged_zero_point) * scale

        nudged_min_expr = relax.const(nudged_min, "float32")
        clamped = relax.op.clip(in_expr, nudged_min, nudged_max)
        clamped_shifted = relax.op.subtract(clamped, nudged_min_expr)

        half = relax.const(0.5, "float32")
        one = relax.const(1.0, "float32")
        scale_expr = relax.const(scale, "float32")
        inv_scale = relax.op.divide(one, scale_expr)
        rounded = relax.op.floor(_op.add(_op.multiply(clamped_shifted, inv_scale), half))
        return relax.op.add(_op.multiply(rounded, scale_expr), nudged_min_expr)

    def convert_real(self, op):
        """Convert TFLite REAL op.

        TFLite complex64 tensors are represented as float32[..., 2] in Relax,
        where index 0 = real part, index 1 = imaginary part along the last axis
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = self.get_expr(input_tensors[0].tensor_idx)
        # slice last axis at index 0, and squeeze to remove the last axis
        real = _op.strided_slice(input_tensor, begin=[0], end=[1], strides=[1], axes=[-1])
        return _op.squeeze(real, axis=[-1])

    def convert_imag(self, op):
        """Convert TFLite IMAG op.

        See convert_real for representation of complex64 tensors in Relax.
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = self.get_expr(input_tensors[0].tensor_idx)
        # slice last axis at index 1, and squeeze to remove the last axis
        imag = _op.strided_slice(input_tensor, begin=[1], end=[2], strides=[1], axes=[-1])
        return _op.squeeze(imag, axis=[-1])

    def convert_complex_abs(self, op):
        """Convert TFLite COMPLEX_ABS op: sqrt(real^2 + imag^2)

        See convert_real for the float32[..., 2] complex representation convention.
        """
        input_tensors = self.get_input_tensors(op)
        assert len(input_tensors) == 1, "input tensors length should be 1"
        input_tensor = self.get_expr(input_tensors[0].tensor_idx)
        real = self.bb.emit(
            _op.strided_slice(input_tensor, begin=[0], end=[1], strides=[1], axes=[-1])
        )
        real = self.bb.emit(_op.squeeze(real, axis=[-1]))
        imag = self.bb.emit(
            _op.strided_slice(input_tensor, begin=[1], end=[2], strides=[1], axes=[-1])
        )
        imag = self.bb.emit(_op.squeeze(imag, axis=[-1]))
        real_sq = self.bb.emit(_op.multiply(real, real))
        imag_sq = self.bb.emit(_op.multiply(imag, imag))
        sum_expr = self.bb.emit(_op.add(real_sq, imag_sq))
        return _op.sqrt(sum_expr)

    def get_expr(self, input_tensor_idx):
        return self.exp_tab.get_expr(get_tensor_name(self.subgraph, input_tensor_idx))

    def has_expr(self, input_tensor_idx):
        return self.exp_tab.has_expr(get_tensor_name(self.subgraph, input_tensor_idx))

    def is_prefetched(self, input_tensor_idx):
        return (
            self.prefetched_nodes.get(get_tensor_name(self.subgraph, input_tensor_idx)) is not None
        )

    def set_prefetched_node(self, input_tensor_idx, value):
        self.prefetched_nodes[get_tensor_name(self.subgraph, input_tensor_idx)] = value

    def get_prefetched_node(self, input_tensor_idx):
        return self.prefetched_nodes[get_tensor_name(self.subgraph, input_tensor_idx)]

    def get_tensor_value_or_prefetched(self, tensor, is_sparse=False):
        if self.is_prefetched(tensor.tensor_idx):
            return self.get_prefetched_node(tensor.tensor_idx)
        return self.get_tensor_value(tensor, is_sparse)

    def get_tensor_expr(self, tensor, is_sparse=False):
        """Return the Relax expr for tensor."""
        if self.has_expr(tensor.tensor_idx):
            return self.get_expr(tensor.tensor_idx)

        type_str = self.get_tensor_type_str(tensor.tensor.Type())
        value = self.get_tensor_value_or_prefetched(tensor, is_sparse)
        # complex64 constants have no native Relax dtype. Reinterpret the
        # interleaved float32 storage as float32[..., 2] to match the
        # convention used for input tensors.
        if type_str == "complex64":
            value = value.view(np.float32).reshape(value.shape + (2,))
            type_str = "float32"
        return self.exp_tab.new_const(value, dtype=type_str, source_name=tensor.tensor.Name())

    def get_tensor_shape(self, tensor_wrapper):
        """Returns the TFLite tensor shape, inferring it if the TFLite shape is empty.

        This returns the *raw TFLite* shape (no pair axis), even for COMPLEX64 tensors.
        It is distinct from ``_get_relax_tensor_shape``, which returns the *Relax*
        representation (TFLite shape with a trailing (2,) axis appended for COMPLEX64).
        Operators that build Relax IR for COMPLEX64 inputs should use
        ``_get_relax_tensor_shape``; operators that need the TFLite shape for validation
        (e.g. comparing against the model's declared output shape) should use this method.
        """
        assert isinstance(tensor_wrapper, TensorWrapper), "Expecting TensorWrapper here"
        return (
            tensor_wrapper.tensor.ShapeAsNumpy()
            if tensor_wrapper.tensor.ShapeLength() > 0
            else self._infer_shape(self.get_tensor_expr(tensor_wrapper))
        )


def _is_power_of_2(n):
    """Return True iff ``n`` is a positive power of 2."""
    return n > 0 and (n & (n - 1)) == 0


def _bit_reversal_swap_pairs(n):
    """Return the (i, j) index pairs (i < j) for the bit-reversal permutation of length n.

    For a Cooley-Tukey radix-2 FFT, the input must be permuted by bit-reversing
    each index in log2(n) bits before the butterfly stages. Precomputing the
    swap pairs as constants is much cheaper in TIR than computing the
    bit-reverse on the fly.
    """
    assert _is_power_of_2(n), f"bit-reversal requires power of 2, got {n}"
    length = n.bit_length() - 1  # log2(n)
    swaps = []
    for i in range(1, n):
        j = 0
        for k in range(length):
            if i & (1 << k):
                j |= 1 << (length - 1 - k)
        if i < j:
            swaps.append((i, j))
    return swaps


def _build_tflite_rfft2d_primfunc(input_shape, output_pair_shape):
    """Build a reference TIR kernel for TFLite RFFT2D.

    The TFLite frontend represents complex tensors as float32 real/imag pairs
    with a trailing dimension of size 2 because TVM does not have a native
    complex64 dtype. This kernel computes the unnormalized 2-D real FFT over
    the last two input dimensions and writes that pair representation.

    All trig and accumulation are in float32, so the result agrees with
    ``np.fft.rfft2`` to about ``1e-5`` absolute tolerance for typical input
    sizes. Higher-precision backends should override this kernel.

    Notes
    -----
    This is a **naive O(B * H * W * H * W) DFT**, not an FFT. For an input of
    spatial shape (H, W) the inner sum runs H*W times per output position, and
    there are H*W' output positions per batch (W' = W // 2 + 1). This is
    intentionally simple for correctness validation against
    ``np.fft.rfft2``; production use cases with large spatial dimensions should
    override the kernel with an FFT-based implementation. The outer
    (batch, out_y, out_x) iteration is structured as S-TIR spatial axes so a
    downstream ``tvm.tir.schedule`` pass can parallelize it.
    """
    from tvm.script.parser import tirx as T

    batch = 1
    for dim in input_shape[:-2]:
        batch *= int(dim)
    height = int(input_shape[-2])
    width = int(input_shape[-1])
    out_width = int(output_pair_shape[-2])
    input_total = batch * height * width
    output_complex_total = batch * height * out_width
    neg_two_pi = np.float32(-2.0 * math.pi)

    @T.prim_func(private=True, s_tir=True, check_well_formed=False)
    def kernel(
        data: T.Buffer(input_shape, "float32"), output: T.Buffer(output_pair_shape, "float32")
    ):
        # Flat 1D aliases of the multi-dim buffers. The kernel is rank-agnostic
        # over the leading batch dimensions, so collapsing the index space
        # avoids special-casing 2D / 3D / 4D input shapes.
        data_flat = T.decl_buffer((input_total,), "float32", data=data.data)
        output_flat = T.decl_buffer((output_complex_total * 2,), "float32", data=output.data)
        neg_two_pi_const = T.float32(neg_two_pi)

        for b_idx, out_y, out_x in T.grid(batch, height, out_width):
            with T.sblock("rfft2d"):
                v_b, v_oy, v_ox = T.axis.remap("SSS", [b_idx, out_y, out_x])
                real_sum = T.float32(0)
                imag_sum = T.float32(0)
                input_base = v_b * height * width
                for in_y, in_x in T.grid(height, width):
                    phase_y = T.Cast("float32", v_oy) * T.Cast("float32", in_y) / T.float32(height)
                    phase_x = T.Cast("float32", v_ox) * T.Cast("float32", in_x) / T.float32(width)
                    angle = neg_two_pi_const * (phase_y + phase_x)
                    value = data_flat[input_base + in_y * width + in_x]
                    real_sum = real_sum + value * T.cos(angle)
                    imag_sum = imag_sum + value * T.sin(angle)
                flat_out_idx = ((v_b * height + v_oy) * out_width + v_ox) * 2
                output_flat[flat_out_idx] = real_sum
                output_flat[flat_out_idx + 1] = imag_sum

    return kernel


def _build_tflite_rfft2d_fft_primfunc(input_shape, output_pair_shape):
    """Build a 2D Cooley-Tukey FFT TIR kernel for TFLite RFFT2D.

    Precondition: both ``input_shape[-2]`` (height) and ``input_shape[-1]``
    (width) must be positive powers of 2. The frontend dispatches to this
    kernel via ``_is_power_of_2`` checks; the DFT reference kernel handles
    the remaining cases (odd / non-power-of-2 sizes).

    Algorithm
    ---------
    1. Copy the real input into a scratch complex buffer (imag = 0) of shape
       ``(B * H * W,)``.
    2. For each batch and each row, run an in-place radix-2 1D FFT of length
       ``W`` along the width axis.
    3. For each batch and each column, run an in-place radix-2 1D FFT of
       length ``H`` along the height axis (with stride ``W``).
    4. Write the first ``W // 2 + 1`` complex bins per row to the output
       pair representation.

    The bit-reversal permutation required by iterative Cooley-Tukey is done
    by emitting the (i, j) swap pairs directly in the TIR source (one
    inlined swap per pair), avoiding the need for runtime index tables.

    Complexity is ``O(B * H * W * (log2(H) + log2(W)))``, vs the DFT
    reference kernel's ``O(B * H * W * H * W)``.
    """
    from tvm.script.parser import tirx as T

    batch = 1
    for dim in input_shape[:-2]:
        batch *= int(dim)
    height = int(input_shape[-2])
    width = int(input_shape[-1])
    out_width = int(output_pair_shape[-2])
    input_total = batch * height * width
    output_complex_total = batch * height * out_width
    # Cast to Python float so the f-string-interpolated repr is a plain number
    # (np.float32's repr is "np.float32(...)", which the TIR parser can't resolve).
    neg_two_pi = float(np.float32(-2.0 * math.pi))
    log2_w = int(math.log2(width))
    log2_h = int(math.log2(height))

    if not (_is_power_of_2(height) and _is_power_of_2(width)):
        raise ValueError(
            f"_build_tflite_rfft2d_fft_primfunc requires power-of-2 height and width, "
            f"got H={height}, W={width}"
        )

    # Precompute the bit-reversal swap pairs at Python level. These are
    # constant for a given FFT length and will be inlined in the TIR source.
    # Each emitted line is indented 16 spaces (4 levels: top → b_idx loop →
    # sblock → row/col loop body) so it lands inside the for loop when
    # concatenated into the primfunc source.
    row_swap_stmts = []
    for i, j in _bit_reversal_swap_pairs(width):
        row_swap_stmts.append(
            f"                i_idx = row_base + {i}\n"
            f"                j_idx = row_base + {j}\n"
            f"                tmp_r = scratch_real[i_idx]\n"
            f"                scratch_real[i_idx] = scratch_real[j_idx]\n"
            f"                scratch_real[j_idx] = tmp_r\n"
            f"                tmp_i = scratch_imag[i_idx]\n"
            f"                scratch_imag[i_idx] = scratch_imag[j_idx]\n"
            f"                scratch_imag[j_idx] = tmp_i\n"
        )
    row_swaps_code = "".join(row_swap_stmts) if row_swap_stmts else "                pass\n"

    col_swap_stmts = []
    for i, j in _bit_reversal_swap_pairs(height):
        col_swap_stmts.append(
            f"                i_idx = col_base + {i * width}\n"
            f"                j_idx = col_base + {j * width}\n"
            f"                tmp_r = scratch_real[i_idx]\n"
            f"                scratch_real[i_idx] = scratch_real[j_idx]\n"
            f"                scratch_real[j_idx] = tmp_r\n"
            f"                tmp_i = scratch_imag[i_idx]\n"
            f"                scratch_imag[i_idx] = scratch_imag[j_idx]\n"
            f"                scratch_imag[j_idx] = tmp_i\n"
        )
    col_swaps_code = "".join(col_swap_stmts) if col_swap_stmts else "                pass\n"

    # Build the per-stage butterfly code with the stage loop fully unrolled
    # at primfunc-construction time. After unrolling, all loop bounds
    # (block_start, k) are compile-time integers, so the TIR parser doesn't
    # need to reason about runtime loop extents and the scheduler can
    # see static twiddle factors instead of runtime trig calls.
    def _stage_stmts(stage_count, length, indent, stride=1, base_expr="row_base"):
        """Generate fully-unrolled Cooley-Tukey butterfly stage bodies.

        ``base_expr`` is the TIR expression holding the base offset of the
        FFT being transformed (e.g. ``"row_base"`` for rows or
        ``"col_base"`` for columns). ``stride`` is the integer distance
        between adjacent butterfly taps: 1 for the row-FFT (contiguous
        elements) and ``width`` for the column-FFT (strided access).
        """
        out = []
        for stage in range(1, stage_count + 1):
            m_val = 1 << stage
            half_val = m_val >> 1
            for block_start in range(0, length, m_val):
                for k in range(half_val):
                    if stride == 1:
                        a_idx = f"{base_expr} + {block_start} + {k}"
                        b_idx_expr = f"{base_expr} + {block_start} + {k} + {half_val}"
                    else:
                        a_idx = f"{base_expr} + ({block_start} + {k}) * {stride}"
                        b_idx_expr = f"{base_expr} + ({block_start} + {k} + {half_val}) * {stride}"
                    angle_val = float(np.float32(neg_two_pi * k / m_val))
                    w_real_val = float(np.float32(math.cos(angle_val)))
                    w_imag_val = float(np.float32(math.sin(angle_val)))
                    out.extend(
                        [
                            f"{indent}a_idx = {a_idx}\n",
                            f"{indent}b_idx_local = {b_idx_expr}\n",
                            f"{indent}t_real = scratch_real[b_idx_local] * T.float32({w_real_val!r}) - scratch_imag[b_idx_local] * T.float32({w_imag_val!r})\n",
                            f"{indent}t_imag = scratch_real[b_idx_local] * T.float32({w_imag_val!r}) + scratch_imag[b_idx_local] * T.float32({w_real_val!r})\n",
                            f"{indent}u_real = scratch_real[a_idx]\n",
                            f"{indent}u_imag = scratch_imag[a_idx]\n",
                            f"{indent}scratch_real[a_idx] = u_real + t_real\n",
                            f"{indent}scratch_imag[a_idx] = u_imag + t_imag\n",
                            f"{indent}scratch_real[b_idx_local] = u_real - t_real\n",
                            f"{indent}scratch_imag[b_idx_local] = u_imag - t_imag\n",
                        ]
                    )
        return "".join(out)

    row_stages_code = _stage_stmts(
        log2_w, width, "                ", stride=1, base_expr="row_base"
    )
    col_stages_code = _stage_stmts(
        log2_h, height, "                ", stride=width, base_expr="col_base"
    )

    # Build the primfunc source. The bit-reversal swaps are inlined (one
    # unconditional block per (i, j) pair) and the butterfly stages are
    # fully unrolled, so the TIR parser sees ordinary statements rather than
    # runtime table lookups or runtime-magnitude loops. The body is wrapped
    # in a single S-TIR block over the batch dimension so the Relax
    # pipeline (which expects an SBlockRealize at the primfunc body) accepts
    # this kernel.
    primfunc_source = (
        "from tvm.script.parser import tirx as T\n"
        "@T.prim_func(private=True, s_tir=True, check_well_formed=False)\n"
        "def kernel(\n"
        f"    data: T.Buffer({tuple(int(x) for x in input_shape)}, 'float32'),\n"
        f"    output: T.Buffer({tuple(int(x) for x in output_pair_shape)}, 'float32'),\n"
        "):\n"
        f"    data_flat = T.decl_buffer(({input_total},), 'float32', data=data.data)\n"
        f"    output_flat = T.decl_buffer(({output_complex_total * 2},), 'float32', data=output.data)\n"
        f"    scratch_real = T.decl_buffer(({input_total},), 'float32')\n"
        f"    scratch_imag = T.decl_buffer(({input_total},), 'float32')\n"
        f"    for b_idx in T.serial({batch}):\n"
        f"        with T.sblock('rfft2d_fft'):\n"
        f"            v_b = T.axis.remap('S', [b_idx])\n"
        f"            # Initialize scratch from real input; imag = 0.\n"
        f"            for i in T.serial({height * width}):\n"
        f"                src = v_b * {height * width} + i\n"
        f"                scratch_real[src] = data_flat[src]\n"
        f"                scratch_imag[src] = T.float32(0)\n"
        f"            # FFT along width axis (one 1D FFT per row).\n"
        f"            for row in T.serial({height}):\n"
        f"                row_base = v_b * {height * width} + row * {width}\n"
        f"                # Bit-reversal permutation (inlined swaps).\n"
        f"{row_swaps_code}"
        f"                # Cooley-Tukey butterfly stages (fully unrolled).\n"
        f"{row_stages_code}"
        f"            # FFT along height axis (one 1D FFT per column, stride = width).\n"
        f"            for col in T.serial({width}):\n"
        f"                col_base = v_b * {height * width} + col\n"
        f"                # Bit-reversal permutation (strided, inlined swaps).\n"
        f"{col_swaps_code}"
        f"                # Cooley-Tukey butterfly stages (strided, fully unrolled).\n"
        f"{col_stages_code}"
        f"            # Write the first out_width complex bins per row to output.\n"
        f"            for row in T.serial({height}):\n"
        f"                for out_x in T.serial({out_width}):\n"
        f"                    src = v_b * {height * width} + row * {width} + out_x\n"
        f"                    dst = ((v_b * {height} + row) * {out_width} + out_x) * 2\n"
        f"                    output_flat[dst] = scratch_real[src]\n"
        f"                    output_flat[dst + 1] = scratch_imag[src]\n"
    )

    namespace = {"T": T, "tirx": T}
    # Register the generated source in linecache so the TIR parser (which calls
    # inspect.getsourcelines) can find it. The fake filename is stable across
    # calls — we include input/output shapes so the cache is self-disambiguating
    # for any callers who want to introspect the generated source.
    import linecache as _linecache

    fake_file = f"<tflite_rfft2d_fft_primfunc H={height} W={width} outW={out_width}>"
    _linecache.cache[fake_file] = (
        len(primfunc_source.splitlines()),
        None,
        [line + "\n" for line in primfunc_source.splitlines()],
        fake_file,
    )
    code = compile(primfunc_source, fake_file, "exec")
    exec(code, namespace)
    return namespace["kernel"]


# Constants for the Random123 counter-based PRNGs used by STABLEHLO_RNG_BIT_GENERATOR,
# matching tensorflow/lite/kernels/rng_util.cc.
_STABLEHLO_RNG_THREEFRY_PARITY = 0x1BD11BDA
_STABLEHLO_RNG_PHILOX_MUL_A = 0xD2511F53
_STABLEHLO_RNG_PHILOX_MUL_B = 0xCD9E8D57
_STABLEHLO_RNG_PHILOX_WEYL_A = 0x9E3779B9
_STABLEHLO_RNG_PHILOX_WEYL_B = 0xBB67AE85


def _build_stablehlo_rng_bit_generator_primfunc(algorithm, state_len, out_dtype, out_shape):
    """Build a bit-exact TIR kernel for STABLEHLO_RNG_BIT_GENERATOR.

    Mirrors the TFLite runtime kernel (tensorflow/lite/kernels/rng_bit_generator.cc),
    implementing the Random123 Threefry2x32 (20 rounds) and Philox4x32 (10 rounds)
    counter-based PRNGs. The kernel reinterprets the uint64 state as uint32 words,
    advances a 64-bit block counter, and packs the generated words into the output
    tensor. The updated state keeps the key unchanged and only advances the counter,
    which is the only behaviour the runtime relies on.
    """
    from tvm.script.parser import tirx as T

    total = 1
    for dim in out_shape:
        total *= int(dim)
    is_64bit = out_dtype in ("int64", "uint64")
    block_words = 2 if algorithm == "threefry" else 4
    out_word_count = total * (2 if is_64bit else 1)
    num_blocks = (out_word_count + block_words - 1) // block_words
    writes_per_block = block_words // (2 if is_64bit else 1)
    parity = _STABLEHLO_RNG_THREEFRY_PARITY
    mul_a, mul_b = _STABLEHLO_RNG_PHILOX_MUL_A, _STABLEHLO_RNG_PHILOX_MUL_B
    weyl_a, weyl_b = _STABLEHLO_RNG_PHILOX_WEYL_A, _STABLEHLO_RNG_PHILOX_WEYL_B

    def _u32(value):
        return T.Cast("uint32", value)

    def _u64(value):
        return T.Cast("uint64", value)

    def _store_value(words, write_index):
        # Pack the generated uint32 words into one output element, reinterpreting
        # the bit pattern into the (possibly signed) output dtype.
        if is_64bit:
            low = _u64(words[2 * write_index])
            high = _u64(words[2 * write_index + 1])
            return T.reinterpret(out_dtype, low | (high << T.uint64(32)))
        return T.reinterpret(out_dtype, words[write_index])

    if algorithm == "threefry":

        @T.prim_func(private=True, s_tir=True)
        def kernel(
            initial_state: T.Buffer((state_len,), "uint64"),
            output_state: T.Buffer((state_len,), "uint64"),
            output: T.Buffer(out_shape, out_dtype),
        ):
            # A single opaque structured block keeps the imperative kernel as a
            # well-formed block-structured PrimFunc, as required by the Relax
            # pipeline (e.g. HasReshapePattern).
            with T.sblock("rng_bit_generator"):
                state_key = initial_state[0]
                state_counter = initial_state[1]
                key_0 = _u32(state_key & T.uint64(0xFFFFFFFF))
                key_1 = _u32(state_key >> T.uint64(32))
                output_state[0] = state_key
                output_state[1] = state_counter + T.uint64(num_blocks)
                out_flat = T.decl_buffer((total,), out_dtype, data=output.data)
                keys = T.decl_buffer((3,), "uint32", scope="local")
                rotations = T.decl_buffer((8,), "uint32", scope="local")
                ctr = T.decl_buffer((2,), "uint32", scope="local")
                keys[0] = key_0
                keys[1] = key_1
                keys[2] = key_0 ^ key_1 ^ T.uint32(parity)
                rotations[0] = T.uint32(13)
                rotations[1] = T.uint32(15)
                rotations[2] = T.uint32(26)
                rotations[3] = T.uint32(6)
                rotations[4] = T.uint32(17)
                rotations[5] = T.uint32(29)
                rotations[6] = T.uint32(16)
                rotations[7] = T.uint32(24)
                for block in T.serial(num_blocks):
                    counter = state_counter + _u64(block)
                    ctr[0] = _u32(counter & T.uint64(0xFFFFFFFF)) + key_0
                    ctr[1] = _u32(counter >> T.uint64(32)) + key_1
                    for group in T.serial(5):
                        for step in T.serial(4):
                            rot = rotations[(group * 4 + step) % 8]
                            ctr[0] = ctr[0] + ctr[1]
                            ctr[1] = (ctr[1] << rot) | (ctr[1] >> (T.uint32(32) - rot))
                            ctr[1] = ctr[1] ^ ctr[0]
                        ctr[0] = ctr[0] + keys[(group + 1) % 3]
                        ctr[1] = ctr[1] + keys[(group + 2) % 3] + _u32(group + 1)
                    for write_index in T.serial(writes_per_block):
                        element = block * writes_per_block + write_index
                        if element < total:
                            out_flat[element] = _store_value(ctr, write_index)

        return kernel

    @T.prim_func(private=True, s_tir=True)
    def kernel(
        initial_state: T.Buffer((state_len,), "uint64"),
        output_state: T.Buffer((state_len,), "uint64"),
        output: T.Buffer(out_shape, out_dtype),
    ):
        with T.sblock("rng_bit_generator"):
            state_key = initial_state[0]
            state_counter = initial_state[1]
            key_0 = _u32(state_key & T.uint64(0xFFFFFFFF))
            key_1 = _u32(state_key >> T.uint64(32))
            output_state[0] = state_key
            output_state[1] = state_counter + T.uint64(num_blocks)
            out_flat = T.decl_buffer((total,), out_dtype, data=output.data)
            ctr = T.decl_buffer((4,), "uint32", scope="local")
            keys = T.decl_buffer((2,), "uint32", scope="local")
            high_ctr = T.decl_buffer((2,), "uint32", scope="local")
            if state_len == 3:
                # PHILOX u64[3]: the third state word feeds the high counter and
                # is passed through to the output state unchanged.
                high_state = initial_state[2]
                output_state[2] = high_state
                high_ctr[0] = _u32(high_state & T.uint64(0xFFFFFFFF))
                high_ctr[1] = _u32(high_state >> T.uint64(32))
            else:
                high_ctr[0] = key_0
                high_ctr[1] = key_1
            for block in T.serial(num_blocks):
                counter = state_counter + _u64(block)
                ctr[0] = _u32(counter & T.uint64(0xFFFFFFFF))
                ctr[1] = _u32(counter >> T.uint64(32))
                ctr[2] = high_ctr[0]
                ctr[3] = high_ctr[1]
                keys[0] = key_0
                keys[1] = key_1
                for _round in T.serial(10):
                    prod_0 = T.uint64(mul_a) * _u64(ctr[0])
                    prod_1 = T.uint64(mul_b) * _u64(ctr[2])
                    new_0 = _u32(prod_1 >> T.uint64(32)) ^ ctr[1] ^ keys[0]
                    new_1 = _u32(prod_1 & T.uint64(0xFFFFFFFF))
                    new_2 = _u32(prod_0 >> T.uint64(32)) ^ ctr[3] ^ keys[1]
                    new_3 = _u32(prod_0 & T.uint64(0xFFFFFFFF))
                    ctr[0] = new_0
                    ctr[1] = new_1
                    ctr[2] = new_2
                    ctr[3] = new_3
                    keys[0] = keys[0] + T.uint32(weyl_a)
                    keys[1] = keys[1] + T.uint32(weyl_b)
                for write_index in T.serial(writes_per_block):
                    element = block * writes_per_block + write_index
                    if element < total:
                        out_flat[element] = _store_value(ctr, write_index)

    return kernel


# pylint: disable=no-else-return
def prepare_dense_matrix_from_sparse(sparse_tensor, sparse_tensor_value, sparse_tensor_type):
    """Prepare sparse indices and dense matrix from TFLite sparse parameters."""
    # The function is implemented based on TFLite sparse parameter specifications
    # Please refer
    # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/schema/schema.fbs#L89
    # for details about each parameters
    sparsity = sparse_tensor.Sparsity()
    dense_shape = sparse_tensor.ShapeAsNumpy()
    orig_rank = len(dense_shape)

    # The traversal order of the dimensions defined in the `shape` field of the to be dense tensor.
    traversal_order = sparsity.TraversalOrderAsNumpy()

    # For an n-dimensional tensor with a k-dimensional block (0 <= k <= n),
    # stores how a block dimension in (dn, ..., dn+k-1) maps to the original
    # tensor dimension in (d0, ..., dn). It's stored in the order of (dn, ..., dn+k-1).
    # If not block-sparse, this field is NULL.
    block_map = sparsity.BlockMapAsNumpy()

    total_rank = sparsity.TraversalOrderLength()
    dense_mat = np.full(shape=dense_shape, fill_value=0, dtype=sparse_tensor_type).flatten()

    from enum import Enum

    # NOTE: Here the Vector term is borrowed from TFLite spec.
    class VectorType(Enum):
        Empty = 0
        Int32 = 1
        Uint16 = 2
        Uint8 = 3

    def _get_vector_flag(v_type):
        if VectorType(v_type) == VectorType.Int32:
            return N.Int32Flags
        elif VectorType(v_type) == VectorType.Uint16:
            return N.Uint16Flags
        elif VectorType(v_type) == VectorType.Uint8:
            return N.Uint8Flags
        else:
            raise tvm.error.OpNotImplemented(f"The provided type {v_type} is not supported")

    def _get_flattened_index(indices, shape):
        index = 0
        sub_elements = 1
        for i in reversed(range(0, len(dense_shape))):
            index += indices[i] * sub_elements
            sub_elements *= shape[i]
        return index

    # DimensionMetadata per dimension: the metadata needed for
    #     each dimension to locate the non-zero values in the original dense tensor
    #     inline with traversal order parameter.
    #
    # sp_format has 2 possible values: {DENSE = 0, SPARSE_CSR = 1}
    # If format = DENSE{0} : DenseSize represents size of that dimension
    # If format = SPARSE_CSR{1} : array_segments represents how to segment the indices array,
    #      each segment corresponds to one element in the previous dimension. array_indices
    #      represents the index of the non-zero elements within this dimension
    #      (as those in the CSR matrix format, where the first array is row pointers
    #       and the second array is column indices).
    sp_format = np.zeros(sparsity.DimMetadataLength())
    dim_metadata = [None] * (2 * sparsity.DimMetadataLength())

    # Below loop will fetch all meta data per dimension based on format type
    # Dense or Sparse and will put it in an agnostic array for easy access
    # while preparing dense buffer or indices.
    for i in range(sparsity.DimMetadataLength()):
        sp_format[i] = sparsity.DimMetadata(i).Format()
        if sp_format[i] == 0:
            dim_metadata[2 * i] = [sparsity.DimMetadata(i).DenseSize()]
        else:
            from flatbuffers import number_types as N

            dim_metadata[2 * i] = (
                sparsity.DimMetadata(i)
                .ArraySegments()
                .GetVectorAsNumpy(
                    flags=_get_vector_flag(sparsity.DimMetadata(i).ArraySegmentsType()), off=4
                )
            )
            dim_metadata[2 * i + 1] = (
                sparsity.DimMetadata(i)
                .ArrayIndices()
                .GetVectorAsNumpy(
                    flags=_get_vector_flag(sparsity.DimMetadata(i).ArrayIndicesType()), off=4
                )
            )

    block_dim = 0
    block_size = np.zeros(sparsity.BlockMapLength())

    # Block size parameter if encoded in BSR format
    for i in range(orig_rank):
        if block_dim < sparsity.BlockMapLength() and block_map[block_dim] == i:
            orig_dim = traversal_order[orig_rank + block_dim]
            block_size[block_dim] = sparsity.DimMetadata(orig_dim).DenseSize()
            block_dim += 1

    indices_list = []

    # Below function iterates through each applicable indices per dimension
    # based on format type specified and finally produce the dense matrix and the NZ indices.
    def _def_prepare_dense_matrix_from_sparse(indices, level, prev_idx):
        if level == len(indices):
            start_pos = 0
            orig_idx = np.zeros(orig_rank, dtype="int32")
            while start_pos < orig_rank:
                orig_idx[traversal_order[start_pos]] = indices[start_pos]
                start_pos += 1
            while start_pos < len(indices):
                block_idx = traversal_order[start_pos] - orig_rank
                orig_dim = block_map[block_idx]
                orig_idx[orig_dim] = orig_idx[orig_dim] * block_size[block_idx] + indices[start_pos]
                start_pos += 1
            indices_list.append(orig_idx)
            nonlocal value_idx
            dense_mat[_get_flattened_index(orig_idx, dense_shape)] = sparse_tensor_value[value_idx]
            value_idx += 1
        else:
            metadata_idx = 2 * level
            if sp_format[level] == 0:
                shape_of_level = dim_metadata[metadata_idx][0]
                for idx in range(shape_of_level):
                    indices[level] = idx
                    _def_prepare_dense_matrix_from_sparse(
                        indices, level + 1, prev_idx * shape_of_level + idx
                    )
            else:
                array_segments = dim_metadata[metadata_idx]
                array_indices = dim_metadata[metadata_idx + 1]
                for idx in range(array_segments[prev_idx], array_segments[prev_idx + 1]):
                    indices[level] = array_indices[idx]
                    _def_prepare_dense_matrix_from_sparse(indices, level + 1, idx)

    indices = np.zeros(total_rank)
    value_idx = 0
    _def_prepare_dense_matrix_from_sparse(indices, 0, 0)
    return np.array(indices_list, dtype="int32"), dense_mat.reshape(dense_shape)


def get_scalar_from_constant(expr):
    """Returns scalar value from Relax constant scalar."""
    assert isinstance(expr, relax.Constant) and not expr.data.shape, (
        "Expr is not a constant scalar."
    )
    value = expr.data.numpy()
    assert value.dtype == np.dtype(np.int32) or value.dtype == np.dtype(np.float32), (
        "value must be float32/int32"
    )
    return value.item(0)


def get_tensor_from_constant(expr):
    """Returns tensor of values from Relax constant node."""
    assert isinstance(expr, relax.const)
    value = expr.data.numpy()
    assert value.dtype == np.dtype(np.int32) or value.dtype == np.dtype(np.float32), (
        "value must be float32/int32"
    )
    return value


def build_str_map(obj):
    """Build string map of TFLite enum int value

    Parameters
    ----------
    obj:
        TFLite class which contains enum int value, such as BuiltInOptions

    Returns
    -------
        String representation map of TFLite class enum int value
    """
    ret = {}
    for field_name in dir(obj):
        if not field_name.startswith("_"):
            field_value = getattr(obj, field_name)
            if isinstance(field_value, int):
                ret[field_value] = field_name
    return ret


# SAME padding: https://www.tensorflow.org/api_guides/python/nn
def get_pad_value(data, kernel, stride):
    """Get the pad tuple of value for SAME padding

    Parameters
    ----------
    data:
        1D input data

    kernel:
        1D input kernel

    stride:
        1D input stride

    Returns
    -------
        pad tuple of value
    """

    out = math.ceil(float(data) / float(stride))
    pad = max(0, (out - 1) * stride + kernel - data)
    pad_before = pad // 2
    pad_after = pad - pad_before
    return pad_before, pad_after


def get_tensor_name(subgraph, tensor_idx):
    """Get the tensor name.

    Parameters
    ----------
    subgraph:
        tflite.Subgraph.Subgraph

    tensor:
        tensor index in subgraph

    Returns
    -------
        tensor name in UTF-8 encoding
    """
    tensor_name = subgraph.Tensors(tensor_idx).Name()
    if tensor_name is not None:
        tensor_name = tensor_name.decode("utf-8")
    else:
        tensor_name = "tvmgen_tensor_" + str(tensor_idx)
    return tensor_name


def _decode_type(n):
    _tflite_m = {
        0: "float32",
        1: "float16",
        2: "int32",
        3: "uint8",
        4: "int64",
        5: "string",
        6: "bool",
        7: "int16",
        8: "complex64",
        9: "int8",
        12: "uint64",
        15: "uint32",
    }
    return _tflite_m[n]


def _input_type(model):
    subgraph_count = model.SubgraphsLength()
    assert subgraph_count > 0
    shape_dict = {}
    dtype_dict = {}
    subgraph = model.Subgraphs(0)
    inputs_count = subgraph.InputsLength()
    # TFLite subgraphs can validly have zero inputs (e.g. constant-only RANGE models).
    for input_index in range(inputs_count):
        input_ = subgraph.Inputs(input_index)
        assert subgraph.TensorsLength() > input_
        tensor = subgraph.Tensors(input_)
        input_shape = tuple(tensor.ShapeAsNumpy())
        tensor_type = tensor.Type()
        input_name = get_tensor_name(subgraph, input_)
        input_dtype = _decode_type(tensor_type)
        shape_dict[input_name] = input_shape
        dtype_dict[input_name] = input_dtype

    return shape_dict, dtype_dict


def from_tflite(
    model,
    shape_dict: dict[str, tuple[int]] | None = None,
    dtype_dict: dict[str, str] | None = None,
    op_converter=OperatorConverter,
):
    """Convert from tflite model into compatible relax Function.

    Parameters
    ----------
    model:
        tflite.Model or tflite.Model.Model (depending on tflite version)

    shape_dict : dict of str to int list/tuple
        Input shapes of the model.

    dtype_dict : dict of str to str
        Input types of the model.

    Returns
    -------
    mod : tvm.IRModule
        The relax module for compilation.

    params : dict of str to tvm.nd.NDArray
        The parameter dict to be used by relax

    Examples
    --------
    Users can use TFLiteConverter to convert a concrete function followed by
    tflite.Model.Model.GetRootAsModel to get a tflite.Model object
    The following codes show how to convert a simple concrete function to a Relax program.

    .. code-block:: python

        import tensorflow as tf
        from tvm.relax.frontend.tflite import from_tflite


        # A concrete function be defined from network architecture or invoking a model
        # from keras applications

        # Define a tf.Module containing a tf.function with network definition
        class Conv2DModule(tf.Module):
            @tf.function(
                input_signature=[
                    tf.TensorSpec(shape=(1, 128, 128, 32), dtype=tf.float32),
                    tf.TensorSpec(shape=(3, 3, 32, 32), dtype=tf.float32),
                ]
            )
            def func(self, data, kernel):
                return tf.nn.conv2d(
                    input=(1, 128, 128, 32),
                    filters=(3, 3, 32, 32),
                    data_format="NHWC",
                    strides=(1, 1, 1, 1),
                    padding="SAME",
                )
        concrete_func = Conv2DModule().func.get_concrete_function()

        # Alternatively you can use tensorflow.keras.application wrapped arounf a tf.function
        class NetworkModule(tf.Module):
            def __init__(self):
                self.model = net(weights="imagenet", include_top=True)

            @tf.function
            def func(self, data):
                return self.model(data, training=False)

        model = NetworkModule()
        concrete_func = model.func.get_concrete_function(
            tf.TensorSpec(shape=shape, dtype=tf.float32)
        )


        # Which ever the way we have created the concrete function make tflite.Model as shown below

        converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
        converter.target_spec.supported_ops = [
            tf.lite.OpsSet.TFLITE_BUILTINS,
            tf.lite.OpsSet.SELECT_TF_OPS,
        ]
        tflite_model = tf.lite.Model.Model.GetRootAsModel(converter.convert(), 0)

        # Call the imported to conver the tflite.Model to relax.Module
        mod = from_tflite(tflite_model)

        # Print out the imported model.
        print(mod.script())

    """
    # TFLite.Model.Model has changed to TFLite.Model from 1.14 to 2.1
    try:
        import tflite

        assert isinstance(model, tflite.Model)
    except ImportError as err:
        raise ImportError(
            "tflite is required by the TFLite frontend. Install it with: pip install tflite"
        ) from err
    except TypeError:
        import tflite.Model

        assert isinstance(model, tflite.Model.Model)

    _shape_dict, _dtype_dict = _input_type(model)
    if shape_dict is not None:
        _shape_dict.update(shape_dict)
    if dtype_dict is not None:
        _dtype_dict.update(dtype_dict)

    # Only Subgraphs(0) is converted into Relax main. Additional subgraphs are
    # region/control-flow bodies referenced by specific TFLite ops and are
    # consumed by those op converters as needed.
    assert model.SubgraphsLength() >= 1, "TFLite model must contain at least one subgraph"
    subgraph = model.Subgraphs(0)

    # model inputs / outputs
    model_inputs = subgraph.InputsAsNumpy()
    model_outputs = subgraph.OutputsAsNumpy()

    bb = relax.BlockBuilder()  # pylint: disable=invalid-name

    with bb.function("main"):
        input_list = []
        with bb.dataflow() as df:  # noqa: F841  # pylint: disable=invalid-name, unused-variable
            exp_tab = ExprTable()
            for model_input in model_inputs:
                model_input_name = get_tensor_name(subgraph, model_input)
                shape = _shape_dict[model_input_name] if model_input_name in _shape_dict else None
                dtype = (
                    _dtype_dict[model_input_name] if model_input_name in _dtype_dict else "float32"
                )
                if dtype == "complex64":
                    dtype = "float32"
                    if shape is not None:
                        shape = tuple(shape) + (2,)
                if dtype == "string":
                    # Relax has no string tensor type, so TFLite TensorType.STRING graph
                    # inputs cannot be represented. This also covers runtime string queries
                    # for ops like HASHTABLE_FIND, whose constant-foldable subset is handled
                    # in the op converter.
                    raise tvm.error.OpNotImplemented(
                        "Relax TFLite frontend does not support TensorType.STRING graph inputs "
                        "(e.g. runtime string queries)"
                    )
                input_var = relax.Var(
                    name_hint=model_input_name,
                    ty=relax.TensorType(shape=shape, dtype=dtype),
                )
                exp_tab.set_expr(model_input_name, input_var)
                input_list.append(input_var)

            # op code in model
            op_converter = op_converter(model, subgraph, exp_tab, bb)
            op_converter.check_unsupported_ops()
            op_converter.convert_op_to_relax()

            # params and outputs
            # Resolve outputs through tensor wrappers so constant/prefetched outputs are handled.
            output_tensors = op_converter.get_tensors(model_outputs)
            outputs = [op_converter.get_tensor_expr(tensor) for tensor in output_tensors]
            outputs = outputs[0] if len(outputs) == 1 else relax.Tuple(outputs)
            output_var = bb.emit_output(outputs)

        bb.emit_func_output(output_var, input_list)

        relax_mod = bb.get()
        # Attach attributes.
        param_value_list = []
        if exp_tab.params:
            _, param_value_list = map(list, zip(*exp_tab.params.values()))
        func_attrs = {}
        func_attrs["num_input"] = len(input_list)
        func_attrs["params"] = [tvm.runtime.tensor(arr) for arr in param_value_list]
        relax_mod["main"] = relax_mod["main"].with_attrs(func_attrs)

        return relax_mod