from __future__ import annotations

import copy
from itertools import permutations
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import onnx
from onnx import numpy_helper

from onnx2tf.utils.common_functions import check_model_has_external_data
from onnx2tf.tflite_builder.dispatcher import dispatch_node
from onnx2tf.tflite_builder.op_registry import (
    NodeValidationError,
    get_custom_op_candidate_ops,
    get_supported_onnx_ops,
    resolve_node_dispatch,
)
from onnx2tf.tflite_builder.ir import (
    ModelIR,
    OperatorIR,
    QuantParamIR,
    TensorIR,
    normalize_onnx_shape,
)
from onnx2tf.tflite_builder.tensor_buffer_builder import tflite_dtype_from_numpy


_ONNX_TYPE_TO_TFLITE_DTYPE = {
    onnx.TensorProto.FLOAT: "FLOAT32",
    onnx.TensorProto.FLOAT16: "FLOAT16",
    onnx.TensorProto.DOUBLE: "FLOAT64",
    onnx.TensorProto.INT8: "INT8",
    onnx.TensorProto.INT16: "INT16",
    onnx.TensorProto.INT32: "INT32",
    onnx.TensorProto.INT64: "INT64",
    onnx.TensorProto.UINT8: "UINT8",
    onnx.TensorProto.UINT16: "UINT16",
    onnx.TensorProto.UINT32: "UINT32",
    onnx.TensorProto.UINT64: "UINT64",
    onnx.TensorProto.BOOL: "BOOL",
    onnx.TensorProto.STRING: "STRING",
}

_TFLITE_DTYPE_TO_NUMPY_DTYPE = {
    "FLOAT16": np.dtype(np.float16),
    "FLOAT32": np.dtype(np.float32),
    "FLOAT64": np.dtype(np.float64),
    "INT8": np.dtype(np.int8),
    "INT16": np.dtype(np.int16),
    "INT32": np.dtype(np.int32),
    "INT64": np.dtype(np.int64),
    "UINT8": np.dtype(np.uint8),
    "UINT16": np.dtype(np.uint16),
    "UINT32": np.dtype(np.uint32),
    "UINT64": np.dtype(np.uint64),
    "BOOL": np.dtype(np.bool_),
    "STRING": np.dtype(np.object_),
}

def _create_progress_bar(
    *,
    total: int,
    desc: str,
    enabled: bool,
):
    if not enabled or int(total) <= 0:
        return None
    try:
        from tqdm.auto import tqdm
    except Exception:
        return None
    return tqdm(
        total=int(total),
        desc=str(desc),
        dynamic_ncols=True,
    )


def _dtype_from_onnx_elem_type(elem_type: Optional[int]) -> str:
    if elem_type is None:
        return "FLOAT32"
    if elem_type not in _ONNX_TYPE_TO_TFLITE_DTYPE:
        raise NotImplementedError(f"Unsupported ONNX dtype in flatbuffer_direct: elem_type={elem_type}")
    return _ONNX_TYPE_TO_TFLITE_DTYPE[elem_type]


def _cast_const_array_to_tflite_dtype(
    values: Any,
    out_dtype: str,
) -> Optional[np.ndarray]:
    dtype_key = str(out_dtype).upper()
    target_dtype = _TFLITE_DTYPE_TO_NUMPY_DTYPE.get(dtype_key, None)
    if target_dtype is None:
        return None
    try:
        src = np.asarray(values)
        if dtype_key == "BOOL":
            return np.asarray(src != 0, dtype=np.bool_)
        return np.asarray(src, dtype=target_dtype)
    except Exception:
        return None


def _collect_constant_arrays(onnx_graph: onnx.ModelProto) -> Dict[str, np.ndarray]:
    constants: Dict[str, np.ndarray] = {}
    for ini in onnx_graph.graph.initializer:
        constants[ini.name] = np.asarray(numpy_helper.to_array(ini))
    for node in onnx_graph.graph.node:
        if str(node.op_type) != "Constant" or len(node.output) < 1:
            continue
        output_name = str(node.output[0])
        if output_name == "":
            continue
        value_attr = None
        for attr in node.attribute:
            if str(attr.name) == "value":
                value_attr = attr
                break
        if value_attr is None:
            continue
        try:
            constants[output_name] = np.asarray(numpy_helper.to_array(value_attr.t))
        except Exception:
            continue
    return constants


def _node_attr_int(node: onnx.NodeProto, attr_name: str, default: int) -> int:
    for attr in node.attribute:
        if str(attr.name) != str(attr_name):
            continue
        if attr.type == onnx.AttributeProto.INT:
            return int(attr.i)
        if attr.type == onnx.AttributeProto.INTS and len(attr.ints) > 0:
            return int(attr.ints[0])
    return int(default)


def _node_attr_ints(node: onnx.NodeProto, attr_name: str) -> List[int]:
    for attr in node.attribute:
        if str(attr.name) != str(attr_name):
            continue
        if attr.type == onnx.AttributeProto.INTS:
            return [int(v) for v in attr.ints]
        if attr.type == onnx.AttributeProto.INT:
            return [int(attr.i)]
    return []


def _infer_missing_tensor_ranks_with_axis_constraints(
    *,
    onnx_graph: onnx.ModelProto,
    shape_map: Dict[str, List[Any]],
) -> None:
    rank_map: Dict[str, int] = {}
    for name, shape in shape_map.items():
        if isinstance(shape, list) and len(shape) > 0:
            rank_map[str(name)] = max(int(rank_map.get(str(name), 0)), int(len(shape)))

    constants = _collect_constant_arrays(onnx_graph)

    def _set_min_rank(tensor_name: str, rank: int) -> bool:
        name = str(tensor_name)
        if name == "" or int(rank) <= 0:
            return False
        prev = int(rank_map.get(name, 0))
        if int(rank) > prev:
            rank_map[name] = int(rank)
            return True
        return False

    def _unify_ranks(names: List[str]) -> bool:
        normalized = [str(v) for v in names if str(v) != ""]
        if len(normalized) == 0:
            return False
        max_rank = max(int(rank_map.get(name, 0)) for name in normalized)
        if max_rank <= 0:
            return False
        changed = False
        for name in normalized:
            changed = _set_min_rank(name, max_rank) or changed
        return changed

    def _slice_axes(node: onnx.NodeProto) -> List[int]:
        if len(node.input) >= 4 and str(node.input[3]) != "":
            axes_arr = constants.get(str(node.input[3]), None)
            if axes_arr is not None:
                return [int(v) for v in np.asarray(axes_arr).reshape(-1).tolist()]
        attr_axes = _node_attr_ints(node, "axes")
        if len(attr_axes) > 0:
            return [int(v) for v in attr_axes]
        if len(node.input) >= 2 and str(node.input[1]) != "":
            starts_arr = constants.get(str(node.input[1]), None)
            if starts_arr is not None:
                starts_len = int(np.asarray(starts_arr).size)
                return [int(v) for v in range(starts_len)]
        return []

    max_iter = max(1, int(len(onnx_graph.graph.node)) * 6)
    for _ in range(max_iter):
        changed = False
        for node in onnx_graph.graph.node:
            op = str(node.op_type)
            inputs = [str(v) for v in node.input if str(v) != ""]
            outputs = [str(v) for v in node.output if str(v) != ""]
            if len(outputs) == 0:
                continue

            if op in {"QuantizeLinear", "DequantizeLinear", "Identity"}:
                if len(inputs) >= 1:
                    changed = _unify_ranks([inputs[0], outputs[0]]) or changed
                continue

            if op in {"Softmax", "LogSoftmax"}:
                if len(inputs) >= 1:
                    changed = _unify_ranks([inputs[0], outputs[0]]) or changed
                    axis = _node_attr_int(node, "axis", 1)
                    if axis >= 0:
                        min_rank = int(axis) + 1
                        changed = _set_min_rank(inputs[0], min_rank) or changed
                        changed = _set_min_rank(outputs[0], min_rank) or changed
                continue

            if op == "Slice":
                if len(inputs) >= 1:
                    changed = _unify_ranks([inputs[0], outputs[0]]) or changed
                    axes = _slice_axes(node)
                    positive_axes = [int(v) for v in axes if int(v) >= 0]
                    if len(positive_axes) > 0:
                        min_rank = max(positive_axes) + 1
                        changed = _set_min_rank(inputs[0], min_rank) or changed
                        changed = _set_min_rank(outputs[0], min_rank) or changed
                continue

            if op == "QLinearConv":
                if len(inputs) >= 1:
                    changed = _set_min_rank(inputs[0], 4) or changed
                changed = _set_min_rank(outputs[0], 4) or changed
                continue

            if op in {"Concat", "QLinearConcat"}:
                if op == "QLinearConcat":
                    data_inputs = [inputs[idx] for idx in range(2, len(inputs), 3)]
                else:
                    data_inputs = list(inputs)
                rank_linked = [str(v) for v in data_inputs if str(v) != ""] + [outputs[0]]
                changed = _unify_ranks(rank_linked) or changed
                axis = _node_attr_int(node, "axis", 1)
                if axis >= 0:
                    min_rank = int(axis) + 1
                    for tensor_name in rank_linked:
                        changed = _set_min_rank(tensor_name, min_rank) or changed
                continue

        if not changed:
            break

    for tensor_name, rank in rank_map.items():
        if int(rank) <= 0:
            continue
        current = shape_map.get(str(tensor_name), None)
        if current is None or len(current) == 0:
            shape_map[str(tensor_name)] = [-1 for _ in range(int(rank))]


def _extract_tensor_info(
    onnx_graph: onnx.ModelProto,
) -> Tuple[Dict[str, List[Any]], Dict[str, str]]:
    shape_map: Dict[str, List[Any]] = {}
    dtype_map: Dict[str, str] = {}

    def _fill_value_info(value_info):
        if not value_info.type.HasField("tensor_type"):
            return
        name = value_info.name
        tensor_type = value_info.type.tensor_type
        dims: List[Any] = []
        if tensor_type.HasField("shape"):
            for d in tensor_type.shape.dim:
                if d.HasField("dim_value") and d.dim_value >= 0:
                    dims.append(int(d.dim_value))
                else:
                    dims.append(-1)
        shape_map[name] = dims
        dtype_map[name] = _dtype_from_onnx_elem_type(tensor_type.elem_type)

    for vi in onnx_graph.graph.input:
        _fill_value_info(vi)
    for vi in onnx_graph.graph.value_info:
        _fill_value_info(vi)
    for vi in onnx_graph.graph.output:
        _fill_value_info(vi)

    for ini in onnx_graph.graph.initializer:
        arr = numpy_helper.to_array(ini)
        shape_map[ini.name] = list(arr.shape)
        dtype_map[ini.name] = tflite_dtype_from_numpy(arr.dtype)

    _infer_missing_tensor_ranks_with_axis_constraints(
        onnx_graph=onnx_graph,
        shape_map=shape_map,
    )

    return shape_map, dtype_map


def _collect_dynamic_boundary_tensor_names(onnx_graph: onnx.ModelProto) -> Dict[str, List[str]]:
    initializer_names = {str(ini.name) for ini in onnx_graph.graph.initializer}

    def _has_dynamic_dim(value_info: Any) -> bool:
        if not value_info.type.HasField("tensor_type"):
            return True
        tensor_type = value_info.type.tensor_type
        if not tensor_type.HasField("shape"):
            return True
        for d in tensor_type.shape.dim:
            if not (d.HasField("dim_value") and int(d.dim_value) >= 0):
                return True
        return False

    dynamic_inputs: List[str] = []
    for vi in onnx_graph.graph.input:
        name = str(vi.name)
        if name in initializer_names:
            continue
        if _has_dynamic_dim(vi):
            dynamic_inputs.append(name)

    dynamic_outputs: List[str] = []
    for vi in onnx_graph.graph.output:
        name = str(vi.name)
        if _has_dynamic_dim(vi):
            dynamic_outputs.append(name)

    return {
        "inputs": dynamic_inputs,
        "outputs": dynamic_outputs,
    }


def _build_onnx_boundary_shape_signature_map(
    *,
    onnx_graph: onnx.ModelProto,
    shape_map: Dict[str, List[Any]],
) -> Dict[str, List[int]]:
    signature_map: Dict[str, List[int]] = {}
    initializer_names = {str(ini.name) for ini in onnx_graph.graph.initializer}
    for value_info in list(onnx_graph.graph.input) + list(onnx_graph.graph.output):
        name = str(value_info.name)
        if name == "" or name in initializer_names:
            continue
        raw_shape = shape_map.get(name, None)
        _, signature = normalize_onnx_shape(raw_shape)
        signature_map[name] = [int(v) for v in list(signature)]
    return signature_map


def _align_boundary_signature_to_current_shape(
    *,
    boundary_signature: Optional[List[int]],
    current_shape: Optional[List[int]],
) -> Optional[List[int]]:
    if boundary_signature is None or current_shape is None:
        return None
    signature = [int(v) for v in list(boundary_signature)]
    shape = [int(v) for v in list(current_shape)]
    if len(signature) == 0 or len(signature) != len(shape):
        return None

    # Fast-path when no layout permutation is observed.
    # Keep boundary axis contracts (including singleton dims) whenever the
    # current static shape is compatible on the same axes.
    static_axes = [int(i) for i, v in enumerate(signature) if int(v) >= 0]
    if len(static_axes) > 0 and all(int(signature[i]) == int(shape[i]) for i in static_axes):
        return [int(v) for v in list(signature)]

    aligned = [-1 for _ in range(len(signature))]
    used_axes: set[int] = set()
    static_values = sorted({int(v) for v in signature if int(v) > 1})
    for value in static_values:
        needed = int(sum(1 for v in signature if int(v) == value))
        if needed <= 0:
            continue
        candidate_axes = [int(i) for i, dim in enumerate(shape) if int(dim) == value and int(i) not in used_axes]
        if len(candidate_axes) < needed:
            continue
        for axis in candidate_axes[:needed]:
            aligned[int(axis)] = int(value)
            used_axes.add(int(axis))
    return [int(v) for v in aligned]


def _graph_has_missing_rank_info(onnx_graph: onnx.ModelProto) -> bool:
    value_infos = (
        list(onnx_graph.graph.input)
        + list(onnx_graph.graph.value_info)
        + list(onnx_graph.graph.output)
    )
    for vi in value_infos:
        if not vi.type.HasField("tensor_type"):
            continue
        if not vi.type.tensor_type.HasField("shape"):
            return True
    return False


def _infer_shapes_with_fallback(onnx_graph: onnx.ModelProto) -> onnx.ModelProto:
    if check_model_has_external_data(onnx_graph):
        return onnx_graph

    inferred_graph = onnx_graph
    try:
        inferred_graph = onnx.shape_inference.infer_shapes(inferred_graph)
    except Exception:
        pass

    if not _graph_has_missing_rank_info(inferred_graph):
        return inferred_graph

    try:
        from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference
    except Exception:
        return inferred_graph

    try:
        return SymbolicShapeInference.infer_shapes(
            inferred_graph,
            auto_merge=True,
            guess_output_rank=True,
        )
    except TypeError:
        try:
            return SymbolicShapeInference.infer_shapes(inferred_graph)
        except Exception:
            return inferred_graph
    except Exception:
        return inferred_graph


class LoweringContext:
    def __init__(
        self,
        model_ir: ModelIR,
        shape_map: Dict[str, List[Any]],
        dtype_map: Dict[str, str],
        constants: Dict[str, np.ndarray],
        onnx_model: Optional[Any] = None,
        allow_custom_ops: bool = False,
        custom_op_allowlist: Optional[List[str]] = None,
        disable_group_convolution: bool = False,
        tensor_consumer_count: Optional[Dict[str, int]] = None,
        graph_output_names: Optional[List[str]] = None,
        output_nms_with_argmax: bool = False,
        switch_nms_version: str = "v4",
        mvn_epsilon: float = 1e-10,
        disable_suppression_flextranspose: bool = False,
        number_of_dimensions_after_flextranspose_compression: int = 6,
        disable_suppression_flexstridedslice: bool = False,
        number_of_dimensions_after_flexstridedslice_compression: int = 5,
        optimization_for_gpu_delegate: bool = False,
        replace_argmax_to_reducemax_and_indices_is_int64: bool = False,
        replace_argmax_to_reducemax_and_indices_is_float32: bool = False,
        replace_argmax_to_fused_argmax_and_indices_is_int64: bool = False,
        replace_argmax_to_fused_argmax_and_indices_is_float32: bool = False,
        fused_argmax_scale_ratio: float = 0.5,
        replace_to_pseudo_operators: Optional[List[str]] = None,
    ):
        self.model_ir = model_ir
        self.shape_map = shape_map
        self.dtype_map = dtype_map
        self.constants = constants
        self.onnx_model = onnx_model
        self.allow_custom_ops = bool(allow_custom_ops)
        self.custom_op_allowlist = (
            list(custom_op_allowlist) if custom_op_allowlist is not None else None
        )
        self.disable_group_convolution = bool(disable_group_convolution)
        self.tensor_consumer_count = (
            dict(tensor_consumer_count)
            if isinstance(tensor_consumer_count, dict)
            else {}
        )
        self.graph_output_names = set(graph_output_names) if graph_output_names is not None else set()
        self.output_nms_with_argmax = bool(output_nms_with_argmax)
        self.mvn_epsilon = float(mvn_epsilon)
        self.disable_suppression_flextranspose = bool(
            disable_suppression_flextranspose
        )
        # TFLite TRANSPOSE kernel supports rank <= 6.
        # Keep user intent (>=2) while clamping to runtime-safe upper bound.
        self.number_of_dimensions_after_flextranspose_compression = int(
            max(2, min(6, int(number_of_dimensions_after_flextranspose_compression)))
        )
        self.disable_suppression_flexstridedslice = bool(
            disable_suppression_flexstridedslice
        )
        # TFLite SLICE kernel supports rank <= 5.
        # Keep user intent (>=1) while clamping to runtime-safe upper bound.
        self.number_of_dimensions_after_flexstridedslice_compression = int(
            max(1, min(5, int(number_of_dimensions_after_flexstridedslice_compression)))
        )
        self.optimization_for_gpu_delegate = bool(optimization_for_gpu_delegate)
        self.replace_argmax_to_reducemax_and_indices_is_int64 = bool(
            replace_argmax_to_reducemax_and_indices_is_int64
        )
        self.replace_argmax_to_reducemax_and_indices_is_float32 = bool(
            replace_argmax_to_reducemax_and_indices_is_float32
        )
        self.replace_argmax_to_fused_argmax_and_indices_is_int64 = bool(
            replace_argmax_to_fused_argmax_and_indices_is_int64
        )
        self.replace_argmax_to_fused_argmax_and_indices_is_float32 = bool(
            replace_argmax_to_fused_argmax_and_indices_is_float32
        )
        self.fused_argmax_scale_ratio = float(fused_argmax_scale_ratio)
        self.argmax_mode = "none"
        if self.replace_argmax_to_reducemax_and_indices_is_int64:
            self.argmax_mode = "reducemax_int64"
        elif self.replace_argmax_to_reducemax_and_indices_is_float32:
            self.argmax_mode = "reducemax_float32"
        elif self.replace_argmax_to_fused_argmax_and_indices_is_int64:
            self.argmax_mode = "fused_int64"
        elif self.replace_argmax_to_fused_argmax_and_indices_is_float32:
            self.argmax_mode = "fused_float32"
        self.fused_argmax_restore_shapes: Dict[str, Dict[str, Any]] = {}
        nms_version = str(switch_nms_version).strip().lower()
        if nms_version not in {"v4", "v5"}:
            raise ValueError(
                "switch_nms_version must be 'v4' or 'v5'. "
                f"got: {switch_nms_version}"
            )
        self.switch_nms_version = nms_version
        self.replace_to_pseudo_operators: set[str] = {
            str(v).strip().lower()
            for v in (replace_to_pseudo_operators or [])
            if str(v).strip() != ""
        }
        self._serial = 0
        self.onnx_tensor_consumers: Dict[str, List[Any]] = {}
        self.onnx_tensor_producers: Dict[str, Any] = {}
        if (
            self.onnx_model is not None
            and hasattr(self.onnx_model, "graph")
            and hasattr(self.onnx_model.graph, "node")
        ):
            for graph_node in self.onnx_model.graph.node:
                for output_name in graph_node.output:
                    key = str(output_name).strip()
                    if key == "":
                        continue
                    self.onnx_tensor_producers[key] = graph_node
                for input_name in graph_node.input:
                    key = str(input_name).strip()
                    if key == "":
                        continue
                    self.onnx_tensor_consumers.setdefault(key, []).append(graph_node)

    def _next_name(self, base: str) -> str:
        self._serial += 1
        return f"{base}_{self._serial}"

    def get_tensor_shape(self, name: str) -> List[int]:
        if name in self.model_ir.tensors:
            return list(self.model_ir.tensors[name].shape)
        shape = self.shape_map.get(name, None)
        norm_shape, _ = normalize_onnx_shape(shape)
        return norm_shape

    def get_tensor_dtype(self, name: str) -> str:
        if name in self.model_ir.tensors:
            return self.model_ir.tensors[name].dtype
        return self.dtype_map.get(name, "FLOAT32")

    def get_constant_array(self, name: str) -> Optional[np.ndarray]:
        if name in self.constants:
            return self.constants[name]
        t = self.model_ir.tensors.get(name, None)
        if t is not None and isinstance(t.data, np.ndarray):
            return t.data
        return None

    def ensure_tensor(self, name: str, dtype: Optional[str] = None, shape: Optional[List[int]] = None) -> str:
        if name == "":
            raise ValueError("Tensor name must not be empty in flatbuffer_direct lowering.")
        if name in self.model_ir.tensors:
            return name
        if dtype is None:
            dtype = self.dtype_map.get(name, "FLOAT32")
        if shape is None:
            shape = self.shape_map.get(name, None)
        shape, signature = normalize_onnx_shape(shape)
        self.model_ir.tensors[name] = TensorIR(
            name=name,
            dtype=dtype,
            shape=list(shape),
            shape_signature=list(signature),
            data=self.constants.get(name, None),
        )
        return name

    def add_const_tensor(self, base_name: str, data: np.ndarray) -> str:
        name = base_name
        if name in self.model_ir.tensors:
            name = self._next_name(base_name)
        data = np.asarray(data)
        dtype = tflite_dtype_from_numpy(data.dtype)
        shape, signature = normalize_onnx_shape(list(data.shape))
        self.model_ir.tensors[name] = TensorIR(
            name=name,
            dtype=dtype,
            shape=shape,
            shape_signature=signature,
            data=data,
        )
        self.constants[name] = data
        return name

    def add_intermediate_tensor(self, base_name: str, dtype: str, shape: List[int]) -> str:
        if base_name == "":
            raise ValueError("Tensor name must not be empty in flatbuffer_direct lowering.")
        name = base_name
        if name in self.model_ir.tensors:
            name = self._next_name(base_name)
        norm_shape, signature = normalize_onnx_shape(shape)
        self.model_ir.tensors[name] = TensorIR(
            name=name,
            dtype=dtype,
            shape=norm_shape,
            shape_signature=signature,
            data=None,
        )
        return name

    def add_operator(self, op: OperatorIR) -> None:
        self.model_ir.operators.append(op)


class _NodeWrap:
    def __init__(
        self,
        n: onnx.NodeProto,
        input_name_remap: Optional[Dict[str, str]] = None,
    ):
        self.name = n.name if n.name else n.op_type
        self.op = n.op_type
        self.attrs = {}
        for a in n.attribute:
            if a.type == onnx.AttributeProto.INT:
                self.attrs[a.name] = int(a.i)
            elif a.type == onnx.AttributeProto.FLOAT:
                self.attrs[a.name] = float(a.f)
            elif a.type == onnx.AttributeProto.INTS:
                self.attrs[a.name] = [int(v) for v in a.ints]
            elif a.type == onnx.AttributeProto.FLOATS:
                self.attrs[a.name] = [float(v) for v in a.floats]
            elif a.type == onnx.AttributeProto.STRING:
                self.attrs[a.name] = a.s.decode("utf-8")
            elif a.type == onnx.AttributeProto.STRINGS:
                self.attrs[a.name] = [v.decode("utf-8") for v in a.strings]
            elif a.type == onnx.AttributeProto.GRAPH:
                self.attrs[a.name] = a.g
            elif a.type == onnx.AttributeProto.GRAPHS:
                self.attrs[a.name] = [g for g in a.graphs]
            else:
                pass
        remap = input_name_remap if isinstance(input_name_remap, dict) else {}
        self.inputs = [
            type(
                "In",
                (),
                {"name": (remap.get(i, i) if i != "" else "")},
            )
            for i in n.input
        ]
        self.outputs = [type("Out", (), {"name": o}) for o in n.output if o != ""]


def _build_tensor_consumer_map(model_ir: ModelIR) -> Dict[str, List[int]]:
    consumers: Dict[str, List[int]] = {}
    for op_idx, op in enumerate(model_ir.operators):
        for input_name in op.inputs:
            if input_name not in consumers:
                consumers[input_name] = []
            consumers[input_name].append(op_idx)
    return consumers


def _read_transpose_perm(model_ir: ModelIR, op: OperatorIR) -> Optional[List[int]]:
    if str(op.op_type) != "TRANSPOSE":
        return None
    if len(op.inputs) < 2:
        return None
    perm_tensor_name = op.inputs[1]
    perm_tensor = model_ir.tensors.get(perm_tensor_name, None)
    if perm_tensor is None or perm_tensor.data is None:
        return None
    perm = np.asarray(perm_tensor.data).reshape(-1)
    if perm.size == 0:
        return None
    return [int(v) for v in perm.tolist()]


def _is_identity_perm(perm: List[int]) -> bool:
    return perm == [int(i) for i in range(len(perm))]


def _is_inverse_perm(perm_a: List[int], perm_b: List[int]) -> bool:
    if len(perm_a) != len(perm_b):
        return False
    rank = len(perm_a)
    if sorted(perm_a) != [int(i) for i in range(rank)]:
        return False
    if sorted(perm_b) != [int(i) for i in range(rank)]:
        return False
    for idx, value in enumerate(perm_a):
        if perm_b[value] != idx:
            return False
    return True


def _invert_perm(perm: List[int]) -> Optional[List[int]]:
    rank = len(perm)
    if sorted(perm) != [int(i) for i in range(rank)]:
        return None
    inv = [0 for _ in range(rank)]
    for idx, value in enumerate(perm):
        inv[int(value)] = int(idx)
    return inv


def _append_tensor_lineage_event(
    *,
    model_ir: ModelIR,
    event: Dict[str, Any],
) -> None:
    if "tensor_lineage_events" not in model_ir.metadata:
        model_ir.metadata["tensor_lineage_events"] = []
    events = model_ir.metadata["tensor_lineage_events"]
    if not isinstance(events, list):
        events = []
        model_ir.metadata["tensor_lineage_events"] = events
    normalized_event = dict(event)
    normalized_event["event_index"] = int(len(events))
    events.append(normalized_event)


def _replace_tensor_inputs(model_ir: ModelIR, src_name: str, dst_name: str) -> None:
    if str(src_name) != str(dst_name):
        _append_tensor_lineage_event(
            model_ir=model_ir,
            event={
                "kind": "replace_input",
                "src_name": str(src_name),
                "dst_name": str(dst_name),
            },
        )
    for op in model_ir.operators:
        if len(op.inputs) > 0:
            op.inputs = [dst_name if input_name == src_name else input_name for input_name in op.inputs]


def _set_operator_inputs(
    *,
    model_ir: ModelIR,
    op: OperatorIR,
    new_inputs: List[str],
) -> None:
    old_inputs = [str(v) for v in list(op.inputs)]
    normalized_new_inputs = [str(v) for v in list(new_inputs)]
    pair_count = int(min(len(old_inputs), len(normalized_new_inputs)))
    for index in range(pair_count):
        old_name = str(old_inputs[index])
        new_name = str(normalized_new_inputs[index])
        if old_name == new_name:
            continue
        _append_tensor_lineage_event(
            model_ir=model_ir,
            event={
                "kind": "replace_input",
                "src_name": old_name,
                "dst_name": new_name,
                "source": "set_operator_inputs",
            },
        )
    op.inputs = normalized_new_inputs


def _set_operator_outputs(
    *,
    model_ir: ModelIR,
    op: OperatorIR,
    new_outputs: List[str],
) -> None:
    old_outputs = [str(v) for v in list(op.outputs)]
    normalized_new_outputs = [str(v) for v in list(new_outputs)]
    pair_count = int(min(len(old_outputs), len(normalized_new_outputs)))
    for index in range(pair_count):
        old_name = str(old_outputs[index])
        new_name = str(normalized_new_outputs[index])
        if old_name == new_name:
            continue
        _append_tensor_lineage_event(
            model_ir=model_ir,
            event={
                "kind": "rename_tensor",
                "old_name": old_name,
                "new_name": new_name,
                "source": "set_operator_outputs",
            },
        )
    op.outputs = normalized_new_outputs


def _get_protected_boundary_tensor_names(model_ir: ModelIR) -> set[str]:
    raw_names = model_ir.metadata.get(
        "protected_boundary_tensor_names",
        [],
    )
    if not isinstance(raw_names, list):
        return set()
    return {
        str(name)
        for name in raw_names
        if str(name).strip() != ""
    }


def _append_model_outputs_preserving_order(
    model_ir: ModelIR,
    output_names: List[str],
) -> None:
    if len(output_names) <= 0:
        return
    merged_outputs = [
        str(name)
        for name in list(model_ir.outputs) + list(output_names)
        if str(name) != ""
    ]
    model_ir.outputs = list(dict.fromkeys(merged_outputs))


def _replace_operator_input_at(
    *,
    model_ir: ModelIR,
    op: OperatorIR,
    input_index: int,
    new_input_name: str,
) -> None:
    if int(input_index) < 0 or int(input_index) >= len(op.inputs):
        return
    old_name = str(op.inputs[int(input_index)])
    new_name = str(new_input_name)
    if old_name != new_name:
        _append_tensor_lineage_event(
            model_ir=model_ir,
            event={
                "kind": "replace_input",
                "src_name": old_name,
                "dst_name": new_name,
                "source": "replace_operator_input_at",
            },
        )
    op.inputs[int(input_index)] = new_name


def _prune_unused_tensors(model_ir: ModelIR) -> None:
    used_tensor_names = set(model_ir.inputs + model_ir.outputs)
    for op in model_ir.operators:
        used_tensor_names.update(op.inputs)
        used_tensor_names.update(op.outputs)
    unused_tensor_names = [name for name in model_ir.tensors.keys() if name not in used_tensor_names]
    if len(unused_tensor_names) > 0:
        _append_tensor_lineage_event(
            model_ir=model_ir,
            event={
                "kind": "prune_unused_tensors",
                "removed_names": [str(v) for v in unused_tensor_names],
            },
        )
    for name in unused_tensor_names:
        del model_ir.tensors[name]


def _prune_dead_operators(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove operators that do not contribute to graph outputs.

    This pass performs a reverse liveness walk from model outputs and keeps only
    operators whose outputs are required by downstream live operators or graph outputs.
    """
    if len(model_ir.operators) == 0:
        return {"removed_dead_operators": 0}

    live_tensors = set(model_ir.outputs)
    keep_flags = [False for _ in model_ir.operators]

    for op_idx in range(len(model_ir.operators) - 1, -1, -1):
        op = model_ir.operators[op_idx]
        outputs = list(op.outputs)
        outputs_live = any(out_name in live_tensors for out_name in outputs)

        # Some kernels (e.g. UNIDIRECTIONAL_SEQUENCE_RNN) mutate variable input
        # tensors in-place. Keep such ops when a live tensor depends on that
        # variable state even if the op's explicit outputs are currently dead.
        mutates_live_variable_input = False
        if not outputs_live:
            for input_name in op.inputs:
                if input_name not in live_tensors:
                    continue
                input_tensor = model_ir.tensors.get(str(input_name), None)
                if input_tensor is not None and bool(input_tensor.is_variable):
                    mutates_live_variable_input = True
                    break

        if outputs_live or mutates_live_variable_input:
            keep_flags[op_idx] = True
            for input_name in op.inputs:
                live_tensors.add(input_name)

    removed = int(sum(1 for keep in keep_flags if not keep))
    if removed == 0:
        return {"removed_dead_operators": 0}

    model_ir.operators = [
        op for idx, op in enumerate(model_ir.operators) if keep_flags[idx]
    ]
    _prune_unused_tensors(model_ir)
    return {"removed_dead_operators": removed}


def _find_unbound_nonconstant_operator_inputs(model_ir: ModelIR) -> List[Dict[str, Any]]:
    """
    Detect operator inputs that are neither:
    - produced by another operator,
    - model inputs, nor
    - constant tensors with embedded data.

    Such tensors become unexpected runtime-fed inputs in TFLite and can trigger
    errors like "Input tensor N lacks data".
    """
    producers = _build_tensor_producer_map(model_ir)
    model_inputs = {str(v) for v in list(model_ir.inputs)}
    issues: List[Dict[str, Any]] = []

    for op_index, op in enumerate(model_ir.operators):
        for input_index, input_name_raw in enumerate(op.inputs):
            input_name = str(input_name_raw)
            if input_name == "":
                continue
            if input_name in model_inputs:
                continue
            if input_name in producers:
                continue
            tensor = model_ir.tensors.get(input_name, None)
            if tensor is not None and tensor.data is not None:
                continue
            issues.append(
                {
                    "op_index": int(op_index),
                    "op_type": str(op.op_type),
                    "input_index": int(input_index),
                    "tensor_name": str(input_name),
                }
            )
    return issues


def _repair_unbound_nonconstant_operator_inputs_with_layout_transpose(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Repair a strict subset of unbound dynamic inputs by inserting a layout transpose.

    Current targets (conservative):
    - Unbound tensor `t` is rank-4 NCHW and used as data input of:
      - RESHAPE input[0], or
      - SHAPE input[0], or
      - SPLIT input[1].
    - There exists an earlier produced rank-4 NHWC tensor `s` with
      shape [N,H,W,C] where `t` has shape [N,C,H,W].
    - Insert `TRANSPOSE(s, [0,3,1,2]) -> t` just before the consumer.
    - Unbound tensor `t` is used by MUL fanout (input[0]) and a nearest
      upstream ADD already materializes the corresponding NHWC tensor under a
      renamed `_input_nhwc` output. Insert `TRANSPOSE(s, [0,3,1,2]) -> t`
      to reconnect dropped alias names.
    """
    repaired = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        serial = 1
        while name in model_ir.tensors:
            name = f"{base}_{serial}"
            serial += 1
        return name

    while True:
        issues = _find_unbound_nonconstant_operator_inputs(model_ir)
        if len(issues) <= 0:
            break

        producers = _build_tensor_producer_map(model_ir)
        consumers = _build_tensor_consumer_map(model_ir)
        changed = False

        for issue in issues:
            op_index = int(issue.get("op_index", -1))
            input_index = int(issue.get("input_index", -1))
            tensor_name = str(issue.get("tensor_name", ""))
            if op_index < 0 or op_index >= len(model_ir.operators):
                continue
            consumer_op = model_ir.operators[int(op_index)]
            consumer_op_type = str(consumer_op.op_type)
            target_input_index = 0
            if consumer_op_type == "SPLIT":
                target_input_index = 1
            if consumer_op_type in ["RESHAPE", "SHAPE", "SPLIT"] and int(input_index) == int(target_input_index):
                if input_index >= len(consumer_op.inputs) or str(consumer_op.inputs[input_index]) != tensor_name:
                    continue

                orphan_tensor = model_ir.tensors.get(tensor_name, None)
                if orphan_tensor is None or not _is_fully_known_positive_shape(orphan_tensor.shape):
                    continue
                orphan_shape = [int(v) for v in list(orphan_tensor.shape)]
                if len(orphan_shape) != 4:
                    continue

                n, c, h, w = [int(v) for v in list(orphan_shape)]
                expected_source_shape = [int(n), int(h), int(w), int(c)]

                best_source_name: Optional[str] = None
                best_source_op_index = -1
                for produced_name, source_op_index in producers.items():
                    source_idx = int(source_op_index)
                    if source_idx >= int(op_index):
                        continue
                    source_tensor = model_ir.tensors.get(str(produced_name), None)
                    if source_tensor is None:
                        continue
                    if source_tensor.data is not None:
                        continue
                    if not _is_fully_known_positive_shape(source_tensor.shape):
                        continue
                    source_shape = [int(v) for v in list(source_tensor.shape)]
                    if source_shape != expected_source_shape:
                        continue
                    if source_idx > int(best_source_op_index):
                        best_source_op_index = int(source_idx)
                        best_source_name = str(produced_name)

                if best_source_name is None:
                    continue

                perm_name = _unique_tensor_name(f"{tensor_name}_repair_perm")
                model_ir.tensors[perm_name] = TensorIR(
                    name=perm_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray([0, 3, 1, 2], dtype=np.int32),
                    is_variable=False,
                )

                model_ir.operators.insert(
                    int(op_index),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(best_source_name), str(perm_name)],
                        outputs=[str(tensor_name)],
                        options={},
                    ),
                )

                repaired += 1
                changed = True
                break

            # Alias-repair fallback: recover dropped ONNX-style intermediate names
            # (e.g. `input.xxx`) when a same-shape NHWC ADD output exists nearby.
            if input_index != 0:
                continue
            if consumer_op_type != "MUL":
                continue
            if input_index >= len(consumer_op.inputs) or str(consumer_op.inputs[input_index]) != tensor_name:
                continue
            if not str(tensor_name).startswith("input."):
                continue
            tensor_consumers = [int(v) for v in consumers.get(str(tensor_name), [])]
            if len(tensor_consumers) == 0:
                continue
            if not all(
                0 <= int(idx) < len(model_ir.operators)
                and str(model_ir.operators[int(idx)].op_type) == "MUL"
                for idx in tensor_consumers
            ):
                continue
            if not all(
                len(model_ir.operators[int(idx)].inputs) > 0
                and str(model_ir.operators[int(idx)].inputs[0]) == str(tensor_name)
                for idx in tensor_consumers
            ):
                continue

            orphan_tensor = model_ir.tensors.get(tensor_name, None)
            if orphan_tensor is None or not _is_fully_known_positive_shape(orphan_tensor.shape):
                continue
            orphan_shape = [int(v) for v in list(orphan_tensor.shape)]
            orphan_dtype = str(orphan_tensor.dtype)
            if len(orphan_shape) != 4:
                continue
            n, c, h, w = [int(v) for v in list(orphan_shape)]
            expected_source_shape = [int(n), int(h), int(w), int(c)]

            best_source_name: Optional[str] = None
            best_source_op_index = -1
            for produced_name, source_op_index in producers.items():
                source_idx = int(source_op_index)
                if source_idx >= int(op_index):
                    continue
                source_name = str(produced_name)
                if not source_name.endswith("_input_nhwc"):
                    continue
                source_op = model_ir.operators[int(source_idx)]
                if str(source_op.op_type) != "ADD":
                    continue
                source_tensor = model_ir.tensors.get(source_name, None)
                if source_tensor is None or source_tensor.data is not None:
                    continue
                if not _is_fully_known_positive_shape(source_tensor.shape):
                    continue
                if [int(v) for v in list(source_tensor.shape)] != expected_source_shape:
                    continue
                if str(source_tensor.dtype) != orphan_dtype:
                    continue
                if int(source_idx) > int(best_source_op_index):
                    best_source_op_index = int(source_idx)
                    best_source_name = str(source_name)

            if best_source_name is None:
                continue

            perm_name = _unique_tensor_name(f"{tensor_name}_repair_perm")
            model_ir.tensors[perm_name] = TensorIR(
                name=perm_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray([0, 3, 1, 2], dtype=np.int32),
                is_variable=False,
            )

            model_ir.operators.insert(
                int(op_index),
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(best_source_name), str(perm_name)],
                    outputs=[str(tensor_name)],
                    options={},
                ),
            )
            source_tensor = model_ir.tensors.get(str(best_source_name), None)
            if source_tensor is not None:
                orphan_tensor.quantization = _clone_quantization(source_tensor.quantization)
                orphan_tensor.shape_signature = (
                    [int(v) for v in list(source_tensor.shape_signature)]
                    if source_tensor.shape_signature is not None
                    else [int(v) for v in list(orphan_shape)]
                )

            repaired += 1
            changed = True
            break

        if not changed:
            break

    if repaired > 0:
        _reconcile_static_tensor_shapes(model_ir)
    return {"repaired_unbound_nonconstant_inputs_with_layout_transpose": int(repaired)}


def _count_ops_by_type(model_ir: ModelIR, op_type: str) -> int:
    target = str(op_type)
    return int(sum(1 for op in model_ir.operators if str(op.op_type) == target))


def _apply_safe_transpose_reduction_lite(model_ir: ModelIR) -> Dict[str, int]:
    """
    Run a conservative transpose-reduction bundle for fallback/non-aggressive paths.

    Safety contract:
    - Apply only a curated set of historically safe transpose/layout passes.
    - Roll back all changes if unbound runtime inputs are introduced.
    """
    before_transpose = _count_ops_by_type(model_ir, "TRANSPOSE")
    if int(before_transpose) <= 0:
        return {
            "safe_transpose_reduction_lite_applied": 0,
            "safe_transpose_reduction_lite_reduced": 0,
            "safe_transpose_reduction_lite_unbound_after": 0,
        }

    snapshot = copy.deepcopy(model_ir)
    pass_sequence = [
        _optimize_transpose_quant_dequant_bridges,
        _optimize_duplicate_transpose_fanout,
        _optimize_singleton_channel_layout_transpose_to_reshape,
        _optimize_transpose_unary_fanout_inverse_post_bridges,
        _optimize_transpose_unary_passthrough_chains,
        _optimize_transpose_elementwise_concat_conv_nhwc_groups,
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter,
        _optimize_transpose_pre_add_nhwc_chains,
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains,
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains,
        _optimize_transposeconv_output_nhwc_passthrough_chains,
    ]

    applied = 0
    for optimize_pass in pass_sequence:
        optimize_pass(model_ir)
        applied += 1

    _prune_dead_operators(model_ir)
    _reconcile_static_tensor_shapes(model_ir)

    unbound_after = _find_unbound_nonconstant_operator_inputs(model_ir)
    if len(unbound_after) > 0:
        model_ir.tensors = snapshot.tensors
        model_ir.operators = snapshot.operators
        model_ir.inputs = snapshot.inputs
        model_ir.outputs = snapshot.outputs
        model_ir.subgraphs = snapshot.subgraphs
        model_ir.metadata = snapshot.metadata
        return {
            "safe_transpose_reduction_lite_applied": 0,
            "safe_transpose_reduction_lite_reduced": 0,
            "safe_transpose_reduction_lite_unbound_after": int(len(unbound_after)),
        }

    after_transpose = _count_ops_by_type(model_ir, "TRANSPOSE")
    reduced = int(before_transpose - after_transpose)
    if reduced <= 0:
        model_ir.tensors = snapshot.tensors
        model_ir.operators = snapshot.operators
        model_ir.inputs = snapshot.inputs
        model_ir.outputs = snapshot.outputs
        model_ir.subgraphs = snapshot.subgraphs
        model_ir.metadata = snapshot.metadata
        return {
            "safe_transpose_reduction_lite_applied": 0,
            "safe_transpose_reduction_lite_reduced": 0,
            "safe_transpose_reduction_lite_unbound_after": 0,
        }

    return {
        "safe_transpose_reduction_lite_applied": int(applied),
        "safe_transpose_reduction_lite_reduced": int(reduced),
        "safe_transpose_reduction_lite_unbound_after": 0,
    }


def _build_tensor_producer_map(model_ir: ModelIR) -> Dict[str, int]:
    producers: Dict[str, int] = {}
    for op_idx, op in enumerate(model_ir.operators):
        for output_name in op.outputs:
            producers[output_name] = op_idx
    return producers


def _topologically_sort_operators(model_ir: ModelIR) -> Dict[str, int]:
    """
    Reorder operators so every producer appears before its consumers.

    This is a safety net for late graph-rewrite passes that may rewrite tensor
    names without physically moving operators.
    """
    op_count = len(model_ir.operators)
    if op_count <= 1:
        return {"reordered_operators": 0, "cycle_detected": 0}

    producers = _build_tensor_producer_map(model_ir)
    dependencies: List[set[int]] = [set() for _ in range(op_count)]
    dependents: List[set[int]] = [set() for _ in range(op_count)]

    for consumer_idx, op in enumerate(model_ir.operators):
        for input_name in op.inputs:
            producer_idx = producers.get(str(input_name), None)
            if producer_idx is None or int(producer_idx) == int(consumer_idx):
                continue
            if int(producer_idx) in dependencies[int(consumer_idx)]:
                continue
            dependencies[int(consumer_idx)].add(int(producer_idx))
            dependents[int(producer_idx)].add(int(consumer_idx))

    indegree = [int(len(v)) for v in dependencies]
    ready = [int(idx) for idx, deg in enumerate(indegree) if int(deg) == 0]
    ready.sort()
    sorted_indices: List[int] = []

    while len(ready) > 0:
        current = int(ready.pop(0))
        sorted_indices.append(current)
        for dependent_idx in sorted(list(dependents[current])):
            indegree[int(dependent_idx)] -= 1
            if int(indegree[int(dependent_idx)]) == 0:
                ready.append(int(dependent_idx))
        ready.sort()

    if len(sorted_indices) != op_count:
        return {"reordered_operators": 0, "cycle_detected": 1}

    if sorted_indices == [int(v) for v in range(op_count)]:
        return {"reordered_operators": 0, "cycle_detected": 0}

    original_ops = list(model_ir.operators)
    model_ir.operators = [original_ops[int(idx)] for idx in sorted_indices]
    return {
        "reordered_operators": int(
            sum(1 for new_idx, old_idx in enumerate(sorted_indices) if int(new_idx) != int(old_idx))
        ),
        "cycle_detected": 0,
    }


def _rename_tensor_globally(
    model_ir: ModelIR,
    old_name: str,
    new_name: str,
) -> None:
    if old_name == new_name:
        return

    for op in model_ir.operators:
        if len(op.inputs) > 0:
            op.inputs = [new_name if input_name == old_name else input_name for input_name in op.inputs]
        if len(op.outputs) > 0:
            op.outputs = [new_name if output_name == old_name else output_name for output_name in op.outputs]

    if len(model_ir.inputs) > 0:
        model_ir.inputs = [new_name if input_name == old_name else input_name for input_name in model_ir.inputs]
    if len(model_ir.outputs) > 0:
        model_ir.outputs = [new_name if output_name == old_name else output_name for output_name in model_ir.outputs]

    old_tensor = model_ir.tensors.get(old_name, None)
    if old_tensor is None:
        if new_name in model_ir.tensors:
            del model_ir.tensors[new_name]
        return
    old_tensor.name = new_name
    if new_name in model_ir.tensors and new_name != old_name:
        del model_ir.tensors[new_name]
    del model_ir.tensors[old_name]
    model_ir.tensors[new_name] = old_tensor
    _append_tensor_lineage_event(
        model_ir=model_ir,
        event={
            "kind": "rename_tensor",
            "old_name": str(old_name),
            "new_name": str(new_name),
        },
    )


def _clone_quantization(quantization: Any) -> Any:
    if quantization is None:
        return None
    if isinstance(quantization, QuantParamIR):
        return QuantParamIR(
            scale=list(quantization.scale),
            zero_point=list(quantization.zero_point),
            quantized_dimension=int(quantization.quantized_dimension),
            min=list(quantization.min) if quantization.min is not None else None,
            max=list(quantization.max) if quantization.max is not None else None,
        )
    return copy.deepcopy(quantization)


def _quantize_prelu_alpha(
    alpha: np.ndarray,
    target_dtype: str,
) -> Tuple[np.ndarray, QuantParamIR]:
    alpha = np.asarray(alpha, dtype=np.float32)
    if str(target_dtype) == "INT8":
        max_abs = float(np.max(np.abs(alpha))) if alpha.size > 0 else 0.0
        scale = max(max_abs / 127.0, 1e-8)
        quantized = np.clip(np.round(alpha / scale), -128, 127).astype(np.int8)
        return quantized, QuantParamIR(
            scale=[float(scale)],
            zero_point=[0],
            quantized_dimension=0,
        )
    if str(target_dtype) == "UINT8":
        mn = float(np.min(alpha)) if alpha.size > 0 else 0.0
        mx = float(np.max(alpha)) if alpha.size > 0 else 0.0
        scale = max((mx - mn) / 255.0, 1e-8)
        zp = int(np.round(-mn / scale))
        zp = int(np.clip(zp, 0, 255))
        quantized = np.clip(np.round(alpha / scale) + zp, 0, 255).astype(np.uint8)
        return quantized, QuantParamIR(
            scale=[float(scale)],
            zero_point=[int(zp)],
            quantized_dimension=0,
        )
    raise NotImplementedError(f"Quantized PRELU alpha requires INT8/UINT8. got={target_dtype}")


def _quantize_tensor_per_tensor(
    values: np.ndarray,
    target_dtype: str,
) -> Tuple[np.ndarray, QuantParamIR]:
    values = np.asarray(values, dtype=np.float32)
    if str(target_dtype) == "INT8":
        max_abs = float(np.max(np.abs(values))) if values.size > 0 else 0.0
        scale = max(max_abs / 127.0, 1e-8)
        quantized = np.clip(np.round(values / scale), -128, 127).astype(np.int8)
        return quantized, QuantParamIR(
            scale=[float(scale)],
            zero_point=[0],
            quantized_dimension=0,
        )
    if str(target_dtype) == "UINT8":
        mn = float(np.min(values)) if values.size > 0 else 0.0
        mx = float(np.max(values)) if values.size > 0 else 0.0
        scale = max((mx - mn) / 255.0, 1e-8)
        zp = int(np.round(-mn / scale))
        zp = int(np.clip(zp, 0, 255))
        quantized = np.clip(np.round(values / scale) + zp, 0, 255).astype(np.uint8)
        return quantized, QuantParamIR(
            scale=[float(scale)],
            zero_point=[int(zp)],
            quantized_dimension=0,
        )
    raise NotImplementedError(f"Per-tensor quantization supports INT8/UINT8 only. got={target_dtype}")


def _get_per_tensor_scale_zero_point(quantization: Any) -> Optional[Tuple[float, int]]:
    if quantization is None:
        return None
    if not _is_per_tensor_quantization(quantization):
        return None
    if isinstance(quantization, QuantParamIR):
        if len(list(quantization.scale)) == 0 or len(list(quantization.zero_point)) == 0:
            return None
        return float(quantization.scale[0]), int(quantization.zero_point[0])
    if isinstance(quantization, dict):
        try:
            scale = np.asarray(quantization.get("scale", []), dtype=np.float32).reshape(-1)
            zero_point = np.asarray(quantization.get("zero_point", []), dtype=np.int64).reshape(-1)
            if scale.size == 0 or zero_point.size == 0:
                return None
            return float(scale[0]), int(zero_point[0])
        except Exception:
            return None
    return None


def _is_same_per_tensor_quantization(
    quantization_a: Any,
    quantization_b: Any,
    *,
    rtol: float = 1e-6,
    atol: float = 1e-8,
) -> bool:
    qparams_a = _get_per_tensor_scale_zero_point(quantization_a)
    qparams_b = _get_per_tensor_scale_zero_point(quantization_b)
    if qparams_a is None or qparams_b is None:
        return False
    scale_a, zp_a = qparams_a
    scale_b, zp_b = qparams_b
    if int(zp_a) != int(zp_b):
        return False
    return bool(np.isclose(float(scale_a), float(scale_b), rtol=rtol, atol=atol))


def _quant_scale_count(quantization: Any) -> int:
    if quantization is None:
        return 0
    if isinstance(quantization, QuantParamIR):
        return len(list(quantization.scale))
    if isinstance(quantization, dict):
        scale = quantization.get("scale", None)
        if scale is None:
            return 0
        if isinstance(scale, (list, tuple)):
            return len(list(scale))
        try:
            return int(np.asarray(scale).size)
        except Exception:
            return 0
    return 0


def _is_per_tensor_quantization(quantization: Any) -> bool:
    count = _quant_scale_count(quantization)
    return count <= 1


def _permute_shape(shape: Optional[List[int]], perm: List[int]) -> Optional[List[int]]:
    if shape is None:
        return None
    rank = len(shape)
    if len(perm) != rank:
        return None
    if sorted(perm) != [int(i) for i in range(rank)]:
        return None
    return [int(shape[int(axis)]) for axis in perm]


def _shapes_equal(shape_a: Optional[List[int]], shape_b: Optional[List[int]]) -> bool:
    if shape_a is None or shape_b is None:
        return False
    if len(shape_a) != len(shape_b):
        return False
    return all(int(a) == int(b) for a, b in zip(shape_a, shape_b))


def _is_unknown_shape(shape: Optional[List[int]]) -> bool:
    if shape is None:
        return True
    if len(shape) == 0:
        return True
    if any(int(dim) < 0 for dim in shape):
        return True
    # Internal placeholder shape used before static shape propagation.
    if len(shape) == 1 and int(shape[0]) == 1:
        return True
    return False


def _shapes_match_if_known(shape_a: Optional[List[int]], shape_b: Optional[List[int]]) -> bool:
    if _is_unknown_shape(shape_a) or _is_unknown_shape(shape_b):
        return True
    return _shapes_equal(shape_a, shape_b)


def _all_per_tensor_quantized(tensors: List[Optional[TensorIR]]) -> bool:
    for tensor in tensors:
        if tensor is None:
            continue
        if tensor.quantization is None:
            continue
        if not _is_per_tensor_quantization(tensor.quantization):
            return False
    return True


def _resolve_reshape_new_shape_from_static_input(
    new_shape: List[int],
    input_signature: Optional[List[int]],
    allow_zero: Optional[bool] = None,
) -> Optional[List[int]]:
    if len(new_shape) == 0:
        return None
    candidate = [int(v) for v in new_shape]

    if any(int(dim) == 0 for dim in candidate):
        # Without allowzero metadata we avoid rewriting zero-dim requests.
        if allow_zero is None:
            return None
        if allow_zero:
            return None
        if input_signature is None or len(input_signature) == 0:
            return None
        for idx, dim in enumerate(candidate):
            if int(dim) != 0:
                continue
            if idx >= len(input_signature):
                return None
            in_dim = int(input_signature[idx])
            if in_dim <= 0:
                return None
            candidate[idx] = int(in_dim)

    if all(int(dim) >= 0 for dim in candidate):
        if input_signature is None or len(input_signature) == 0:
            return list(candidate)
        if any(int(dim) <= 0 for dim in input_signature):
            return list(candidate)

        input_product = int(np.prod(np.asarray(input_signature, dtype=np.int64)))
        if input_product <= 0:
            return list(candidate)
        target_product = int(np.prod(np.asarray(candidate, dtype=np.int64)))
        if target_product == input_product:
            return list(candidate)

        # Stale reshape constants can appear after late layout rewrites.
        # If exactly one dimension can be re-inferred from static input size,
        # repair that dimension instead of keeping an invalid constant.
        inferable_dims: List[Tuple[int, int]] = []
        for idx in range(len(candidate)):
            other_product = 1
            valid = True
            for j, dim in enumerate(candidate):
                if j == idx:
                    continue
                d = int(dim)
                if d <= 0:
                    valid = False
                    break
                other_product *= d
            if not valid or other_product <= 0:
                continue
            if input_product % other_product != 0:
                continue
            inferred_dim = int(input_product // other_product)
            if inferred_dim <= 0:
                continue
            inferable_dims.append((int(idx), int(inferred_dim)))

        if len(inferable_dims) == 1:
            dim_idx, inferred_dim = inferable_dims[0]
            repaired = [int(v) for v in candidate]
            repaired[dim_idx] = int(inferred_dim)
            return repaired

        return None

    minus_one_indices = [idx for idx, dim in enumerate(candidate) if int(dim) == -1]
    if len(minus_one_indices) != 1:
        return None
    if input_signature is None or len(input_signature) == 0:
        return None
    if any(int(dim) <= 0 for dim in input_signature):
        return None

    known_product = 1
    for dim in candidate:
        if int(dim) == -1:
            continue
        if int(dim) <= 0:
            return None
        known_product *= int(dim)
    if known_product <= 0:
        return None

    input_product = int(np.prod(np.asarray(input_signature, dtype=np.int64)))
    if input_product <= 0 or input_product % known_product != 0:
        return None
    inferred = int(input_product // known_product)
    if inferred <= 0:
        return None
    candidate[minus_one_indices[0]] = inferred
    return candidate


def _resolve_dynamic_reshape_shapes(
    model_ir: ModelIR,
    prefer_runtime_inferable_from_onnx_raw: bool = False,
) -> Dict[str, int]:
    def _sanitize_reshape_template(
        *,
        template: List[int],
        input_dims: List[int],
    ) -> List[int]:
        sanitized = [int(v) for v in list(template)]
        for idx, dim in enumerate(list(sanitized)):
            if int(dim) == 0:
                copied = (
                    int(input_dims[idx])
                    if idx < len(input_dims) and int(input_dims[idx]) > 0
                    else None
                )
                if copied is not None:
                    sanitized[idx] = int(copied)
                else:
                    sanitized[idx] = -1
            elif int(dim) < -1:
                sanitized[idx] = 1

        minus_one_indices = [idx for idx, dim in enumerate(sanitized) if int(dim) == -1]
        if len(minus_one_indices) > 1:
            keep_idx = int(minus_one_indices[0])
            for idx in minus_one_indices[1:]:
                copied = (
                    int(input_dims[idx])
                    if idx < len(input_dims) and int(input_dims[idx]) > 0
                    else 1
                )
                sanitized[idx] = int(copied)
            sanitized[keep_idx] = -1

        return [int(v) for v in list(sanitized)]

    resolved_count = 0
    for op in model_ir.operators:
        if str(op.op_type) != "RESHAPE":
            continue
        if bool(op.options.get("preserveDynamicShape", False)):
            continue
        if len(op.inputs) < 1 or len(op.outputs) != 1:
            continue

        has_onnx_raw_new_shape = "onnxRawNewShape" in op.options
        raw_new_shape = op.options.get("onnxRawNewShape", op.options.get("newShape", []))
        try:
            new_shape = [int(v) for v in np.asarray(raw_new_shape).reshape(-1).tolist()]
        except Exception:
            continue
        shape_tensor_new_shape: List[int] = []
        if len(op.inputs) >= 2:
            shape_tensor = model_ir.tensors.get(op.inputs[1], None)
            if shape_tensor is not None and shape_tensor.data is not None:
                try:
                    shape_tensor_new_shape = [
                        int(v) for v in np.asarray(shape_tensor.data).reshape(-1).tolist()
                    ]
                except Exception:
                    shape_tensor_new_shape = []
        if len(new_shape) == 0 and len(shape_tensor_new_shape) > 0:
            new_shape = [int(v) for v in list(shape_tensor_new_shape)]
        existing_new_shape = op.options.get("newShape", [])
        try:
            existing_new_shape_list = [
                int(v) for v in np.asarray(existing_new_shape).reshape(-1).tolist()
            ]
        except Exception:
            existing_new_shape_list = []
        if len(new_shape) == 0:
            # Runtime shape-driven RESHAPE (e.g. SHAPE/GATHER/CONCAT pipelines)
            # must remain dynamic. Do not force-resolve from empty templates.
            continue

        input_tensor = model_ir.tensors.get(op.inputs[0], None)
        if input_tensor is None:
            continue
        input_signature = (
            list(input_tensor.shape_signature)
            if input_tensor.shape_signature is not None
            else list(input_tensor.shape)
        )
        input_shape_fallback = (
            list(input_tensor.shape)
            if input_tensor.shape is not None
            else []
        )
        effective_input_signature: List[int] = []
        if len(input_signature) > 0:
            for idx, dim in enumerate(input_signature):
                sig_dim = int(dim)
                if sig_dim > 0:
                    effective_input_signature.append(int(sig_dim))
                    continue
                if idx < len(input_shape_fallback) and int(input_shape_fallback[idx]) > 0:
                    effective_input_signature.append(int(input_shape_fallback[idx]))
                else:
                    effective_input_signature.append(int(sig_dim))
        else:
            effective_input_signature = [int(v) for v in list(input_shape_fallback)]
        has_zero_dim = any(int(dim) == 0 for dim in new_shape)
        has_multi_minus_one = sum(1 for dim in new_shape if int(dim) == -1) > 1
        signature_for_resolve = (
            [int(v) for v in list(effective_input_signature)]
            if has_zero_dim or has_multi_minus_one
            else [int(v) for v in list(input_signature)]
        )
        if has_onnx_raw_new_shape and len(new_shape) > 0:
            # ONNX-exported raw reshape constants are semantically authoritative.
            # Do not re-infer them from possibly stale intermediate static metadata.
            # If a concrete newShape already exists, keep it as the most stable form.
            raw_has_minus_one = any(int(dim) == -1 for dim in new_shape)
            raw_has_zero_dim = any(int(dim) == 0 for dim in new_shape)
            if (
                bool(prefer_runtime_inferable_from_onnx_raw)
                and raw_has_minus_one
                and not raw_has_zero_dim
                and len(new_shape) >= 5
            ):
                # Final-stage safety mode:
                # keep ONNX's runtime-inferable `-1` instead of stale concretized values.
                resolved_shape = _sanitize_reshape_template(
                    template=new_shape,
                    input_dims=signature_for_resolve,
                )
            elif (
                len(existing_new_shape_list) > 0
                and all(int(dim) > 0 for dim in existing_new_shape_list)
            ):
                resolved_shape = [int(v) for v in existing_new_shape_list]
            elif all(int(dim) > 0 for dim in new_shape):
                resolved_shape = [int(v) for v in new_shape]
            elif (
                any(int(dim) == -1 for dim in new_shape)
                and any(int(dim) <= 0 for dim in input_signature)
                and not any(int(dim) == 0 for dim in new_shape)
            ):
                # Keep runtime "-1" inference dimensions when input signature is dynamic.
                resolved_shape = [int(v) for v in new_shape]
            else:
                resolved_shape = _resolve_reshape_new_shape_from_static_input(
                    new_shape=_sanitize_reshape_template(
                        template=new_shape,
                        input_dims=signature_for_resolve,
                    ),
                    input_signature=signature_for_resolve,
                    allow_zero=(
                        bool(op.options.get("allowZero"))
                        if "allowZero" in op.options
                        else None
                    ),
                )
        else:
            resolved_shape = _resolve_reshape_new_shape_from_static_input(
                new_shape=_sanitize_reshape_template(
                    template=new_shape,
                    input_dims=signature_for_resolve,
                ),
                input_signature=signature_for_resolve,
                allow_zero=(
                    bool(op.options.get("allowZero"))
                    if "allowZero" in op.options
                    else None
                ),
            )
        if resolved_shape is None and len(new_shape) > 0:
            fallback_shape = _sanitize_reshape_template(
                template=new_shape,
                input_dims=signature_for_resolve,
            )
            if (
                sum(1 for dim in fallback_shape if int(dim) == -1) <= 1
                and all(int(dim) >= -1 for dim in fallback_shape)
            ):
                resolved_shape = [int(v) for v in list(fallback_shape)]
        if resolved_shape is None:
            continue
        if (
            has_onnx_raw_new_shape
            and len(new_shape) == len(resolved_shape)
            and any(int(dim) == 0 for dim in new_shape)
            and sum(1 for dim in new_shape if int(dim) == -1) == 1
            and any(int(dim) < 0 for dim in input_signature)
            and all(int(dim) > 0 for dim in resolved_shape)
        ):
            zero_axes = [int(i) for i, dim in enumerate(new_shape) if int(dim) == 0]
            inferred_axis = next(
                (int(i) for i, dim in enumerate(new_shape) if int(dim) == -1),
                None,
            )
            if inferred_axis is not None and len(zero_axes) > 0:
                dynamic_candidate_axes = [
                    int(axis) for axis in zero_axes if int(resolved_shape[axis]) > 1
                ]
                if len(dynamic_candidate_axes) == 0:
                    dynamic_candidate_axes = [int(v) for v in list(zero_axes)]
                dynamic_axis = int(dynamic_candidate_axes[0])
                if int(dynamic_axis) != int(inferred_axis):
                    relaxed_shape = [int(v) for v in list(resolved_shape)]
                    relaxed_shape[dynamic_axis] = -1
                    resolved_shape = [int(v) for v in list(relaxed_shape)]

        changed = False
        should_update_newshape_option = True
        if (
            len(existing_new_shape_list) == 0
            and any(int(dim) < 0 for dim in resolved_shape)
        ):
            # Keep runtime-driven dynamic reshape intent (newShape=[]) when the
            # resolved template still contains unknown dimensions.
            should_update_newshape_option = False
        if should_update_newshape_option:
            if existing_new_shape_list != [int(v) for v in resolved_shape]:
                op.options["newShape"] = [int(v) for v in resolved_shape]
                changed = True

        if len(op.inputs) >= 2:
            shape_tensor = model_ir.tensors.get(op.inputs[1], None)
            if shape_tensor is not None:
                existing_shape_data = None
                if shape_tensor.data is not None:
                    existing_shape_data = [
                        int(v) for v in np.asarray(shape_tensor.data).reshape(-1).tolist()
                    ]
                if existing_shape_data != [int(v) for v in resolved_shape]:
                    shape_tensor.data = np.asarray(resolved_shape, dtype=np.int32)
                    shape_tensor.dtype = "INT32"
                    shape_tensor.shape = [int(len(resolved_shape))]
                    shape_tensor.shape_signature = [int(len(resolved_shape))]
                    changed = True

        output_tensor = model_ir.tensors.get(op.outputs[0], None)
        if output_tensor is not None:
            output_shape = list(output_tensor.shape)
            output_signature = (
                list(output_tensor.shape_signature)
                if output_tensor.shape_signature is not None
                else None
            )
            materialized_shape = [int(v) if int(v) > 0 else 1 for v in resolved_shape]
            resolved_signature = [int(v) for v in resolved_shape]
            if output_shape != materialized_shape or output_signature != resolved_signature:
                output_tensor.shape = [int(v) for v in materialized_shape]
                output_tensor.shape_signature = [int(v) for v in resolved_shape]
                changed = True

        if changed:
            resolved_count += 1

    return {"resolved_dynamic_reshape_shapes": int(resolved_count)}


def _sanitize_hardswish_tensor_shapes(model_ir: ModelIR) -> Dict[str, int]:
    """
    HARD_SWISH must preserve input shape exactly.
    Late transpose/layout rewrites can occasionally leave stale output metadata.
    """
    fixed = 0
    for op in model_ir.operators:
        if str(op.op_type) != "HARD_SWISH":
            continue
        if len(op.inputs) < 1 or len(op.outputs) != 1:
            continue

        input_name = str(op.inputs[0])
        output_name = str(op.outputs[0])
        input_tensor = model_ir.tensors.get(input_name, None)
        output_tensor = model_ir.tensors.get(output_name, None)
        if input_tensor is None or output_tensor is None:
            continue
        if input_tensor.shape is None or len(list(input_tensor.shape)) == 0:
            continue

        input_shape = [int(v) for v in list(input_tensor.shape)]
        input_signature_current = (
            [int(v) for v in list(input_tensor.shape_signature)]
            if input_tensor.shape_signature is not None
            else []
        )
        if _is_fully_known_positive_shape(input_shape):
            input_signature = [int(v) for v in list(input_shape)]
        else:
            input_signature = (
                [int(v) for v in list(input_signature_current)]
                if len(input_signature_current) > 0
                else [int(v) for v in list(input_shape)]
            )
        output_shape = (
            [int(v) for v in list(output_tensor.shape)]
            if output_tensor.shape is not None
            else []
        )
        output_signature = (
            [int(v) for v in list(output_tensor.shape_signature)]
            if output_tensor.shape_signature is not None
            else [int(v) for v in list(output_shape)]
        )

        if (
            input_signature_current == input_signature
            and output_shape == input_shape
            and output_signature == input_signature
        ):
            continue

        input_tensor.shape_signature = [int(v) for v in list(input_signature)]
        output_tensor.shape = [int(v) for v in list(input_shape)]
        output_tensor.shape_signature = [int(v) for v in list(input_signature)]
        fixed += 1

    return {"sanitized_hardswish_tensor_shapes": int(fixed)}


def _sanitize_squeeze_axes_with_static_input_shapes(model_ir: ModelIR) -> Dict[str, int]:
    """
    Keep SQUEEZE axes consistent with concrete input metadata.

    Late rewrite passes can leave a SQUEEZE axis targeting a non-singleton
    dimension, which crashes LiteRT during `allocate_tensors()`.
    This pass normalizes axes and repairs static metadata conservatively:
    - invalid / duplicated axes are removed
    - non-singleton axes are forced to singleton on non-constant inputs
    - for constant inputs whose payload shape proves non-singleton axes,
      those axes are dropped
    """
    sanitized_ops = 0
    repaired_input_dims = 0
    updated_output_shapes = 0

    for op in model_ir.operators:
        if str(op.op_type) != "SQUEEZE":
            continue
        if len(op.inputs) < 1 or len(op.outputs) != 1:
            continue

        input_name = str(op.inputs[0])
        output_name = str(op.outputs[0])
        input_tensor = model_ir.tensors.get(input_name, None)
        output_tensor = model_ir.tensors.get(output_name, None)
        if input_tensor is None or output_tensor is None:
            continue
        if not _is_fully_known_positive_shape(input_tensor.shape):
            continue

        input_shape = [int(v) for v in list(input_tensor.shape)]
        rank = len(input_shape)
        input_signature = (
            [int(v) for v in list(input_tensor.shape_signature)]
            if input_tensor.shape_signature is not None
            and len(list(input_tensor.shape_signature)) == rank
            else [int(v) for v in list(input_shape)]
        )

        options = dict(op.options) if isinstance(op.options, dict) else {}
        raw_axes = _parse_axes_option(options.get("squeezeDims", []))

        normalized_axes: List[int] = []
        axes_modified = False
        for axis in raw_axes:
            a = int(axis)
            if a < 0:
                a += int(rank)
            if a < 0 or a >= int(rank):
                axes_modified = True
                continue
            if int(a) in normalized_axes:
                axes_modified = True
                continue
            normalized_axes.append(int(a))

        data_shape: Optional[List[int]] = None
        if input_tensor.data is not None:
            try:
                data_arr = np.asarray(input_tensor.data)
                if int(data_arr.ndim) == int(rank):
                    data_shape = [int(v) for v in list(data_arr.shape)]
            except Exception:
                data_shape = None

        sanitized_axes: List[int] = []
        repaired_shape = [int(v) for v in list(input_shape)]
        repaired_signature = [int(v) for v in list(input_signature)]
        for axis in normalized_axes:
            current_dim = int(repaired_shape[axis])
            if current_dim == 1:
                sanitized_axes.append(int(axis))
                if int(repaired_signature[axis]) != 1:
                    repaired_signature[axis] = 1
                continue

            can_force_singleton = (
                data_shape is None
                or int(axis) >= len(data_shape)
                or int(data_shape[axis]) == 1
            )
            if can_force_singleton:
                repaired_shape[axis] = 1
                repaired_signature[axis] = 1
                sanitized_axes.append(int(axis))
                repaired_input_dims += 1
                continue

            axes_modified = True

        input_changed = (
            repaired_shape != [int(v) for v in list(input_tensor.shape)]
            or (
                input_tensor.shape_signature is None
                or [int(v) for v in list(input_tensor.shape_signature)] != repaired_signature
            )
        )
        if input_changed:
            input_tensor.shape = [int(v) for v in list(repaired_shape)]
            input_tensor.shape_signature = [int(v) for v in list(repaired_signature)]

        if axes_modified or sanitized_axes != raw_axes:
            options["squeezeDims"] = [int(v) for v in sanitized_axes]
            op.options = options
            sanitized_ops += 1

        out_shape, out_signature = _infer_squeeze_output_shape_and_signature(
            input_shape=[int(v) for v in list(input_tensor.shape)],
            input_signature=(
                [int(v) for v in list(input_tensor.shape_signature)]
                if input_tensor.shape_signature is not None
                else [int(v) for v in list(input_tensor.shape)]
            ),
            squeeze_axes=[int(v) for v in list(sanitized_axes)],
        )
        if out_shape is None or out_signature is None:
            continue
        if not _is_fully_known_positive_shape(out_shape):
            continue

        if (
            [int(v) for v in list(output_tensor.shape)] != [int(v) for v in list(out_shape)]
            or output_tensor.shape_signature is None
            or [int(v) for v in list(output_tensor.shape_signature)] != [int(v) for v in list(out_signature)]
        ):
            output_tensor.shape = [int(v) for v in list(out_shape)]
            output_tensor.shape_signature = [int(v) for v in list(out_signature)]
            updated_output_shapes += 1

    return {
        "sanitized_squeeze_axes_with_static_input_shapes": int(sanitized_ops),
        "repaired_squeeze_input_singleton_dims": int(repaired_input_dims),
        "updated_squeeze_output_shapes": int(updated_output_shapes),
    }


def _replace_expand_dims_and_squeeze_with_reshape(model_ir: ModelIR) -> Dict[str, int]:
    """
    Replace EXPAND_DIMS/SQUEEZE with RESHAPE for LiteRT.js WebGPU compatibility.

    The replacement uses output tensor metadata as target shape. When possible,
    a single unknown dimension from shape_signature is preserved as -1.
    """
    rewritten = 0
    shape_tensors_created = 0
    pre_ops_by_index: Dict[int, List[OperatorIR]] = {}

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 0
        while name in model_ir.tensors:
            suffix += 1
            name = f"{base}_{suffix}"
        return name

    def _pick_dynamic_axis_for_reshape_target(
        *,
        op: OperatorIR,
        input_name: str,
        output_shape: List[int],
    ) -> Optional[int]:
        if len(output_shape) == 0:
            return None
        input_tensor = model_ir.tensors.get(input_name, None)
        if input_tensor is None:
            return None
        input_signature = (
            [int(v) for v in list(input_tensor.shape_signature)]
            if input_tensor.shape_signature is not None
            else [int(v) for v in list(input_tensor.shape)]
        )
        dynamic_input_axes = [int(i) for i, dim in enumerate(input_signature) if int(dim) < 0]
        if len(dynamic_input_axes) == 0:
            return None

        op_type = str(op.op_type)
        if op_type == "EXPAND_DIMS" and len(op.inputs) >= 2:
            axis_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
            axis_values = _read_const_ints_from_tensor(axis_tensor)
            if axis_values is not None and len(axis_values) > 0:
                out_rank = len(output_shape)
                axis = int(axis_values[0])
                if axis < 0:
                    axis += int(out_rank)
                if 0 <= axis < out_rank:
                    mapped_axes: List[int] = []
                    for in_axis in dynamic_input_axes:
                        out_axis = int(in_axis) if int(in_axis) < axis else int(in_axis) + 1
                        if 0 <= out_axis < out_rank:
                            mapped_axes.append(int(out_axis))
                    if len(mapped_axes) > 0:
                        for out_axis in mapped_axes:
                            if int(output_shape[out_axis]) > 1:
                                return int(out_axis)
                        return int(mapped_axes[0])

        for axis, dim in enumerate(output_shape):
            if int(dim) > 1:
                return int(axis)
        return 0

    for op_index, op in enumerate(model_ir.operators):
        op_type = str(op.op_type)
        if op_type not in {"EXPAND_DIMS", "SQUEEZE"}:
            continue
        if len(op.inputs) < 1 or len(op.outputs) != 1:
            continue

        input_name = str(op.inputs[0])
        output_name = str(op.outputs[0])
        input_tensor = model_ir.tensors.get(input_name, None)
        output_tensor = model_ir.tensors.get(output_name, None)
        if input_tensor is None or output_tensor is None or output_tensor.shape is None:
            continue

        input_signature = (
            [int(v) for v in list(input_tensor.shape_signature)]
            if input_tensor.shape_signature is not None
            else [int(v) for v in list(input_tensor.shape)]
        )
        output_shape = [int(v) for v in list(output_tensor.shape)]
        output_signature = (
            [int(v) for v in list(output_tensor.shape_signature)]
            if output_tensor.shape_signature is not None
            else [int(v) for v in list(output_shape)]
        )

        if op_type == "SQUEEZE" and len(output_shape) > 0 and all(int(v) == 1 for v in output_shape):
            # Conservative guard:
            # all-ones squeeze outputs are frequently metadata-ambiguous in
            # dynamic coordinate paths (e.g. GridSample auxiliaries). Rewriting
            # to RESHAPE can hard-code "1" and break runtime element counts.
            continue
        if op_type == "EXPAND_DIMS" and (
            any(int(v) < 0 for v in input_signature)
            or any(int(v) < 0 for v in output_signature)
        ):
            # EXPAND_DIMS can preserve multiple dynamic axes, while static
            # RESHAPE targets can represent at most one unknown axis reliably.
            # Keep EXPAND_DIMS to avoid hard-coding dynamic extents.
            continue
        if op_type == "SQUEEZE" and (
            any(int(v) < 0 for v in input_signature)
            or any(int(v) < 0 for v in output_signature)
        ):
            # Dynamic SQUEEZE can still be rewritten with runtime shape plumbing:
            #   SHAPE(x) -> GATHER(kept_axes) -> RESHAPE(x, gathered_shape)
            squeeze_dims_raw = op.options.get("squeezeDims", [])
            squeeze_dims = (
                [int(v) for v in np.asarray(squeeze_dims_raw).reshape(-1).tolist()]
                if squeeze_dims_raw is not None
                else []
            )
            input_rank = len(list(input_signature))
            if input_rank <= 0 or len(squeeze_dims) == 0:
                # squeezeDims=[] semantics are data-dependent and cannot be
                # represented as a rank-stable RESHAPE.
                continue
            normalized_axes: List[int] = []
            valid_axes = True
            for axis in squeeze_dims:
                norm_axis = int(axis)
                if norm_axis < 0:
                    norm_axis += int(input_rank)
                if norm_axis < 0 or norm_axis >= int(input_rank):
                    valid_axes = False
                    break
                if norm_axis not in normalized_axes:
                    normalized_axes.append(int(norm_axis))
            if not valid_axes:
                continue

            kept_axes = [
                int(axis)
                for axis in range(input_rank)
                if int(axis) not in set(int(v) for v in normalized_axes)
            ]
            runtime_shape_name = _unique_tensor_name(f"{output_name}_runtime_shape")
            runtime_shape_filtered_name = _unique_tensor_name(
                f"{output_name}_runtime_shape_filtered"
            )
            kept_axes_name = _unique_tensor_name(f"{output_name}_squeeze_kept_axes")
            model_ir.tensors[runtime_shape_name] = TensorIR(
                name=runtime_shape_name,
                dtype="INT32",
                shape=[int(input_rank)],
                shape_signature=[int(input_rank)],
                data=None,
                is_variable=False,
                quantization=None,
            )
            model_ir.tensors[kept_axes_name] = TensorIR(
                name=kept_axes_name,
                dtype="INT32",
                shape=[int(len(kept_axes))],
                shape_signature=[int(len(kept_axes))],
                data=np.asarray([int(v) for v in list(kept_axes)], dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            model_ir.tensors[runtime_shape_filtered_name] = TensorIR(
                name=runtime_shape_filtered_name,
                dtype="INT32",
                shape=[int(len(kept_axes))],
                shape_signature=[int(len(kept_axes))],
                data=None,
                is_variable=False,
                quantization=None,
            )
            pre_ops = pre_ops_by_index.setdefault(int(op_index), [])
            pre_ops.append(
                OperatorIR(
                    op_type="SHAPE",
                    inputs=[input_name],
                    outputs=[runtime_shape_name],
                    options={"outType": "INT32"},
                )
            )
            pre_ops.append(
                OperatorIR(
                    op_type="GATHER",
                    inputs=[runtime_shape_name, kept_axes_name],
                    outputs=[runtime_shape_filtered_name],
                    options={"axis": 0, "batchDims": 0},
                )
            )
            op.op_type = "RESHAPE"
            op.inputs = [input_name, runtime_shape_filtered_name]
            op.options = {"newShape": []}
            rewritten += 1
            shape_tensors_created += 1
            continue

        if len(output_signature) == 0:
            reshape_target = []
        else:
            reshape_target = [int(v) for v in list(output_shape)]
            if len(output_signature) == len(output_shape):
                if all(int(v) > 0 for v in output_signature):
                    reshape_target = [int(v) for v in list(output_signature)]
                else:
                    negative_count = sum(1 for dim in output_signature if int(dim) < 0)
                    if negative_count == 1:
                        reshape_target = [
                            int(dim) if int(dim) > 0 else -1 for dim in output_signature
                        ]

        preserve_dynamic_shape = False
        if len(reshape_target) > 0 and all(int(v) > 0 for v in reshape_target):
            dynamic_axis = _pick_dynamic_axis_for_reshape_target(
                op=op,
                input_name=input_name,
                output_shape=reshape_target,
            )
            if dynamic_axis is not None and 0 <= int(dynamic_axis) < len(reshape_target):
                reshape_target[int(dynamic_axis)] = -1
                preserve_dynamic_shape = True
                if output_tensor is not None and len(output_tensor.shape) == len(reshape_target):
                    output_tensor.shape_signature = [int(v) for v in list(reshape_target)]

        shape_name = _unique_tensor_name(f"{output_name}_reshape_shape")
        model_ir.tensors[shape_name] = TensorIR(
            name=shape_name,
            dtype="INT32",
            shape=[int(len(reshape_target))],
            shape_signature=[int(len(reshape_target))],
            data=np.asarray(reshape_target, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )
        op.op_type = "RESHAPE"
        op.inputs = [input_name, shape_name]
        op.options = {"newShape": [int(v) for v in list(reshape_target)]}
        if preserve_dynamic_shape:
            op.options["preserveDynamicShape"] = True
        rewritten += 1
        shape_tensors_created += 1

    if len(pre_ops_by_index) > 0:
        rebuilt_ops: List[OperatorIR] = []
        for idx, op in enumerate(model_ir.operators):
            if idx in pre_ops_by_index:
                rebuilt_ops.extend(pre_ops_by_index[idx])
            rebuilt_ops.append(op)
        model_ir.operators = rebuilt_ops

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {
        "replaced_expand_dims_and_squeeze_with_reshape": int(rewritten),
        "expand_dims_squeeze_rewrite_shape_tensors": int(shape_tensors_created),
    }


def _sanitize_wrong_way_nchw_to_nhwc_transpose_before_conv(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove invalid NCHW->NHWC transposes that are applied to already-NHWC tensors
    and break downstream CONV_2D channel alignment.
    """
    removed = 0
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    model_outputs = set(str(v) for v in model_ir.outputs)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE" or len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, op) != perm_nchw_to_nhwc:
                continue

            in_name = str(op.inputs[0])
            out_name = str(op.outputs[0])
            if out_name in model_outputs:
                continue

            in_tensor = model_ir.tensors.get(in_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if (
                in_tensor is None
                or out_tensor is None
                or len(list(in_tensor.shape)) != 4
                or len(list(out_tensor.shape)) != 4
            ):
                continue

            user_indices = [int(v) for v in consumers.get(out_name, [])]
            if len(user_indices) == 0:
                continue

            remove_this = True
            in_c = int(in_tensor.shape[3])
            out_c = int(out_tensor.shape[3])
            for user_idx in user_indices:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "CONV_2D"
                    or len(user_op.inputs) < 2
                    or str(user_op.inputs[0]) != out_name
                ):
                    remove_this = False
                    break
                filter_tensor = model_ir.tensors.get(str(user_op.inputs[1]), None)
                if filter_tensor is None or len(list(filter_tensor.shape)) != 4:
                    remove_this = False
                    break
                expected_c = int(filter_tensor.shape[3])
                if not (in_c == expected_c and out_c != expected_c):
                    remove_this = False
                    break

            if not remove_this:
                continue

            _replace_tensor_inputs(model_ir, out_name, in_name)
            del model_ir.operators[int(op_idx)]
            removed += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"sanitized_wrong_way_nchw_to_nhwc_transpose_before_conv": int(removed)}


def _repair_rank4_binary_layout_mismatch_with_transpose_adapter(model_ir: ModelIR) -> Dict[str, int]:
    """
    Repair rank-4 layout mismatches for elementwise binary ops by inserting a
    minimal transpose adapter on one input when shapes are pure NHWC/NCHW
    permutations of each other.
    """
    inserted = 0
    binary_ops = {"ADD", "MUL", "SUB", "DIV", "MAXIMUM", "MINIMUM"}
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    def _make_unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        for op_idx, op in enumerate(model_ir.operators):
            op_type = str(op.op_type)
            if op_type not in binary_ops or len(op.inputs) != 2 or len(op.outputs) != 1:
                continue

            in0_name = str(op.inputs[0])
            in1_name = str(op.inputs[1])
            t0 = model_ir.tensors.get(in0_name, None)
            t1 = model_ir.tensors.get(in1_name, None)
            if t0 is None or t1 is None:
                continue
            s0 = [int(v) for v in list(t0.shape)]
            s1 = [int(v) for v in list(t1.shape)]
            if len(s0) != 4 or len(s1) != 4:
                continue
            if any(int(v) <= 0 for v in s0 + s1):
                continue
            if s0 == s1:
                continue

            perm_to_use: Optional[List[int]] = None
            if s0 == [int(s1[idx]) for idx in perm_nchw_to_nhwc]:
                perm_to_use = [int(v) for v in list(perm_nchw_to_nhwc)]
            elif s0 == [int(s1[idx]) for idx in perm_nhwc_to_nchw]:
                perm_to_use = [int(v) for v in list(perm_nhwc_to_nchw)]
            else:
                continue

            perm_name = _make_unique_tensor_name(f"{in1_name}_layout_fix_perm")
            adapted_name = _make_unique_tensor_name(f"{in1_name}_layout_fix")
            model_ir.tensors[perm_name] = TensorIR(
                name=perm_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray(perm_to_use, dtype=np.int32),
                is_variable=False,
            )
            model_ir.tensors[adapted_name] = TensorIR(
                name=adapted_name,
                dtype=str(t1.dtype),
                shape=[int(v) for v in list(s0)],
                shape_signature=[int(v) for v in list(s0)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(t1.quantization),
            )
            model_ir.operators.insert(
                int(op_idx),
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[in1_name, perm_name],
                    outputs=[adapted_name],
                ),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=op,
                input_index=1,
                new_input_name=adapted_name,
            )
            inserted += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"inserted_rank4_binary_layout_fix_transpose": int(inserted)}


def _sanitize_static_shape_signature_consistency(model_ir: ModelIR) -> Dict[str, int]:
    """
    For fully static tensors, keep shape_signature consistent with runtime shape.

    Notes:
    - Keep dynamic `-1` only for ONNX graph-boundary tensors that were dynamic
      in the source model (input/output contract).
    - Preserve dynamic `-1` on internal tensors only when they are reachable
      from those dynamic graph boundaries through operator dependencies.
    - Internal tensors are auto-completed: if runtime `shape` is fully static,
      `shape_signature` is overwritten to the same static shape.
    """
    fixed = 0
    preserved_dynamic_boundary = 0
    preserved_dynamic_leading_axis = 0
    preserved_dynamic_lineage = 0
    dynamic_boundary_names = set()
    for key in (
        "onnx_dynamic_input_tensor_names",
        "onnx_dynamic_output_tensor_names",
    ):
        names = model_ir.metadata.get(key, [])
        if isinstance(names, list):
            dynamic_boundary_names.update(str(v) for v in names if str(v) != "")
    dynamic_boundary_signature_map = model_ir.metadata.get(
        "dynamic_boundary_shape_signature_map", {}
    )
    if not isinstance(dynamic_boundary_signature_map, dict):
        dynamic_boundary_signature_map = {}
    graph_output_names = set(str(v) for v in list(model_ir.outputs))
    producer_map = _build_tensor_producer_map(model_ir)
    dynamic_lineage_root_names = set(dynamic_boundary_names)
    for tensor_name, producer_idx in producer_map.items():
        if int(producer_idx) < 0 or int(producer_idx) >= len(model_ir.operators):
            continue
        producer_op = model_ir.operators[int(producer_idx)]
        producer_op_type = str(producer_op.op_type)
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            continue
        signature = (
            [int(v) for v in list(tensor.shape_signature)]
            if tensor.shape_signature is not None
            else [int(v) for v in list(tensor.shape)]
        )
        if producer_op_type == "WHERE" and len(signature) >= 1 and int(signature[0]) < 0:
            dynamic_lineage_root_names.add(str(tensor_name))
            continue
        if producer_op_type == "RANGE" and any(int(v) < 0 for v in signature):
            # RANGE output length is runtime-dependent when `limit` is non-const.
            # Keep this tensor as a dynamic-lineage root so downstream GATHER /
            # RESHAPE chains preserve leading -1 signatures.
            dynamic_lineage_root_names.add(str(tensor_name))
            continue
        if producer_op_type == "RESHAPE":
            reshape_target: List[int] = []
            try:
                reshape_target = [
                    int(v)
                    for v in np.asarray(
                        producer_op.options.get("newShape", [])
                    ).reshape(-1).tolist()
                ]
            except Exception:
                reshape_target = []
            if len(reshape_target) == 0 and len(producer_op.inputs) >= 2:
                shape_tensor = model_ir.tensors.get(str(producer_op.inputs[1]), None)
                shape_vals = _read_const_ints_from_tensor(shape_tensor)
                if shape_vals is not None:
                    reshape_target = [int(v) for v in list(shape_vals)]
            if any(int(v) < 0 for v in reshape_target):
                dynamic_lineage_root_names.add(str(tensor_name))
                continue
        # TOPK_V2 output extent depends on runtime k. Preserve internal -1 axes
        # even when dynamic lineage is not anchored to graph boundaries.
        if producer_op_type == "TOPK_V2" and any(int(v) < 0 for v in signature):
            dynamic_lineage_root_names.add(str(tensor_name))
    dynamic_lineage_cache: Dict[str, bool] = {}
    dynamic_lineage_visiting: set[str] = set()

    def _has_dynamic_boundary_lineage(tensor_name: str) -> bool:
        key = str(tensor_name)
        cached = dynamic_lineage_cache.get(key, None)
        if cached is not None:
            return bool(cached)
        if key in dynamic_lineage_root_names:
            dynamic_lineage_cache[key] = True
            return True
        if key in dynamic_lineage_visiting:
            dynamic_lineage_cache[key] = False
            return False
        tensor = model_ir.tensors.get(key, None)
        if tensor is None:
            dynamic_lineage_cache[key] = False
            return False
        if tensor.data is not None:
            dynamic_lineage_cache[key] = False
            return False
        producer_idx = producer_map.get(key, None)
        if producer_idx is None:
            dynamic_lineage_cache[key] = False
            return False
        dynamic_lineage_visiting.add(key)
        producer = model_ir.operators[int(producer_idx)]
        has_dynamic_ancestor = False
        for input_name in producer.inputs:
            parent_name = str(input_name)
            if parent_name == "":
                continue
            if _has_dynamic_boundary_lineage(parent_name):
                has_dynamic_ancestor = True
                break
        dynamic_lineage_visiting.discard(key)
        dynamic_lineage_cache[key] = bool(has_dynamic_ancestor)
        return bool(has_dynamic_ancestor)

    for tensor in model_ir.tensors.values():
        if tensor.shape is None:
            continue
        shape = [int(v) for v in list(tensor.shape)]
        if len(shape) == 0:
            if tensor.shape_signature is None:
                tensor.shape_signature = []
                fixed += 1
            continue
        if not _is_fully_known_positive_shape(shape):
            continue

        boundary_signature = dynamic_boundary_signature_map.get(str(tensor.name), None)
        if isinstance(boundary_signature, list) and len(boundary_signature) == len(shape):
            normalized_boundary_signature = [
                int(-1 if int(v) < 0 else shape[idx])
                for idx, v in enumerate(boundary_signature)
            ]
            if any(int(v) < 0 for v in normalized_boundary_signature):
                if tensor.shape_signature != normalized_boundary_signature:
                    tensor.shape_signature = list(normalized_boundary_signature)
                    fixed += 1
                preserved_dynamic_boundary += 1
                continue

        signature = (
            [int(v) for v in list(tensor.shape_signature)]
            if tensor.shape_signature is not None
            else None
        )
        if signature is None:
            tensor.shape_signature = [int(v) for v in list(shape)]
            fixed += 1
            continue
        if len(signature) != len(shape):
            tensor.shape_signature = [int(v) for v in list(shape)]
            fixed += 1
            continue
        if any(int(v) < 0 for v in signature):
            negative_axes = [int(i) for i, v in enumerate(signature) if int(v) < 0]
            trailing_axes_static = all(
                int(signature[idx]) > 0 for idx in range(1, len(signature))
            )
            producer_idx = producer_map.get(str(tensor.name), None)
            producer_op_type = (
                str(model_ir.operators[int(producer_idx)].op_type)
                if producer_idx is not None and int(producer_idx) < len(model_ir.operators)
                else ""
            )
            if str(tensor.name) in dynamic_lineage_root_names:
                preserved_dynamic_boundary += 1
                continue
            if (
                producer_op_type == "WHERE"
                and negative_axes == [0]
                and trailing_axes_static
            ):
                preserved_dynamic_leading_axis += 1
                continue
            has_dynamic_lineage = _has_dynamic_boundary_lineage(str(tensor.name))
            if negative_axes == [0] and trailing_axes_static:
                if str(tensor.name) in graph_output_names:
                    preserved_dynamic_leading_axis += 1
                    continue
                if has_dynamic_lineage:
                    preserved_dynamic_leading_axis += 1
                    continue
            if has_dynamic_lineage:
                if negative_axes == [0] and trailing_axes_static:
                    preserved_dynamic_leading_axis += 1
                else:
                    preserved_dynamic_lineage += 1
                continue
            tensor.shape_signature = [int(v) for v in list(shape)]
            fixed += 1
            continue
        if signature != shape:
            tensor.shape_signature = [int(v) for v in list(shape)]
            fixed += 1

    return {
        "sanitized_static_shape_signature_consistency": int(fixed),
        "preserved_dynamic_boundary_shape_signature": int(preserved_dynamic_boundary),
        "preserved_dynamic_leading_axis_shape_signature": int(
            preserved_dynamic_leading_axis
        ),
        "preserved_dynamic_lineage_shape_signature": int(
            preserved_dynamic_lineage
        ),
    }


def _realign_dynamic_boundary_shape_signature_map(model_ir: ModelIR) -> Dict[str, int]:
    signature_map = model_ir.metadata.get("dynamic_boundary_shape_signature_map", {})
    if not isinstance(signature_map, dict):
        return {"realigned_dynamic_boundary_shape_signature_map": 0}

    updated = 0
    for tensor_name, boundary_signature in list(signature_map.items()):
        if not isinstance(boundary_signature, list):
            continue
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.shape is None:
            continue
        aligned = _align_boundary_signature_to_current_shape(
            boundary_signature=[int(v) for v in list(boundary_signature)],
            current_shape=[int(v) for v in list(tensor.shape)],
        )
        if aligned is None:
            continue
        if [int(v) for v in list(boundary_signature)] != [int(v) for v in list(aligned)]:
            signature_map[str(tensor_name)] = [int(v) for v in list(aligned)]
            updated += 1

    model_ir.metadata["dynamic_boundary_shape_signature_map"] = signature_map
    return {"realigned_dynamic_boundary_shape_signature_map": int(updated)}


def _is_fully_known_positive_shape(shape: Optional[List[int]]) -> bool:
    if shape is None or len(shape) == 0:
        return False
    return all(int(dim) > 0 for dim in shape)


def _read_const_ints_from_tensor(tensor: Optional[TensorIR]) -> Optional[List[int]]:
    if tensor is None or tensor.data is None:
        return None
    try:
        return [int(v) for v in np.asarray(tensor.data).reshape(-1).tolist()]
    except Exception:
        return None


def _write_const_ints_to_tensor(tensor: Optional[TensorIR], values: List[int]) -> bool:
    if tensor is None:
        return False
    current = _read_const_ints_from_tensor(tensor)
    normalized = [int(v) for v in list(values)]
    if current == normalized:
        return False
    np_dtype = np.int32
    if tensor.data is not None:
        try:
            np_dtype = np.asarray(tensor.data).dtype
        except Exception:
            np_dtype = np.int32
    tensor.data = np.asarray(normalized, dtype=np_dtype)
    tensor.shape = [int(len(normalized))]
    tensor.shape_signature = [int(len(normalized))]
    return True


def _replace_unsupported_split_with_slice(model_ir: ModelIR) -> Dict[str, int]:
    """
    Replace SPLIT ops whose input dtype is unsupported by LiteRT with SLICE chains.

    LiteRT SPLIT currently accepts only:
      FLOAT32, UINT8, INT8, INT16, INT32, INT64

    For unsupported input dtypes (e.g. FLOAT16), rewrite:
      SPLIT(axis, src) -> out_i
    to:
      [optional CAST(src -> target_dtype)]
      SLICE(src_or_cast, begin_i, size_i) -> out_i
    """
    split_supported_input_dtypes = {
        "FLOAT32",
        "UINT8",
        "INT8",
        "INT16",
        "INT32",
        "INT64",
    }

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        serial = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{serial}"
            serial += 1
        return candidate

    rewritten = 0
    new_operators: List[OperatorIR] = []

    for op in model_ir.operators:
        if str(op.op_type) != "SPLIT" or len(op.inputs) < 2 or len(op.outputs) <= 0:
            new_operators.append(op)
            continue

        axis_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
        source_name = str(op.inputs[1])
        source_tensor = model_ir.tensors.get(source_name, None)
        if source_tensor is None:
            new_operators.append(op)
            continue

        source_dtype = str(source_tensor.dtype).upper()
        if source_dtype in split_supported_input_dtypes:
            new_operators.append(op)
            continue

        axis_values = _read_const_ints_from_tensor(axis_tensor)
        if axis_values is None or len(axis_values) == 0:
            new_operators.append(op)
            continue

        source_shape = (
            [int(v) for v in list(source_tensor.shape)]
            if source_tensor.shape is not None
            else []
        )
        rank = int(len(source_shape))
        if rank <= 0:
            new_operators.append(op)
            continue

        axis = int(axis_values[0])
        if axis < 0:
            axis += int(rank)
        if axis < 0 or axis >= int(rank):
            new_operators.append(op)
            continue

        outputs = [str(v) for v in list(op.outputs)]
        num_splits = int(op.options.get("numSplits", len(outputs)))
        if num_splits <= 0 or len(outputs) != int(num_splits):
            new_operators.append(op)
            continue

        # Prefer explicit output tensor metadata for chunk sizes.
        split_sizes: List[int] = []
        can_derive_from_outputs = True
        output_dtypes: List[str] = []
        for output_name in outputs:
            output_tensor = model_ir.tensors.get(output_name, None)
            if output_tensor is None or output_tensor.shape is None:
                can_derive_from_outputs = False
                break
            output_shape = [int(v) for v in list(output_tensor.shape)]
            if len(output_shape) != int(rank):
                can_derive_from_outputs = False
                break
            split_dim = int(output_shape[int(axis)])
            if split_dim <= 0:
                can_derive_from_outputs = False
                break
            split_sizes.append(int(split_dim))
            output_dtypes.append(str(output_tensor.dtype).upper())

        if not can_derive_from_outputs:
            axis_dim = int(source_shape[int(axis)]) if int(axis) < len(source_shape) else -1
            if axis_dim <= 0 or axis_dim % int(num_splits) != 0:
                new_operators.append(op)
                continue
            each = int(axis_dim // int(num_splits))
            split_sizes = [int(each) for _ in range(int(num_splits))]
            output_dtypes = [
                str(model_ir.tensors.get(output_name, source_tensor).dtype).upper()
                for output_name in outputs
            ]

        slice_source_name = str(source_name)
        unique_output_dtypes = sorted(set(output_dtypes))
        if len(unique_output_dtypes) == 1:
            target_dtype = str(unique_output_dtypes[0]).upper()
            if target_dtype != "" and target_dtype != str(source_dtype).upper():
                cast_output_name = _unique_tensor_name(f"{source_name}_split_cast")
                source_signature = (
                    [int(v) for v in list(source_tensor.shape_signature)]
                    if source_tensor.shape_signature is not None
                    else [int(v) for v in list(source_shape)]
                )
                model_ir.tensors[cast_output_name] = TensorIR(
                    name=cast_output_name,
                    dtype=target_dtype,
                    shape=[int(v) for v in list(source_shape)],
                    shape_signature=[int(v) for v in list(source_signature)],
                    data=None,
                    is_variable=False,
                    quantization=None,
                )
                new_operators.append(
                    OperatorIR(
                        op_type="CAST",
                        inputs=[str(source_name)],
                        outputs=[cast_output_name],
                        options={"outDataType": target_dtype},
                    )
                )
                slice_source_name = str(cast_output_name)

        offset = 0
        for output_index, output_name in enumerate(outputs):
            begin = [0 for _ in range(int(rank))]
            begin[int(axis)] = int(offset)
            size = [-1 for _ in range(int(rank))]
            size[int(axis)] = int(split_sizes[int(output_index)])
            offset += int(split_sizes[int(output_index)])

            begin_name = _unique_tensor_name(f"{output_name}_split_fallback_begin")
            size_name = _unique_tensor_name(f"{output_name}_split_fallback_size")

            model_ir.tensors[begin_name] = TensorIR(
                name=begin_name,
                dtype="INT32",
                shape=[int(rank)],
                shape_signature=[int(rank)],
                data=np.asarray(begin, dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            model_ir.tensors[size_name] = TensorIR(
                name=size_name,
                dtype="INT32",
                shape=[int(rank)],
                shape_signature=[int(rank)],
                data=np.asarray(size, dtype=np.int32),
                is_variable=False,
                quantization=None,
            )

            new_operators.append(
                OperatorIR(
                    op_type="SLICE",
                    inputs=[str(slice_source_name), begin_name, size_name],
                    outputs=[output_name],
                    options={},
                )
            )

        rewritten += 1

    if rewritten > 0:
        model_ir.operators = new_operators
        _prune_unused_tensors(model_ir)

    return {"replaced_unsupported_split_with_slice": int(rewritten)}


def _infer_slice_output_shape_and_resolved_params(
    input_shape: Optional[List[int]],
    begin_vals: Optional[List[int]],
    size_vals: Optional[List[int]],
) -> Tuple[Optional[List[int]], Optional[List[int]], Optional[List[int]]]:
    if not _is_fully_known_positive_shape(input_shape):
        return None, None, None
    if begin_vals is None or size_vals is None:
        return None, None, None

    in_shape = [int(v) for v in list(input_shape)]
    begin = [int(v) for v in list(begin_vals)]
    size = [int(v) for v in list(size_vals)]
    rank = len(in_shape)
    if len(begin) != rank or len(size) != rank:
        return None, None, None

    resolved_begin: List[int] = []
    resolved_size: List[int] = []
    out_shape: List[int] = []
    for axis, dim in enumerate(in_shape):
        b = int(begin[axis])
        if b < 0:
            b += int(dim)
        b = max(0, min(int(b), int(dim)))
        remain = max(int(dim) - int(b), 0)

        raw_size = int(size[axis])
        if raw_size == -1:
            s = int(remain)
        elif raw_size >= 0:
            s = int(raw_size)
        else:
            return None, None, None

        out_dim = max(min(int(s), int(remain)), 0)
        resolved_begin.append(int(b))
        resolved_size.append(int(s))
        out_shape.append(int(out_dim))

    return out_shape, resolved_begin, resolved_size


def _infer_slice_output_signature(
    *,
    input_shape: Optional[List[int]],
    input_signature: Optional[List[int]],
    begin_vals: Optional[List[int]],
    size_vals: Optional[List[int]],
) -> Optional[List[int]]:
    if not _is_fully_known_positive_shape(input_shape):
        return None
    if begin_vals is None or size_vals is None:
        return None

    in_shape = [int(v) for v in list(input_shape)]
    rank = len(in_shape)
    begin = [int(v) for v in list(begin_vals)]
    size = [int(v) for v in list(size_vals)]
    if len(begin) != rank or len(size) != rank:
        return None

    if input_signature is not None and len(list(input_signature)) == rank:
        in_signature = [int(v) for v in list(input_signature)]
    else:
        in_signature = [int(v) for v in list(in_shape)]

    out_signature: List[int] = []
    for axis, dim in enumerate(in_shape):
        b = int(begin[axis])
        if b < 0:
            b += int(dim)
        b = max(0, min(int(b), int(dim)))
        remain = max(int(dim) - int(b), 0)

        raw_size = int(size[axis])
        if raw_size == -1:
            if int(in_signature[axis]) < 0:
                out_signature.append(-1)
            else:
                out_signature.append(int(remain))
            continue
        if raw_size < 0:
            return None
        if int(in_signature[axis]) < 0:
            # Preserve explicit static slice sizes on dynamic axes. Clamping to
            # placeholder `input_shape` (often 1) can corrupt downstream shape
            # metadata and break kernels that validate state widths statically.
            out_signature.append(int(max(int(raw_size), 0)))
        else:
            out_signature.append(int(max(min(int(raw_size), int(remain)), 0)))

    return out_signature


def _broadcast_static_shapes(
    shape_a: Optional[List[int]],
    shape_b: Optional[List[int]],
) -> Optional[List[int]]:
    if not _is_fully_known_positive_shape(shape_a) or not _is_fully_known_positive_shape(shape_b):
        return None
    a = [int(v) for v in list(shape_a)]
    b = [int(v) for v in list(shape_b)]
    rank = max(len(a), len(b))
    a = [1] * (rank - len(a)) + a
    b = [1] * (rank - len(b)) + b
    out: List[int] = []
    for dim_a, dim_b in zip(a, b):
        if int(dim_a) == int(dim_b):
            out.append(int(dim_a))
            continue
        if int(dim_a) == 1:
            out.append(int(dim_b))
            continue
        if int(dim_b) == 1:
            out.append(int(dim_a))
            continue
        return None
    return out


def _broadcast_shape_signatures(
    signature_a: Optional[List[int]],
    signature_b: Optional[List[int]],
) -> Optional[List[int]]:
    if signature_a is None or signature_b is None:
        return None
    a = [int(v) for v in list(signature_a)]
    b = [int(v) for v in list(signature_b)]
    rank = max(len(a), len(b))
    a = [1] * (rank - len(a)) + a
    b = [1] * (rank - len(b)) + b
    out: List[int] = []
    for dim_a, dim_b in zip(a, b):
        if int(dim_a) == int(dim_b):
            out.append(int(dim_a))
            continue
        if int(dim_a) == 1:
            out.append(-1 if int(dim_b) < 0 else int(dim_b))
            continue
        if int(dim_b) == 1:
            out.append(-1 if int(dim_a) < 0 else int(dim_a))
            continue
        if int(dim_a) < 0 and int(dim_b) < 0:
            out.append(-1)
            continue
        if int(dim_a) < 0 and int(dim_b) > 1:
            out.append(int(dim_b))
            continue
        if int(dim_b) < 0 and int(dim_a) > 1:
            out.append(int(dim_a))
            continue
        return None
    return out


def _infer_batch_matmul_output_shape_and_signature(
    shape_a: Optional[List[int]],
    shape_b: Optional[List[int]],
    signature_a: Optional[List[int]],
    signature_b: Optional[List[int]],
    adj_x: bool,
    adj_y: bool,
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
    if not _is_fully_known_positive_shape(shape_a) or not _is_fully_known_positive_shape(shape_b):
        return None, None
    a_shape = [int(v) for v in list(shape_a)]
    b_shape = [int(v) for v in list(shape_b)]
    if len(a_shape) < 2 or len(b_shape) < 2:
        return None, None

    a_sig = (
        [int(v) for v in list(signature_a)]
        if signature_a is not None and len(list(signature_a)) == len(a_shape)
        else list(a_shape)
    )
    b_sig = (
        [int(v) for v in list(signature_b)]
        if signature_b is not None and len(list(signature_b)) == len(b_shape)
        else list(b_shape)
    )

    batch_shape = _broadcast_static_shapes(a_shape[:-2], b_shape[:-2])
    if batch_shape is None:
        return None, None

    a_m_idx = -1 if bool(adj_x) else -2
    a_k_idx = -2 if bool(adj_x) else -1
    b_k_idx = -1 if bool(adj_y) else -2
    b_n_idx = -2 if bool(adj_y) else -1

    a_m = int(a_shape[a_m_idx])
    a_k = int(a_shape[a_k_idx])
    b_k = int(b_shape[b_k_idx])
    b_n = int(b_shape[b_n_idx])

    a_k_sig = int(a_sig[a_k_idx])
    b_k_sig = int(b_sig[b_k_idx])
    if a_k != b_k and a_k_sig >= 0 and b_k_sig >= 0:
        return None, None

    out_shape = [int(v) for v in list(batch_shape)] + [int(a_m), int(b_n)]

    batch_sig = _broadcast_shape_signatures(a_sig[:-2], b_sig[:-2])
    if batch_sig is None:
        batch_sig = [int(v) for v in list(batch_shape)]
    out_sig_m = -1 if int(a_sig[a_m_idx]) < 0 else int(a_m)
    out_sig_n = -1 if int(b_sig[b_n_idx]) < 0 else int(b_n)
    out_signature = [int(v) for v in list(batch_sig)] + [int(out_sig_m), int(out_sig_n)]

    return out_shape, out_signature


def _infer_rank4_signature_from_input(
    *,
    input_signature: Optional[List[int]],
    output_shape: Optional[List[int]],
    existing_output_signature: Optional[List[int]] = None,
    propagate_channel: bool = False,
) -> Optional[List[int]]:
    if output_shape is None or len(list(output_shape)) != 4:
        return None
    signature = [int(v) for v in list(output_shape)]

    normalized_input_signature = (
        [int(v) for v in list(input_signature)]
        if input_signature is not None and len(list(input_signature)) == 4
        else None
    )
    dynamic_from_input = [False, False, False, False]
    if normalized_input_signature is not None:
        dynamic_from_input[0] = int(normalized_input_signature[0]) < 0
        dynamic_from_input[1] = int(normalized_input_signature[1]) < 0
        dynamic_from_input[2] = int(normalized_input_signature[2]) < 0
        dynamic_from_input[3] = bool(propagate_channel) and int(normalized_input_signature[3]) < 0

    # Preserve existing dynamic marks only when input-derived dynamics can explain
    # that axis. This avoids stale all-unknown placeholders dominating new
    # operator-local inferences (e.g. static channel axis from filters).
    if (
        existing_output_signature is not None
        and len(list(existing_output_signature)) == 4
    ):
        if normalized_input_signature is None:
            for axis in range(4):
                if int(existing_output_signature[axis]) < 0:
                    signature[axis] = -1
        else:
            for axis in range(4):
                if int(existing_output_signature[axis]) < 0 and bool(dynamic_from_input[axis]):
                    signature[axis] = -1

    if normalized_input_signature is not None:
        if int(normalized_input_signature[0]) < 0:
            signature[0] = -1
        if int(normalized_input_signature[1]) < 0:
            signature[1] = -1
        if int(normalized_input_signature[2]) < 0:
            signature[2] = -1
        if bool(propagate_channel) and int(normalized_input_signature[3]) < 0:
            signature[3] = -1
    return signature


def _normalize_reduce_axes_for_rank(
    axes: Optional[List[int]],
    rank: int,
) -> Optional[List[int]]:
    if axes is None:
        return None
    if rank < 0:
        return None
    normalized: List[int] = []
    for axis in list(axes):
        a = int(axis)
        if a < 0:
            a += int(rank)
        if a < 0 or a >= int(rank):
            return None
        if int(a) not in normalized:
            normalized.append(int(a))
    return normalized


def _infer_reduce_output_shape_and_signature(
    *,
    input_shape: Optional[List[int]],
    input_signature: Optional[List[int]],
    axes: Optional[List[int]],
    keep_dims: bool,
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
    if input_shape is None:
        return None, None
    in_shape = [int(v) for v in list(input_shape)]
    rank = len(in_shape)
    if rank == 0:
        return list(in_shape), list(in_shape)
    normalized_axes = _normalize_reduce_axes_for_rank(axes, rank)
    if normalized_axes is None:
        return None, None

    if input_signature is not None and len(list(input_signature)) == rank:
        in_signature = [int(v) for v in list(input_signature)]
    else:
        in_signature = [int(v) for v in list(in_shape)]

    if bool(keep_dims):
        out_shape = [int(v) for v in list(in_shape)]
        out_signature = [int(v) for v in list(in_signature)]
        for axis in normalized_axes:
            out_shape[int(axis)] = 1
            out_signature[int(axis)] = 1
        return out_shape, out_signature

    reduced_axes = set(int(v) for v in normalized_axes)
    out_shape = [int(in_shape[idx]) for idx in range(rank) if idx not in reduced_axes]
    out_signature = [int(in_signature[idx]) for idx in range(rank) if idx not in reduced_axes]
    return out_shape, out_signature


def _parse_axes_option(raw_axes: Any) -> List[int]:
    if raw_axes is None:
        return []
    if isinstance(raw_axes, (list, tuple, np.ndarray)):
        return [int(v) for v in list(raw_axes)]
    return [int(raw_axes)]


def _normalize_squeeze_axes_for_rank(
    axes: List[int],
    rank: int,
) -> Optional[List[int]]:
    normalized: List[int] = []
    for axis in list(axes):
        a = int(axis)
        if a < 0:
            a += int(rank)
        if a < 0 or a >= int(rank):
            return None
        if int(a) not in normalized:
            normalized.append(int(a))
    return normalized


def _infer_squeeze_output_shape_and_signature(
    *,
    input_shape: Optional[List[int]],
    input_signature: Optional[List[int]],
    squeeze_axes: Optional[List[int]],
) -> Tuple[Optional[List[int]], Optional[List[int]]]:
    if input_shape is None:
        return None, None
    in_shape = [int(v) for v in list(input_shape)]
    rank = len(in_shape)
    if input_signature is not None and len(list(input_signature)) == rank:
        in_signature = [int(v) for v in list(input_signature)]
    else:
        in_signature = [int(v) for v in list(in_shape)]

    axes_list = [int(v) for v in list(squeeze_axes)] if squeeze_axes is not None else []
    normalized_axes = _normalize_squeeze_axes_for_rank(axes_list, rank)
    if normalized_axes is None:
        return None, None
    if len(normalized_axes) == 0:
        normalized_axes = [
            int(idx)
            for idx in range(rank)
            if idx < len(in_signature) and int(in_signature[idx]) == 1
        ]
        if len(normalized_axes) == 0:
            normalized_axes = [
                int(idx)
                for idx, dim in enumerate(in_shape)
                if int(dim) == 1
            ]

    remove_axes = set(int(v) for v in normalized_axes)
    out_shape = [int(in_shape[idx]) for idx in range(rank) if idx not in remove_axes]
    out_signature = [int(in_signature[idx]) for idx in range(rank) if idx not in remove_axes]
    return out_shape, out_signature


def _infer_conv_out_dim(
    in_size: int,
    kernel_size: int,
    stride: int,
    dilation: int,
    padding: str,
) -> Optional[int]:
    if any(int(v) <= 0 for v in [in_size, kernel_size, stride, dilation]):
        return None
    effective_kernel = int((int(kernel_size) - 1) * int(dilation) + 1)
    mode = str(padding).upper()
    if mode == "SAME":
        return int((int(in_size) + int(stride) - 1) // int(stride))
    if mode == "VALID":
        return int((int(in_size) - int(effective_kernel)) // int(stride) + 1)
    return None


def _reconcile_static_tensor_shapes(model_ir: ModelIR) -> Dict[str, int]:
    """
    Recompute static tensor shapes after aggressive graph rewrites.

    Some transpose-bridge optimizations intentionally relax local shape guards and can
    leave stale static metadata. This pass performs a conservative forward fixed-point
    shape propagation and syncs `shape` / `shape_signature` for common TFLite ops.
    """
    updated_tensors = 0

    def _update_tensor_shape(
        tensor_name: str,
        new_shape: Optional[List[int]],
        new_shape_signature: Optional[List[int]] = None,
    ) -> bool:
        nonlocal updated_tensors
        if new_shape is None:
            return False
        if not _is_fully_known_positive_shape(new_shape):
            return False
        tensor = model_ir.tensors.get(tensor_name, None)
        if tensor is None:
            return False
        normalized = [int(v) for v in list(new_shape)]
        if new_shape_signature is not None and len(list(new_shape_signature)) == len(normalized):
            signature = [int(v) for v in list(new_shape_signature)]
        else:
            existing_signature = (
                list(tensor.shape_signature)
                if tensor.shape_signature is not None and len(list(tensor.shape_signature)) == len(normalized)
                else None
            )
            if existing_signature is None:
                signature = [int(v) for v in list(normalized)]
            else:
                # Preserve unknown dims inferred upstream to avoid collapsing dynamic axes
                # into placeholder static 1s during conservative shape reconciliation.
                signature = [
                    int(existing_signature[idx]) if int(existing_signature[idx]) < 0 else int(normalized[idx])
                    for idx in range(len(normalized))
                ]
        if tensor.shape == normalized and tensor.shape_signature == signature:
            return False
        tensor.shape = normalized
        tensor.shape_signature = signature
        updated_tensors += 1
        return True

    max_passes = 32
    for _ in range(max_passes):
        changed = False
        for op in model_ir.operators:
            op_type = str(op.op_type)
            inputs = [str(v) for v in list(op.inputs)]
            outputs = [str(v) for v in list(op.outputs)]
            if len(outputs) == 0:
                continue

            # Pass-through ops: output shape == first input shape.
            if op_type in {
                "QUANTIZE",
                "DEQUANTIZE",
                "SOFTMAX",
                "LOGISTIC",
                "TANH",
                "RELU",
                "RELU6",
                "RELU_0_TO_1",
                "PRELU",
                "HARD_SWISH",
                "LEAKY_RELU",
                "GELU",
                "CAST",
                "NEG",
                "FLOOR",
            }:
                if len(inputs) >= 1:
                    in_tensor = model_ir.tensors.get(inputs[0], None)
                    if in_tensor is not None and _is_fully_known_positive_shape(in_tensor.shape):
                        in_signature = (
                            list(in_tensor.shape_signature)
                            if in_tensor.shape_signature is not None
                            else list(in_tensor.shape)
                        )
                        changed |= _update_tensor_shape(
                            outputs[0],
                            list(in_tensor.shape),
                            in_signature,
                        )
                continue

            if op_type == "TRANSPOSE" and len(inputs) >= 2 and len(outputs) == 1:
                in_tensor = model_ir.tensors.get(inputs[0], None)
                out_tensor = model_ir.tensors.get(outputs[0], None)
                perm_tensor = model_ir.tensors.get(inputs[1], None)
                perm = _read_const_ints_from_tensor(perm_tensor)
                if perm is None:
                    continue
                perm = [int(v) for v in perm]
                if in_tensor is not None and _is_fully_known_positive_shape(in_tensor.shape):
                    out_shape = _permute_shape(list(in_tensor.shape), perm)
                    in_signature = (
                        list(in_tensor.shape_signature)
                        if in_tensor.shape_signature is not None
                        else list(in_tensor.shape)
                    )
                    out_signature = _permute_shape(in_signature, perm)
                    changed |= _update_tensor_shape(
                        outputs[0],
                        out_shape,
                        out_signature,
                    )
                if out_tensor is not None and _is_fully_known_positive_shape(out_tensor.shape):
                    inv_perm = _invert_perm(perm)
                    if inv_perm is not None:
                        in_shape = _permute_shape(list(out_tensor.shape), inv_perm)
                        out_signature = (
                            list(out_tensor.shape_signature)
                            if out_tensor.shape_signature is not None
                            else list(out_tensor.shape)
                        )
                        in_signature = _permute_shape(out_signature, inv_perm)
                        changed |= _update_tensor_shape(
                            inputs[0],
                            in_shape,
                            in_signature,
                        )
                continue

            if op_type == "GATHER" and len(inputs) >= 2 and len(outputs) == 1:
                params_tensor = model_ir.tensors.get(inputs[0], None)
                indices_tensor = model_ir.tensors.get(inputs[1], None)
                if params_tensor is None or indices_tensor is None:
                    continue
                if not _is_fully_known_positive_shape(list(params_tensor.shape)):
                    continue
                params_shape = [int(v) for v in list(params_tensor.shape)]
                params_signature = (
                    [int(v) for v in list(params_tensor.shape_signature)]
                    if params_tensor.shape_signature is not None
                    else list(params_shape)
                )
                indices_shape = [int(v) for v in list(indices_tensor.shape)]
                indices_signature = (
                    [int(v) for v in list(indices_tensor.shape_signature)]
                    if indices_tensor.shape_signature is not None
                    else list(indices_shape)
                )
                axis = int(op.options.get("axis", 0))
                batch_dims = int(op.options.get("batchDims", 0))
                if batch_dims != 0:
                    continue
                rank = len(params_shape)
                if rank <= 0:
                    continue
                if axis < 0:
                    axis += rank
                if axis < 0 or axis >= rank:
                    continue
                out_shape = (
                    [int(v) for v in params_shape[:axis]]
                    + [int(v) for v in indices_shape]
                    + [int(v) for v in params_shape[axis + 1:]]
                )
                out_signature = (
                    [int(v) for v in params_signature[:axis]]
                    + [int(v) for v in indices_signature]
                    + [int(v) for v in params_signature[axis + 1:]]
                )
                if len(out_shape) == 0:
                    out_shape = [1]
                if len(out_signature) == 0:
                    out_signature = [1]
                changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type == "GATHER_ND" and len(inputs) >= 2 and len(outputs) == 1:
                params_tensor = model_ir.tensors.get(inputs[0], None)
                indices_tensor = model_ir.tensors.get(inputs[1], None)
                if params_tensor is None or indices_tensor is None:
                    continue
                if not _is_fully_known_positive_shape(list(params_tensor.shape)):
                    continue
                params_shape = [int(v) for v in list(params_tensor.shape)]
                params_signature = (
                    [int(v) for v in list(params_tensor.shape_signature)]
                    if params_tensor.shape_signature is not None
                    else list(params_shape)
                )
                indices_shape = [int(v) for v in list(indices_tensor.shape)]
                indices_signature = (
                    [int(v) for v in list(indices_tensor.shape_signature)]
                    if indices_tensor.shape_signature is not None
                    else list(indices_shape)
                )
                if len(indices_shape) < 1:
                    continue
                gather_dims = int(indices_shape[-1])
                if len(indices_signature) > 0 and int(indices_signature[-1]) > 0:
                    gather_dims = int(indices_signature[-1])
                if gather_dims <= 0 or gather_dims > len(params_shape):
                    continue
                out_shape = (
                    [int(v) for v in indices_shape[:-1]]
                    + [int(v) for v in params_shape[gather_dims:]]
                )
                out_signature = (
                    [int(v) for v in indices_signature[:-1]]
                    + [int(v) for v in params_signature[gather_dims:]]
                )
                if len(out_shape) == 0:
                    out_shape = [1]
                if len(out_signature) == 0:
                    out_signature = [1]
                changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type in {"BIDIRECTIONAL_SEQUENCE_LSTM", "UNIDIRECTIONAL_SEQUENCE_LSTM"} \
                and len(inputs) >= 1 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                output_tensor = model_ir.tensors.get(outputs[0], None)
                if (
                    input_tensor is None
                    or output_tensor is None
                    or len(list(input_tensor.shape)) != 3
                    or len(list(output_tensor.shape)) != 3
                    or not _is_fully_known_positive_shape(list(input_tensor.shape))
                ):
                    continue
                input_shape = [int(v) for v in list(input_tensor.shape)]
                input_signature = (
                    [int(v) for v in list(input_tensor.shape_signature)]
                    if input_tensor.shape_signature is not None
                    else list(input_shape)
                )
                output_shape = [int(v) for v in list(output_tensor.shape)]
                output_signature_existing = (
                    [int(v) for v in list(output_tensor.shape_signature)]
                    if output_tensor.shape_signature is not None
                    else list(output_shape)
                )
                time_major = bool(op.options.get("timeMajor", True))
                if time_major:
                    inferred_output_shape = [int(input_shape[0]), int(input_shape[1]), int(output_shape[2])]
                    inferred_output_signature = [
                        int(input_signature[0]),
                        int(input_signature[1]),
                        int(output_signature_existing[2]),
                    ]
                else:
                    inferred_output_shape = [int(input_shape[1]), int(input_shape[0]), int(output_shape[2])]
                    inferred_output_signature = [
                        int(input_signature[1]),
                        int(input_signature[0]),
                        int(output_signature_existing[2]),
                    ]
                changed |= _update_tensor_shape(
                    outputs[0],
                    inferred_output_shape,
                    inferred_output_signature,
                )
                continue

            if op_type == "SPLIT" and len(inputs) >= 2 and len(outputs) >= 1:
                axis_tensor = model_ir.tensors.get(inputs[0], None)
                input_tensor = model_ir.tensors.get(inputs[1], None)
                axis_values = _read_const_ints_from_tensor(axis_tensor)
                if (
                    axis_values is None
                    or len(axis_values) == 0
                    or input_tensor is None
                    or not _is_fully_known_positive_shape(list(input_tensor.shape))
                ):
                    continue
                input_shape = [int(v) for v in list(input_tensor.shape)]
                rank = len(input_shape)
                axis = int(axis_values[0])
                if axis < 0:
                    axis += rank
                if axis < 0 or axis >= rank:
                    continue
                num_splits = int(op.options.get("numSplits", len(outputs)))
                if num_splits <= 0 or len(outputs) != num_splits:
                    continue
                if int(input_shape[axis]) % num_splits != 0:
                    continue
                split_dim = int(input_shape[axis] // num_splits)
                input_signature = (
                    [int(v) for v in list(input_tensor.shape_signature)]
                    if input_tensor.shape_signature is not None
                    else list(input_shape)
                )
                split_signature_dim = -1
                if axis < len(input_signature):
                    sig_dim = int(input_signature[axis])
                    if sig_dim > 0 and sig_dim % num_splits == 0:
                        split_signature_dim = int(sig_dim // num_splits)
                for output_name in outputs:
                    out_shape = [int(v) for v in list(input_shape)]
                    out_shape[axis] = int(split_dim)
                    out_signature = [int(v) for v in list(input_signature)]
                    out_signature[axis] = int(split_signature_dim if split_signature_dim > 0 else -1)
                    changed |= _update_tensor_shape(
                        output_name,
                        out_shape,
                        out_signature,
                    )
                continue

            if op_type == "EXPAND_DIMS" and len(inputs) >= 2 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                axis_tensor = model_ir.tensors.get(inputs[1], None)
                axis_values = _read_const_ints_from_tensor(axis_tensor)
                if (
                    input_tensor is None
                    or axis_values is None
                    or len(axis_values) == 0
                    or not _is_fully_known_positive_shape(list(input_tensor.shape))
                ):
                    continue
                input_shape = [int(v) for v in list(input_tensor.shape)]
                input_signature = (
                    [int(v) for v in list(input_tensor.shape_signature)]
                    if input_tensor.shape_signature is not None
                    else list(input_shape)
                )
                out_rank = len(input_shape) + 1
                axis = int(axis_values[0])
                if axis < 0:
                    axis += out_rank
                if axis < 0 or axis >= out_rank:
                    continue
                out_shape = [int(v) for v in list(input_shape)]
                out_signature = [int(v) for v in list(input_signature)]
                out_shape.insert(axis, 1)
                out_signature.insert(axis, 1)
                changed |= _update_tensor_shape(
                    outputs[0],
                    out_shape,
                    out_signature,
                )
                continue

            if op_type in {"ADD", "SUB", "MUL", "DIV", "FLOOR_MOD", "MAXIMUM", "MINIMUM"} and len(inputs) >= 2 and len(outputs) == 1:
                in0 = model_ir.tensors.get(inputs[0], None)
                in1 = model_ir.tensors.get(inputs[1], None)
                shape0 = list(in0.shape) if in0 is not None else None
                shape1 = list(in1.shape) if in1 is not None else None
                sig0 = (
                    list(in0.shape_signature)
                    if in0 is not None and in0.shape_signature is not None
                    else (list(shape0) if shape0 is not None else None)
                )
                sig1 = (
                    list(in1.shape_signature)
                    if in1 is not None and in1.shape_signature is not None
                    else (list(shape1) if shape1 is not None else None)
                )
                out_shape = _broadcast_static_shapes(shape0, shape1)
                out_signature = _broadcast_shape_signatures(sig0, sig1)
                if out_shape is not None:
                    changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type == "BATCH_MATMUL" and len(inputs) >= 2 and len(outputs) == 1:
                in0 = model_ir.tensors.get(inputs[0], None)
                in1 = model_ir.tensors.get(inputs[1], None)
                shape0 = list(in0.shape) if in0 is not None else None
                shape1 = list(in1.shape) if in1 is not None else None
                sig0 = (
                    list(in0.shape_signature)
                    if in0 is not None and in0.shape_signature is not None
                    else (list(shape0) if shape0 is not None else None)
                )
                sig1 = (
                    list(in1.shape_signature)
                    if in1 is not None and in1.shape_signature is not None
                    else (list(shape1) if shape1 is not None else None)
                )
                out_shape, out_signature = _infer_batch_matmul_output_shape_and_signature(
                    shape_a=shape0,
                    shape_b=shape1,
                    signature_a=sig0,
                    signature_b=sig1,
                    adj_x=bool(op.options.get("adjX", False)),
                    adj_y=bool(op.options.get("adjY", False)),
                )
                if out_shape is not None:
                    changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type == "CONCATENATION" and len(inputs) >= 1 and len(outputs) == 1:
                axis = op.options.get("axis", None)
                if axis is None:
                    continue
                in_shapes: List[List[int]] = []
                in_signatures: List[List[int]] = []
                ranks: List[int] = []
                valid = True
                for input_name in inputs:
                    t = model_ir.tensors.get(input_name, None)
                    if t is None or not _is_fully_known_positive_shape(t.shape):
                        valid = False
                        break
                    shape = [int(v) for v in list(t.shape)]
                    in_shapes.append(shape)
                    in_signatures.append(
                        [int(v) for v in list(t.shape_signature)]
                        if t.shape_signature is not None
                        else list(shape)
                    )
                    ranks.append(len(shape))
                if not valid or len(in_shapes) == 0 or len(set(ranks)) != 1:
                    continue
                rank = int(ranks[0])
                axis_new = int(axis)
                if axis_new < 0:
                    axis_new += rank
                if axis_new < 0 or axis_new >= rank:
                    continue
                out_shape = list(in_shapes[0])
                compatible = True
                for shape in in_shapes[1:]:
                    for dim_idx in range(rank):
                        if dim_idx == axis_new:
                            continue
                        if int(shape[dim_idx]) != int(out_shape[dim_idx]):
                            compatible = False
                            break
                    if not compatible:
                        break
                    out_shape[axis_new] += int(shape[axis_new])
                if compatible:
                    out_signature = [int(v) for v in list(out_shape)]
                    if len(in_signatures) == len(in_shapes):
                        for dim_idx in range(rank):
                            dim_values = [int(sig[dim_idx]) for sig in in_signatures]
                            if dim_idx == axis_new:
                                if any(int(v) < 0 for v in dim_values):
                                    out_signature[dim_idx] = -1
                                else:
                                    out_signature[dim_idx] = int(sum(int(v) for v in dim_values))
                            else:
                                if any(int(v) < 0 for v in dim_values):
                                    out_signature[dim_idx] = -1
                    changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type == "RESHAPE" and len(inputs) >= 1 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                if input_tensor is None:
                    continue
                input_signature = (
                    list(input_tensor.shape_signature)
                    if input_tensor.shape_signature is not None
                    else list(input_tensor.shape)
                )
                raw_new_shape = op.options.get("newShape", [])
                try:
                    new_shape = [int(v) for v in np.asarray(raw_new_shape).reshape(-1).tolist()]
                except Exception:
                    new_shape = []
                has_onnx_raw_new_shape = "onnxRawNewShape" in op.options
                onnx_raw_shape = op.options.get("onnxRawNewShape", [])
                try:
                    onnx_raw_shape_list = [
                        int(v) for v in np.asarray(onnx_raw_shape).reshape(-1).tolist()
                    ]
                except Exception:
                    onnx_raw_shape_list = []
                if (
                    has_onnx_raw_new_shape
                    and len(onnx_raw_shape_list) > 0
                    and all(int(dim) > 0 for dim in onnx_raw_shape_list)
                ):
                    resolved = [int(v) for v in onnx_raw_shape_list]
                else:
                    resolved = _resolve_reshape_new_shape_from_static_input(
                        new_shape=new_shape,
                        input_signature=input_signature,
                        allow_zero=(
                            bool(op.options.get("allowZero"))
                            if "allowZero" in op.options
                            else None
                        ),
                    )
                if resolved is not None and _is_fully_known_positive_shape(resolved):
                    changed |= _update_tensor_shape(outputs[0], resolved)
                continue

            if op_type == "SQUEEZE" and len(inputs) >= 1 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                if input_tensor is None or not _is_fully_known_positive_shape(input_tensor.shape):
                    continue
                raw_axes = op.options.get("squeezeDims", [])
                squeeze_axes = _parse_axes_option(raw_axes)
                input_signature = (
                    list(input_tensor.shape_signature)
                    if input_tensor.shape_signature is not None
                    else list(input_tensor.shape)
                )
                out_shape, out_signature = _infer_squeeze_output_shape_and_signature(
                    input_shape=list(input_tensor.shape),
                    input_signature=input_signature,
                    squeeze_axes=squeeze_axes,
                )
                if out_shape is None or not _is_fully_known_positive_shape(out_shape):
                    continue
                changed |= _update_tensor_shape(
                    outputs[0],
                    out_shape,
                    out_signature,
                )
                continue

            if op_type == "MEAN" and len(inputs) >= 2 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                axes_tensor = model_ir.tensors.get(inputs[1], None)
                if input_tensor is None or not _is_fully_known_positive_shape(input_tensor.shape):
                    continue
                axes_vals = _read_const_ints_from_tensor(axes_tensor)
                keep_dims = bool(op.options.get("keepDims", False))
                input_signature = (
                    list(input_tensor.shape_signature)
                    if input_tensor.shape_signature is not None
                    else list(input_tensor.shape)
                )
                out_shape, out_signature = _infer_reduce_output_shape_and_signature(
                    input_shape=list(input_tensor.shape),
                    input_signature=input_signature,
                    axes=axes_vals,
                    keep_dims=keep_dims,
                )
                if out_shape is None or not _is_fully_known_positive_shape(out_shape):
                    continue
                changed |= _update_tensor_shape(
                    outputs[0],
                    out_shape,
                    out_signature,
                )
                continue

            if op_type == "SLICE" and len(inputs) >= 3 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                begin_tensor = model_ir.tensors.get(inputs[1], None)
                size_tensor = model_ir.tensors.get(inputs[2], None)
                if input_tensor is None:
                    continue
                input_signature = (
                    list(input_tensor.shape_signature)
                    if input_tensor.shape_signature is not None
                    else list(input_tensor.shape)
                )
                has_dynamic_input_dim = any(int(v) < 0 for v in input_signature)
                begin_vals = _read_const_ints_from_tensor(begin_tensor)
                size_vals = _read_const_ints_from_tensor(size_tensor)
                out_shape, resolved_begin, resolved_size = _infer_slice_output_shape_and_resolved_params(
                    input_shape=list(input_tensor.shape),
                    begin_vals=begin_vals,
                    size_vals=size_vals,
                )
                if out_shape is None or resolved_begin is None or resolved_size is None:
                    continue
                out_signature = _infer_slice_output_signature(
                    input_shape=list(input_tensor.shape),
                    input_signature=input_signature,
                    begin_vals=begin_vals,
                    size_vals=size_vals,
                )
                shape_for_update = [int(v) for v in list(out_shape)]
                if (
                    has_dynamic_input_dim
                    and out_signature is not None
                    and len(out_signature) == len(shape_for_update)
                ):
                    existing_output_shape = []
                    output_tensor = model_ir.tensors.get(outputs[0], None)
                    if output_tensor is not None:
                        existing_output_shape = [int(v) for v in list(output_tensor.shape)]
                    merged_shape: List[int] = []
                    for idx, dim in enumerate(shape_for_update):
                        sig_dim = int(out_signature[idx])
                        if sig_dim > 0:
                            merged_shape.append(int(sig_dim))
                            continue
                        if idx < len(existing_output_shape) and int(existing_output_shape[idx]) > 0:
                            merged_shape.append(int(existing_output_shape[idx]))
                            continue
                        merged_shape.append(int(dim) if int(dim) > 0 else 1)
                    shape_for_update = [int(v) for v in list(merged_shape)]
                out_shape_is_fully_known_positive = _is_fully_known_positive_shape(shape_for_update)
                if out_shape_is_fully_known_positive:
                    changed |= _update_tensor_shape(outputs[0], shape_for_update, out_signature)
                # Keep runtime-driven SLICE semantics when input dimensions are dynamic.
                # Also avoid mutating begin/size when inferred output dims include 0 because
                # stale static metadata can collapse valid slicing ranges into empty outputs.
                if (
                    out_shape_is_fully_known_positive
                    and not has_dynamic_input_dim
                    and not bool(op.options.get("preserveDynamicShape", False))
                ):
                    changed |= _write_const_ints_to_tensor(begin_tensor, resolved_begin)
                    changed |= _write_const_ints_to_tensor(size_tensor, resolved_size)
                continue

            if op_type in {"PAD", "MIRROR_PAD"} and len(inputs) >= 2 and len(outputs) == 1:
                input_tensor = model_ir.tensors.get(inputs[0], None)
                pads_tensor = model_ir.tensors.get(inputs[1], None)
                if (
                    input_tensor is None
                    or pads_tensor is None
                    or not _is_fully_known_positive_shape(input_tensor.shape)
                ):
                    continue
                in_shape = [int(v) for v in list(input_tensor.shape)]
                rank = len(in_shape)
                pads_vals = _read_const_ints_from_tensor(pads_tensor)
                if pads_vals is None or len(pads_vals) != int(rank * 2):
                    continue
                pads_shape = (
                    [int(v) for v in list(pads_tensor.shape)]
                    if pads_tensor.shape is not None
                    else []
                )

                pad_pairs: List[Tuple[int, int]] = []
                if len(pads_shape) == 2 and int(pads_shape[0]) == int(rank) and int(pads_shape[1]) == 2:
                    for axis in range(rank):
                        before = int(pads_vals[int(axis * 2)])
                        after = int(pads_vals[int(axis * 2 + 1)])
                        pad_pairs.append((before, after))
                elif len(pads_shape) == 2 and int(pads_shape[0]) == 2 and int(pads_shape[1]) == int(rank):
                    for axis in range(rank):
                        before = int(pads_vals[int(axis)])
                        after = int(pads_vals[int(rank + axis)])
                        pad_pairs.append((before, after))
                else:
                    for axis in range(rank):
                        before = int(pads_vals[int(axis * 2)])
                        after = int(pads_vals[int(axis * 2 + 1)])
                        pad_pairs.append((before, after))

                if any(int(before) < 0 or int(after) < 0 for before, after in pad_pairs):
                    continue

                out_shape: List[int] = []
                valid_out = True
                for axis in range(rank):
                    before, after = pad_pairs[axis]
                    dim = int(in_shape[axis]) + int(before) + int(after)
                    if int(dim) <= 0:
                        valid_out = False
                        break
                    out_shape.append(int(dim))
                if not valid_out:
                    continue

                input_signature = (
                    list(input_tensor.shape_signature)
                    if input_tensor.shape_signature is not None
                    else list(in_shape)
                )
                out_signature: List[int] = []
                for axis in range(rank):
                    sig_dim = int(input_signature[axis]) if axis < len(input_signature) else int(in_shape[axis])
                    if int(sig_dim) < 0:
                        out_signature.append(-1)
                    else:
                        out_signature.append(int(out_shape[axis]))

                changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type in {"CONV_2D", "DEPTHWISE_CONV_2D"} and len(inputs) >= 2 and len(outputs) == 1:
                in_tensor = model_ir.tensors.get(inputs[0], None)
                filter_tensor = model_ir.tensors.get(inputs[1], None)
                out_tensor = model_ir.tensors.get(outputs[0], None)
                if (
                    in_tensor is None
                    or filter_tensor is None
                    or not _is_fully_known_positive_shape(in_tensor.shape)
                    or not _is_fully_known_positive_shape(filter_tensor.shape)
                ):
                    continue
                in_shape = [int(v) for v in list(in_tensor.shape)]
                filter_shape = [int(v) for v in list(filter_tensor.shape)]
                if len(in_shape) != 4 or len(filter_shape) != 4:
                    continue
                padding = str(op.options.get("padding", "SAME"))
                stride_h = int(op.options.get("strideH", 1))
                stride_w = int(op.options.get("strideW", 1))
                dilation_h = int(op.options.get("dilationHFactor", 1))
                dilation_w = int(op.options.get("dilationWFactor", 1))
                kernel_h = int(filter_shape[1])
                kernel_w = int(filter_shape[2])
                out_h = _infer_conv_out_dim(in_shape[1], kernel_h, stride_h, dilation_h, padding)
                out_w = _infer_conv_out_dim(in_shape[2], kernel_w, stride_w, dilation_w, padding)
                if out_h is None or out_w is None or int(out_h) <= 0 or int(out_w) <= 0:
                    continue
                if op_type == "CONV_2D":
                    out_c = int(filter_shape[0])
                else:
                    # TFLite depthwise filter layout: [1, KH, KW, OC]
                    out_c = int(filter_shape[3])
                out_shape = [int(in_shape[0]), int(out_h), int(out_w), int(out_c)]
                input_signature = (
                    list(in_tensor.shape_signature)
                    if in_tensor.shape_signature is not None
                    else list(in_shape)
                )
                existing_out_signature = (
                    list(out_tensor.shape_signature)
                    if out_tensor is not None and out_tensor.shape_signature is not None
                    else None
                )
                out_signature = _infer_rank4_signature_from_input(
                    input_signature=input_signature,
                    output_shape=out_shape,
                    existing_output_signature=existing_out_signature,
                    propagate_channel=False,
                )
                changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type in {"AVERAGE_POOL_2D", "MAX_POOL_2D"} and len(inputs) >= 1 and len(outputs) == 1:
                in_tensor = model_ir.tensors.get(inputs[0], None)
                out_tensor = model_ir.tensors.get(outputs[0], None)
                if in_tensor is None or not _is_fully_known_positive_shape(in_tensor.shape):
                    continue
                in_shape = [int(v) for v in list(in_tensor.shape)]
                if len(in_shape) != 4:
                    continue

                padding = str(op.options.get("padding", "SAME")).upper()
                stride_h = int(op.options.get("strideH", 1))
                stride_w = int(op.options.get("strideW", 1))
                filter_h = int(op.options.get("filterHeight", 1))
                filter_w = int(op.options.get("filterWidth", 1))

                out_h = _infer_conv_out_dim(in_shape[1], filter_h, stride_h, 1, padding)
                out_w = _infer_conv_out_dim(in_shape[2], filter_w, stride_w, 1, padding)

                # GlobalAveragePool lowering can become stale after aggressive transpose/layout
                # rewrites. Recover invalid VALID+stride=1 average-pool metadata by snapping
                # pool kernel to the current input spatial size.
                if (
                    op_type == "AVERAGE_POOL_2D"
                    and (out_h is None or out_w is None or int(out_h) <= 0 or int(out_w) <= 0)
                    and padding == "VALID"
                    and int(stride_h) == 1
                    and int(stride_w) == 1
                ):
                    if int(filter_h) > int(in_shape[1]) or int(filter_w) > int(in_shape[2]):
                        filter_h = int(in_shape[1])
                        filter_w = int(in_shape[2])
                        op.options["filterHeight"] = int(filter_h)
                        op.options["filterWidth"] = int(filter_w)
                        out_h = _infer_conv_out_dim(in_shape[1], filter_h, stride_h, 1, padding)
                        out_w = _infer_conv_out_dim(in_shape[2], filter_w, stride_w, 1, padding)

                if out_h is None or out_w is None or int(out_h) <= 0 or int(out_w) <= 0:
                    continue
                out_shape = [int(in_shape[0]), int(out_h), int(out_w), int(in_shape[3])]
                input_signature = (
                    list(in_tensor.shape_signature)
                    if in_tensor.shape_signature is not None
                    else list(in_shape)
                )
                existing_out_signature = (
                    list(out_tensor.shape_signature)
                    if out_tensor is not None and out_tensor.shape_signature is not None
                    else None
                )
                out_signature = _infer_rank4_signature_from_input(
                    input_signature=input_signature,
                    output_shape=out_shape,
                    existing_output_signature=existing_out_signature,
                    propagate_channel=True,
                )
                changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

            if op_type in {"RESIZE_BILINEAR", "RESIZE_NEAREST_NEIGHBOR"} and len(inputs) >= 2 and len(outputs) == 1:
                in_tensor = model_ir.tensors.get(inputs[0], None)
                size_tensor = model_ir.tensors.get(inputs[1], None)
                out_tensor = model_ir.tensors.get(outputs[0], None)
                if in_tensor is None or not _is_fully_known_positive_shape(in_tensor.shape):
                    continue
                in_shape = [int(v) for v in list(in_tensor.shape)]
                if len(in_shape) != 4:
                    continue

                out_h = None
                out_w = None
                size_source = "unknown"
                size_vals = _read_const_ints_from_tensor(size_tensor)
                size_tensor_is_const = bool(size_tensor is not None and size_tensor.data is not None)
                onnx_sizes_hw = op.options.get("onnxSizesHW", None)
                if isinstance(onnx_sizes_hw, (list, tuple)) and len(onnx_sizes_hw) >= 2:
                    try:
                        out_h = int(onnx_sizes_hw[0])
                        out_w = int(onnx_sizes_hw[1])
                        size_source = "onnx_sizes"
                    except Exception:
                        out_h = None
                        out_w = None
                        size_source = "unknown"
                if out_h is None or out_w is None:
                    onnx_scales_hw = op.options.get("onnxScalesHW", None)
                    if isinstance(onnx_scales_hw, (list, tuple)) and len(onnx_scales_hw) >= 2:
                        try:
                            out_h = int(round(float(in_shape[1]) * float(onnx_scales_hw[0])))
                            out_w = int(round(float(in_shape[2]) * float(onnx_scales_hw[1])))
                            size_source = "onnx_scales"
                        except Exception:
                            out_h = None
                            out_w = None
                            size_source = "unknown"
                if out_h is None or out_w is None:
                    if size_vals is None or len(size_vals) < 2:
                        continue
                    out_h = int(size_vals[0])
                    out_w = int(size_vals[1])
                    size_source = "const_size"
                if out_h <= 0 or out_w <= 0:
                    continue
                # Synchronize resize-size constants with ONNX-derived hints so downstream
                # prepare-time shape inference follows the reconciled NHWC metadata.
                if size_tensor_is_const and size_source in {"onnx_sizes", "onnx_scales"}:
                    changed |= _write_const_ints_to_tensor(size_tensor, [int(out_h), int(out_w)])
                out_shape = [int(in_shape[0]), int(out_h), int(out_w), int(in_shape[3])]
                input_signature = (
                    list(in_tensor.shape_signature)
                    if in_tensor.shape_signature is not None
                    else list(in_shape)
                )
                existing_out_signature = (
                    list(out_tensor.shape_signature)
                    if out_tensor is not None and out_tensor.shape_signature is not None
                    else None
                )
                out_signature = _infer_rank4_signature_from_input(
                    input_signature=input_signature,
                    output_shape=out_shape,
                    existing_output_signature=existing_out_signature,
                    propagate_channel=True,
                )
                if len(out_signature) == 4 and len(input_signature) == 4 and size_source == "onnx_scales":
                    if int(input_signature[1]) < 0:
                        out_signature[1] = -1
                    if int(input_signature[2]) < 0:
                        out_signature[2] = -1
                changed |= _update_tensor_shape(outputs[0], out_shape, out_signature)
                continue

        if not changed:
            break

    return {"reconciled_static_tensor_shapes": int(updated_tensors)}


def _optimize_transpose_quant_dequant_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose-quantize/dequantize-transpose bridges when both transposes cancel.

    Target patterns:
      X --Transpose(P)--> A --QUANTIZE--> B --Transpose(inv(P))--> Y
      X --Transpose(P)--> A --DEQUANTIZE--> B --Transpose(inv(P))--> Y

    This is safe for per-tensor quantization only.
    """
    removed_bridge_pairs = 0
    rewritten_add_qdq_residual_bridges = 0
    rewritten_mixed_add_qdq_residual_bridges = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        # Pattern A:
        #   X -T(P)-> A -Q-> B -DQ-> C -T(invP)-> Y
        # -> X -Q-> B -DQ-> Y
        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            bridge_a = pre_op.outputs[0]
            q_users = consumers.get(bridge_a, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if q_op.inputs[0] != bridge_a:
                continue

            bridge_b = q_op.outputs[0]
            dq_users = consumers.get(bridge_b, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if dq_op.inputs[0] != bridge_b:
                continue

            bridge_c = dq_op.outputs[0]
            post_users = consumers.get(bridge_c, [])
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if post_op.inputs[0] != bridge_c:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_pre is None or perm_post is None:
                continue
            if not _is_inverse_perm(perm_pre, perm_post):
                continue

            q_out_tensor = model_ir.tensors.get(bridge_b, None)
            if q_out_tensor is None or not _is_per_tensor_quantization(q_out_tensor.quantization):
                continue

            if bridge_a in model_ir.outputs or bridge_b in model_ir.outputs or bridge_c in model_ir.outputs:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=q_op,
                new_inputs=[pre_op.inputs[0]],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=dq_op,
                new_outputs=[post_op.outputs[0]],
            )

            new_dq_out = model_ir.tensors.get(post_op.outputs[0], None)
            old_dq_out = model_ir.tensors.get(bridge_c, None)
            if new_dq_out is not None and old_dq_out is not None:
                new_dq_out.dtype = str(old_dq_out.dtype)

            for remove_idx in sorted([pre_idx, post_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_bridge_pairs += 1
            changed = True
            break

        if changed:
            continue

        # Pattern B:
        #   X -T(P)-> A -(Q|DQ)-> B -T(invP)-> Y
        # -> X -(Q|DQ)-> Y
        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in {"QUANTIZE", "DEQUANTIZE"}:
                continue
            if len(mid_op.inputs) != 1 or len(mid_op.outputs) != 1:
                continue

            bridge_in = mid_op.inputs[0]
            bridge_out = mid_op.outputs[0]

            pre_idx = producers.get(bridge_in, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[pre_idx]
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.outputs) != 1 or len(pre_op.inputs) < 2:
                continue
            if pre_op.outputs[0] != bridge_in:
                continue
            in_users = set(consumers.get(bridge_in, []))
            if int(mid_idx) not in in_users:
                continue
            can_remove_pre = in_users == {int(mid_idx)}

            post_users = [int(v) for v in consumers.get(bridge_out, []) if int(v) != int(mid_idx)]
            if len(post_users) == 0:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            valid_posts = True
            for post_idx in post_users:
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or post_op.inputs[0] != bridge_out
                ):
                    valid_posts = False
                    break
                perm_post = _read_transpose_perm(model_ir, post_op)
                if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                    valid_posts = False
                    break
                post_indices.append(int(post_idx))
                post_output_names.append(str(post_op.outputs[0]))
            if not valid_posts or len(post_indices) == 0:
                continue

            # Keep visible output names stable.
            if bridge_out in model_ir.outputs:
                continue
            if can_remove_pre and bridge_in in model_ir.outputs:
                continue
            single_post = len(post_indices) == 1
            if (not single_post) and any(
                post_output_name in model_ir.outputs for post_output_name in post_output_names
            ):
                continue

            pre_input_name = pre_op.inputs[0]
            representative_post_output_name = post_output_names[0]

            if str(mid_op.op_type) == "QUANTIZE":
                bridge_out_tensor = model_ir.tensors.get(bridge_out, None)
                if bridge_out_tensor is None:
                    continue
                if not _is_per_tensor_quantization(bridge_out_tensor.quantization):
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_inputs=[pre_input_name],
                )
                if single_post:
                    post_output_name = representative_post_output_name
                    _set_operator_outputs(
                        model_ir=model_ir,
                        op=mid_op,
                        new_outputs=[post_output_name],
                    )
                    post_output_tensor = model_ir.tensors.get(post_output_name, None)
                    if post_output_tensor is not None:
                        post_output_tensor.dtype = str(bridge_out_tensor.dtype)
                        post_output_tensor.quantization = _clone_quantization(
                            bridge_out_tensor.quantization
                        )
                else:
                    rep_post_tensor = model_ir.tensors.get(representative_post_output_name, None)
                    if rep_post_tensor is not None:
                        bridge_out_tensor.shape = list(rep_post_tensor.shape)
                        bridge_out_tensor.shape_signature = (
                            list(rep_post_tensor.shape_signature)
                            if rep_post_tensor.shape_signature is not None
                            else list(rep_post_tensor.shape)
                        )
                        bridge_out_tensor.dtype = str(rep_post_tensor.dtype)
                    bridge_out_tensor.quantization = _clone_quantization(
                        bridge_out_tensor.quantization
                    )
            else:
                bridge_in_tensor = model_ir.tensors.get(bridge_in, None)
                if bridge_in_tensor is None:
                    continue
                if not _is_per_tensor_quantization(bridge_in_tensor.quantization):
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_inputs=[pre_input_name],
                )
                if single_post:
                    post_output_name = representative_post_output_name
                    _set_operator_outputs(
                        model_ir=model_ir,
                        op=mid_op,
                        new_outputs=[post_output_name],
                    )
                    pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
                    if pre_input_tensor is not None and pre_input_tensor.quantization is None:
                        pre_input_tensor.quantization = _clone_quantization(
                            bridge_in_tensor.quantization
                        )
                    post_output_tensor = model_ir.tensors.get(post_output_name, None)
                    bridge_out_tensor = model_ir.tensors.get(bridge_out, None)
                    if post_output_tensor is not None and bridge_out_tensor is not None:
                        post_output_tensor.dtype = str(bridge_out_tensor.dtype)
                else:
                    pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
                    if pre_input_tensor is not None and pre_input_tensor.quantization is None:
                        pre_input_tensor.quantization = _clone_quantization(
                            bridge_in_tensor.quantization
                        )
                    bridge_out_tensor = model_ir.tensors.get(bridge_out, None)
                    rep_post_tensor = model_ir.tensors.get(representative_post_output_name, None)
                    if bridge_out_tensor is not None and rep_post_tensor is not None:
                        bridge_out_tensor.shape = list(rep_post_tensor.shape)
                        bridge_out_tensor.shape_signature = (
                            list(rep_post_tensor.shape_signature)
                            if rep_post_tensor.shape_signature is not None
                            else list(rep_post_tensor.shape)
                        )
                        bridge_out_tensor.dtype = str(rep_post_tensor.dtype)

            if not single_post:
                for post_output_name in post_output_names:
                    _replace_tensor_inputs(model_ir, post_output_name, bridge_out)

            remove_indices = list(post_indices)
            if can_remove_pre:
                remove_indices.append(pre_idx)
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[remove_idx]
            removed_bridge_pairs += int(1 if single_post else max(1, len(post_indices)))
            changed = True
            break

        if changed:
            continue

        # Pattern C:
        #   a_nhwc -T(n2c)-> a_nchw -Q-> aq -DQ-> adq -(optional T(c2n))*->
        #   b_nhwc -T(n2c)-> b_nchw -Q-> bq -DQ-> bdq
        #   ADD(adq, bdq) -> sum_nchw -Q-> sum_q -DQ-> sum_dq -T(c2n)-> sum_nhwc
        # -> Remove the NCHW adapters by keeping Q/DQ and ADD in NHWC.
        #
        # Notes:
        # - This intentionally targets strict residual branches where ADD inputs are
        #   Q/DQ-dequantized tensors produced from pre-transpose wrappers.
        # - Side fanout from branch DQ outputs is allowed only via inverse post-transpose.
        model_outputs = set(str(v) for v in model_ir.outputs)
        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            branch_specs: List[Dict[str, Any]] = []
            common_perm_post: Optional[List[int]] = None
            valid_add_branches = True

            for add_input_name in [str(v) for v in list(add_op.inputs)]:
                if add_input_name in model_outputs:
                    valid_add_branches = False
                    break

                dq_idx = producers.get(add_input_name, None)
                if dq_idx is None:
                    valid_add_branches = False
                    break
                dq_op = model_ir.operators[int(dq_idx)]
                if (
                    str(dq_op.op_type) != "DEQUANTIZE"
                    or len(dq_op.inputs) != 1
                    or len(dq_op.outputs) != 1
                    or str(dq_op.outputs[0]) != add_input_name
                ):
                    valid_add_branches = False
                    break

                q_out_name = str(dq_op.inputs[0])
                q_idx = producers.get(q_out_name, None)
                if q_idx is None:
                    valid_add_branches = False
                    break
                q_op = model_ir.operators[int(q_idx)]
                if (
                    str(q_op.op_type) != "QUANTIZE"
                    or len(q_op.inputs) != 1
                    or len(q_op.outputs) != 1
                    or str(q_op.outputs[0]) != q_out_name
                ):
                    valid_add_branches = False
                    break

                q_out_tensor = model_ir.tensors.get(q_out_name, None)
                if q_out_tensor is None or not _is_per_tensor_quantization(q_out_tensor.quantization):
                    valid_add_branches = False
                    break

                pre_in_name = str(q_op.inputs[0])
                pre_idx = producers.get(pre_in_name, None)
                if pre_idx is None:
                    valid_add_branches = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != pre_in_name
                ):
                    valid_add_branches = False
                    break

                perm_pre = _read_transpose_perm(model_ir, pre_op)
                if perm_pre is None:
                    valid_add_branches = False
                    break
                perm_post_expected = _invert_perm(perm_pre)
                if perm_post_expected is None:
                    valid_add_branches = False
                    break
                if common_perm_post is None:
                    common_perm_post = [int(v) for v in list(perm_post_expected)]
                elif list(common_perm_post) != list(perm_post_expected):
                    valid_add_branches = False
                    break

                pre_output_name = str(pre_op.outputs[0])
                pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
                can_remove_pre = set(pre_users) == {int(q_idx)}

                branch_post_indices: List[int] = []
                branch_post_output_names: List[str] = []
                branch_other_users = [
                    int(v)
                    for v in consumers.get(add_input_name, [])
                    if int(v) != int(add_idx)
                ]
                for user_idx in branch_other_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) != "TRANSPOSE"
                        or len(user_op.inputs) < 2
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != add_input_name
                    ):
                        valid_add_branches = False
                        break
                    user_perm = _read_transpose_perm(model_ir, user_op)
                    if user_perm is None or list(user_perm) != list(perm_post_expected):
                        valid_add_branches = False
                        break
                    branch_post_out_name = str(user_op.outputs[0])
                    if branch_post_out_name in model_outputs:
                        valid_add_branches = False
                        break
                    branch_post_indices.append(int(user_idx))
                    branch_post_output_names.append(branch_post_out_name)
                if not valid_add_branches:
                    break

                branch_specs.append(
                    {
                        "pre_idx": int(pre_idx),
                        "q_idx": int(q_idx),
                        "dq_idx": int(dq_idx),
                        "pre_input_name": str(pre_op.inputs[0]),
                        "q_out_name": str(q_out_name),
                        "dq_out_name": str(add_input_name),
                        "perm_post": [int(v) for v in list(perm_post_expected)],
                        "post_indices": [int(v) for v in branch_post_indices],
                        "post_output_names": [str(v) for v in branch_post_output_names],
                        "can_remove_pre": bool(can_remove_pre),
                    }
                )

            if not valid_add_branches or len(branch_specs) != 2 or common_perm_post is None:
                continue

            add_users = [
                int(v)
                for v in consumers.get(add_out_name, [])
                if int(v) != int(add_idx)
            ]
            if len(add_users) != 1:
                continue

            out_q_idx = int(add_users[0])
            out_q_op = model_ir.operators[int(out_q_idx)]
            if str(out_q_op.op_type) != "QUANTIZE" or len(out_q_op.inputs) != 1 or len(out_q_op.outputs) != 1:
                continue
            if str(out_q_op.inputs[0]) != add_out_name:
                continue

            out_q_name = str(out_q_op.outputs[0])
            out_q_tensor = model_ir.tensors.get(out_q_name, None)
            if out_q_tensor is None or not _is_per_tensor_quantization(out_q_tensor.quantization):
                continue
            if out_q_name in model_outputs:
                continue

            out_q_users = [int(v) for v in consumers.get(out_q_name, []) if int(v) != int(out_q_idx)]
            if len(out_q_users) != 1:
                continue

            out_dq_idx = int(out_q_users[0])
            out_dq_op = model_ir.operators[int(out_dq_idx)]
            if (
                str(out_dq_op.op_type) != "DEQUANTIZE"
                or len(out_dq_op.inputs) != 1
                or len(out_dq_op.outputs) != 1
                or str(out_dq_op.inputs[0]) != out_q_name
            ):
                continue

            out_dq_name = str(out_dq_op.outputs[0])
            if out_dq_name in model_outputs:
                continue

            out_post_indices: List[int] = []
            out_post_output_names: List[str] = []
            out_legacy_user_indices: List[int] = []
            invalid_output_tail = False
            out_dq_users = [
                int(v)
                for v in consumers.get(out_dq_name, [])
                if int(v) != int(out_dq_idx)
            ]
            for user_idx in out_dq_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == out_dq_name
                ):
                    user_perm = _read_transpose_perm(model_ir, user_op)
                    if user_perm is not None and list(user_perm) == list(common_perm_post):
                        out_post_name = str(user_op.outputs[0])
                        if out_post_name in model_outputs:
                            invalid_output_tail = True
                            break
                        out_post_indices.append(int(user_idx))
                        out_post_output_names.append(out_post_name)
                        continue
                out_legacy_user_indices.append(int(user_idx))

            if invalid_output_tail:
                continue
            if len(out_legacy_user_indices) > 0 and len(out_post_indices) == 0:
                continue

            remove_indices: set[int] = set()
            removed_transposes_local = 0
            out_dq_tensor = model_ir.tensors.get(out_dq_name, None)
            out_dq_shape_before_permute = (
                [int(v) for v in list(out_dq_tensor.shape)]
                if out_dq_tensor is not None and out_dq_tensor.shape is not None
                else None
            )
            out_dq_signature_before_permute = (
                [int(v) for v in list(out_dq_tensor.shape_signature)]
                if out_dq_tensor is not None and out_dq_tensor.shape_signature is not None
                else (
                    [int(v) for v in list(out_dq_tensor.shape)]
                    if out_dq_tensor is not None and out_dq_tensor.shape is not None
                    else None
                )
            )

            for branch_spec in branch_specs:
                q_op = model_ir.operators[int(branch_spec["q_idx"])]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=q_op,
                    new_inputs=[str(branch_spec["pre_input_name"])],
                )

                q_tensor = model_ir.tensors.get(str(branch_spec["q_out_name"]), None)
                dq_tensor = model_ir.tensors.get(str(branch_spec["dq_out_name"]), None)
                post_output_names = [str(v) for v in list(branch_spec["post_output_names"])]
                if len(post_output_names) > 0:
                    representative_post_tensor = model_ir.tensors.get(post_output_names[0], None)
                    if dq_tensor is not None and representative_post_tensor is not None:
                        dq_tensor.shape = [int(v) for v in list(representative_post_tensor.shape)]
                        dq_tensor.shape_signature = (
                            [int(v) for v in list(representative_post_tensor.shape_signature)]
                            if representative_post_tensor.shape_signature is not None
                            else [int(v) for v in list(representative_post_tensor.shape)]
                        )
                    if q_tensor is not None and dq_tensor is not None:
                        q_tensor.shape = [int(v) for v in list(dq_tensor.shape)]
                        q_tensor.shape_signature = (
                            [int(v) for v in list(dq_tensor.shape_signature)]
                            if dq_tensor.shape_signature is not None
                            else [int(v) for v in list(dq_tensor.shape)]
                        )
                else:
                    _permute_tensor_metadata_if_rank_matches(
                        q_tensor,
                        [int(v) for v in list(branch_spec["perm_post"])],
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        dq_tensor,
                        [int(v) for v in list(branch_spec["perm_post"])],
                    )

                for post_output_name in post_output_names:
                    _replace_tensor_inputs(model_ir, post_output_name, str(branch_spec["dq_out_name"]))
                for post_idx in [int(v) for v in list(branch_spec["post_indices"])]:
                    remove_indices.add(int(post_idx))
                    removed_transposes_local += 1

                if bool(branch_spec["can_remove_pre"]):
                    remove_indices.add(int(branch_spec["pre_idx"]))
                    removed_transposes_local += 1

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_out_name, None),
                [int(v) for v in list(common_perm_post)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(out_q_name, None),
                [int(v) for v in list(common_perm_post)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(out_dq_name, None),
                [int(v) for v in list(common_perm_post)],
            )

            keep_post_idx: Optional[int] = None
            keep_post_output_name: Optional[str] = None
            perm_nhwc_to_nchw = _invert_perm(list(common_perm_post))
            if len(out_legacy_user_indices) > 0:
                keep_post_idx = int(out_post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_post_output_name = str(keep_post_op.outputs[0])
                if perm_nhwc_to_nchw is None:
                    keep_post_idx = None
                    keep_post_output_name = None
                else:
                    keep_perm_name = str(keep_post_op.inputs[1])
                    keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                    if keep_perm_tensor is not None:
                        keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=keep_post_op,
                        new_inputs=[out_dq_name, keep_perm_name],
                    )
                    _set_operator_outputs(
                        model_ir=model_ir,
                        op=keep_post_op,
                        new_outputs=[keep_post_output_name],
                    )
                    if keep_post_output_name is not None:
                        keep_post_tensor = model_ir.tensors.get(keep_post_output_name, None)
                        if (
                            keep_post_tensor is not None
                            and out_dq_shape_before_permute is not None
                            and out_dq_signature_before_permute is not None
                        ):
                            keep_post_tensor.shape = [int(v) for v in list(out_dq_shape_before_permute)]
                            keep_post_tensor.shape_signature = [
                                int(v) for v in list(out_dq_signature_before_permute)
                            ]
            if len(out_legacy_user_indices) > 0 and keep_post_output_name is None:
                continue

            for out_post_output_name in out_post_output_names:
                _replace_tensor_inputs(model_ir, out_post_output_name, out_dq_name)
            for legacy_user_idx in out_legacy_user_indices:
                if keep_post_output_name is None:
                    continue
                legacy_user_op = model_ir.operators[int(legacy_user_idx)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=legacy_user_op,
                    new_inputs=[
                        keep_post_output_name if str(inp) == out_dq_name else str(inp)
                        for inp in list(legacy_user_op.inputs)
                    ],
                )
            for out_post_idx in out_post_indices:
                if keep_post_idx is not None and int(out_post_idx) == int(keep_post_idx):
                    continue
                remove_indices.add(int(out_post_idx))
                removed_transposes_local += 1

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            removed_bridge_pairs += int(max(1, removed_transposes_local // 2))
            rewritten_add_qdq_residual_bridges += 1
            changed = True
            break

        if changed:
            continue

        # Pattern D:
        #   skip_nhwc --T(P)------------------------------> skip_nchw
        #   proj_nhwc --T(P)--> QUANTIZE -> DEQUANTIZE --> proj_nchw
        #   ADD(proj_nchw, skip_nchw) -> y_nchw -> QUANTIZE -> DEQUANTIZE -> T(invP) -> y_nhwc
        # -> Keep Q/DQ and ADD in NHWC by bypassing skip/proj pre-transposes and final post-transpose.
        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            post_input_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if post_input_name in model_ir.outputs:
                continue

            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None:
                continue
            perm_pre_expected = _invert_perm(perm_post)
            if perm_pre_expected is None:
                continue

            dq_out_name = post_input_name
            dq_idx = producers.get(dq_out_name, None)
            if dq_idx is None:
                continue
            dq_op = model_ir.operators[int(dq_idx)]
            if (
                str(dq_op.op_type) != "DEQUANTIZE"
                or len(dq_op.inputs) != 1
                or len(dq_op.outputs) != 1
                or str(dq_op.outputs[0]) != dq_out_name
            ):
                continue

            dq_out_users = [int(v) for v in consumers.get(dq_out_name, []) if int(v) != int(dq_idx)]
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_user_indices: List[int] = []
            valid_post_fanout = True
            for user_idx in dq_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == dq_out_name
                ):
                    user_perm = _read_transpose_perm(model_ir, user_op)
                    if user_perm is not None and list(user_perm) == list(perm_post):
                        post_indices.append(int(user_idx))
                        post_output_names.append(str(user_op.outputs[0]))
                        continue
                legacy_user_indices.append(int(user_idx))
            if not valid_post_fanout:
                continue
            if int(post_idx) not in set(int(v) for v in post_indices):
                continue
            if len(post_indices) == 0:
                continue

            out_q_name = str(dq_op.inputs[0])
            out_q_idx = producers.get(out_q_name, None)
            if out_q_idx is None:
                continue
            out_q_op = model_ir.operators[int(out_q_idx)]
            if (
                str(out_q_op.op_type) != "QUANTIZE"
                or len(out_q_op.inputs) != 1
                or len(out_q_op.outputs) != 1
                or str(out_q_op.outputs[0]) != out_q_name
            ):
                continue

            out_q_tensor = model_ir.tensors.get(out_q_name, None)
            if out_q_tensor is None or not _is_per_tensor_quantization(out_q_tensor.quantization):
                continue
            add_out_name = str(out_q_op.inputs[0])
            if add_out_name in model_ir.outputs:
                continue

            add_idx = producers.get(add_out_name, None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            if str(add_op.outputs[0]) != add_out_name:
                continue

            proj_branch: Optional[Dict[str, Any]] = None
            skip_branch: Optional[Dict[str, Any]] = None
            valid_inputs = True

            for add_input_name in [str(v) for v in list(add_op.inputs)]:
                add_input_prod_idx = producers.get(add_input_name, None)
                if add_input_prod_idx is None:
                    valid_inputs = False
                    break
                add_input_prod_op = model_ir.operators[int(add_input_prod_idx)]
                add_input_prod_type = str(add_input_prod_op.op_type)

                if add_input_prod_type == "DEQUANTIZE":
                    if (
                        len(add_input_prod_op.inputs) != 1
                        or len(add_input_prod_op.outputs) != 1
                        or str(add_input_prod_op.outputs[0]) != add_input_name
                    ):
                        valid_inputs = False
                        break
                    in_q_name = str(add_input_prod_op.inputs[0])
                    in_q_idx = producers.get(in_q_name, None)
                    if in_q_idx is None:
                        valid_inputs = False
                        break
                    in_q_op = model_ir.operators[int(in_q_idx)]
                    if (
                        str(in_q_op.op_type) != "QUANTIZE"
                        or len(in_q_op.inputs) != 1
                        or len(in_q_op.outputs) != 1
                        or str(in_q_op.outputs[0]) != in_q_name
                    ):
                        valid_inputs = False
                        break
                    in_q_tensor = model_ir.tensors.get(in_q_name, None)
                    if in_q_tensor is None or not _is_per_tensor_quantization(in_q_tensor.quantization):
                        valid_inputs = False
                        break

                    pre_name = str(in_q_op.inputs[0])
                    pre_idx = producers.get(pre_name, None)
                    if pre_idx is None:
                        valid_inputs = False
                        break
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != pre_name
                    ):
                        valid_inputs = False
                        break
                    pre_perm = _read_transpose_perm(model_ir, pre_op)
                    if pre_perm is None or list(pre_perm) != list(perm_pre_expected):
                        valid_inputs = False
                        break

                    dq_input_users = [int(v) for v in consumers.get(add_input_name, [])]
                    if set(dq_input_users) != {int(add_idx)}:
                        valid_inputs = False
                        break

                    pre_output_name = str(pre_op.outputs[0])
                    pre_output_users = [int(v) for v in consumers.get(pre_output_name, [])]
                    can_remove_pre = set(pre_output_users) == {int(in_q_idx)}

                    if proj_branch is not None:
                        valid_inputs = False
                        break
                    proj_branch = {
                        "pre_idx": int(pre_idx),
                        "q_idx": int(in_q_idx),
                        "dq_name": str(add_input_name),
                        "pre_input_name": str(pre_op.inputs[0]),
                        "can_remove_pre": bool(can_remove_pre),
                    }
                    continue

                if add_input_prod_type == "TRANSPOSE":
                    if (
                        len(add_input_prod_op.inputs) < 2
                        or len(add_input_prod_op.outputs) != 1
                        or str(add_input_prod_op.outputs[0]) != add_input_name
                    ):
                        valid_inputs = False
                        break
                    pre_perm = _read_transpose_perm(model_ir, add_input_prod_op)
                    if pre_perm is None or list(pre_perm) != list(perm_pre_expected):
                        valid_inputs = False
                        break

                    skip_users = [int(v) for v in consumers.get(add_input_name, [])]
                    if set(skip_users) != {int(add_idx)}:
                        valid_inputs = False
                        break

                    if skip_branch is not None:
                        valid_inputs = False
                        break
                    skip_branch = {
                        "pre_idx": int(add_input_prod_idx),
                        "pre_input_name": str(add_input_prod_op.inputs[0]),
                        "skip_name": str(add_input_name),
                    }
                    continue

                valid_inputs = False
                break

            if (not valid_inputs) or proj_branch is None or skip_branch is None:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=model_ir.operators[int(proj_branch["q_idx"])],
                new_inputs=[str(proj_branch["pre_input_name"])],
            )

            new_add_inputs = [
                str(skip_branch["pre_input_name"]) if str(v) == str(skip_branch["skip_name"]) else str(v)
                for v in list(add_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=new_add_inputs,
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(proj_branch["dq_name"]), None),
                list(perm_post),
            )
            proj_q_op = model_ir.operators[int(proj_branch["q_idx"])]
            if len(proj_q_op.outputs) == 1:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(proj_q_op.outputs[0]), None),
                    list(perm_post),
                )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_out_name, None),
                list(perm_post),
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(out_q_name, None),
                list(perm_post),
            )
            canonical_post_idx = int(post_idx)
            canonical_post_output_name = str(post_output_name)

            _set_operator_outputs(
                model_ir=model_ir,
                op=dq_op,
                new_outputs=[canonical_post_output_name],
            )
            new_dq_out_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            old_dq_out_tensor = model_ir.tensors.get(dq_out_name, None)
            if new_dq_out_tensor is not None and old_dq_out_tensor is not None:
                new_dq_out_tensor.dtype = str(old_dq_out_tensor.dtype)

            remove_indices = {int(skip_branch["pre_idx"])}
            if bool(proj_branch["can_remove_pre"]):
                remove_indices.add(int(proj_branch["pre_idx"]))
            if len(legacy_user_indices) > 0:
                keep_post_op = model_ir.operators[int(canonical_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(list(perm_pre_expected), dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[dq_out_name],
                )
                for post_output in post_output_names:
                    if str(post_output) == canonical_post_output_name:
                        continue
                    _replace_tensor_inputs(model_ir, str(post_output), canonical_post_output_name)
                for idx_rm in post_indices:
                    if int(idx_rm) == int(canonical_post_idx):
                        continue
                    remove_indices.add(int(idx_rm))
            else:
                for post_output in post_output_names:
                    if str(post_output) == canonical_post_output_name:
                        continue
                    _replace_tensor_inputs(model_ir, str(post_output), canonical_post_output_name)
                for idx_rm in post_indices:
                    remove_indices.add(int(idx_rm))

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten_mixed_add_qdq_residual_bridges += 1
            removed_bridge_pairs += int(1)
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_quantize_dequantize_bridges": int(removed_bridge_pairs),
        "rewritten_add_qdq_residual_transpose_bridges": int(rewritten_add_qdq_residual_bridges),
        "rewritten_mixed_add_qdq_residual_transpose_bridges": int(
            rewritten_mixed_add_qdq_residual_bridges
        ),
    }


def _optimize_duplicate_transpose_fanout(model_ir: ModelIR) -> Dict[str, int]:
    """
    Deduplicate fan-out TRANSPOSE nodes with identical input and permutation.

    Target pattern:
      X --TRANSPOSE(P)--> Y0
      X --TRANSPOSE(P)--> Y1
      ...

    Rewritten:
      X --TRANSPOSE(P)--> Y0
      (all uses of Y1, ... are rewired to Y0; duplicate TRANSPOSE nodes removed)
    """
    removed_duplicates = 0

    while True:
        changed = False
        canonical_by_key: Dict[Tuple[str, Tuple[int, ...]], int] = {}

        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE":
                continue
            if len(op.inputs) < 2 or len(op.outputs) != 1:
                continue

            input_name = str(op.inputs[0])
            output_name = str(op.outputs[0])
            perm = _read_transpose_perm(model_ir, op)
            if perm is None:
                continue

            key = (input_name, tuple(int(v) for v in perm))
            canonical_idx = canonical_by_key.get(key, None)
            if canonical_idx is None:
                canonical_by_key[key] = int(op_idx)
                continue

            if output_name in model_ir.outputs:
                # Preserve user-visible graph output names.
                continue

            canonical_op = model_ir.operators[int(canonical_idx)]
            if len(canonical_op.outputs) != 1:
                continue
            canonical_output = str(canonical_op.outputs[0])
            if canonical_output == output_name:
                continue

            canonical_tensor = model_ir.tensors.get(canonical_output, None)
            duplicate_tensor = model_ir.tensors.get(output_name, None)
            if canonical_tensor is not None and duplicate_tensor is not None:
                if canonical_tensor.shape == [1] and duplicate_tensor.shape != [1]:
                    canonical_tensor.shape = [int(v) for v in list(duplicate_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(duplicate_tensor.shape_signature)]
                        if duplicate_tensor.shape_signature is not None
                        else [int(v) for v in list(duplicate_tensor.shape)]
                    )
                if canonical_tensor.quantization is None and duplicate_tensor.quantization is not None:
                    canonical_tensor.quantization = _clone_quantization(duplicate_tensor.quantization)
                if str(canonical_tensor.dtype) == "FLOAT32" and str(duplicate_tensor.dtype) != "FLOAT32":
                    canonical_tensor.dtype = str(duplicate_tensor.dtype)

            _replace_tensor_inputs(model_ir, output_name, canonical_output)
            del model_ir.operators[int(op_idx)]
            removed_duplicates += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_duplicate_transpose_fanout": int(removed_duplicates),
    }


def _optimize_duplicate_reshape_fanout(model_ir: ModelIR) -> Dict[str, int]:
    """
    Deduplicate fan-out RESHAPE nodes with identical input and target shape.

    Target pattern:
      X --RESHAPE(S)--> Y0
      X --RESHAPE(S)--> Y1
      ...

    Rewritten:
      X --RESHAPE(S)--> Y0
      (all uses of Y1, ... are rewired to Y0; duplicate RESHAPE nodes removed)
    """
    removed_duplicates = 0

    def _read_reshape_target_shape(op: OperatorIR) -> Optional[List[int]]:
        if str(op.op_type) != "RESHAPE":
            return None
        if isinstance(op.options, dict):
            new_shape = op.options.get("newShape", None)
            if isinstance(new_shape, list) and len(new_shape) > 0:
                try:
                    return [int(v) for v in list(new_shape)]
                except Exception:
                    pass
        if len(op.inputs) >= 2:
            shape_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
            shape_values = _read_const_ints_from_tensor(shape_tensor)
            if shape_values is not None and len(shape_values) > 0:
                return [int(v) for v in list(shape_values)]
        return None

    while True:
        changed = False
        canonical_by_key: Dict[Tuple[str, Tuple[int, ...]], int] = {}

        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "RESHAPE":
                continue
            if len(op.inputs) < 2 or len(op.outputs) != 1:
                continue

            input_name = str(op.inputs[0])
            output_name = str(op.outputs[0])
            target_shape = _read_reshape_target_shape(op)
            if target_shape is None:
                continue

            key = (input_name, tuple(int(v) for v in list(target_shape)))
            canonical_idx = canonical_by_key.get(key, None)
            if canonical_idx is None:
                canonical_by_key[key] = int(op_idx)
                continue

            if output_name in model_ir.outputs:
                # Preserve user-visible graph output names.
                continue

            canonical_op = model_ir.operators[int(canonical_idx)]
            if len(canonical_op.outputs) != 1:
                continue
            canonical_output = str(canonical_op.outputs[0])
            if canonical_output == output_name:
                continue

            canonical_tensor = model_ir.tensors.get(canonical_output, None)
            duplicate_tensor = model_ir.tensors.get(output_name, None)
            if canonical_tensor is not None and duplicate_tensor is not None:
                if canonical_tensor.shape == [1] and duplicate_tensor.shape != [1]:
                    canonical_tensor.shape = [int(v) for v in list(duplicate_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(duplicate_tensor.shape_signature)]
                        if duplicate_tensor.shape_signature is not None
                        else [int(v) for v in list(duplicate_tensor.shape)]
                    )
                if canonical_tensor.quantization is None and duplicate_tensor.quantization is not None:
                    canonical_tensor.quantization = _clone_quantization(duplicate_tensor.quantization)
                if str(canonical_tensor.dtype) == "FLOAT32" and str(duplicate_tensor.dtype) != "FLOAT32":
                    canonical_tensor.dtype = str(duplicate_tensor.dtype)

            _replace_tensor_inputs(model_ir, output_name, canonical_output)
            del model_ir.operators[int(op_idx)]
            removed_duplicates += 1
            changed = True
            break

        if not changed:
            break

    if removed_duplicates > 0:
        _prune_unused_tensors(model_ir)
    return {
        "removed_duplicate_reshape_fanout": int(removed_duplicates),
    }


def _optimize_transpose_dequant_prelu_quantize_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around DEQUANTIZE->PRELU->QUANTIZE chains.

    Target pattern:
      Xq --Transpose(P)--> Aq --DEQUANTIZE--> A --PRELU(alpha)--> B --QUANTIZE--> Bq --Transpose(inv(P))--> Yq

    Rewritten:
      Xq --DEQUANTIZE--> A --PRELU(alpha')--> B --QUANTIZE--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - Pre/Post transpose permutations are exact inverses
    - Quantized tensors use per-tensor quantization only
    - PRELU alpha tensor is constant if rank remap is required
    """
    removed_prelu_bridges = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            bridge_q_in = str(pre_op.outputs[0])
            dq_users = consumers.get(bridge_q_in, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != bridge_q_in:
                continue

            bridge_f_in = str(dq_op.outputs[0])
            prelu_users = consumers.get(bridge_f_in, [])
            if len(prelu_users) != 1:
                continue
            prelu_idx = int(prelu_users[0])
            prelu_op = model_ir.operators[prelu_idx]
            if str(prelu_op.op_type) != "PRELU" or len(prelu_op.inputs) != 2 or len(prelu_op.outputs) != 1:
                continue
            if str(prelu_op.inputs[0]) != bridge_f_in:
                continue

            bridge_f_out = str(prelu_op.outputs[0])
            q_users = consumers.get(bridge_f_out, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != bridge_f_out:
                continue

            bridge_q_out = str(q_op.outputs[0])
            post_users = consumers.get(bridge_q_out, [])
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != bridge_q_out:
                continue

            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                continue

            # Keep user-visible output names stable and avoid breaking observable intermediates.
            if (
                bridge_q_in in model_ir.outputs
                or bridge_f_in in model_ir.outputs
                or bridge_f_out in model_ir.outputs
                or bridge_q_out in model_ir.outputs
            ):
                continue

            q_src_name = str(pre_op.inputs[0])
            q_dst_name = str(post_op.outputs[0])
            if q_src_name in model_ir.outputs:
                continue

            q_src_tensor = model_ir.tensors.get(q_src_name, None)
            q_mid_in_tensor = model_ir.tensors.get(bridge_q_in, None)
            q_mid_out_tensor = model_ir.tensors.get(bridge_q_out, None)
            q_dst_tensor = model_ir.tensors.get(q_dst_name, None)
            if not _all_per_tensor_quantized([q_src_tensor, q_mid_in_tensor, q_mid_out_tensor, q_dst_tensor]):
                continue

            # PRELU alpha layout follows data layout. When alpha rank matches the transposed rank,
            # remap alpha using post permutation so PRELU can run directly on non-transposed data.
            alpha_name = str(prelu_op.inputs[1])
            alpha_tensor = model_ir.tensors.get(alpha_name, None)
            if alpha_tensor is not None and isinstance(alpha_tensor.data, np.ndarray):
                alpha_data = np.asarray(alpha_tensor.data)
                if alpha_data.ndim == len(perm_post):
                    alpha_data = np.transpose(alpha_data, axes=perm_post)
                    alpha_tensor.data = alpha_data
                    alpha_tensor.shape = [int(v) for v in alpha_data.shape]
                    alpha_tensor.shape_signature = [int(v) for v in alpha_data.shape]

            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_src_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_op,
                new_outputs=[q_dst_name],
            )

            # Update bridge tensor metadata to the non-transposed layout.
            q_src_shape = list(q_src_tensor.shape) if q_src_tensor is not None else None
            q_src_signature = (
                list(q_src_tensor.shape_signature)
                if q_src_tensor is not None and q_src_tensor.shape_signature is not None
                else q_src_shape
            )
            if q_src_shape is not None:
                dq_out_tensor = model_ir.tensors.get(bridge_f_in, None)
                prelu_out_tensor = model_ir.tensors.get(bridge_f_out, None)
                if dq_out_tensor is not None:
                    dq_out_tensor.shape = [int(v) for v in q_src_shape]
                    dq_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if prelu_out_tensor is not None:
                    prelu_out_tensor.shape = [int(v) for v in q_src_shape]
                    prelu_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )

            if q_dst_tensor is not None and q_mid_out_tensor is not None:
                q_dst_tensor.dtype = str(q_mid_out_tensor.dtype)
                q_dst_tensor.quantization = _clone_quantization(q_mid_out_tensor.quantization)

            for remove_idx in sorted([pre_idx, post_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_prelu_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_dequant_prelu_quantize_bridges": int(removed_prelu_bridges),
    }


def _optimize_transpose_dequant_relu_quantize_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around DEQUANTIZE->(RELU|RELU6)->QUANTIZE chains.

    Target pattern:
      Xq --Transpose(P)--> Aq --DEQUANTIZE--> A --(RELU|RELU6)--> B --QUANTIZE--> Bq --Transpose(inv(P))--> Yq

    Rewritten:
      Xq --DEQUANTIZE--> A --(RELU|RELU6)--> B --QUANTIZE--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - Pre/Post transpose permutations are exact inverses
    - Quantized tensors use per-tensor quantization only
    """
    removed_relu_bridges = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            bridge_q_in = str(pre_op.outputs[0])
            dq_users = consumers.get(bridge_q_in, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != bridge_q_in:
                continue

            bridge_f_in = str(dq_op.outputs[0])
            relu_users = consumers.get(bridge_f_in, [])
            if len(relu_users) != 1:
                continue
            relu_idx = int(relu_users[0])
            relu_op = model_ir.operators[relu_idx]
            if str(relu_op.op_type) not in {"RELU", "RELU6"} or len(relu_op.inputs) != 1 or len(relu_op.outputs) != 1:
                continue
            if str(relu_op.inputs[0]) != bridge_f_in:
                continue

            bridge_f_out = str(relu_op.outputs[0])
            q_users = consumers.get(bridge_f_out, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != bridge_f_out:
                continue

            bridge_q_out = str(q_op.outputs[0])
            post_users = consumers.get(bridge_q_out, [])
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != bridge_q_out:
                continue

            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                continue

            # Keep user-visible output names stable and avoid breaking observable intermediates.
            if (
                bridge_q_in in model_ir.outputs
                or bridge_f_in in model_ir.outputs
                or bridge_f_out in model_ir.outputs
                or bridge_q_out in model_ir.outputs
            ):
                continue

            q_src_name = str(pre_op.inputs[0])
            q_dst_name = str(post_op.outputs[0])
            if q_src_name in model_ir.outputs:
                continue

            q_src_tensor = model_ir.tensors.get(q_src_name, None)
            q_mid_in_tensor = model_ir.tensors.get(bridge_q_in, None)
            q_mid_out_tensor = model_ir.tensors.get(bridge_q_out, None)
            q_dst_tensor = model_ir.tensors.get(q_dst_name, None)
            if not _all_per_tensor_quantized([q_src_tensor, q_mid_in_tensor, q_mid_out_tensor, q_dst_tensor]):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_src_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_op,
                new_outputs=[q_dst_name],
            )

            # Update bridge tensor metadata to the non-transposed layout.
            q_src_shape = list(q_src_tensor.shape) if q_src_tensor is not None else None
            q_src_signature = (
                list(q_src_tensor.shape_signature)
                if q_src_tensor is not None and q_src_tensor.shape_signature is not None
                else q_src_shape
            )
            if q_src_shape is not None:
                dq_out_tensor = model_ir.tensors.get(bridge_f_in, None)
                relu_out_tensor = model_ir.tensors.get(bridge_f_out, None)
                if dq_out_tensor is not None:
                    dq_out_tensor.shape = [int(v) for v in q_src_shape]
                    dq_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if relu_out_tensor is not None:
                    relu_out_tensor.shape = [int(v) for v in q_src_shape]
                    relu_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )

            if q_dst_tensor is not None and q_mid_out_tensor is not None:
                q_dst_tensor.dtype = str(q_mid_out_tensor.dtype)
                q_dst_tensor.quantization = _clone_quantization(q_mid_out_tensor.quantization)

            for remove_idx in sorted([pre_idx, post_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_relu_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_dequant_relu_quantize_bridges": int(removed_relu_bridges),
    }


def _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around DEQUANTIZE->HardSigmoid(expanded)->QUANTIZE chains.

    Target pattern:
      Xq --Transpose(P)--> Aq --DEQUANTIZE--> A
         --MUL(c1)--> B --ADD(c2)--> C --(RELU_0_TO_1 | MAXIMUM(c3)->MINIMUM(c4))--> D
         --QUANTIZE--> Dq --Transpose(inv(P))--> Yq

    Rewritten:
      Xq --DEQUANTIZE--> A --MUL(c1')--> B --ADD(c2')--> C --(RELU_0_TO_1 | MAXIMUM(c3')->MINIMUM(c4'))--> D
         --QUANTIZE--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - Pre/Post transpose permutations are exact inverses
    - Quantized tensors use per-tensor quantization only
    - Side inputs are constants; rank-matched constants are layout-remapped
    """
    removed_hardsigmoid_bridges = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            bridge_q_in = str(pre_op.outputs[0])
            dq_users = consumers.get(bridge_q_in, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != bridge_q_in:
                continue

            bridge_f_in = str(dq_op.outputs[0])
            mul_users = consumers.get(bridge_f_in, [])
            if len(mul_users) != 1:
                continue
            mul_idx = int(mul_users[0])
            mul_op = model_ir.operators[mul_idx]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_in0 = str(mul_op.inputs[0])
            mul_in1 = str(mul_op.inputs[1])
            if mul_in0 == bridge_f_in:
                mul_const_input_index = 1
            elif mul_in1 == bridge_f_in:
                mul_const_input_index = 0
            else:
                continue

            bridge_mul_out = str(mul_op.outputs[0])
            add_users = consumers.get(bridge_mul_out, [])
            if len(add_users) != 1:
                continue
            add_idx = int(add_users[0])
            add_op = model_ir.operators[add_idx]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_in0 = str(add_op.inputs[0])
            add_in1 = str(add_op.inputs[1])
            if add_in0 == bridge_mul_out:
                add_const_input_index = 1
            elif add_in1 == bridge_mul_out:
                add_const_input_index = 0
            else:
                continue

            bridge_add_out = str(add_op.outputs[0])
            hs_users = consumers.get(bridge_add_out, [])
            if len(hs_users) != 1:
                continue

            hs_terminal_idx: Optional[int] = None
            hs_terminal_op: Optional[OperatorIR] = None
            bridge_hs_out: Optional[str] = None
            const_inputs_to_remap: List[Tuple[int, OperatorIR, int]] = [
                (int(mul_idx), mul_op, int(mul_const_input_index)),
                (int(add_idx), add_op, int(add_const_input_index)),
            ]
            hs_intermediate_names: List[str] = [bridge_mul_out, bridge_add_out]

            hs_candidate_idx = int(hs_users[0])
            hs_candidate_op = model_ir.operators[hs_candidate_idx]
            if (
                str(hs_candidate_op.op_type) == "RELU_0_TO_1"
                and len(hs_candidate_op.inputs) == 1
                and len(hs_candidate_op.outputs) == 1
                and str(hs_candidate_op.inputs[0]) == bridge_add_out
            ):
                hs_terminal_idx = int(hs_candidate_idx)
                hs_terminal_op = hs_candidate_op
                bridge_hs_out = str(hs_candidate_op.outputs[0])
            elif (
                str(hs_candidate_op.op_type) == "MAXIMUM"
                and len(hs_candidate_op.inputs) == 2
                and len(hs_candidate_op.outputs) == 1
            ):
                max_inputs = [str(v) for v in list(hs_candidate_op.inputs)]
                if max_inputs[0] == bridge_add_out:
                    max_const_input_index = 1
                elif max_inputs[1] == bridge_add_out:
                    max_const_input_index = 0
                else:
                    continue
                max_out_name = str(hs_candidate_op.outputs[0])
                min_users = consumers.get(max_out_name, [])
                if len(min_users) != 1:
                    continue
                min_idx = int(min_users[0])
                min_op = model_ir.operators[min_idx]
                if str(min_op.op_type) != "MINIMUM" or len(min_op.inputs) != 2 or len(min_op.outputs) != 1:
                    continue
                min_inputs = [str(v) for v in list(min_op.inputs)]
                if min_inputs[0] == max_out_name:
                    min_const_input_index = 1
                elif min_inputs[1] == max_out_name:
                    min_const_input_index = 0
                else:
                    continue
                hs_terminal_idx = int(min_idx)
                hs_terminal_op = min_op
                bridge_hs_out = str(min_op.outputs[0])
                hs_intermediate_names.append(max_out_name)
                const_inputs_to_remap.extend(
                    [
                        (int(hs_candidate_idx), hs_candidate_op, int(max_const_input_index)),
                        (int(min_idx), min_op, int(min_const_input_index)),
                    ]
                )
            else:
                continue

            if hs_terminal_idx is None or hs_terminal_op is None or bridge_hs_out is None:
                continue

            q_users = consumers.get(str(bridge_hs_out), [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != str(bridge_hs_out):
                continue

            bridge_q_out = str(q_op.outputs[0])
            post_users = consumers.get(bridge_q_out, [])
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != bridge_q_out:
                continue

            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                continue

            # Keep user-visible output names stable and avoid breaking observable intermediates.
            if (
                bridge_q_in in model_ir.outputs
                or bridge_f_in in model_ir.outputs
                or bridge_mul_out in model_ir.outputs
                or bridge_add_out in model_ir.outputs
                or str(bridge_hs_out) in model_ir.outputs
                or bridge_q_out in model_ir.outputs
            ):
                continue

            q_src_name = str(pre_op.inputs[0])
            q_dst_name = str(post_op.outputs[0])
            if q_src_name in model_ir.outputs:
                continue

            q_src_tensor = model_ir.tensors.get(q_src_name, None)
            q_mid_in_tensor = model_ir.tensors.get(bridge_q_in, None)
            q_mid_out_tensor = model_ir.tensors.get(bridge_q_out, None)
            q_dst_tensor = model_ir.tensors.get(q_dst_name, None)
            if not _all_per_tensor_quantized([q_src_tensor, q_mid_in_tensor, q_mid_out_tensor, q_dst_tensor]):
                continue

            def _remap_constant_input_for_op(
                *,
                op: OperatorIR,
                op_idx: int,
                input_index: int,
            ) -> bool:
                const_name = str(op.inputs[input_index])
                const_tensor = model_ir.tensors.get(const_name, None)
                if const_tensor is None or const_tensor.data is None:
                    return False
                const_data = np.asarray(const_tensor.data)
                if const_data.ndim != len(perm_post):
                    return True
                transposed_data = np.transpose(const_data, axes=perm_post)
                const_users = consumers.get(const_name, [])
                if len(const_users) == 1 and int(const_users[0]) == int(op_idx):
                    const_tensor.data = np.asarray(transposed_data)
                    const_tensor.shape = [int(v) for v in transposed_data.shape]
                    const_tensor.shape_signature = [int(v) for v in transposed_data.shape]
                    return True
                new_name = _unique_tensor_name(f"{const_name}_nhwc")
                model_ir.tensors[new_name] = TensorIR(
                    name=new_name,
                    dtype=str(const_tensor.dtype),
                    shape=[int(v) for v in transposed_data.shape],
                    shape_signature=[int(v) for v in transposed_data.shape],
                    data=np.asarray(transposed_data),
                    is_variable=False,
                    quantization=_clone_quantization(const_tensor.quantization),
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=op,
                    input_index=int(input_index),
                    new_input_name=str(new_name),
                )
                return True

            remap_ok = True
            for remap_idx, remap_op, remap_input_index in const_inputs_to_remap:
                if not _remap_constant_input_for_op(
                    op=remap_op,
                    op_idx=int(remap_idx),
                    input_index=int(remap_input_index),
                ):
                    remap_ok = False
                    break
            if not remap_ok:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_src_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_op,
                new_outputs=[q_dst_name],
            )

            # Update bridge tensor metadata to the non-transposed layout.
            q_src_shape = list(q_src_tensor.shape) if q_src_tensor is not None else None
            q_src_signature = (
                list(q_src_tensor.shape_signature)
                if q_src_tensor is not None and q_src_tensor.shape_signature is not None
                else q_src_shape
            )
            if q_src_shape is not None:
                dq_out_tensor = model_ir.tensors.get(bridge_f_in, None)
                mul_out_tensor = model_ir.tensors.get(bridge_mul_out, None)
                add_out_tensor = model_ir.tensors.get(bridge_add_out, None)
                hs_out_tensor = model_ir.tensors.get(str(bridge_hs_out), None)
                if dq_out_tensor is not None:
                    dq_out_tensor.shape = [int(v) for v in q_src_shape]
                    dq_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if mul_out_tensor is not None:
                    mul_out_tensor.shape = [int(v) for v in q_src_shape]
                    mul_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if add_out_tensor is not None:
                    add_out_tensor.shape = [int(v) for v in q_src_shape]
                    add_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                for hs_intermediate_name in hs_intermediate_names:
                    hs_mid_tensor = model_ir.tensors.get(str(hs_intermediate_name), None)
                    if hs_mid_tensor is None:
                        continue
                    hs_mid_tensor.shape = [int(v) for v in q_src_shape]
                    hs_mid_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if hs_out_tensor is not None:
                    hs_out_tensor.shape = [int(v) for v in q_src_shape]
                    hs_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )

            if q_dst_tensor is not None and q_mid_out_tensor is not None:
                q_dst_tensor.dtype = str(q_mid_out_tensor.dtype)
                q_dst_tensor.quantization = _clone_quantization(q_mid_out_tensor.quantization)

            for remove_idx in sorted([pre_idx, post_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_hardsigmoid_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_dequant_hardsigmoid_quantize_bridges": int(removed_hardsigmoid_bridges),
    }


def _optimize_transpose_dequant_mul_add_prelu_quantize_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around DEQUANTIZE->MUL->ADD->PRELU->QUANTIZE chains.

    Target pattern:
      Xq --Transpose(P)--> Aq --DEQUANTIZE--> A --MUL(c1)--> B --ADD(c2)--> C
         --PRELU(alpha)--> D --QUANTIZE--> Dq --Transpose(inv(P))--> Yq

    Rewritten:
      Xq --DEQUANTIZE--> A --MUL(c1')--> B --ADD(c2')--> C --PRELU(alpha')--> D --QUANTIZE--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - Pre/Post transpose permutations are exact inverses
    - Quantized tensors use per-tensor quantization only
    - Side inputs c1/c2/alpha are constants; rank-matched constants are layout-remapped
    """
    removed_mul_add_prelu_bridges = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            bridge_q_in = str(pre_op.outputs[0])
            dq_users = consumers.get(bridge_q_in, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != bridge_q_in:
                continue

            bridge_f_in = str(dq_op.outputs[0])
            mul_users = consumers.get(bridge_f_in, [])
            if len(mul_users) != 1:
                continue
            mul_idx = int(mul_users[0])
            mul_op = model_ir.operators[mul_idx]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_in0 = str(mul_op.inputs[0])
            mul_in1 = str(mul_op.inputs[1])
            if mul_in0 == bridge_f_in:
                mul_const_input_index = 1
            elif mul_in1 == bridge_f_in:
                mul_const_input_index = 0
            else:
                continue

            bridge_mul_out = str(mul_op.outputs[0])
            add_users = consumers.get(bridge_mul_out, [])
            if len(add_users) != 1:
                continue
            add_idx = int(add_users[0])
            add_op = model_ir.operators[add_idx]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_in0 = str(add_op.inputs[0])
            add_in1 = str(add_op.inputs[1])
            if add_in0 == bridge_mul_out:
                add_const_input_index = 1
            elif add_in1 == bridge_mul_out:
                add_const_input_index = 0
            else:
                continue

            bridge_add_out = str(add_op.outputs[0])
            prelu_users = consumers.get(bridge_add_out, [])
            if len(prelu_users) != 1:
                continue
            prelu_idx = int(prelu_users[0])
            prelu_op = model_ir.operators[prelu_idx]
            if str(prelu_op.op_type) != "PRELU" or len(prelu_op.inputs) != 2 or len(prelu_op.outputs) != 1:
                continue
            if str(prelu_op.inputs[0]) != bridge_add_out:
                continue

            bridge_prelu_out = str(prelu_op.outputs[0])
            q_users = consumers.get(bridge_prelu_out, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != bridge_prelu_out:
                continue

            bridge_q_out = str(q_op.outputs[0])
            post_users = consumers.get(bridge_q_out, [])
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != bridge_q_out:
                continue

            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                continue

            # Keep user-visible output names stable and avoid breaking observable intermediates.
            if (
                bridge_q_in in model_ir.outputs
                or bridge_f_in in model_ir.outputs
                or bridge_mul_out in model_ir.outputs
                or bridge_add_out in model_ir.outputs
                or bridge_prelu_out in model_ir.outputs
                or bridge_q_out in model_ir.outputs
            ):
                continue

            q_src_name = str(pre_op.inputs[0])
            q_dst_name = str(post_op.outputs[0])
            if q_src_name in model_ir.outputs:
                continue

            q_src_tensor = model_ir.tensors.get(q_src_name, None)
            q_mid_in_tensor = model_ir.tensors.get(bridge_q_in, None)
            q_mid_out_tensor = model_ir.tensors.get(bridge_q_out, None)
            q_dst_tensor = model_ir.tensors.get(q_dst_name, None)
            if not _all_per_tensor_quantized([q_src_tensor, q_mid_in_tensor, q_mid_out_tensor, q_dst_tensor]):
                continue

            mul_const_name = str(mul_op.inputs[mul_const_input_index])
            add_const_name = str(add_op.inputs[add_const_input_index])
            alpha_name = str(prelu_op.inputs[1])
            mul_const_tensor = model_ir.tensors.get(mul_const_name, None)
            add_const_tensor = model_ir.tensors.get(add_const_name, None)
            alpha_tensor = model_ir.tensors.get(alpha_name, None)
            if (
                mul_const_tensor is None
                or add_const_tensor is None
                or alpha_tensor is None
                or not isinstance(mul_const_tensor.data, np.ndarray)
                or not isinstance(add_const_tensor.data, np.ndarray)
                or not isinstance(alpha_tensor.data, np.ndarray)
            ):
                continue

            def _remap_constant_input_for_op(
                *,
                op: OperatorIR,
                op_idx: int,
                input_index: int,
            ) -> bool:
                const_name = str(op.inputs[input_index])
                const_tensor = model_ir.tensors.get(const_name, None)
                if const_tensor is None or not isinstance(const_tensor.data, np.ndarray):
                    return False
                const_data = np.asarray(const_tensor.data)
                if const_data.ndim != len(perm_post):
                    return True
                transposed_data = np.transpose(const_data, axes=perm_post)
                const_users = consumers.get(const_name, [])
                if len(const_users) == 1 and int(const_users[0]) == int(op_idx):
                    const_tensor.data = np.asarray(transposed_data)
                    const_tensor.shape = [int(v) for v in transposed_data.shape]
                    const_tensor.shape_signature = [int(v) for v in transposed_data.shape]
                    return True
                new_name = _unique_tensor_name(f"{const_name}_nhwc")
                model_ir.tensors[new_name] = TensorIR(
                    name=new_name,
                    dtype=str(const_tensor.dtype),
                    shape=[int(v) for v in transposed_data.shape],
                    shape_signature=[int(v) for v in transposed_data.shape],
                    data=np.asarray(transposed_data),
                    is_variable=False,
                    quantization=_clone_quantization(const_tensor.quantization),
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=op,
                    input_index=int(input_index),
                    new_input_name=new_name,
                )
                return True

            if not _remap_constant_input_for_op(
                op=mul_op,
                op_idx=mul_idx,
                input_index=mul_const_input_index,
            ):
                continue
            if not _remap_constant_input_for_op(
                op=add_op,
                op_idx=add_idx,
                input_index=add_const_input_index,
            ):
                continue
            if not _remap_constant_input_for_op(
                op=prelu_op,
                op_idx=prelu_idx,
                input_index=1,
            ):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_src_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_op,
                new_outputs=[q_dst_name],
            )

            # Update bridge tensor metadata to the non-transposed layout.
            q_src_shape = list(q_src_tensor.shape) if q_src_tensor is not None else None
            q_src_signature = (
                list(q_src_tensor.shape_signature)
                if q_src_tensor is not None and q_src_tensor.shape_signature is not None
                else q_src_shape
            )
            if q_src_shape is not None:
                dq_out_tensor = model_ir.tensors.get(bridge_f_in, None)
                mul_out_tensor = model_ir.tensors.get(bridge_mul_out, None)
                add_out_tensor = model_ir.tensors.get(bridge_add_out, None)
                prelu_out_tensor = model_ir.tensors.get(bridge_prelu_out, None)
                if dq_out_tensor is not None:
                    dq_out_tensor.shape = [int(v) for v in q_src_shape]
                    dq_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if mul_out_tensor is not None:
                    mul_out_tensor.shape = [int(v) for v in q_src_shape]
                    mul_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if add_out_tensor is not None:
                    add_out_tensor.shape = [int(v) for v in q_src_shape]
                    add_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if prelu_out_tensor is not None:
                    prelu_out_tensor.shape = [int(v) for v in q_src_shape]
                    prelu_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )

            if q_dst_tensor is not None and q_mid_out_tensor is not None:
                q_dst_tensor.dtype = str(q_mid_out_tensor.dtype)
                q_dst_tensor.quantization = _clone_quantization(q_mid_out_tensor.quantization)

            for remove_idx in sorted([pre_idx, post_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_mul_add_prelu_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_dequant_mul_add_prelu_quantize_bridges": int(removed_mul_add_prelu_bridges),
    }


def _optimize_transpose_dequant_prelu_transpose_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around DEQUANTIZE->PRELU chains.

    Target pattern:
      Xq --Transpose(P)--> Aq --DEQUANTIZE--> A --PRELU(alpha)--> B --Transpose(inv(P))--> Y

    Rewritten:
      Xq --DEQUANTIZE--> A' --PRELU(alpha')--> Y

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - Pre/Post transpose permutations are exact inverses
    - Quantized tensors on DEQUANTIZE input path are per-tensor
    """
    removed_prelu_transpose_bridges = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            bridge_q_in = str(pre_op.outputs[0])
            dq_users = consumers.get(bridge_q_in, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != bridge_q_in:
                continue

            bridge_f_in = str(dq_op.outputs[0])
            prelu_users = consumers.get(bridge_f_in, [])
            if len(prelu_users) != 1:
                continue
            prelu_idx = int(prelu_users[0])
            prelu_op = model_ir.operators[prelu_idx]
            if str(prelu_op.op_type) != "PRELU" or len(prelu_op.inputs) != 2 or len(prelu_op.outputs) != 1:
                continue
            if str(prelu_op.inputs[0]) != bridge_f_in:
                continue

            bridge_f_out = str(prelu_op.outputs[0])
            post_users = consumers.get(bridge_f_out, [])
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != bridge_f_out:
                continue

            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                continue

            if (
                bridge_q_in in model_ir.outputs
                or bridge_f_in in model_ir.outputs
                or bridge_f_out in model_ir.outputs
                or post_op.outputs[0] in model_ir.outputs
            ):
                continue

            q_src_name = str(pre_op.inputs[0])
            y_name = str(post_op.outputs[0])
            q_src_tensor = model_ir.tensors.get(q_src_name, None)
            q_mid_tensor = model_ir.tensors.get(bridge_q_in, None)
            if not _all_per_tensor_quantized([q_src_tensor, q_mid_tensor]):
                continue

            alpha_name = str(prelu_op.inputs[1])
            alpha_tensor = model_ir.tensors.get(alpha_name, None)
            if alpha_tensor is not None:
                alpha_data = alpha_tensor.data
                if isinstance(alpha_data, np.ndarray) and alpha_data.ndim == len(perm_post):
                    transposed_alpha = np.transpose(alpha_data, axes=perm_post)
                    alpha_users = consumers.get(alpha_name, [])
                    if len(alpha_users) == 1 and int(alpha_users[0]) == int(prelu_idx):
                        alpha_tensor.data = transposed_alpha
                        alpha_tensor.shape = [int(v) for v in transposed_alpha.shape]
                        alpha_tensor.shape_signature = [int(v) for v in transposed_alpha.shape]
                    else:
                        new_alpha_name = _unique_tensor_name(f"{alpha_name}_nhwc")
                        model_ir.tensors[new_alpha_name] = TensorIR(
                            name=new_alpha_name,
                            dtype=str(alpha_tensor.dtype),
                            shape=[int(v) for v in transposed_alpha.shape],
                            shape_signature=[int(v) for v in transposed_alpha.shape],
                            data=np.asarray(transposed_alpha),
                            is_variable=False,
                            quantization=_clone_quantization(alpha_tensor.quantization),
                        )
                        _replace_operator_input_at(
                            model_ir=model_ir,
                            op=prelu_op,
                            input_index=1,
                            new_input_name=new_alpha_name,
                        )

            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_src_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu_op,
                new_outputs=[y_name],
            )

            q_src_shape = list(q_src_tensor.shape) if q_src_tensor is not None else None
            q_src_signature = (
                list(q_src_tensor.shape_signature)
                if q_src_tensor is not None and q_src_tensor.shape_signature is not None
                else q_src_shape
            )
            if q_src_shape is not None:
                dq_out_tensor = model_ir.tensors.get(bridge_f_in, None)
                y_tensor = model_ir.tensors.get(y_name, None)
                if dq_out_tensor is not None:
                    dq_out_tensor.shape = [int(v) for v in q_src_shape]
                    dq_out_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )
                if y_tensor is not None:
                    y_tensor.shape = [int(v) for v in q_src_shape]
                    y_tensor.shape_signature = (
                        [int(v) for v in q_src_signature]
                        if q_src_signature is not None
                        else [int(v) for v in q_src_shape]
                    )

            for remove_idx in sorted([pre_idx, post_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_prelu_transpose_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_dequant_prelu_transpose_bridges": int(removed_prelu_transpose_bridges),
    }


def _optimize_transpose_dequant_logistic_mul_quantize_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around quantized logistic-gate MUL bridges.

    Target pattern:
      Xq --T(P)--> Xq' --DQ--> A --LOGISTIC--> S --Q--> Sq --DQ--> S'
                    \\--DQ--> B ------------------------------/
      MUL(B, S') -> M --Q--> Mq --T(inv(P))--> Yq

    Rewritten:
      Xq --DQ--> A' --LOGISTIC--> S' --Q--> Sq' --DQ--> S''
        \\--DQ--> B' ------------------------------/
      MUL(B', S'') -> M' --Q--> Yq

    Safety:
    - Pre/Post transpose permutations are exact NHWC<->NCHW inverses.
    - Matched branch is linear (single consumer at each bridge tensor).
    - All post users of `Mq` are inverse transposes (fanout aliases supported).
    - Involved quantized tensors are per-tensor quantized.
    """
    removed_logistic_mul_bridges = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            q_mul_out_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if q_mul_out_name in model_outputs or post_output_name in model_outputs:
                continue

            q_mul_idx = producers.get(q_mul_out_name, None)
            if q_mul_idx is None:
                continue
            q_mul_op = model_ir.operators[int(q_mul_idx)]
            if str(q_mul_op.op_type) != "QUANTIZE" or len(q_mul_op.inputs) != 1 or len(q_mul_op.outputs) != 1:
                continue
            if str(q_mul_op.outputs[0]) != q_mul_out_name:
                continue
            mul_out_name = str(q_mul_op.inputs[0])
            if mul_out_name in model_outputs:
                continue

            mul_idx = producers.get(mul_out_name, None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(mul_op.outputs[0]) != mul_out_name:
                continue

            dq_gate_idx: Optional[int] = None
            dq_gate_op: Optional[OperatorIR] = None
            dq_data_idx: Optional[int] = None
            dq_data_op: Optional[OperatorIR] = None
            q_gate_idx: Optional[int] = None
            q_gate_op: Optional[OperatorIR] = None
            logistic_idx: Optional[int] = None
            logistic_op: Optional[OperatorIR] = None
            dq_log_idx: Optional[int] = None
            dq_log_op: Optional[OperatorIR] = None

            for mul_input_name in [str(v) for v in list(mul_op.inputs)]:
                candidate_dq_idx = producers.get(str(mul_input_name), None)
                if candidate_dq_idx is None:
                    continue
                candidate_dq_op = model_ir.operators[int(candidate_dq_idx)]
                if (
                    str(candidate_dq_op.op_type) != "DEQUANTIZE"
                    or len(candidate_dq_op.inputs) != 1
                    or len(candidate_dq_op.outputs) != 1
                    or str(candidate_dq_op.outputs[0]) != str(mul_input_name)
                ):
                    continue

                dq_input_name = str(candidate_dq_op.inputs[0])
                candidate_q_gate_idx = producers.get(dq_input_name, None)
                if candidate_q_gate_idx is not None:
                    candidate_q_gate_op = model_ir.operators[int(candidate_q_gate_idx)]
                    if (
                        str(candidate_q_gate_op.op_type) == "QUANTIZE"
                        and len(candidate_q_gate_op.inputs) == 1
                        and len(candidate_q_gate_op.outputs) == 1
                        and str(candidate_q_gate_op.outputs[0]) == dq_input_name
                    ):
                        candidate_logistic_out = str(candidate_q_gate_op.inputs[0])
                        candidate_logistic_idx = producers.get(candidate_logistic_out, None)
                        if candidate_logistic_idx is not None:
                            candidate_logistic_op = model_ir.operators[int(candidate_logistic_idx)]
                            if (
                                str(candidate_logistic_op.op_type) == "LOGISTIC"
                                and len(candidate_logistic_op.inputs) == 1
                                and len(candidate_logistic_op.outputs) == 1
                                and str(candidate_logistic_op.outputs[0]) == candidate_logistic_out
                            ):
                                candidate_dq_log_out = str(candidate_logistic_op.inputs[0])
                                candidate_dq_log_idx = producers.get(candidate_dq_log_out, None)
                                if candidate_dq_log_idx is not None:
                                    candidate_dq_log_op = model_ir.operators[int(candidate_dq_log_idx)]
                                    if (
                                        str(candidate_dq_log_op.op_type) == "DEQUANTIZE"
                                        and len(candidate_dq_log_op.inputs) == 1
                                        and len(candidate_dq_log_op.outputs) == 1
                                        and str(candidate_dq_log_op.outputs[0]) == candidate_dq_log_out
                                    ):
                                        if dq_gate_idx is not None:
                                            dq_gate_idx = None
                                            break
                                        dq_gate_idx = int(candidate_dq_idx)
                                        dq_gate_op = candidate_dq_op
                                        q_gate_idx = int(candidate_q_gate_idx)
                                        q_gate_op = candidate_q_gate_op
                                        logistic_idx = int(candidate_logistic_idx)
                                        logistic_op = candidate_logistic_op
                                        dq_log_idx = int(candidate_dq_log_idx)
                                        dq_log_op = candidate_dq_log_op
                                        continue

                if dq_data_idx is not None:
                    dq_data_idx = None
                    break
                dq_data_idx = int(candidate_dq_idx)
                dq_data_op = candidate_dq_op

            if (
                dq_gate_idx is None
                or dq_gate_op is None
                or dq_data_idx is None
                or dq_data_op is None
                or q_gate_idx is None
                or q_gate_op is None
                or logistic_idx is None
                or logistic_op is None
                or dq_log_idx is None
                or dq_log_op is None
            ):
                continue

            pre_out_name = str(dq_data_op.inputs[0])
            if str(dq_log_op.inputs[0]) != pre_out_name:
                continue
            if pre_out_name in model_outputs:
                continue

            pre_idx = producers.get(pre_out_name, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != pre_out_name
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_in_name = str(pre_op.inputs[0])
            if pre_in_name in model_outputs:
                continue

            # Keep this pass strict so pre/post transposes can be safely deleted.
            if set(int(v) for v in consumers.get(pre_out_name, [])) != {int(dq_data_idx), int(dq_log_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(dq_data_op.outputs[0]), [])) != {int(mul_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(dq_log_op.outputs[0]), [])) != {int(logistic_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(logistic_op.outputs[0]), [])) != {int(q_gate_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(q_gate_op.outputs[0]), [])) != {int(dq_gate_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(dq_gate_op.outputs[0]), [])) != {int(mul_idx)}:
                continue
            if set(int(v) for v in consumers.get(mul_out_name, [])) != {int(q_mul_idx)}:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            q_mul_users = [int(v) for v in consumers.get(q_mul_out_name, [])]
            if len(q_mul_users) == 0:
                continue
            valid_posts = True
            for user_idx in q_mul_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != q_mul_out_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                ):
                    valid_posts = False
                    break
                cand_out = str(user_op.outputs[0])
                if cand_out in model_outputs:
                    valid_posts = False
                    break
                post_indices.append(int(user_idx))
                post_output_names.append(cand_out)
            if not valid_posts or len(post_indices) == 0:
                continue

            q_src_tensor = model_ir.tensors.get(pre_in_name, None)
            q_mid_tensor = model_ir.tensors.get(pre_out_name, None)
            q_gate_out_tensor = model_ir.tensors.get(str(q_gate_op.outputs[0]), None)
            q_mul_out_tensor = model_ir.tensors.get(q_mul_out_name, None)
            if not _all_per_tensor_quantized([q_src_tensor, q_mid_tensor, q_gate_out_tensor, q_mul_out_tensor]):
                continue

            # Bypass pre-transpose by rewiring both DEQUANTIZE inputs to the
            # original quantized NHWC tensor.
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_data_op,
                new_inputs=[pre_in_name],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_log_op,
                new_inputs=[pre_in_name],
            )

            # Elementwise subgraph is layout-agnostic; update metadata to NHWC.
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(dq_data_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(dq_log_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(logistic_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(q_gate_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(dq_gate_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mul_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(q_mul_out_name, None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_mul_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if canonical_tensor is not None and q_mul_out_tensor is not None:
                canonical_tensor.dtype = str(q_mul_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(q_mul_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(q_mul_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(q_mul_out_tensor.shape_signature)]
                    if q_mul_out_tensor.shape_signature is not None
                    else [int(v) for v in list(q_mul_out_tensor.shape)]
                )

            remove_indices = sorted(
                set([int(pre_idx)] + [int(v) for v in list(post_indices)]),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            removed_logistic_mul_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_dequant_logistic_mul_quantize_bridges": int(removed_logistic_mul_bridges),
    }


def _optimize_transpose_swish_qdq_nhwc_islands(
    model_ir: ModelIR,
    *,
    min_spatial_stage: int = 160,
    require_concat_closure: bool = False,
) -> Dict[str, int]:
    """
    Rewrite transpose-wrapped quantized Swish islands from NCHW back to NHWC.

    Strictly matches branches of:
      T(0,3,1,2) -> DQ -> LOGISTIC -> Q -> DQ
                   DQ ----------------------> MUL -> (Q)?

    and then propagates NHWC metadata through layout-agnostic quant/elementwise
    operators, upgrades rank-4 CONCAT axis 1->3 when inputs are NHWC-aligned,
    and removes inverse post-transposes fed by rewritten tensors.
    """
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    model_outputs = set(str(v) for v in model_ir.outputs)

    rewritten_branches = 0
    removed_pre = 0
    propagated_tensors = 0
    rewritten_concat_axis = 0
    removed_post = 0

    rewritten_tensors: set[str] = set()

    def _copy_shape_signature(dst_name: str, src_name: str) -> bool:
        dst = model_ir.tensors.get(str(dst_name), None)
        src = model_ir.tensors.get(str(src_name), None)
        if dst is None or src is None:
            return False
        changed_local = False
        src_shape = [int(v) for v in list(src.shape)]
        src_sig = (
            [int(v) for v in list(src.shape_signature)]
            if src.shape_signature is not None
            else [int(v) for v in list(src.shape)]
        )
        if [int(v) for v in list(dst.shape)] != src_shape:
            dst.shape = [int(v) for v in list(src_shape)]
            changed_local = True
        dst_sig = (
            [int(v) for v in list(dst.shape_signature)]
            if dst.shape_signature is not None
            else [int(v) for v in list(dst.shape)]
        )
        if dst.shape_signature is None or dst_sig != src_sig:
            dst.shape_signature = [int(v) for v in list(src_sig)]
            changed_local = True
        return changed_local

    def _is_swish_quantized_output(tensor_name: str) -> bool:
        producers = _build_tensor_producer_map(model_ir)
        producer_idx = producers.get(str(tensor_name), None)
        if producer_idx is None:
            return False
        producer_op = model_ir.operators[int(producer_idx)]
        if (
            str(producer_op.op_type) != "QUANTIZE"
            or len(producer_op.inputs) != 1
            or len(producer_op.outputs) != 1
            or str(producer_op.outputs[0]) != str(tensor_name)
        ):
            return False
        mul_name = str(producer_op.inputs[0])
        mul_idx = producers.get(mul_name, None)
        if mul_idx is None:
            return False
        mul_op = model_ir.operators[int(mul_idx)]
        return (
            str(mul_op.op_type) == "MUL"
            and len(mul_op.inputs) == 2
            and len(mul_op.outputs) == 1
            and str(mul_op.outputs[0]) == mul_name
        )

    def _concat_has_quantize_transpose_tail(concat_out_name: str, consumers: Dict[str, List[int]]) -> bool:
        concat_users = [int(v) for v in consumers.get(str(concat_out_name), [])]
        if len(concat_users) != 1:
            return False
        q_idx = int(concat_users[0])
        q_op = model_ir.operators[int(q_idx)]
        if (
            str(q_op.op_type) != "QUANTIZE"
            or len(q_op.inputs) != 1
            or len(q_op.outputs) != 1
            or str(q_op.inputs[0]) != str(concat_out_name)
        ):
            return False
        q_out_name = str(q_op.outputs[0])
        q_users = [int(v) for v in consumers.get(q_out_name, [])]
        if len(q_users) == 0:
            return False
        for user_idx in q_users:
            user_op = model_ir.operators[int(user_idx)]
            if (
                str(user_op.op_type) != "TRANSPOSE"
                or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
            ):
                return False
        return True

    def _has_concat_closure_from_tensor(
        tensor_name: str,
        consumers: Dict[str, List[int]],
    ) -> bool:
        # Traverse downstream through a conservative allowlist and accept only
        # paths that eventually land on CONCAT(axis=1)->QUANTIZE->TRANSPOSE.
        allowed_passthrough_ops = {
            "DEQUANTIZE",
            "QUANTIZE",
            "LOGISTIC",
            "MUL",
            "ADD",
            "SUB",
            "DIV",
            "MAXIMUM",
            "MINIMUM",
            "TRANSPOSE",
            "CONV_2D",
            "DEPTHWISE_CONV_2D",
            "MAX_POOL_2D",
            "AVERAGE_POOL_2D",
            "RESIZE_NEAREST_NEIGHBOR",
            "RESIZE_BILINEAR",
            "SHAPE",
            "SLICE",
            "STRIDED_SLICE",
            "RELU",
            "RELU6",
            "LEAKY_RELU",
            "HARD_SWISH",
            "CAST",
        }
        queue: List[Tuple[str, int]] = [(str(tensor_name), 0)]
        visited: set[str] = set()
        max_hops = 192

        while len(queue) > 0:
            cur_name, depth = queue.pop(0)
            if str(cur_name) in visited:
                continue
            visited.add(str(cur_name))
            if int(depth) > int(max_hops):
                continue

            for user_idx in [int(v) for v in consumers.get(str(cur_name), [])]:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)

                if user_type == "CONCATENATION" and len(user_op.outputs) == 1:
                    axis = int(user_op.options.get("axis", 1))
                    if axis < 0:
                        axis += 4
                    if axis == 1 and str(cur_name) in set(str(v) for v in list(user_op.inputs)):
                        if _concat_has_quantize_transpose_tail(str(user_op.outputs[0]), consumers):
                            return True

                if len(user_op.outputs) != 1:
                    continue
                if user_type not in allowed_passthrough_ops:
                    continue
                next_name = str(user_op.outputs[0])
                if str(next_name) in visited:
                    continue
                queue.append((str(next_name), int(depth) + 1))

        return False

    # 1) Rewire Swish branches fed by NHWC->NCHW transposes.
    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue
            pre_tensor = model_ir.tensors.get(pre_out_name, None)
            if pre_tensor is None:
                continue
            pre_shape = [int(v) for v in list(pre_tensor.shape)]
            if len(pre_shape) != 4:
                continue
            if int(min_spatial_stage) > 0 and min(int(pre_shape[2]), int(pre_shape[3])) < int(min_spatial_stage):
                continue

            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) < 2:
                continue
            dq_candidates = [
                int(idx)
                for idx in pre_users
                if str(model_ir.operators[int(idx)].op_type) == "DEQUANTIZE"
                and len(model_ir.operators[int(idx)].inputs) == 1
                and len(model_ir.operators[int(idx)].outputs) == 1
                and str(model_ir.operators[int(idx)].inputs[0]) == pre_out_name
            ]
            if len(dq_candidates) < 2:
                continue

            used_dq: set[int] = set()
            local_rewritten = 0

            for dq_log_idx in dq_candidates:
                if int(dq_log_idx) in used_dq:
                    continue
                dq_log_op = model_ir.operators[int(dq_log_idx)]
                dq_log_out = str(dq_log_op.outputs[0])
                if dq_log_out in model_outputs:
                    continue

                log_users = [int(v) for v in consumers.get(dq_log_out, [])]
                if len(log_users) != 1:
                    continue
                logistic_idx = int(log_users[0])
                logistic_op = model_ir.operators[int(logistic_idx)]
                if (
                    str(logistic_op.op_type) != "LOGISTIC"
                    or len(logistic_op.inputs) != 1
                    or len(logistic_op.outputs) != 1
                    or str(logistic_op.inputs[0]) != dq_log_out
                ):
                    continue
                logistic_out = str(logistic_op.outputs[0])
                if logistic_out in model_outputs:
                    continue

                q_users = [int(v) for v in consumers.get(logistic_out, [])]
                if len(q_users) != 1:
                    continue
                q_gate_idx = int(q_users[0])
                q_gate_op = model_ir.operators[int(q_gate_idx)]
                if (
                    str(q_gate_op.op_type) != "QUANTIZE"
                    or len(q_gate_op.inputs) != 1
                    or len(q_gate_op.outputs) != 1
                    or str(q_gate_op.inputs[0]) != logistic_out
                ):
                    continue
                q_gate_out = str(q_gate_op.outputs[0])
                if q_gate_out in model_outputs:
                    continue

                dq_gate_users = [int(v) for v in consumers.get(q_gate_out, [])]
                if len(dq_gate_users) != 1:
                    continue
                dq_gate_idx = int(dq_gate_users[0])
                dq_gate_op = model_ir.operators[int(dq_gate_idx)]
                if (
                    str(dq_gate_op.op_type) != "DEQUANTIZE"
                    or len(dq_gate_op.inputs) != 1
                    or len(dq_gate_op.outputs) != 1
                    or str(dq_gate_op.inputs[0]) != q_gate_out
                ):
                    continue
                dq_gate_out = str(dq_gate_op.outputs[0])
                if dq_gate_out in model_outputs:
                    continue

                matched_data_idx: Optional[int] = None
                matched_mul_idx: Optional[int] = None
                for dq_data_idx in dq_candidates:
                    if int(dq_data_idx) == int(dq_log_idx) or int(dq_data_idx) in used_dq:
                        continue
                    dq_data_op = model_ir.operators[int(dq_data_idx)]
                    dq_data_out = str(dq_data_op.outputs[0])
                    mul_users = [int(v) for v in consumers.get(dq_data_out, [])]
                    if len(mul_users) != 1:
                        continue
                    cand_mul_idx = int(mul_users[0])
                    cand_mul_op = model_ir.operators[int(cand_mul_idx)]
                    if (
                        str(cand_mul_op.op_type) != "MUL"
                        or len(cand_mul_op.inputs) != 2
                        or len(cand_mul_op.outputs) != 1
                    ):
                        continue
                    cand_inputs = [str(v) for v in list(cand_mul_op.inputs)]
                    if dq_data_out not in cand_inputs or dq_gate_out not in cand_inputs:
                        continue
                    gate_users = [int(v) for v in consumers.get(dq_gate_out, [])]
                    if len(gate_users) != 1 or int(gate_users[0]) != int(cand_mul_idx):
                        continue
                    matched_data_idx = int(dq_data_idx)
                    matched_mul_idx = int(cand_mul_idx)
                    break

                if matched_data_idx is None or matched_mul_idx is None:
                    continue

                dq_data_op = model_ir.operators[int(matched_data_idx)]
                dq_data_out = str(dq_data_op.outputs[0])
                mul_op = model_ir.operators[int(matched_mul_idx)]
                mul_out = str(mul_op.outputs[0])
                if mul_out in model_outputs:
                    continue

                # Guard: only rewrite branches with safe downstream users.
                mul_users = [int(v) for v in consumers.get(mul_out, [])]
                if len(mul_users) == 0:
                    continue

                has_q_mul = False
                if len(mul_users) == 1:
                    q_mul_idx = int(mul_users[0])
                    q_mul_op = model_ir.operators[int(q_mul_idx)]
                    if (
                        str(q_mul_op.op_type) == "QUANTIZE"
                        and len(q_mul_op.inputs) == 1
                        and len(q_mul_op.outputs) == 1
                        and str(q_mul_op.inputs[0]) == mul_out
                    ):
                        has_q_mul = True
                        q_mul_out = str(q_mul_op.outputs[0])
                        q_mul_users = [int(v) for v in consumers.get(q_mul_out, [])]
                        if len(q_mul_users) == 0:
                            continue
                        safe_q_users = True
                        for user_idx in q_mul_users:
                            user_op = model_ir.operators[int(user_idx)]
                            user_type = str(user_op.op_type)
                            if user_type == "TRANSPOSE":
                                if (
                                    _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                                    or str(user_op.outputs[0]) in model_outputs
                                ):
                                    safe_q_users = False
                                    break
                                continue
                            if (
                                bool(require_concat_closure)
                                and user_type == "DEQUANTIZE"
                                and len(user_op.inputs) == 1
                                and len(user_op.outputs) == 1
                                and str(user_op.inputs[0]) == q_mul_out
                                and _has_concat_closure_from_tensor(str(user_op.outputs[0]), consumers)
                            ):
                                continue
                            if user_type not in {"ADD", "MUL", "SUB", "DIV", "MAXIMUM", "MINIMUM"}:
                                safe_q_users = False
                                break
                            other_inputs = [str(v) for v in list(user_op.inputs) if str(v) != q_mul_out]
                            if len(other_inputs) != 1:
                                safe_q_users = False
                                break
                            other_name = str(other_inputs[0])
                            if other_name in rewritten_tensors:
                                continue
                            other_tensor = model_ir.tensors.get(other_name, None)
                            if other_tensor is not None and other_tensor.data is not None:
                                continue
                            if _is_swish_quantized_output(other_name):
                                continue
                            q_tensor = model_ir.tensors.get(q_mul_out, None)
                            if q_tensor is not None and other_tensor is not None:
                                q_shape = [int(v) for v in list(q_tensor.shape)]
                                if bool(require_concat_closure) and [int(v) for v in list(other_tensor.shape)] == [int(v) for v in list(q_shape)]:
                                    continue
                                permuted_q_shape = (
                                    [int(q_shape[idx]) for idx in perm_nchw_to_nhwc]
                                    if len(q_shape) == len(perm_nchw_to_nhwc)
                                    else None
                                )
                                if (
                                    permuted_q_shape is not None
                                    and [int(v) for v in list(other_tensor.shape)] == [int(v) for v in list(permuted_q_shape)]
                                ):
                                    continue
                            safe_q_users = False
                            break
                        if not safe_q_users:
                            continue
                if not has_q_mul:
                    safe_float_users = True
                    for user_idx in mul_users:
                        user_op = model_ir.operators[int(user_idx)]
                        user_type = str(user_op.op_type)
                        if user_type in {"ADD", "MUL", "SUB", "DIV", "MAXIMUM", "MINIMUM"}:
                            continue
                        if user_type == "CONCATENATION":
                            axis = int(user_op.options.get("axis", 1))
                            if axis < 0:
                                axis += 4
                            if axis != 1:
                                safe_float_users = False
                                break
                            continue
                        if user_type == "TRANSPOSE":
                            if (
                                _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                                or str(user_op.outputs[0]) in model_outputs
                            ):
                                safe_float_users = False
                                break
                            continue
                        safe_float_users = False
                        break
                    if not safe_float_users:
                        continue

                if bool(require_concat_closure):
                    is_concat_member = False
                    if has_q_mul:
                        is_concat_member = _has_concat_closure_from_tensor(
                            str(q_mul_out),
                            consumers,
                        )
                    else:
                        is_concat_member = _has_concat_closure_from_tensor(
                            str(mul_out),
                            consumers,
                        )
                    if not is_concat_member:
                        continue

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=dq_log_op,
                    new_inputs=[pre_in_name],
                )
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=dq_data_op,
                    new_inputs=[pre_in_name],
                )

                for tensor_name in [
                    dq_log_out,
                    logistic_out,
                    q_gate_out,
                    dq_gate_out,
                    dq_data_out,
                    mul_out,
                ]:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(tensor_name), None),
                        perm_nchw_to_nhwc,
                    )
                    rewritten_tensors.add(str(tensor_name))

                if len(mul_users) == 1:
                    q_mul_idx = int(mul_users[0])
                    q_mul_op = model_ir.operators[int(q_mul_idx)]
                    if (
                        str(q_mul_op.op_type) == "QUANTIZE"
                        and len(q_mul_op.inputs) == 1
                        and len(q_mul_op.outputs) == 1
                        and str(q_mul_op.inputs[0]) == mul_out
                    ):
                        q_mul_out = str(q_mul_op.outputs[0])
                        _permute_tensor_metadata_if_rank_matches(
                            model_ir.tensors.get(str(q_mul_out), None),
                            perm_nchw_to_nhwc,
                        )
                        rewritten_tensors.add(str(q_mul_out))

                used_dq.add(int(dq_log_idx))
                used_dq.add(int(matched_data_idx))
                rewritten_branches += 1
                local_rewritten += 1
                changed = True

            if local_rewritten <= 0:
                continue

            consumers_after = _build_tensor_consumer_map(model_ir)
            if len([int(v) for v in consumers_after.get(pre_out_name, [])]) == 0:
                del model_ir.operators[int(pre_idx)]
                removed_pre += 1
                changed = True
                break

        if not changed:
            break

    # 2) Propagate NHWC metadata through layout-agnostic quant/elementwise ops.
    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for op in model_ir.operators:
            op_type = str(op.op_type)
            if len(op.outputs) != 1:
                continue
            out_name = str(op.outputs[0])
            if out_name in model_outputs:
                continue

            if op_type in {"DEQUANTIZE", "QUANTIZE", "LOGISTIC"} and len(op.inputs) == 1:
                in_name = str(op.inputs[0])
                if in_name not in rewritten_tensors:
                    continue
                if _copy_shape_signature(out_name, in_name):
                    changed = True
                if out_name not in rewritten_tensors:
                    rewritten_tensors.add(out_name)
                    changed = True
                continue

            if op_type in {"ADD", "MUL", "SUB", "DIV", "MAXIMUM", "MINIMUM"} and len(op.inputs) == 2:
                in0 = str(op.inputs[0])
                in1 = str(op.inputs[1])
                if in0 not in rewritten_tensors or in1 not in rewritten_tensors:
                    continue
                t0 = model_ir.tensors.get(in0, None)
                t1 = model_ir.tensors.get(in1, None)
                tout = model_ir.tensors.get(out_name, None)
                if t0 is None or t1 is None or tout is None:
                    continue

                out_shape = _broadcast_static_shapes(t0.shape, t1.shape)
                sig0 = (
                    [int(v) for v in list(t0.shape_signature)]
                    if t0.shape_signature is not None
                    else [int(v) for v in list(t0.shape)]
                )
                sig1 = (
                    [int(v) for v in list(t1.shape_signature)]
                    if t1.shape_signature is not None
                    else [int(v) for v in list(t1.shape)]
                )
                out_sig = _broadcast_shape_signatures(sig0, sig1)

                if out_shape is None:
                    out_shape = [int(v) for v in list(t0.shape)]
                if out_sig is None:
                    out_sig = [int(v) for v in list(sig0)]

                if [int(v) for v in list(tout.shape)] != [int(v) for v in list(out_shape)]:
                    tout.shape = [int(v) for v in list(out_shape)]
                    changed = True
                current_sig = (
                    [int(v) for v in list(tout.shape_signature)]
                    if tout.shape_signature is not None
                    else [int(v) for v in list(tout.shape)]
                )
                if tout.shape_signature is None or current_sig != [int(v) for v in list(out_sig)]:
                    tout.shape_signature = [int(v) for v in list(out_sig)]
                    changed = True
                if out_name not in rewritten_tensors:
                    rewritten_tensors.add(out_name)
                    changed = True
                continue

            if op_type in {"MAX_POOL_2D", "AVERAGE_POOL_2D", "RESIZE_NEAREST_NEIGHBOR", "RESIZE_BILINEAR"} and len(op.inputs) >= 1:
                in_name = str(op.inputs[0])
                if in_name not in rewritten_tensors:
                    continue
                in_tensor = model_ir.tensors.get(in_name, None)
                out_tensor = model_ir.tensors.get(out_name, None)
                if in_tensor is None or out_tensor is None:
                    continue
                in_shape = [int(v) for v in list(in_tensor.shape)]
                out_shape = [int(v) for v in list(out_tensor.shape)]
                if len(in_shape) != 4 or len(out_shape) != 4:
                    continue
                in_c = int(in_shape[3])
                out_c = int(out_shape[3])
                if in_c >= 0 and out_c >= 0 and in_c != out_c:
                    continue
                if out_name not in rewritten_tensors:
                    rewritten_tensors.add(out_name)
                    changed = True
                continue

            if op_type == "CONCATENATION" and len(op.inputs) >= 2:
                axis = int(op.options.get("axis", 1))
                if axis < 0:
                    axis += 4
                if axis != 1:
                    continue
                input_names = [str(v) for v in list(op.inputs)]
                if any(name not in rewritten_tensors for name in input_names):
                    continue
                input_tensors = [model_ir.tensors.get(name, None) for name in input_names]
                if any(t is None for t in input_tensors):
                    continue
                input_shapes = [[int(v) for v in list(t.shape)] for t in input_tensors if t is not None]
                if any(len(shape) != 4 for shape in input_shapes):
                    continue
                if not all(
                    int(shape[0]) == int(input_shapes[0][0])
                    and int(shape[1]) == int(input_shapes[0][1])
                    and int(shape[2]) == int(input_shapes[0][2])
                    for shape in input_shapes[1:]
                ):
                    continue

                # Require downstream quantize->transpose bridge to keep this strict.
                concat_users = [int(v) for v in consumers.get(out_name, [])]
                if len(concat_users) != 1:
                    continue
                q_idx = int(concat_users[0])
                q_op = model_ir.operators[int(q_idx)]
                if (
                    str(q_op.op_type) != "QUANTIZE"
                    or len(q_op.inputs) != 1
                    or len(q_op.outputs) != 1
                    or str(q_op.inputs[0]) != out_name
                ):
                    continue
                q_out = str(q_op.outputs[0])
                q_users = [int(v) for v in consumers.get(q_out, [])]
                if len(q_users) == 0:
                    continue
                if any(
                    str(model_ir.operators[int(user_idx)].op_type) != "TRANSPOSE"
                    or _read_transpose_perm(model_ir, model_ir.operators[int(user_idx)]) != perm_nchw_to_nhwc
                    for user_idx in q_users
                ):
                    continue

                op.options["axis"] = 3
                rewritten_concat_axis += 1

                concat_tensor = model_ir.tensors.get(out_name, None)
                q_tensor = model_ir.tensors.get(q_out, None)
                if concat_tensor is not None:
                    out_shape = [int(v) for v in list(input_shapes[0])]
                    out_shape[3] = int(sum(int(shape[3]) for shape in input_shapes))
                    concat_tensor.shape = [int(v) for v in list(out_shape)]
                    concat_tensor.shape_signature = [int(v) for v in list(out_shape)]
                    rewritten_tensors.add(out_name)
                if q_tensor is not None and concat_tensor is not None:
                    q_tensor.shape = [int(v) for v in list(concat_tensor.shape)]
                    q_tensor.shape_signature = (
                        [int(v) for v in list(concat_tensor.shape_signature)]
                        if concat_tensor.shape_signature is not None
                        else [int(v) for v in list(concat_tensor.shape)]
                    )
                    rewritten_tensors.add(q_out)
                changed = True

        if not changed:
            break

    # 3) Remove inverse post-transposes fed by rewritten NHWC tensors.
    while True:
        changed = False
        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE" or len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, op) != perm_nchw_to_nhwc:
                continue
            in_name = str(op.inputs[0])
            out_name = str(op.outputs[0])
            if out_name in model_outputs:
                continue
            if in_name not in rewritten_tensors:
                continue

            _replace_tensor_inputs(model_ir, out_name, in_name)
            del model_ir.operators[int(op_idx)]
            removed_post += 1
            changed = True
            break
        if not changed:
            break

    # 3b) Late concat-axis rewrite after post-transpose cleanup.
    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        for op in model_ir.operators:
            if str(op.op_type) != "CONCATENATION" or len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            out_name = str(op.outputs[0])
            if out_name in model_outputs:
                continue

            axis = int(op.options.get("axis", 1))
            if axis < 0:
                axis += 4
            if axis != 1:
                continue

            input_names = [str(v) for v in list(op.inputs)]
            normalized_inputs: List[str] = []
            normalized_shapes: List[List[int]] = []
            bypass_transpose_indices: List[int] = []
            rewire_dequant_inputs: List[Tuple[int, str, str]] = []

            for input_name in input_names:
                input_tensor = model_ir.tensors.get(input_name, None)
                if input_tensor is None:
                    normalized_inputs = []
                    break
                input_shape = [int(v) for v in list(input_tensor.shape)]
                if len(input_shape) != 4:
                    normalized_inputs = []
                    break

                producer_idx = producers.get(input_name, None)
                if producer_idx is not None:
                    producer_op = model_ir.operators[int(producer_idx)]
                    if (
                        str(producer_op.op_type) == "DEQUANTIZE"
                        and len(producer_op.inputs) == 1
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == input_name
                    ):
                        dq_in_name = str(producer_op.inputs[0])
                        dq_in_producer_idx = producers.get(dq_in_name, None)
                        if dq_in_producer_idx is not None:
                            dq_in_producer_op = model_ir.operators[int(dq_in_producer_idx)]
                            if (
                                str(dq_in_producer_op.op_type) == "TRANSPOSE"
                                and len(dq_in_producer_op.inputs) >= 2
                                and len(dq_in_producer_op.outputs) == 1
                                and str(dq_in_producer_op.outputs[0]) == dq_in_name
                                and _read_transpose_perm(model_ir, dq_in_producer_op) == perm_nhwc_to_nchw
                            ):
                                pre_name = str(dq_in_producer_op.inputs[0])
                                if pre_name in model_outputs:
                                    normalized_inputs = []
                                    break
                                pre_tensor = model_ir.tensors.get(pre_name, None)
                                if pre_tensor is None:
                                    normalized_inputs = []
                                    break
                                pre_shape = [int(v) for v in list(pre_tensor.shape)]
                                if len(pre_shape) != 4:
                                    normalized_inputs = []
                                    break
                                normalized_inputs.append(input_name)
                                normalized_shapes.append(pre_shape)
                                rewire_dequant_inputs.append((int(producer_idx), pre_name, input_name))
                                bypass_transpose_indices.append(int(dq_in_producer_idx))
                                continue

                    if (
                        str(producer_op.op_type) == "TRANSPOSE"
                        and len(producer_op.inputs) >= 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == input_name
                        and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                    ):
                        pre_name = str(producer_op.inputs[0])
                        if pre_name in model_outputs:
                            normalized_inputs = []
                            break
                        pre_tensor = model_ir.tensors.get(pre_name, None)
                        if pre_tensor is None:
                            normalized_inputs = []
                            break
                        pre_shape = [int(v) for v in list(pre_tensor.shape)]
                        if len(pre_shape) != 4:
                            normalized_inputs = []
                            break
                        normalized_inputs.append(pre_name)
                        normalized_shapes.append(pre_shape)
                        bypass_transpose_indices.append(int(producer_idx))
                        continue

                normalized_inputs.append(input_name)
                normalized_shapes.append(input_shape)

            if len(normalized_inputs) != len(input_names):
                continue
            if len(bypass_transpose_indices) == 0:
                continue
            if not all(
                int(shape[0]) == int(normalized_shapes[0][0])
                and int(shape[1]) == int(normalized_shapes[0][1])
                and int(shape[2]) == int(normalized_shapes[0][2])
                for shape in normalized_shapes[1:]
            ):
                continue

            concat_users = [int(v) for v in consumers.get(out_name, [])]
            if len(concat_users) != 1:
                continue
            q_idx = int(concat_users[0])
            q_op = model_ir.operators[int(q_idx)]
            if (
                str(q_op.op_type) != "QUANTIZE"
                or len(q_op.inputs) != 1
                or len(q_op.outputs) != 1
                or str(q_op.inputs[0]) != out_name
            ):
                continue
            q_out = str(q_op.outputs[0])
            q_users = [int(v) for v in consumers.get(q_out, [])]
            if len(q_users) == 0:
                continue
            if any(
                str(model_ir.operators[int(user_idx)].op_type) != "TRANSPOSE"
                or _read_transpose_perm(model_ir, model_ir.operators[int(user_idx)]) != perm_nchw_to_nhwc
                for user_idx in q_users
            ):
                continue

            op.options["axis"] = 3
            rewritten_concat_axis += 1

            concat_tensor = model_ir.tensors.get(out_name, None)
            q_tensor = model_ir.tensors.get(q_out, None)
            if concat_tensor is not None:
                out_shape = [int(v) for v in list(normalized_shapes[0])]
                out_shape[3] = int(sum(int(shape[3]) for shape in normalized_shapes))
                concat_tensor.shape = [int(v) for v in list(out_shape)]
                concat_tensor.shape_signature = [int(v) for v in list(out_shape)]
                rewritten_tensors.add(out_name)
            if q_tensor is not None and concat_tensor is not None:
                q_tensor.shape = [int(v) for v in list(concat_tensor.shape)]
                q_tensor.shape_signature = (
                    [int(v) for v in list(concat_tensor.shape_signature)]
                    if concat_tensor.shape_signature is not None
                    else [int(v) for v in list(concat_tensor.shape)]
                )
                rewritten_tensors.add(q_out)

            _set_operator_inputs(
                model_ir=model_ir,
                op=op,
                new_inputs=[str(v) for v in list(normalized_inputs)],
            )
            for dequant_idx, dequant_new_input, dequant_out_name in rewire_dequant_inputs:
                if dequant_idx < 0 or dequant_idx >= len(model_ir.operators):
                    continue
                dequant_op = model_ir.operators[int(dequant_idx)]
                if (
                    str(dequant_op.op_type) != "DEQUANTIZE"
                    or len(dequant_op.inputs) != 1
                    or len(dequant_op.outputs) != 1
                    or str(dequant_op.outputs[0]) != str(dequant_out_name)
                ):
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=dequant_op,
                    new_inputs=[str(dequant_new_input)],
                )
                _copy_shape_signature(str(dequant_out_name), str(dequant_new_input))
            for rewritten_input in normalized_inputs:
                rewritten_tensors.add(str(rewritten_input))

            consumers_after = _build_tensor_consumer_map(model_ir)
            for transpose_idx in sorted(set(int(v) for v in bypass_transpose_indices), reverse=True):
                if transpose_idx < 0 or transpose_idx >= len(model_ir.operators):
                    continue
                transpose_op = model_ir.operators[int(transpose_idx)]
                if str(transpose_op.op_type) != "TRANSPOSE" or len(transpose_op.outputs) != 1:
                    continue
                transpose_out = str(transpose_op.outputs[0])
                if len([int(v) for v in consumers_after.get(transpose_out, [])]) != 0:
                    continue
                del model_ir.operators[int(transpose_idx)]
                removed_post += 1

            changed = True

        if not changed:
            break

    # 3c) Remove post-transposes newly unlocked by late concat-axis rewrite.
    while True:
        changed = False
        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE" or len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, op) != perm_nchw_to_nhwc:
                continue
            in_name = str(op.inputs[0])
            out_name = str(op.outputs[0])
            if out_name in model_outputs:
                continue
            if in_name not in rewritten_tensors:
                continue

            _replace_tensor_inputs(model_ir, out_name, in_name)
            del model_ir.operators[int(op_idx)]
            removed_post += 1
            changed = True
            break
        if not changed:
            break

    # 4) Safety valve: remove wrong-way NCHW->NHWC transposes that would break
    # CONV_2D channel alignment after NHWC propagation.
    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE" or len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, op) != perm_nchw_to_nhwc:
                continue
            in_name = str(op.inputs[0])
            out_name = str(op.outputs[0])
            if out_name in model_outputs:
                continue

            in_tensor = model_ir.tensors.get(in_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if (
                in_tensor is None
                or out_tensor is None
                or len(list(in_tensor.shape)) != 4
                or len(list(out_tensor.shape)) != 4
            ):
                continue

            user_indices = [int(v) for v in consumers.get(out_name, [])]
            if len(user_indices) == 0:
                continue
            all_conv_users = True
            should_bypass = True
            for user_idx in user_indices:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "CONV_2D"
                    or len(user_op.inputs) < 2
                    or str(user_op.inputs[0]) != out_name
                ):
                    all_conv_users = False
                    break
                filter_tensor = model_ir.tensors.get(str(user_op.inputs[1]), None)
                if filter_tensor is None or len(list(filter_tensor.shape)) != 4:
                    should_bypass = False
                    break
                expected_c = int(filter_tensor.shape[3])
                in_c = int(in_tensor.shape[3])
                out_c = int(out_tensor.shape[3])
                if not (in_c == expected_c and out_c != expected_c):
                    should_bypass = False
                    break

            if not all_conv_users or not should_bypass:
                continue

            _replace_tensor_inputs(model_ir, out_name, in_name)
            del model_ir.operators[int(op_idx)]
            removed_post += 1
            changed = True
            break

        if not changed:
            break

    propagated_tensors = int(len(rewritten_tensors))
    _prune_unused_tensors(model_ir)
    return {
        "rewritten_transpose_swish_branches_to_nhwc": int(rewritten_branches),
        "removed_transpose_swish_pre": int(removed_pre),
        "propagated_nhwc_tensor_metadata": int(propagated_tensors),
        "rewritten_concat_axis_to_nhwc": int(rewritten_concat_axis),
        "removed_transpose_swish_post": int(removed_post),
    }


def _optimize_transpose_swish_residual_concat_closure_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Rewrite Swish + Residual + Concat islands as a single closure.

    Unlike the legacy swish-qdq pass, this closure pass is spatial-size agnostic
    and only rewrites branches that are proven to terminate at
    CONCAT(axis=1) -> QUANTIZE -> TRANSPOSE(0,2,3,1) tails.
    """
    stats = _optimize_transpose_swish_qdq_nhwc_islands(
        model_ir,
        min_spatial_stage=0,
        require_concat_closure=True,
    )
    return {
        "optimized_transpose_swish_residual_concat_closure_nhwc_chains": int(
            stats.get("rewritten_transpose_swish_branches_to_nhwc", 0)
        ),
        "removed_transpose_swish_residual_concat_closure_pre": int(
            stats.get("removed_transpose_swish_pre", 0)
        ),
        "rewritten_transpose_swish_residual_concat_closure_axis_to_nhwc": int(
            stats.get("rewritten_concat_axis_to_nhwc", 0)
        ),
        "removed_transpose_swish_residual_concat_closure_post": int(
            stats.get("removed_transpose_swish_post", 0)
        ),
    }


def _optimize_transpose_binary_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold symmetric transpose wrappers around elementwise binary ops.

    Target pattern:
      A --Transpose(P)--> A'
      B --Transpose(P)--> B'
      A',B' --(ADD|SUB|MUL|DIV)--> C'
      C' --Transpose(inv(P))--> C

    Constraints:
    - Both pre-transposed tensors are consumed only by the binary op.
    - Binary output is consumed only by the post transpose.
    - pre/post permutations are exact inverses.
    """
    removed_binary_bridges = 0
    removed_binary_asymmetric_bridges = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}
    enable_fanout_pattern_c = False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = mid_op.inputs[0]
            in1_name = mid_op.inputs[1]
            out_name = mid_op.outputs[0]

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            if pre0_idx is None or pre1_idx is None:
                continue
            pre0_op = model_ir.operators[pre0_idx]
            pre1_op = model_ir.operators[pre1_idx]
            if str(pre0_op.op_type) != "TRANSPOSE" or str(pre1_op.op_type) != "TRANSPOSE":
                continue
            if len(pre0_op.inputs) < 2 or len(pre1_op.inputs) < 2:
                continue
            if len(pre0_op.outputs) != 1 or len(pre1_op.outputs) != 1:
                continue
            if pre0_op.outputs[0] != in0_name or pre1_op.outputs[0] != in1_name:
                continue

            in0_users = set(consumers.get(in0_name, []))
            in1_users = set(consumers.get(in1_name, []))
            if in0_users != {int(mid_idx)} or in1_users != {int(mid_idx)}:
                continue

            perm0 = _read_transpose_perm(model_ir, pre0_op)
            perm1 = _read_transpose_perm(model_ir, pre1_op)
            if perm0 is None or perm1 is None:
                continue
            if perm0 != perm1:
                continue

            out_users = {int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)}
            if len(out_users) == 0:
                continue

            post_candidates: List[int] = []
            for user_idx in sorted(list(out_users)):
                post_candidate = model_ir.operators[int(user_idx)]
                if (
                    str(post_candidate.op_type) != "TRANSPOSE"
                    or len(post_candidate.inputs) < 2
                    or len(post_candidate.outputs) != 1
                    or post_candidate.inputs[0] != out_name
                ):
                    continue
                perm_post = _read_transpose_perm(model_ir, post_candidate)
                if perm_post is None or not _is_inverse_perm(perm0, perm_post):
                    continue
                post_candidates.append(int(user_idx))

            synthesize_post = False
            if len(post_candidates) == 0:
                if any(str(model_ir.operators[int(user_idx)].op_type) == "TRANSPOSE" for user_idx in out_users):
                    continue
                synthesize_post = True
            if len(post_candidates) > 1:
                continue
            post_idx = None
            post_op = None
            post_perm_name = None
            keep_post_for_fanout = False
            if not synthesize_post:
                post_idx = int(post_candidates[0])
                post_op = model_ir.operators[post_idx]
                post_perm_name = post_op.inputs[1]
                extra_users = set(out_users) - {int(post_idx)}
                # Extra users may remain on out_name, but must not be transpose producers.
                if any(str(model_ir.operators[int(user_idx)].op_type) == "TRANSPOSE" for user_idx in extra_users):
                    continue
                keep_post_for_fanout = len(extra_users) > 0

            # Keep user-visible output names stable.
            if in0_name in model_ir.outputs or in1_name in model_ir.outputs or out_name in model_ir.outputs:
                continue

            original_in0 = pre0_op.inputs[0]
            original_in1 = pre1_op.inputs[0]
            final_out = (
                post_op.outputs[0]
                if post_op is not None
                else out_name
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=[original_in0, original_in1],
            )
            if synthesize_post:
                out_tensor = model_ir.tensors.get(out_name, None)
                if out_tensor is None:
                    continue
                if not _is_per_tensor_quantization(out_tensor.quantization):
                    continue
                inv_perm = _invert_perm(perm0)
                if inv_perm is None:
                    continue
                raw_shape = _permute_shape(list(out_tensor.shape), inv_perm)
                raw_signature = _permute_shape(
                    list(out_tensor.shape_signature)
                    if out_tensor.shape_signature is not None
                    else list(out_tensor.shape),
                    inv_perm,
                )
                if raw_shape is None:
                    continue
                if raw_signature is None:
                    raw_signature = list(raw_shape)
                raw_out_name = f"{out_name}__raw"
                serial = 0
                while raw_out_name in model_ir.tensors:
                    serial += 1
                    raw_out_name = f"{out_name}__raw_{serial}"
                model_ir.tensors[raw_out_name] = TensorIR(
                    name=raw_out_name,
                    dtype=str(out_tensor.dtype),
                    shape=[int(v) for v in raw_shape],
                    shape_signature=[int(v) for v in raw_signature],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(out_tensor.quantization),
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_outputs=[raw_out_name],
                )
            else:
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_outputs=[final_out],
                )

            final_tensor = model_ir.tensors.get(final_out, None)
            old_out_tensor = model_ir.tensors.get(out_name, None)
            if final_tensor is not None and old_out_tensor is not None:
                final_tensor.dtype = str(old_out_tensor.dtype)
                if final_tensor.quantization is None and old_out_tensor.quantization is not None:
                    final_tensor.quantization = _clone_quantization(old_out_tensor.quantization)

            if keep_post_for_fanout and post_op is not None and post_perm_name is not None:
                perm_tensor = model_ir.tensors.get(post_perm_name, None)
                if perm_tensor is not None:
                    perm_tensor.data = np.asarray(perm0, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=post_op,
                    new_inputs=[final_out, post_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=post_op,
                    new_outputs=[out_name],
                )

            if synthesize_post:
                reuse_perm_name = pre0_op.inputs[1]
                insert_at = int(min(out_users))
                remove_indices = sorted(list({int(pre0_idx), int(pre1_idx)}), reverse=True)
                shift = sum(1 for idx in remove_indices if int(idx) < insert_at)
                for remove_idx in remove_indices:
                    del model_ir.operators[remove_idx]
                adjusted_insert_at = int(insert_at - shift)
                model_ir.operators.insert(
                    adjusted_insert_at,
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[raw_out_name, reuse_perm_name],
                        outputs=[out_name],
                    ),
                )
            else:
                remove_indices = (
                    sorted(list({int(pre0_idx), int(pre1_idx)}), reverse=True)
                    if keep_post_for_fanout
                    else sorted(list({int(pre0_idx), int(pre1_idx), int(post_idx)}), reverse=True)
                )
                for remove_idx in remove_indices:
                    del model_ir.operators[remove_idx]
            removed_binary_bridges += 1
            changed = True
            break

        if changed:
            continue

        # Pattern B:
        #   A,B' --(ADD|SUB|MUL|DIV)--> C' --T(invP)--> C
        # where B' is produced by T(P) from B.
        # Rewrite to:
        #   A' --T(invP)--> A'
        #   A',B --(ADD|SUB|MUL|DIV)--> C
        # This removes one transpose around the binary op.
        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = mid_op.inputs[0]
            in1_name = mid_op.inputs[1]
            out_name = mid_op.outputs[0]

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            pre0_is_transpose = False
            pre1_is_transpose = False
            pre0_op = None
            pre1_op = None

            if pre0_idx is not None:
                pre0_op = model_ir.operators[int(pre0_idx)]
                pre0_is_transpose = (
                    str(pre0_op.op_type) == "TRANSPOSE"
                    and len(pre0_op.inputs) >= 2
                    and len(pre0_op.outputs) == 1
                    and pre0_op.outputs[0] == in0_name
                )
            if pre1_idx is not None:
                pre1_op = model_ir.operators[int(pre1_idx)]
                pre1_is_transpose = (
                    str(pre1_op.op_type) == "TRANSPOSE"
                    and len(pre1_op.inputs) >= 2
                    and len(pre1_op.outputs) == 1
                    and pre1_op.outputs[0] == in1_name
                )

            # Asymmetric-only pass: exactly one input is wrapped by pre-transpose.
            if pre0_is_transpose == pre1_is_transpose:
                continue

            if pre0_is_transpose:
                pre_idx = int(pre0_idx)
                pre_op = pre0_op
                transposed_input_name = in0_name
                plain_input_name = in1_name
                transpose_on_lhs = True
            else:
                pre_idx = int(pre1_idx)
                pre_op = pre1_op
                transposed_input_name = in1_name
                plain_input_name = in0_name
                transpose_on_lhs = False

            if pre_op is None:
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            transposed_users = set(consumers.get(transposed_input_name, []))
            if transposed_users != {int(mid_idx)}:
                continue

            out_users = set(consumers.get(out_name, []))
            if len(out_users) != 1:
                continue
            post_idx = int(list(out_users)[0])
            if post_idx == mid_idx:
                continue
            post_op = model_ir.operators[post_idx]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if post_op.inputs[0] != out_name:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_pre is None or perm_post is None:
                continue
            if not _is_inverse_perm(perm_pre, perm_post):
                continue

            # Keep user-visible output names stable.
            if transposed_input_name in model_ir.outputs or out_name in model_ir.outputs:
                continue

            raw_from_pre_name = pre_op.inputs[0]
            final_out_name = post_op.outputs[0]

            transposed_tensor = model_ir.tensors.get(transposed_input_name, None)
            plain_tensor = model_ir.tensors.get(plain_input_name, None)
            raw_tensor = model_ir.tensors.get(raw_from_pre_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            final_out_tensor = model_ir.tensors.get(final_out_name, None)

            if not _all_per_tensor_quantized(
                [transposed_tensor, plain_tensor, raw_tensor, out_tensor, final_out_tensor]
            ):
                continue

            transposed_input_shape = (
                list(transposed_tensor.shape) if transposed_tensor is not None else None
            )
            plain_input_shape = list(plain_tensor.shape) if plain_tensor is not None else None
            plain_input_signature = (
                list(plain_tensor.shape_signature)
                if plain_tensor is not None and plain_tensor.shape_signature is not None
                else plain_input_shape
            )
            raw_pre_shape = list(raw_tensor.shape) if raw_tensor is not None else None
            final_out_shape = list(final_out_tensor.shape) if final_out_tensor is not None else None

            # Require no broadcasting ambiguity.
            if not _shapes_equal(transposed_input_shape, plain_input_shape):
                continue
            expected_raw_shape = _permute_shape(plain_input_shape, perm_post)
            expected_raw_signature = _permute_shape(plain_input_signature, perm_post)
            if expected_raw_shape is None:
                continue
            if expected_raw_signature is None:
                expected_raw_signature = list(expected_raw_shape)
            if raw_pre_shape is not None and not _shapes_equal(expected_raw_shape, raw_pre_shape):
                continue
            if final_out_shape is not None and not _shapes_equal(expected_raw_shape, final_out_shape):
                continue

            # Reuse pre-transpose node by moving it to the other input with inverse perm.
            _set_operator_inputs(
                model_ir=model_ir,
                op=pre_op,
                new_inputs=[plain_input_name, post_op.inputs[1]],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=pre_op,
                new_outputs=[transposed_input_name],
            )

            if transpose_on_lhs:
                # op(T(A), B) then T(invP) -> op(A, T(invP)(B))
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_inputs=[raw_from_pre_name, transposed_input_name],
                )
            else:
                # op(A, T(B)) then T(invP) -> op(T(invP)(A), B)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_inputs=[transposed_input_name, raw_from_pre_name],
                )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[final_out_name],
            )

            if transposed_tensor is not None:
                transposed_tensor.shape = list(expected_raw_shape)
                transposed_tensor.shape_signature = list(expected_raw_signature)
                if plain_tensor is not None:
                    transposed_tensor.dtype = str(plain_tensor.dtype)
                    if transposed_tensor.quantization is None and plain_tensor.quantization is not None:
                        transposed_tensor.quantization = _clone_quantization(plain_tensor.quantization)

            if final_out_tensor is not None and out_tensor is not None:
                final_out_tensor.dtype = str(out_tensor.dtype)
                if final_out_tensor.quantization is None and out_tensor.quantization is not None:
                    final_out_tensor.quantization = _clone_quantization(out_tensor.quantization)

            del model_ir.operators[post_idx]
            removed_binary_asymmetric_bridges += 1
            changed = True
            break

        if changed:
            continue

        # Pattern C:
        # Fanout-aware binary bridge optimization.
        # Allows output consumers to include:
        # - inverse TRANSPOSE nodes (removed)
        # - downstream binary ops (kept, enabling chained rewrites)
        if not enable_fanout_pattern_c:
            break
        binary_chain_cache: Dict[int, bool] = {}

        def _is_rewritable_binary_chain(
            start_idx: int,
            base_perm_local: List[int],
            visiting: Optional[set[int]] = None,
        ) -> bool:
            if start_idx in binary_chain_cache:
                return bool(binary_chain_cache[start_idx])
            if visiting is None:
                visiting = set()
            if int(start_idx) in visiting:
                binary_chain_cache[int(start_idx)] = False
                return False
            if start_idx < 0 or start_idx >= len(model_ir.operators):
                binary_chain_cache[int(start_idx)] = False
                return False
            op_local = model_ir.operators[int(start_idx)]
            if str(op_local.op_type) not in binary_ops or len(op_local.outputs) != 1:
                binary_chain_cache[int(start_idx)] = False
                return False
            out_local = str(op_local.outputs[0])
            if out_local in model_ir.outputs:
                binary_chain_cache[int(start_idx)] = False
                return False

            downstream = [int(v) for v in consumers.get(out_local, []) if int(v) != int(start_idx)]
            if len(downstream) == 0:
                binary_chain_cache[int(start_idx)] = False
                return False

            visiting.add(int(start_idx))
            for down_idx in downstream:
                down_op = model_ir.operators[int(down_idx)]
                if (
                    str(down_op.op_type) == "TRANSPOSE"
                    and len(down_op.inputs) >= 2
                    and len(down_op.outputs) == 1
                    and down_op.inputs[0] == out_local
                ):
                    down_perm = _read_transpose_perm(model_ir, down_op)
                    if down_perm is None or not _is_inverse_perm(base_perm_local, down_perm):
                        visiting.remove(int(start_idx))
                        binary_chain_cache[int(start_idx)] = False
                        return False
                    if str(down_op.outputs[0]) in model_ir.outputs:
                        visiting.remove(int(start_idx))
                        binary_chain_cache[int(start_idx)] = False
                        return False
                    continue
                if str(down_op.op_type) in binary_ops and out_local in set(down_op.inputs):
                    if not _is_rewritable_binary_chain(
                        int(down_idx),
                        base_perm_local,
                        visiting,
                    ):
                        visiting.remove(int(start_idx))
                        binary_chain_cache[int(start_idx)] = False
                        return False
                    continue
                visiting.remove(int(start_idx))
                binary_chain_cache[int(start_idx)] = False
                return False

            visiting.remove(int(start_idx))
            binary_chain_cache[int(start_idx)] = True
            return True

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = mid_op.inputs[0]
            in1_name = mid_op.inputs[1]
            out_name = mid_op.outputs[0]

            if out_name in model_ir.outputs:
                continue

            def _input_info(input_name: str) -> Tuple[str, Optional[int], Optional[List[int]]]:
                pre_idx_local = producers.get(input_name, None)
                if pre_idx_local is None:
                    return input_name, None, None
                pre_op_local = model_ir.operators[int(pre_idx_local)]
                if (
                    str(pre_op_local.op_type) != "TRANSPOSE"
                    or len(pre_op_local.inputs) < 2
                    or len(pre_op_local.outputs) != 1
                    or pre_op_local.outputs[0] != input_name
                ):
                    return input_name, None, None
                if set(consumers.get(input_name, [])) != {int(mid_idx)}:
                    return input_name, None, None
                perm_local = _read_transpose_perm(model_ir, pre_op_local)
                if perm_local is None:
                    return input_name, None, None
                return pre_op_local.inputs[0], int(pre_idx_local), perm_local

            raw_in0_name, pre0_idx, perm0 = _input_info(in0_name)
            raw_in1_name, pre1_idx, perm1 = _input_info(in1_name)

            if perm0 is None and perm1 is None:
                continue
            if in0_name in model_ir.outputs or in1_name in model_ir.outputs:
                continue

            perm_candidates = []
            if perm0 is not None:
                perm_candidates.append(list(perm0))
            if perm1 is not None:
                perm_candidates.append(list(perm1))
            if len(perm_candidates) == 0:
                continue
            base_perm = perm_candidates[0]
            if any(perm != base_perm for perm in perm_candidates[1:]):
                continue

            raw_in0_tensor = model_ir.tensors.get(raw_in0_name, None)
            raw_in1_tensor = model_ir.tensors.get(raw_in1_name, None)
            if raw_in0_tensor is None or raw_in1_tensor is None:
                continue
            raw_shape0 = list(raw_in0_tensor.shape)
            raw_shape1 = list(raw_in1_tensor.shape)
            if not _shapes_equal(raw_shape0, raw_shape1):
                continue
            raw_signature0 = (
                list(raw_in0_tensor.shape_signature)
                if raw_in0_tensor.shape_signature is not None
                else list(raw_shape0)
            )

            out_users = [int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)]
            if len(out_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            users_valid = True

            for user_idx in out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and user_op.inputs[0] == out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is None or not _is_inverse_perm(base_perm, perm_post):
                        users_valid = False
                        break
                    post_output_name = str(user_op.outputs[0])
                    if post_output_name in model_ir.outputs:
                        users_valid = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                    continue

                if str(user_op.op_type) in binary_ops and out_name in set(user_op.inputs):
                    if len(user_op.inputs) != 2:
                        users_valid = False
                        break
                    if user_op.inputs[0] == out_name:
                        other_input_name = user_op.inputs[1]
                    elif user_op.inputs[1] == out_name:
                        other_input_name = user_op.inputs[0]
                    else:
                        users_valid = False
                        break

                    other_ok = False
                    other_prod_idx = producers.get(other_input_name, None)
                    if other_prod_idx is not None:
                        other_prod_op = model_ir.operators[int(other_prod_idx)]
                        if (
                            str(other_prod_op.op_type) == "TRANSPOSE"
                            and len(other_prod_op.inputs) >= 2
                            and len(other_prod_op.outputs) == 1
                            and other_prod_op.outputs[0] == other_input_name
                            and set(consumers.get(other_input_name, [])) == {int(user_idx)}
                        ):
                            other_perm = _read_transpose_perm(model_ir, other_prod_op)
                            if other_perm is not None and other_perm == base_perm:
                                other_ok = True
                    if not other_ok:
                        other_tensor = model_ir.tensors.get(other_input_name, None)
                        if other_tensor is not None and _shapes_equal(list(other_tensor.shape), raw_shape0):
                            other_ok = True
                    if not other_ok:
                        users_valid = False
                        break

                    if not _is_rewritable_binary_chain(int(user_idx), base_perm):
                        users_valid = False
                        break

                    continue

                users_valid = False
                break

            if not users_valid:
                continue

            removable_pre_indices = [idx for idx in [pre0_idx, pre1_idx] if idx is not None]
            if len(removable_pre_indices) == 0 and len(post_indices) == 0:
                continue

            # Rewire binary inputs to raw layout.
            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=[raw_in0_name, raw_in1_name],
            )

            out_tensor = model_ir.tensors.get(out_name, None)
            if out_tensor is not None:
                out_tensor.shape = list(raw_shape0)
                out_tensor.shape_signature = list(raw_signature0)

            for post_output_name in post_output_names:
                _replace_tensor_inputs(model_ir, post_output_name, out_name)

            remove_indices = sorted(list(set(removable_pre_indices + post_indices)), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[remove_idx]

            removed_binary_bridges += int(max(1, len(post_indices) + len(removable_pre_indices)))
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_transpose_binary_bridges": int(removed_binary_bridges),
        "removed_transpose_binary_asymmetric_bridges": int(removed_binary_asymmetric_bridges),
    }


def _optimize_transpose_binary_asymmetric_fanout_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold asymmetric transpose wrappers around binary ops when output has fanout.

    Target pattern:
      T(P)(X) + Y -> Z
      T(inv(P))(Z) -> Z_nhwc
      (Z may have extra non-transpose consumers)

    Rewritten:
      X + T(inv(P))(Y) -> Z_nhwc
      (optional adapter) T(P)(Z_nhwc) -> Z

    This reduces one transpose in fanout cases where the simple asymmetric bridge
    optimization cannot apply.
    """
    rewritten = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = str(mid_op.inputs[0])
            in1_name = str(mid_op.inputs[1])
            out_name = str(mid_op.outputs[0])

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            pre0_is_transpose = False
            pre1_is_transpose = False
            pre0_op = None
            pre1_op = None

            if pre0_idx is not None:
                pre0_op = model_ir.operators[int(pre0_idx)]
                pre0_is_transpose = (
                    str(pre0_op.op_type) == "TRANSPOSE"
                    and len(pre0_op.inputs) >= 2
                    and len(pre0_op.outputs) == 1
                    and str(pre0_op.outputs[0]) == in0_name
                )
            if pre1_idx is not None:
                pre1_op = model_ir.operators[int(pre1_idx)]
                pre1_is_transpose = (
                    str(pre1_op.op_type) == "TRANSPOSE"
                    and len(pre1_op.inputs) >= 2
                    and len(pre1_op.outputs) == 1
                    and str(pre1_op.outputs[0]) == in1_name
                )

            # Exactly one pre-transposed input.
            if pre0_is_transpose == pre1_is_transpose:
                continue

            if pre0_is_transpose:
                pre_idx = int(pre0_idx)
                pre_op = pre0_op
                transposed_input_name = in0_name
                plain_input_name = in1_name
                transpose_on_lhs = True
            else:
                pre_idx = int(pre1_idx)
                pre_op = pre1_op
                transposed_input_name = in1_name
                plain_input_name = in0_name
                transpose_on_lhs = False

            if pre_op is None:
                continue

            transposed_users = [int(v) for v in consumers.get(transposed_input_name, [])]
            if transposed_users != [int(mid_idx)] and set(transposed_users) != {int(mid_idx)}:
                continue

            out_users = [int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)]
            if len(out_users) == 0:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            inv_perm = _invert_perm(perm_pre)
            if inv_perm is None:
                continue

            post_idx = None
            post_op = None
            for user_idx in out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is not None and perm_post == inv_perm:
                        post_idx = int(user_idx)
                        post_op = user_op
                        break
            if post_idx is None or post_op is None:
                continue

            post_output_name = str(post_op.outputs[0])
            if post_output_name in model_ir.outputs and out_name in model_ir.outputs:
                # Avoid ambiguous observable output rewrites.
                continue

            # Find an existing transpose(plain_input, inv_perm) to reuse.
            plain_users = [int(v) for v in consumers.get(plain_input_name, [])]
            plain_transpose_idx = None
            plain_transpose_out = None
            for plain_user_idx in plain_users:
                plain_user_op = model_ir.operators[int(plain_user_idx)]
                if (
                    str(plain_user_op.op_type) == "TRANSPOSE"
                    and len(plain_user_op.inputs) >= 2
                    and len(plain_user_op.outputs) == 1
                    and str(plain_user_op.inputs[0]) == plain_input_name
                ):
                    perm_plain = _read_transpose_perm(model_ir, plain_user_op)
                    if perm_plain is not None and perm_plain == inv_perm:
                        plain_transpose_idx = int(plain_user_idx)
                        plain_transpose_out = str(plain_user_op.outputs[0])
                        break
            if plain_transpose_idx is None or plain_transpose_out is None:
                continue

            raw_from_pre_name = str(pre_op.inputs[0])
            if transpose_on_lhs:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_inputs=[raw_from_pre_name, plain_transpose_out],
                )
            else:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mid_op,
                    new_inputs=[plain_transpose_out, raw_from_pre_name],
                )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[post_output_name],
            )

            out_tensor = model_ir.tensors.get(out_name, None)
            post_out_tensor = model_ir.tensors.get(post_output_name, None)
            if out_tensor is not None and post_out_tensor is not None:
                post_out_tensor.dtype = str(out_tensor.dtype)
                if post_out_tensor.quantization is None and out_tensor.quantization is not None:
                    post_out_tensor.quantization = _clone_quantization(out_tensor.quantization)

            extra_users = [int(v) for v in out_users if int(v) != int(post_idx)]
            adapter_needed = len(extra_users) > 0 or out_name in model_ir.outputs
            if adapter_needed:
                adapter = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[post_output_name, str(pre_op.inputs[1])],
                    outputs=[out_name],
                )

            remove_indices = sorted(list({int(pre_idx), int(post_idx)}), reverse=True)
            shift = sum(1 for ridx in remove_indices if int(ridx) < int(mid_idx))
            for ridx in remove_indices:
                del model_ir.operators[int(ridx)]
            mid_new_idx = int(mid_idx - shift)
            if adapter_needed:
                model_ir.operators.insert(int(mid_new_idx) + 1, adapter)

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "rewritten_transpose_binary_asymmetric_fanout_bridges": int(rewritten),
    }


def _optimize_transpose_binary_full_post_fanout_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold a strict transpose-binary-transpose fanout pattern.

    Target:
      T(P)(A) + T(P)(B) -> C
      C -> T(inv(P)) -> C0
      C -> T(inv(P)) -> C1
      ...

    Rewrite:
      A + B -> C0
      (all uses of C1... are rewired to C0)

    Safety:
    - Binary op has exactly 2 inputs / 1 output (ADD|SUB|MUL|DIV)
    - Both inputs are produced by TRANSPOSE with identical permutation
    - Both pre-transpose outputs are consumed only by the binary op
    - All consumers of binary output are inverse TRANSPOSE ops
    - No broadcasting ambiguity: all participating tensor shapes match exactly
    - Intermediate tensors are not graph outputs
    """
    rewritten = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = str(mid_op.inputs[0])
            in1_name = str(mid_op.inputs[1])
            out_name = str(mid_op.outputs[0])

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            if pre0_idx is None or pre1_idx is None:
                continue
            # This rewrite targets two independent symmetric transpose branches.
            # Shared-producer binary (e.g. MUL(x, x)) is out-of-scope and can
            # cause duplicate index deletion if not excluded.
            if int(pre0_idx) == int(pre1_idx):
                continue
            pre0_op = model_ir.operators[int(pre0_idx)]
            pre1_op = model_ir.operators[int(pre1_idx)]
            if str(pre0_op.op_type) != "TRANSPOSE" or str(pre1_op.op_type) != "TRANSPOSE":
                continue
            if len(pre0_op.inputs) < 2 or len(pre1_op.inputs) < 2:
                continue
            if len(pre0_op.outputs) != 1 or len(pre1_op.outputs) != 1:
                continue
            if str(pre0_op.outputs[0]) != in0_name or str(pre1_op.outputs[0]) != in1_name:
                continue

            # Pre-transposed tensors must be exclusive to the binary op.
            if set(consumers.get(in0_name, [])) != {int(mid_idx)}:
                continue
            if set(consumers.get(in1_name, [])) != {int(mid_idx)}:
                continue

            perm_pre0 = _read_transpose_perm(model_ir, pre0_op)
            perm_pre1 = _read_transpose_perm(model_ir, pre1_op)
            if perm_pre0 is None or perm_pre1 is None:
                continue
            if perm_pre0 != perm_pre1:
                continue

            out_users = [int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)]
            if len(out_users) < 2:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            for user_idx in out_users:
                post_op = model_ir.operators[int(user_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != out_name
                ):
                    post_indices = []
                    break
                perm_post = _read_transpose_perm(model_ir, post_op)
                if perm_post is None or not _is_inverse_perm(perm_pre0, perm_post):
                    post_indices = []
                    break
                post_indices.append(int(user_idx))
                post_output_names.append(str(post_op.outputs[0]))
            if len(post_indices) != len(out_users):
                continue

            if (
                in0_name in model_ir.outputs
                or in1_name in model_ir.outputs
                or out_name in model_ir.outputs
                or any(name in model_ir.outputs for name in post_output_names)
            ):
                continue

            raw_in0_name = str(pre0_op.inputs[0])
            raw_in1_name = str(pre1_op.inputs[0])
            raw_in0_tensor = model_ir.tensors.get(raw_in0_name, None)
            raw_in1_tensor = model_ir.tensors.get(raw_in1_name, None)
            in0_tensor = model_ir.tensors.get(in0_name, None)
            in1_tensor = model_ir.tensors.get(in1_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if not _all_per_tensor_quantized(
                [raw_in0_tensor, raw_in1_tensor, in0_tensor, in1_tensor, out_tensor]
            ):
                continue

            # Conservative shape guards with broadcast-aware validation.
            raw_shape0 = list(raw_in0_tensor.shape) if raw_in0_tensor is not None else None
            raw_shape1 = list(raw_in1_tensor.shape) if raw_in1_tensor is not None else None
            raw_broadcast_shape = _broadcast_static_shapes(raw_shape0, raw_shape1)
            if raw_broadcast_shape is None:
                continue
            mid_expected_shape = _permute_shape(list(raw_broadcast_shape), perm_pre0)
            if mid_expected_shape is None:
                continue

            in_shape0 = list(in0_tensor.shape) if in0_tensor is not None else None
            in_shape1 = list(in1_tensor.shape) if in1_tensor is not None else None
            mid_broadcast_shape = _broadcast_static_shapes(in_shape0, in_shape1)
            if mid_broadcast_shape is None:
                continue
            if not _shapes_match_if_known(mid_broadcast_shape, mid_expected_shape):
                continue
            out_shape = list(out_tensor.shape) if out_tensor is not None else None
            if not _shapes_match_if_known(out_shape, mid_expected_shape):
                continue

            valid_post_shapes = True
            for post_output_name in post_output_names:
                post_tensor = model_ir.tensors.get(post_output_name, None)
                post_shape = list(post_tensor.shape) if post_tensor is not None else None
                if not _shapes_match_if_known(post_shape, raw_broadcast_shape):
                    valid_post_shapes = False
                    break
            if not valid_post_shapes:
                continue

            canonical_post_output = str(post_output_names[0])
            canonical_tensor = model_ir.tensors.get(canonical_post_output, None)

            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=[raw_in0_name, raw_in1_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[canonical_post_output],
            )

            if canonical_tensor is not None:
                canonical_tensor.shape = [int(v) for v in list(raw_broadcast_shape)]
                raw_sig0 = (
                    list(raw_in0_tensor.shape_signature)
                    if raw_in0_tensor is not None and raw_in0_tensor.shape_signature is not None
                    else list(raw_shape0) if raw_shape0 is not None else None
                )
                raw_sig1 = (
                    list(raw_in1_tensor.shape_signature)
                    if raw_in1_tensor is not None and raw_in1_tensor.shape_signature is not None
                    else list(raw_shape1) if raw_shape1 is not None else None
                )
                raw_broadcast_signature = _broadcast_static_shapes(raw_sig0, raw_sig1)
                if raw_broadcast_signature is None:
                    raw_broadcast_signature = [int(v) for v in list(raw_broadcast_shape)]
                canonical_tensor.shape_signature = [int(v) for v in list(raw_broadcast_signature)]
                if out_tensor is not None:
                    canonical_tensor.dtype = str(out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(out_tensor.quantization)

            for post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, post_output_name, canonical_post_output)

            remove_indices = sorted(
                list(set([int(pre0_idx), int(pre1_idx)] + [int(idx) for idx in post_indices])),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[remove_idx]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_binary_full_post_fanout_bridges": int(rewritten)}


def _optimize_transpose_binary_single_post_bridges_safe(model_ir: ModelIR) -> Dict[str, int]:
    """
    Conservative transpose-binary bridge folding with exactly one inverse post-transpose.

    Target:
      T(P)(A), T(P)(B) -> C
      C -> T(inv(P)) -> C_raw
      (C may also have extra non-transpose consumers)

    Rewrite:
      A, B -> C_raw
      If C has extra consumers:
        keep adapter T(P): C_raw -> C
      else:
        remove post-transpose.

    This pass intentionally uses strict shape guards to avoid layout corruption.
    """
    rewritten = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}
    preserve_layout_boundary_marker = "__preserve_layout_boundary__"

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = str(mid_op.inputs[0])
            in1_name = str(mid_op.inputs[1])
            out_name = str(mid_op.outputs[0])

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            if pre0_idx is None or pre1_idx is None:
                continue
            # This rewrite assumes two independent transpose producers.
            # Shared-producer binary (e.g. MUL(x, x)) can invalidate index removal
            # ordering and leave dangling edges.
            if int(pre0_idx) == int(pre1_idx):
                continue
            pre0_op = model_ir.operators[int(pre0_idx)]
            pre1_op = model_ir.operators[int(pre1_idx)]
            if str(pre0_op.op_type) != "TRANSPOSE" or str(pre1_op.op_type) != "TRANSPOSE":
                continue
            if len(pre0_op.inputs) < 2 or len(pre1_op.inputs) < 2:
                continue
            if len(pre0_op.outputs) != 1 or len(pre1_op.outputs) != 1:
                continue
            if str(pre0_op.outputs[0]) != in0_name or str(pre1_op.outputs[0]) != in1_name:
                continue

            if set(consumers.get(in0_name, [])) != {int(mid_idx)}:
                continue
            if set(consumers.get(in1_name, [])) != {int(mid_idx)}:
                continue

            perm_pre0 = _read_transpose_perm(model_ir, pre0_op)
            perm_pre1 = _read_transpose_perm(model_ir, pre1_op)
            if perm_pre0 is None or perm_pre1 is None:
                continue
            if perm_pre0 != perm_pre1:
                continue

            out_users = {int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)}
            if len(out_users) == 0:
                continue

            post_candidates: List[int] = []
            for user_idx in sorted(list(out_users)):
                post_op = model_ir.operators[int(user_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != out_name
                ):
                    continue
                perm_post = _read_transpose_perm(model_ir, post_op)
                if perm_post is None or not _is_inverse_perm(perm_pre0, perm_post):
                    continue
                post_candidates.append(int(user_idx))

            if len(post_candidates) != 1:
                continue

            post_idx = int(post_candidates[0])
            post_op = model_ir.operators[post_idx]
            post_out_name = str(post_op.outputs[0])
            post_perm_name = str(post_op.inputs[1])
            extra_users = set(out_users) - {int(post_idx)}

            # Avoid ambiguous fanout rewrites if any extra fanout is another transpose.
            if any(str(model_ir.operators[int(v)].op_type) == "TRANSPOSE" for v in extra_users):
                continue

            if in0_name in model_ir.outputs or in1_name in model_ir.outputs:
                continue

            raw_in0_name = str(pre0_op.inputs[0])
            raw_in1_name = str(pre1_op.inputs[0])

            raw_in0_tensor = model_ir.tensors.get(raw_in0_name, None)
            raw_in1_tensor = model_ir.tensors.get(raw_in1_name, None)
            in0_tensor = model_ir.tensors.get(in0_name, None)
            in1_tensor = model_ir.tensors.get(in1_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            post_out_tensor = model_ir.tensors.get(post_out_name, None)

            if not _all_per_tensor_quantized(
                [raw_in0_tensor, raw_in1_tensor, in0_tensor, in1_tensor, out_tensor, post_out_tensor]
            ):
                continue

            raw_shape0 = list(raw_in0_tensor.shape) if raw_in0_tensor is not None else None
            raw_shape1 = list(raw_in1_tensor.shape) if raw_in1_tensor is not None else None
            # NOTE:
            # Some optimized branches can retain stale static shapes in IR while runtime
            # shapes are valid. For strict single-post bridge chains, rely on topology
            # constraints and quantization compatibility instead of hard shape equality.

            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=[raw_in0_name, raw_in1_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[post_out_name],
            )

            if post_out_tensor is not None and out_tensor is not None:
                post_out_tensor.dtype = str(out_tensor.dtype)
                post_out_tensor.quantization = _clone_quantization(out_tensor.quantization)
                if not _is_unknown_shape(raw_shape0):
                    post_out_tensor.shape = [int(v) for v in list(raw_shape0)]
                    raw_sig0 = (
                        list(raw_in0_tensor.shape_signature)
                        if raw_in0_tensor is not None and raw_in0_tensor.shape_signature is not None
                        else list(raw_shape0)
                    )
                    post_out_tensor.shape_signature = [int(v) for v in list(raw_sig0)]

            if len(extra_users) > 0:
                post_opts = dict(post_op.options) if isinstance(post_op.options, dict) else {}
                post_opts[preserve_layout_boundary_marker] = True
                post_op.options = post_opts
                perm_tensor = model_ir.tensors.get(post_perm_name, None)
                if perm_tensor is not None:
                    perm_tensor.data = np.asarray(perm_pre0, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=post_op,
                    new_inputs=[post_out_name, post_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=post_op,
                    new_outputs=[out_name],
                )
                remove_indices = sorted(list({int(pre0_idx), int(pre1_idx)}), reverse=True)
            else:
                remove_indices = sorted(list({int(pre0_idx), int(pre1_idx), int(post_idx)}), reverse=True)

            for remove_idx in remove_indices:
                del model_ir.operators[remove_idx]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_binary_single_post_bridges_safe": int(rewritten)}


def _optimize_transpose_binary_mixed_fanout_bridges_safe(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose-binary pattern when output fanout mixes:
    - inverse TRANSPOSE consumers (layout adapters)
    - non-TRANSPOSE consumers (legacy layout path)

    Rewrite keeps one transpose adapter for legacy consumers and removes the
    inverse transpose fanout consumers.
    """
    rewritten = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}
    preserve_layout_boundary_marker = "__preserve_layout_boundary__"

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = str(mid_op.inputs[0])
            in1_name = str(mid_op.inputs[1])
            out_name = str(mid_op.outputs[0])

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            if pre0_idx is None or pre1_idx is None:
                continue
            # This rewrite assumes two independent transpose producers.
            # Shared-producer binary (e.g. MUL(x, x)) can invalidate index removal
            # ordering and leave dangling edges.
            if int(pre0_idx) == int(pre1_idx):
                continue
            pre0_op = model_ir.operators[int(pre0_idx)]
            pre1_op = model_ir.operators[int(pre1_idx)]
            if str(pre0_op.op_type) != "TRANSPOSE" or str(pre1_op.op_type) != "TRANSPOSE":
                continue
            if len(pre0_op.inputs) < 2 or len(pre1_op.inputs) < 2:
                continue
            if len(pre0_op.outputs) != 1 or len(pre1_op.outputs) != 1:
                continue
            if str(pre0_op.outputs[0]) != in0_name or str(pre1_op.outputs[0]) != in1_name:
                continue

            if set(consumers.get(in0_name, [])) != {int(mid_idx)}:
                continue
            if set(consumers.get(in1_name, [])) != {int(mid_idx)}:
                continue

            perm_pre0 = _read_transpose_perm(model_ir, pre0_op)
            perm_pre1 = _read_transpose_perm(model_ir, pre1_op)
            if perm_pre0 is None or perm_pre1 is None or perm_pre0 != perm_pre1:
                continue

            out_users = [int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)]
            if len(out_users) < 2:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid = True
            for user_idx in out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is None or not _is_inverse_perm(perm_pre0, perm_post):
                        valid = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))

            if not valid:
                continue
            if len(post_indices) == 0 or len(legacy_users) == 0:
                continue

            if (
                in0_name in model_ir.outputs
                or in1_name in model_ir.outputs
                or out_name in model_ir.outputs
                or any(name in model_ir.outputs for name in post_output_names)
            ):
                continue

            raw_in0_name = str(pre0_op.inputs[0])
            raw_in1_name = str(pre1_op.inputs[0])
            raw_in0_tensor = model_ir.tensors.get(raw_in0_name, None)
            raw_in1_tensor = model_ir.tensors.get(raw_in1_name, None)
            in0_tensor = model_ir.tensors.get(in0_name, None)
            in1_tensor = model_ir.tensors.get(in1_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if not _all_per_tensor_quantized(
                [raw_in0_tensor, raw_in1_tensor, in0_tensor, in1_tensor, out_tensor]
            ):
                continue

            raw_shape0 = list(raw_in0_tensor.shape) if raw_in0_tensor is not None else None
            raw_shape1 = list(raw_in1_tensor.shape) if raw_in1_tensor is not None else None
            if not _shapes_match_if_known(raw_shape0, raw_shape1):
                continue
            in_shape0 = list(in0_tensor.shape) if in0_tensor is not None else None
            in_shape1 = list(in1_tensor.shape) if in1_tensor is not None else None
            if not _shapes_match_if_known(in_shape0, in_shape1):
                continue
            if not _is_unknown_shape(raw_shape0) and not _is_unknown_shape(in_shape0):
                expected_trans_shape = _permute_shape(raw_shape0, perm_pre0)
                if expected_trans_shape is None or not _shapes_equal(expected_trans_shape, in_shape0):
                    continue

            valid_post_shapes = True
            for post_output_name in post_output_names:
                post_tensor = model_ir.tensors.get(post_output_name, None)
                post_shape = list(post_tensor.shape) if post_tensor is not None else None
                if not _shapes_match_if_known(post_shape, raw_shape0):
                    valid_post_shapes = False
                    break
            if not valid_post_shapes:
                continue

            keep_post_idx = int(post_indices[0])
            keep_post_op = model_ir.operators[int(keep_post_idx)]
            canonical_raw_name = str(keep_post_op.outputs[0])
            keep_perm_name = str(keep_post_op.inputs[1])
            canonical_raw_tensor = model_ir.tensors.get(canonical_raw_name, None)

            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=[raw_in0_name, raw_in1_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[canonical_raw_name],
            )

            if canonical_raw_tensor is not None and out_tensor is not None:
                canonical_raw_tensor.dtype = str(out_tensor.dtype)
                canonical_raw_tensor.quantization = _clone_quantization(out_tensor.quantization)
                if not _is_unknown_shape(raw_shape0):
                    canonical_raw_tensor.shape = [int(v) for v in list(raw_shape0)]
                    raw_sig0 = (
                        list(raw_in0_tensor.shape_signature)
                        if raw_in0_tensor is not None and raw_in0_tensor.shape_signature is not None
                        else list(raw_shape0)
                    )
                    canonical_raw_tensor.shape_signature = [int(v) for v in list(raw_sig0)]

            for post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, post_output_name, canonical_raw_name)

            # Keep one adapter transpose for legacy consumers: raw -> original out layout.
            keep_post_opts = dict(keep_post_op.options) if isinstance(keep_post_op.options, dict) else {}
            keep_post_opts[preserve_layout_boundary_marker] = True
            keep_post_op.options = keep_post_opts
            keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
            if keep_perm_tensor is not None:
                keep_perm_tensor.data = np.asarray(perm_pre0, dtype=np.int32)
            _set_operator_inputs(
                model_ir=model_ir,
                op=keep_post_op,
                new_inputs=[canonical_raw_name, keep_perm_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=keep_post_op,
                new_outputs=[out_name],
            )

            remove_indices = sorted(
                list(set([int(pre0_idx), int(pre1_idx)] + [int(idx) for idx in post_indices[1:]])),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[remove_idx]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_binary_mixed_fanout_bridges_safe": int(rewritten)}


def _optimize_transpose_binary_symmetric_legacy_only_bridges_safe(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold symmetric pre-transpose binary pattern when output has only legacy consumers.

    Pattern:
      T(P)(A), T(P)(B) -> C
      C -> legacy consumers only (no inverse-post transpose)

    Rewrite:
      A, B -> C_raw
      T(P)(C_raw) -> C

    Net effect: remove 2 pre-transposes and insert 1 adapter transpose.
    """
    rewritten = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}
    preserve_layout_boundary_marker = "__preserve_layout_boundary__"

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = str(mid_op.inputs[0])
            in1_name = str(mid_op.inputs[1])
            out_name = str(mid_op.outputs[0])

            pre0_idx = producers.get(in0_name, None)
            pre1_idx = producers.get(in1_name, None)
            if pre0_idx is None or pre1_idx is None:
                continue
            # This rewrite assumes two independent transpose producers.
            # Shared-producer binary (e.g. MUL(x, x)) can invalidate index removal
            # ordering and leave dangling edges.
            if int(pre0_idx) == int(pre1_idx):
                continue
            pre0_op = model_ir.operators[int(pre0_idx)]
            pre1_op = model_ir.operators[int(pre1_idx)]
            if str(pre0_op.op_type) != "TRANSPOSE" or str(pre1_op.op_type) != "TRANSPOSE":
                continue
            if len(pre0_op.inputs) < 2 or len(pre1_op.inputs) < 2:
                continue
            if len(pre0_op.outputs) != 1 or len(pre1_op.outputs) != 1:
                continue
            if str(pre0_op.outputs[0]) != in0_name or str(pre1_op.outputs[0]) != in1_name:
                continue

            if set(consumers.get(in0_name, [])) != {int(mid_idx)}:
                continue
            if set(consumers.get(in1_name, [])) != {int(mid_idx)}:
                continue

            perm_pre0 = _read_transpose_perm(model_ir, pre0_op)
            perm_pre1 = _read_transpose_perm(model_ir, pre1_op)
            if perm_pre0 is None or perm_pre1 is None or perm_pre0 != perm_pre1:
                continue

            out_users = [int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)]
            if len(out_users) == 0:
                continue

            # This pass targets legacy-only fanout; inverse post-transpose is handled elsewhere.
            has_inverse_post = False
            for user_idx in out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is not None and _is_inverse_perm(perm_pre0, perm_post):
                        has_inverse_post = True
                        break
            if has_inverse_post:
                continue

            if in0_name in model_ir.outputs or in1_name in model_ir.outputs:
                continue

            raw_in0_name = str(pre0_op.inputs[0])
            raw_in1_name = str(pre1_op.inputs[0])

            raw_in0_tensor = model_ir.tensors.get(raw_in0_name, None)
            raw_in1_tensor = model_ir.tensors.get(raw_in1_name, None)
            in0_tensor = model_ir.tensors.get(in0_name, None)
            in1_tensor = model_ir.tensors.get(in1_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if not _all_per_tensor_quantized(
                [raw_in0_tensor, raw_in1_tensor, in0_tensor, in1_tensor, out_tensor]
            ):
                continue

            raw_shape0 = list(raw_in0_tensor.shape) if raw_in0_tensor is not None else None
            raw_shape1 = list(raw_in1_tensor.shape) if raw_in1_tensor is not None else None
            if not _shapes_match_if_known(raw_shape0, raw_shape1):
                continue

            # Prepare raw output tensor.
            raw_out_name = f"{out_name}__raw"
            suffix = 1
            while raw_out_name in model_ir.tensors:
                raw_out_name = f"{out_name}__raw_{suffix}"
                suffix += 1

            raw_sig0 = (
                list(raw_in0_tensor.shape_signature)
                if raw_in0_tensor is not None and raw_in0_tensor.shape_signature is not None
                else raw_shape0
            )
            model_ir.tensors[raw_out_name] = TensorIR(
                name=raw_out_name,
                dtype=str(out_tensor.dtype) if out_tensor is not None else "INT8",
                shape=[int(v) for v in list(raw_shape0)] if not _is_unknown_shape(raw_shape0) else [1],
                shape_signature=(
                    [int(v) for v in list(raw_sig0)]
                    if raw_sig0 is not None and not _is_unknown_shape(raw_sig0)
                    else [1]
                ),
                data=None,
                is_variable=False,
                quantization=_clone_quantization(out_tensor.quantization if out_tensor is not None else None),
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=[raw_in0_name, raw_in1_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[raw_out_name],
            )

            insert_at = min(out_users)
            remove_indices = sorted(list({int(pre0_idx), int(pre1_idx)}), reverse=True)
            shift = sum(1 for idx in remove_indices if idx < insert_at)
            for idx in remove_indices:
                del model_ir.operators[idx]
            adjusted_insert_at = int(insert_at - shift)
            model_ir.operators.insert(
                adjusted_insert_at,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[raw_out_name, str(pre0_op.inputs[1])],
                    outputs=[out_name],
                    options={preserve_layout_boundary_marker: True},
                ),
            )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_binary_symmetric_legacy_only_bridges_safe": int(rewritten)}


def _is_singleton_constant_tensor(
    model_ir: ModelIR,
    tensor_name: str,
) -> bool:
    tensor = model_ir.tensors.get(str(tensor_name), None)
    if tensor is None or tensor.data is None:
        return False
    try:
        array = np.asarray(tensor.data)
    except Exception:
        return False
    return int(array.size) == 1


def _is_scalar_like_tensor(
    model_ir: ModelIR,
    tensor_name: str,
) -> bool:
    tensor = model_ir.tensors.get(str(tensor_name), None)
    if tensor is None:
        return False
    if tensor.data is not None:
        try:
            return int(np.asarray(tensor.data).size) == 1
        except Exception:
            return False
    shape = list(tensor.shape) if tensor.shape is not None else []
    if len(shape) == 0:
        return False
    # Require fully-known singleton shape so NHWC/NCHW rewrite keeps broadcast semantics.
    if any(int(dim) <= 0 for dim in shape):
        return False
    return all(int(dim) == 1 for dim in shape)


def _read_singleton_constant_float(
    model_ir: ModelIR,
    tensor_name: str,
) -> Optional[float]:
    tensor = model_ir.tensors.get(str(tensor_name), None)
    if tensor is None or tensor.data is None:
        return None
    try:
        array = np.asarray(tensor.data)
    except Exception:
        return None
    if int(array.size) != 1:
        return None
    value = float(array.reshape(-1)[0])
    if not np.isfinite(value):
        return None
    return float(value)


def _optimize_maximum_minimum_relu0to1_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Replace clamp chains MAXIMUM(0.0) -> MINIMUM(1.0) with RELU_0_TO_1.

    Target:
      X --MAXIMUM(0.0)--> M --MINIMUM(1.0)--> Y

    Rewrite:
      X --RELU_0_TO_1--> Y

    Safety:
    - MAXIMUM and MINIMUM side inputs must be singleton constants.
    - MAXIMUM output must be consumed only by the matched MINIMUM.
    """
    rewritten = 0
    atol = 1e-6

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for min_idx, min_op in enumerate(model_ir.operators):
            if str(min_op.op_type) != "MINIMUM" or len(min_op.inputs) != 2 or len(min_op.outputs) != 1:
                continue

            min_input0 = str(min_op.inputs[0])
            min_input1 = str(min_op.inputs[1])
            if _is_singleton_constant_tensor(model_ir, min_input0):
                min_data_name = str(min_input1)
                min_const_name = str(min_input0)
            elif _is_singleton_constant_tensor(model_ir, min_input1):
                min_data_name = str(min_input0)
                min_const_name = str(min_input1)
            else:
                continue
            min_const_value = _read_singleton_constant_float(model_ir, min_const_name)
            if min_const_value is None or not np.isclose(float(min_const_value), 1.0, atol=atol):
                continue
            if min_data_name in model_outputs:
                continue

            max_idx = producers.get(min_data_name, None)
            if max_idx is None:
                continue
            max_op = model_ir.operators[int(max_idx)]
            if str(max_op.op_type) != "MAXIMUM" or len(max_op.inputs) != 2 or len(max_op.outputs) != 1:
                continue
            if str(max_op.outputs[0]) != str(min_data_name):
                continue

            max_input0 = str(max_op.inputs[0])
            max_input1 = str(max_op.inputs[1])
            if _is_singleton_constant_tensor(model_ir, max_input0):
                max_data_name = str(max_input1)
                max_const_name = str(max_input0)
            elif _is_singleton_constant_tensor(model_ir, max_input1):
                max_data_name = str(max_input0)
                max_const_name = str(max_input1)
            else:
                continue
            max_const_value = _read_singleton_constant_float(model_ir, max_const_name)
            if max_const_value is None or not np.isclose(float(max_const_value), 0.0, atol=atol):
                continue

            max_users = [int(v) for v in consumers.get(str(min_data_name), [])]
            if len(max_users) != 1 or int(max_users[0]) != int(min_idx):
                continue

            min_op.op_type = "RELU_0_TO_1"
            min_op.version = 1
            _set_operator_inputs(
                model_ir=model_ir,
                op=min_op,
                new_inputs=[str(max_data_name)],
            )
            min_op.options = {}

            del model_ir.operators[int(max_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_maximum_minimum_relu0to1_chains": int(rewritten)}


def _optimize_maximum_with_zero_input2_to_relu(model_ir: ModelIR) -> Dict[str, int]:
    """
    Rewrite MAXIMUM(x, 0) into RELU(x) when the second input is singleton zero.

    Target:
      X --MAXIMUM(X, 0.0)--> Y

    Rewrite:
      X --RELU--> Y

    Notes:
    - This pass intentionally checks input2 only to match ONNX Max patterns that
      lower as MAXIMUM(data, const_zero).
    - Restrict to FLOAT16/FLOAT32 tensors for safety.
    """
    rewritten = 0
    atol = 1e-6

    for op in model_ir.operators:
        if str(op.op_type) != "MAXIMUM" or len(op.inputs) != 2 or len(op.outputs) != 1:
            continue

        data_name = str(op.inputs[0])
        const_name = str(op.inputs[1])
        if not _is_singleton_constant_tensor(model_ir, const_name):
            continue
        const_value = _read_singleton_constant_float(model_ir, const_name)
        if const_value is None or not np.isclose(float(const_value), 0.0, atol=atol):
            continue

        data_tensor = model_ir.tensors.get(data_name, None)
        if data_tensor is None:
            continue
        data_dtype = str(data_tensor.dtype).upper()
        if data_dtype not in {"FLOAT16", "FLOAT32"}:
            continue

        op.op_type = "RELU"
        op.version = 1
        _set_operator_inputs(
            model_ir=model_ir,
            op=op,
            new_inputs=[data_name],
        )
        op.options = {}
        rewritten += 1

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_maximum_with_zero_input2_to_relu": int(rewritten)}


def _optimize_fuse_pseudo_leakyrelu_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fuse pseudo-op-expanded LeakyReLU chain into single LEAKY_RELU op.

    Target:
      x --NEG--> n --RELU--> nr --MUL(alpha)--> m
      x --RELU--> p
      p --SUB(m)--> y
    """
    fused = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for sub_idx, sub_op in enumerate(model_ir.operators):
            if str(sub_op.op_type) != "SUB" or len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue

            relu_pos_out_name = str(sub_op.inputs[0])
            mul_out_name = str(sub_op.inputs[1])
            output_name = str(sub_op.outputs[0])

            relu_pos_idx = producers.get(relu_pos_out_name, None)
            mul_idx = producers.get(mul_out_name, None)
            if relu_pos_idx is None or mul_idx is None:
                continue
            relu_pos_op = model_ir.operators[int(relu_pos_idx)]
            mul_op = model_ir.operators[int(mul_idx)]
            if str(relu_pos_op.op_type) != "RELU" or len(relu_pos_op.inputs) != 1 or len(relu_pos_op.outputs) != 1:
                continue
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(relu_pos_op.outputs[0]) != relu_pos_out_name:
                continue
            if str(mul_op.outputs[0]) != mul_out_name:
                continue

            x_name = str(relu_pos_op.inputs[0])
            mul_input0 = str(mul_op.inputs[0])
            mul_input1 = str(mul_op.inputs[1])
            alpha_name: Optional[str] = None
            relu_neg_out_name: Optional[str] = None
            if _is_singleton_constant_tensor(model_ir, mul_input0):
                alpha_name = mul_input0
                relu_neg_out_name = mul_input1
            elif _is_singleton_constant_tensor(model_ir, mul_input1):
                alpha_name = mul_input1
                relu_neg_out_name = mul_input0
            if alpha_name is None or relu_neg_out_name is None:
                continue
            if relu_neg_out_name == relu_pos_out_name:
                continue

            alpha = _read_singleton_constant_float(model_ir, alpha_name)
            if alpha is None:
                continue

            relu_neg_idx = producers.get(relu_neg_out_name, None)
            if relu_neg_idx is None:
                continue
            relu_neg_op = model_ir.operators[int(relu_neg_idx)]
            if str(relu_neg_op.op_type) != "RELU" or len(relu_neg_op.inputs) != 1 or len(relu_neg_op.outputs) != 1:
                continue
            if str(relu_neg_op.outputs[0]) != relu_neg_out_name:
                continue

            neg_out_name = str(relu_neg_op.inputs[0])
            neg_idx = producers.get(neg_out_name, None)
            if neg_idx is None:
                continue
            neg_op = model_ir.operators[int(neg_idx)]
            if str(neg_op.op_type) != "NEG" or len(neg_op.inputs) != 1 or len(neg_op.outputs) != 1:
                continue
            if str(neg_op.outputs[0]) != neg_out_name:
                continue
            if str(neg_op.inputs[0]) != x_name:
                continue

            if any(
                name in model_outputs
                for name in [
                    relu_pos_out_name,
                    mul_out_name,
                    relu_neg_out_name,
                    neg_out_name,
                ]
            ):
                continue

            if [int(v) for v in consumers.get(relu_pos_out_name, [])] != [int(sub_idx)]:
                continue
            if [int(v) for v in consumers.get(mul_out_name, [])] != [int(sub_idx)]:
                continue
            if [int(v) for v in consumers.get(relu_neg_out_name, [])] != [int(mul_idx)]:
                continue
            if [int(v) for v in consumers.get(neg_out_name, [])] != [int(relu_neg_idx)]:
                continue

            x_tensor = model_ir.tensors.get(x_name, None)
            y_tensor = model_ir.tensors.get(output_name, None)
            if x_tensor is not None and y_tensor is not None:
                x_dtype = str(x_tensor.dtype).upper()
                y_dtype = str(y_tensor.dtype).upper()
                if not x_dtype.startswith("FLOAT") or not y_dtype.startswith("FLOAT"):
                    continue

            model_ir.operators[int(sub_idx)] = OperatorIR(
                op_type="LEAKY_RELU",
                inputs=[x_name],
                outputs=[output_name],
                options={"alpha": float(alpha)},
            )

            for remove_idx in sorted(
                [int(neg_idx), int(relu_neg_idx), int(mul_idx), int(relu_pos_idx)],
                reverse=True,
            ):
                del model_ir.operators[int(remove_idx)]

            fused += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"fused_pseudo_leakyrelu_chains": int(fused)}


def _optimize_yolo_decode_mul_square_anchor_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold YOLO decode MUL chains with square expansion.

    Target:
      X --MUL(c)--> A --MUL(A,A)--> B --MUL(anchor_const)--> C --MUL(scale_const)--> Y

    Rewrite:
      X --MUL(X,X)--> B --MUL(anchor_const * (c*c) * scale_const)--> Y

    Safety:
    - All side inputs must be constants.
    - `A`, `B`, and `C` are strict single-consumer intermediates on this chain.
    - `c` must be scalar-like (singleton constant).
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _consumed_only_by(tensor_name: str, op_index: int) -> bool:
        users = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(users) == 0:
            return False
        unique_users = {int(v) for v in users}
        return unique_users == {int(op_index)}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for mul_scale_idx, mul_scale_op in enumerate(model_ir.operators):
            if str(mul_scale_op.op_type) != "MUL" or len(mul_scale_op.inputs) != 2 or len(mul_scale_op.outputs) != 1:
                continue

            scale_input0 = str(mul_scale_op.inputs[0])
            scale_input1 = str(mul_scale_op.inputs[1])
            scale_const_name: Optional[str] = None
            anchor_mul_out_name: Optional[str] = None
            if model_ir.tensors.get(scale_input0, None) is not None and model_ir.tensors[scale_input0].data is not None:
                scale_const_name = scale_input0
                anchor_mul_out_name = scale_input1
            elif model_ir.tensors.get(scale_input1, None) is not None and model_ir.tensors[scale_input1].data is not None:
                scale_const_name = scale_input1
                anchor_mul_out_name = scale_input0
            if scale_const_name is None or anchor_mul_out_name is None:
                continue

            if not _consumed_only_by(anchor_mul_out_name, int(mul_scale_idx)):
                continue
            mul_anchor_idx = producers.get(anchor_mul_out_name, None)
            if mul_anchor_idx is None:
                continue

            mul_anchor_op = model_ir.operators[int(mul_anchor_idx)]
            if str(mul_anchor_op.op_type) != "MUL" or len(mul_anchor_op.inputs) != 2 or len(mul_anchor_op.outputs) != 1:
                continue
            if str(mul_anchor_op.outputs[0]) != str(anchor_mul_out_name):
                continue

            anchor_input0 = str(mul_anchor_op.inputs[0])
            anchor_input1 = str(mul_anchor_op.inputs[1])
            anchor_const_name: Optional[str] = None
            square_out_name: Optional[str] = None
            anchor_const_input_index = -1
            if model_ir.tensors.get(anchor_input0, None) is not None and model_ir.tensors[anchor_input0].data is not None:
                anchor_const_name = anchor_input0
                square_out_name = anchor_input1
                anchor_const_input_index = 0
            elif model_ir.tensors.get(anchor_input1, None) is not None and model_ir.tensors[anchor_input1].data is not None:
                anchor_const_name = anchor_input1
                square_out_name = anchor_input0
                anchor_const_input_index = 1
            if anchor_const_name is None or square_out_name is None:
                continue
            if not _consumed_only_by(square_out_name, int(mul_anchor_idx)):
                continue

            mul_square_idx = producers.get(square_out_name, None)
            if mul_square_idx is None:
                continue
            mul_square_op = model_ir.operators[int(mul_square_idx)]
            if str(mul_square_op.op_type) != "MUL" or len(mul_square_op.inputs) != 2 or len(mul_square_op.outputs) != 1:
                continue
            if str(mul_square_op.outputs[0]) != str(square_out_name):
                continue
            square_input0 = str(mul_square_op.inputs[0])
            square_input1 = str(mul_square_op.inputs[1])
            if square_input0 != square_input1:
                continue

            scaled_x_name = str(square_input0)
            if not _consumed_only_by(scaled_x_name, int(mul_square_idx)):
                continue

            mul_pre_idx = producers.get(scaled_x_name, None)
            if mul_pre_idx is None:
                continue
            mul_pre_op = model_ir.operators[int(mul_pre_idx)]
            if str(mul_pre_op.op_type) != "MUL" or len(mul_pre_op.inputs) != 2 or len(mul_pre_op.outputs) != 1:
                continue
            if str(mul_pre_op.outputs[0]) != str(scaled_x_name):
                continue

            pre_input0 = str(mul_pre_op.inputs[0])
            pre_input1 = str(mul_pre_op.inputs[1])
            pre_const_name: Optional[str] = None
            x_name: Optional[str] = None
            if _is_singleton_constant_tensor(model_ir, pre_input0):
                pre_const_name = pre_input0
                x_name = pre_input1
            elif _is_singleton_constant_tensor(model_ir, pre_input1):
                pre_const_name = pre_input1
                x_name = pre_input0
            if pre_const_name is None or x_name is None:
                continue

            pre_scale = _read_singleton_constant_float(model_ir, pre_const_name)
            if pre_scale is None or not np.isfinite(pre_scale):
                continue

            anchor_tensor = model_ir.tensors.get(anchor_const_name, None)
            scale_tensor = model_ir.tensors.get(scale_const_name, None)
            if anchor_tensor is None or anchor_tensor.data is None or scale_tensor is None or scale_tensor.data is None:
                continue

            anchor_data = np.asarray(anchor_tensor.data)
            scale_data = np.asarray(scale_tensor.data)
            if not np.issubdtype(anchor_data.dtype, np.floating):
                continue
            if not np.issubdtype(scale_data.dtype, np.floating):
                continue

            try:
                fused_anchor_data = (
                    anchor_data.astype(np.float32, copy=False)
                    * float(pre_scale) * float(pre_scale)
                    * scale_data.astype(np.float32, copy=False)
                )
            except Exception:
                continue
            if not np.all(np.isfinite(fused_anchor_data)):
                continue
            fused_anchor_data = fused_anchor_data.astype(anchor_data.dtype, copy=False)

            fused_const_name = _unique_tensor_name(f"{anchor_const_name}_mulsq_fused")
            fused_shape, fused_signature = normalize_onnx_shape(list(fused_anchor_data.shape))
            model_ir.tensors[fused_const_name] = TensorIR(
                name=fused_const_name,
                dtype=tflite_dtype_from_numpy(fused_anchor_data.dtype),
                shape=[int(v) for v in fused_shape],
                shape_signature=[int(v) for v in fused_signature],
                data=fused_anchor_data,
                is_variable=False,
                quantization=_clone_quantization(anchor_tensor.quantization),
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_square_op,
                new_inputs=[str(x_name), str(x_name)],
            )

            anchor_inputs = [str(v) for v in list(mul_anchor_op.inputs)]
            anchor_inputs[int(anchor_const_input_index)] = str(fused_const_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_anchor_op,
                new_inputs=anchor_inputs,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_anchor_op,
                new_outputs=[str(mul_scale_op.outputs[0])],
            )

            for remove_idx in sorted({int(mul_scale_idx), int(mul_pre_idx)}, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_yolo_decode_mul_square_anchor_chains": int(rewritten)}


def _optimize_fold_consecutive_mul_constants_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold consecutive floating-point MUL chains with constant side inputs.

    Target:
      X --MUL(c0)--> A --MUL(c1)--> Y

    Rewrite:
      X --MUL(c0 * c1)--> Y

    Safety:
    - Both MULs must be strict binary single-output operators.
    - `A` must be consumed only by the second MUL and must not be a model output.
    - Constant tensors must be floating-point and broadcastable.
    - Data/output tensors on the path must be floating-point.
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _is_constant_tensor(tensor_name: str) -> bool:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        return tensor is not None and tensor.data is not None

    def _is_float_tensor_name(tensor_name: str) -> bool:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            return False
        dtype = str(tensor.dtype).upper()
        return dtype.startswith("FLOAT")

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for second_mul_idx, second_mul_op in enumerate(model_ir.operators):
            if str(second_mul_op.op_type) != "MUL" or len(second_mul_op.inputs) != 2 or len(second_mul_op.outputs) != 1:
                continue

            second_in0 = str(second_mul_op.inputs[0])
            second_in1 = str(second_mul_op.inputs[1])
            second_const_name: Optional[str] = None
            mid_name: Optional[str] = None
            second_const_input_index = -1
            mid_input_index = -1
            if _is_constant_tensor(second_in0):
                second_const_name = second_in0
                mid_name = second_in1
                second_const_input_index = 0
                mid_input_index = 1
            elif _is_constant_tensor(second_in1):
                second_const_name = second_in1
                mid_name = second_in0
                second_const_input_index = 1
                mid_input_index = 0
            if second_const_name is None or mid_name is None:
                continue
            if mid_name in model_outputs:
                continue

            mid_users = sorted({int(v) for v in consumers.get(mid_name, [])})
            if mid_users != [int(second_mul_idx)]:
                continue

            first_mul_idx = producers.get(mid_name, None)
            if first_mul_idx is None:
                continue
            first_mul_op = model_ir.operators[int(first_mul_idx)]
            if str(first_mul_op.op_type) != "MUL" or len(first_mul_op.inputs) != 2 or len(first_mul_op.outputs) != 1:
                continue
            if str(first_mul_op.outputs[0]) != str(mid_name):
                continue

            first_in0 = str(first_mul_op.inputs[0])
            first_in1 = str(first_mul_op.inputs[1])
            first_const_name: Optional[str] = None
            data_name: Optional[str] = None
            if _is_constant_tensor(first_in0):
                first_const_name = first_in0
                data_name = first_in1
            elif _is_constant_tensor(first_in1):
                first_const_name = first_in1
                data_name = first_in0
            if first_const_name is None or data_name is None:
                continue

            second_output_name = str(second_mul_op.outputs[0])
            if not _is_float_tensor_name(data_name):
                continue
            if not _is_float_tensor_name(mid_name):
                continue
            if not _is_float_tensor_name(second_output_name):
                continue

            first_const_tensor = model_ir.tensors.get(first_const_name, None)
            second_const_tensor = model_ir.tensors.get(second_const_name, None)
            if first_const_tensor is None or first_const_tensor.data is None:
                continue
            if second_const_tensor is None or second_const_tensor.data is None:
                continue
            first_const_data = np.asarray(first_const_tensor.data)
            second_const_data = np.asarray(second_const_tensor.data)
            if not np.issubdtype(first_const_data.dtype, np.floating):
                continue
            if not np.issubdtype(second_const_data.dtype, np.floating):
                continue

            fused_dtype = np.result_type(first_const_data.dtype, second_const_data.dtype)
            if not np.issubdtype(fused_dtype, np.floating):
                continue
            try:
                fused_const_data = (
                    first_const_data.astype(fused_dtype, copy=False)
                    * second_const_data.astype(fused_dtype, copy=False)
                )
            except Exception:
                continue
            if not np.all(np.isfinite(fused_const_data)):
                continue

            fused_const_name = _unique_tensor_name(f"{first_const_name}_mulfused")
            fused_shape, fused_signature = normalize_onnx_shape(list(fused_const_data.shape))
            model_ir.tensors[fused_const_name] = TensorIR(
                name=fused_const_name,
                dtype=tflite_dtype_from_numpy(fused_const_data.dtype),
                shape=[int(v) for v in fused_shape],
                shape_signature=[int(v) for v in fused_signature],
                data=fused_const_data,
                is_variable=False,
                quantization=_clone_quantization(
                    first_const_tensor.quantization
                    if first_const_tensor.quantization is not None
                    else second_const_tensor.quantization
                ),
            )

            new_inputs = [str(v) for v in list(second_mul_op.inputs)]
            new_inputs[int(mid_input_index)] = str(data_name)
            new_inputs[int(second_const_input_index)] = str(fused_const_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=second_mul_op,
                new_inputs=new_inputs,
            )

            del model_ir.operators[int(first_mul_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_fold_consecutive_mul_constants_chains": int(rewritten)}


def _permute_tensor_metadata_if_rank_matches(
    tensor: Optional[TensorIR],
    perm: List[int],
) -> None:
    if tensor is None:
        return
    shape_src = list(tensor.shape) if tensor.shape is not None else None
    if shape_src is not None and len(shape_src) == len(perm):
        new_shape = _permute_shape(shape_src, perm)
        if new_shape is not None:
            tensor.shape = [int(v) for v in list(new_shape)]
    signature_src = (
        list(tensor.shape_signature)
        if tensor.shape_signature is not None
        else (list(tensor.shape) if tensor.shape is not None else None)
    )
    if signature_src is not None and len(signature_src) == len(perm):
        new_signature = _permute_shape(signature_src, perm)
        if new_signature is not None:
            tensor.shape_signature = [int(v) for v in list(new_signature)]


def _optimize_leading_input_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold leading input-boundary transpose chains through layout-agnostic ops.

    Target:
      X_in(NHWC) --T(P)--> X_ncx --(CAST|SUB|ADD|MUL|DIV|MAXIMUM|MINIMUM|ATAN2|QUANTIZE|DEQUANTIZE)*--> Y_ncx --T(inv(P))--> Y_nhwc
      X_in(NHWC) --T(P)--> X_ncx --(same chain)*--> Y_ncx (Y_ncx is model output)

    Rewrite:
      X_in(NHWC) --(same chain)*--> Y_nhwc

    Safety:
    - Leading transpose input must be a model input tensor.
    - Chain must be strictly linear (single-consumer on the main path).
    - Binary ops in the chain must use constant side inputs.
    - Quantization in the chain must remain per-tensor only.
    """
    rewritten = 0
    unary_passthrough_ops = {"CAST", "QUANTIZE", "DEQUANTIZE"}
    binary_passthrough_ops = {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM", "ATAN2"}

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(model_ir.inputs)
        model_outputs = set(model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            if pre_input_name not in model_inputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            # Build a strict linear passthrough chain.
            chain_indices: List[int] = []
            chain_ops: List[OperatorIR] = []
            chain_output_names: List[str] = []
            chain_main_input_index: Dict[int, int] = {}
            chain_binary_side_const_by_op: Dict[int, str] = {}
            current_tensor = pre_output_name
            while True:
                current_users = [int(v) for v in consumers.get(current_tensor, [])]
                if len(current_users) != 1:
                    break
                op_idx = int(current_users[0])
                op = model_ir.operators[op_idx]
                op_type = str(op.op_type)

                if op_type in unary_passthrough_ops:
                    if len(op.inputs) != 1 or len(op.outputs) != 1:
                        break
                    if str(op.inputs[0]) != current_tensor:
                        break
                    out_name = str(op.outputs[0])
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        break
                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    current_tensor = out_name
                    if out_name in model_outputs:
                        break
                    continue

                if op_type in binary_passthrough_ops:
                    if len(op.inputs) != 2 or len(op.outputs) != 1:
                        break
                    input_0 = str(op.inputs[0])
                    input_1 = str(op.inputs[1])
                    if input_0 == current_tensor:
                        side_input_name = input_1
                    elif input_1 == current_tensor:
                        side_input_name = input_0
                    else:
                        break
                    side_tensor = model_ir.tensors.get(side_input_name, None)
                    if side_tensor is None or side_tensor.data is None:
                        break
                    side_array = np.asarray(side_tensor.data)
                    if int(side_array.size) != 1 and int(side_array.ndim) != int(len(perm_pre)):
                        break
                    if side_tensor is not None and not _is_per_tensor_quantization(side_tensor.quantization):
                        break
                    out_name = str(op.outputs[0])
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        break
                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    chain_binary_side_const_by_op[int(op_idx)] = str(side_input_name)
                    current_tensor = out_name
                    if out_name in model_outputs:
                        break
                    continue

                break

            if len(chain_ops) == 0:
                continue

            # Chain must end with one inverse transpose or a model output.
            tail_users = [int(v) for v in consumers.get(current_tensor, [])]
            terminal_mode = ""
            post_idx: Optional[int] = None
            post_op: Optional[OperatorIR] = None
            post_output_name = ""
            if len(tail_users) == 1:
                post_idx = int(tail_users[0])
                if post_idx in set(chain_indices):
                    continue
                post_op = model_ir.operators[post_idx]
                if str(post_op.op_type) == "TRANSPOSE" and len(post_op.inputs) >= 2 and len(post_op.outputs) == 1:
                    if str(post_op.inputs[0]) == current_tensor:
                        perm_post = _read_transpose_perm(model_ir, post_op)
                        if perm_post is not None and perm_post == perm_post_expected:
                            terminal_mode = "inverse_post_transpose"
                            post_output_name = str(post_op.outputs[0])
            if terminal_mode == "" and current_tensor in model_outputs and len(tail_users) == 0:
                terminal_mode = "model_output"
                post_output_name = str(current_tensor)
            if terminal_mode == "":
                continue

            # Ensure chain topology remains linear on the main path.
            linear_ok = True
            previous_tensor_name = pre_output_name
            chain_index_to_pos = {int(idx): pos for pos, idx in enumerate(chain_indices)}
            for op_idx, op in zip(chain_indices, chain_ops):
                if previous_tensor_name not in set(str(v) for v in op.inputs):
                    linear_ok = False
                    break
                out_name = str(op.outputs[0]) if len(op.outputs) > 0 else ""
                if out_name == "":
                    linear_ok = False
                    break
                expected_users = [int(v) for v in consumers.get(out_name, [])]
                pos = int(chain_index_to_pos[int(op_idx)])
                if pos < len(chain_indices) - 1:
                    if expected_users != [int(chain_indices[pos + 1])]:
                        linear_ok = False
                        break
                else:
                    if terminal_mode == "inverse_post_transpose":
                        if post_idx is None or expected_users != [int(post_idx)]:
                            linear_ok = False
                            break
                    else:
                        if len(expected_users) != 0:
                            linear_ok = False
                            break
                previous_tensor_name = out_name
            if not linear_ok:
                continue

            old_last_name = str(chain_ops[-1].outputs[0])
            if terminal_mode == "inverse_post_transpose" and old_last_name in model_outputs:
                continue

            # Rewire chain head: consume model input directly.
            first_op = chain_ops[0]
            first_input_names = [str(v) for v in first_op.inputs]
            if first_input_names[0] == pre_output_name:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=[pre_input_name] + first_input_names[1:],
                )
            elif len(first_input_names) > 1 and first_input_names[1] == pre_output_name:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=[first_input_names[0], pre_input_name],
                )
            else:
                continue

            # Permute non-scalar side constants that carry NCX axis order.
            chain_index_set = set(int(v) for v in chain_indices)
            for op_idx, side_name in chain_binary_side_const_by_op.items():
                side_tensor = model_ir.tensors.get(str(side_name), None)
                if side_tensor is None or side_tensor.data is None:
                    continue
                side_array = np.asarray(side_tensor.data)
                if int(side_array.size) == 1 or int(side_array.ndim) != int(len(perm_post_expected)):
                    continue
                transposed_const = np.transpose(side_array, axes=perm_post_expected)
                side_users = [int(v) for v in consumers.get(str(side_name), [])]
                unique_to_chain = all(int(v) in chain_index_set for v in side_users)
                replacement_name = str(side_name)
                if unique_to_chain:
                    side_tensor.data = transposed_const
                    shape, signature = normalize_onnx_shape(list(transposed_const.shape))
                    side_tensor.shape = [int(v) for v in list(shape)]
                    side_tensor.shape_signature = [int(v) for v in list(signature)]
                else:
                    replacement_name = _unique_tensor_name(f"{side_name}_nhwc")
                    shape, signature = normalize_onnx_shape(list(transposed_const.shape))
                    model_ir.tensors[replacement_name] = TensorIR(
                        name=replacement_name,
                        dtype=str(side_tensor.dtype),
                        shape=[int(v) for v in list(shape)],
                        shape_signature=[int(v) for v in list(signature)],
                        data=transposed_const,
                        is_variable=False,
                        quantization=_clone_quantization(side_tensor.quantization),
                    )
                if replacement_name != str(side_name):
                    target_op = model_ir.operators[int(op_idx)]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=target_op,
                        new_inputs=[
                            replacement_name if str(inp) == str(side_name) else str(inp)
                            for inp in list(target_op.inputs)
                        ],
                    )

            # Chain tail now produces post-transpose output name directly.
            if terminal_mode == "inverse_post_transpose":
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=chain_ops[-1],
                    new_outputs=[post_output_name],
                )

            # Update metadata to NHWC-side layout.
            for out_name in chain_output_names[:-1]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_post_expected,
                )
            if terminal_mode == "inverse_post_transpose":
                old_last_tensor = model_ir.tensors.get(old_last_name, None)
                post_output_tensor = model_ir.tensors.get(post_output_name, None)
                if post_output_tensor is not None and old_last_tensor is not None:
                    post_output_tensor.dtype = str(old_last_tensor.dtype)
                    post_output_tensor.quantization = _clone_quantization(old_last_tensor.quantization)
                    _permute_tensor_metadata_if_rank_matches(
                        post_output_tensor,
                        perm_post_expected,
                    )
            else:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(post_output_name, None),
                    perm_post_expected,
                )

            # Remove boundary transposes.
            remove_set = {int(pre_idx)}
            if terminal_mode == "inverse_post_transpose" and post_idx is not None:
                remove_set.add(int(post_idx))
            remove_indices = sorted(list(remove_set), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_leading_input_transpose_passthrough_chains": int(rewritten)}


def _optimize_boundary_input_transpose_mul_sum_reshape_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove boundary NHWC->NCHW input transpose for a strict normalization prelude.

    Target:
      input_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --MUL(const[1,C,1,1]|scalar)--> y_nchw
      y_nchw --SUM(axis includes channel=1, keepDims=True)--> z_nchw
      z_nchw --RESHAPE([N,H,W,C'])--> z_nhwc

    Rewrite:
      input_nhwc --MUL(const_nhwc|scalar)--> y_nhwc
      y_nhwc --SUM(mapped_axes, keepDims=True)--> z_nhwc
      z_nhwc --RESHAPE(...)--> z_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _normalize_axis(axis: int, rank: int) -> Optional[int]:
        try:
            axis_int = int(axis)
        except Exception:
            return None
        if axis_int < 0:
            axis_int += int(rank)
        if axis_int < 0 or axis_int >= int(rank):
            return None
        return int(axis_int)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(str(v) for v in model_ir.inputs)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
            ):
                continue

            input_name = str(pre_op.inputs[0])
            internal_name = str(pre_op.outputs[0])
            if input_name not in model_inputs:
                continue
            if input_name in model_outputs or internal_name in model_outputs:
                continue
            if not str(internal_name).endswith("_onnx_ncx_internal"):
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            internal_users = [int(v) for v in consumers.get(internal_name, [])]
            if len(internal_users) != 1:
                continue
            mul_idx = int(internal_users[0])
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            if str(mul_inputs[0]) == internal_name:
                mul_data_input_index = 0
                mul_const_name = str(mul_inputs[1])
            elif str(mul_inputs[1]) == internal_name:
                mul_data_input_index = 1
                mul_const_name = str(mul_inputs[0])
            else:
                continue

            mul_const_tensor = model_ir.tensors.get(str(mul_const_name), None)
            if mul_const_tensor is None or mul_const_tensor.data is None:
                continue
            mul_const_array = np.asarray(mul_const_tensor.data)
            if int(mul_const_array.size) != 1:
                if mul_const_array.ndim != 4:
                    continue
                const_shape = [int(v) for v in list(mul_const_array.shape)]
                if not (
                    int(const_shape[0]) == 1
                    and int(const_shape[2]) == 1
                    and int(const_shape[3]) == 1
                ):
                    continue

            mul_out_name = str(mul_op.outputs[0])
            if mul_out_name in model_outputs:
                continue
            mul_out_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_out_users) != 1:
                continue
            sum_idx = int(mul_out_users[0])
            sum_op = model_ir.operators[int(sum_idx)]
            if str(sum_op.op_type) != "SUM" or len(sum_op.inputs) < 2 or len(sum_op.outputs) != 1:
                continue
            if str(sum_op.inputs[0]) != mul_out_name:
                continue
            if not bool(sum_op.options.get("keepDims", False)):
                continue
            sum_axes_tensor = model_ir.tensors.get(str(sum_op.inputs[1]), None)
            sum_axes_vals = _read_const_ints_from_tensor(sum_axes_tensor)
            if sum_axes_vals is None or len(sum_axes_vals) == 0:
                continue

            sum_out_name = str(sum_op.outputs[0])
            if sum_out_name in model_outputs:
                continue
            sum_out_users = [int(v) for v in consumers.get(sum_out_name, [])]
            if len(sum_out_users) != 1:
                continue
            reshape_idx = int(sum_out_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != sum_out_name
            ):
                continue

            mapped_axes: List[int] = []
            axes_ok = True
            includes_channel_axis = False
            for axis in list(sum_axes_vals):
                norm_axis = _normalize_axis(int(axis), 4)
                if norm_axis is None:
                    axes_ok = False
                    break
                if int(norm_axis) == 1:
                    includes_channel_axis = True
                mapped_axes.append(int(perm_nhwc_to_nchw[int(norm_axis)]))
            if not axes_ok or not includes_channel_axis:
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=int(mul_data_input_index),
                new_input_name=str(input_name),
            )

            if int(mul_const_array.size) != 1:
                transposed_const = np.transpose(mul_const_array, axes=perm_nchw_to_nhwc)
                mul_const_tensor.data = np.asarray(transposed_const)
                const_shape, const_signature = normalize_onnx_shape(list(transposed_const.shape))
                mul_const_tensor.shape = [int(v) for v in list(const_shape)]
                mul_const_tensor.shape_signature = [int(v) for v in list(const_signature)]

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(sum_out_name), None),
                perm_nchw_to_nhwc,
            )

            _write_const_ints_to_tensor(sum_axes_tensor, [int(v) for v in list(mapped_axes)])
            if isinstance(sum_op.options, dict):
                sum_opts = dict(sum_op.options)
                for key in ["axis", "axes", "onnxRawAxes"]:
                    value = sum_opts.get(key, None)
                    if isinstance(value, list) and len(value) == len(mapped_axes):
                        sum_opts[key] = [int(v) for v in list(mapped_axes)]
                sum_op.options = sum_opts

            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_boundary_input_transpose_mul_sum_reshape_nhwc_chains": int(rewritten)}


def _optimize_boundary_input_transpose_batchmatmul_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Elide input-boundary transpose wrappers that only feed BATCH_MATMUL.

    Target:
      X_in(external layout) --TRANSPOSE(external->internal)--> X_internal
      BATCH_MATMUL(..., X_internal, ...)

    Rewrite:
      BATCH_MATMUL(..., X_in, ...)

    Safety:
    - Leading transpose input must be a model input tensor.
    - Model input tensor must be consumed only by that transpose.
    - Transpose output must be consumed only by BATCH_MATMUL ops.
    - Transpose must be one of NCX boundary permutations inserted by input
      layout adaptation (rank 3/4/5).
    """
    rewritten = 0
    ncx_boundary_perms = {
        (0, 2, 1),
        (0, 3, 1, 2),
        (0, 4, 1, 2, 3),
    }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(str(v) for v in model_ir.inputs)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
            ):
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name not in model_inputs:
                continue
            if pre_output_name in model_outputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            if tuple(int(v) for v in list(perm_pre)) not in ncx_boundary_perms:
                continue

            pre_input_users = [int(v) for v in consumers.get(pre_input_name, [])]
            if set(pre_input_users) != {int(pre_idx)}:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue
            if not all(
                str(model_ir.operators[int(user_idx)].op_type) == "BATCH_MATMUL"
                for user_idx in pre_users
            ):
                continue

            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=user_op,
                    new_inputs=[
                        pre_input_name if str(v) == pre_output_name else str(v)
                        for v in list(user_op.inputs)
                    ],
                )

            input_tensor = model_ir.tensors.get(pre_input_name, None)
            internal_tensor = model_ir.tensors.get(pre_output_name, None)
            if input_tensor is not None and internal_tensor is not None:
                if internal_tensor.shape is not None:
                    input_tensor.shape = [int(v) for v in list(internal_tensor.shape)]
                if internal_tensor.shape_signature is not None:
                    input_tensor.shape_signature = [
                        int(v) for v in list(internal_tensor.shape_signature)
                    ]

            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_boundary_input_transpose_batchmatmul_chains": int(rewritten)}


def _optimize_asin_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold leading input TRANSPOSE through Asin/Acos decomposition chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --MUL(x_t, x_t)--> x2
      c1 --SUB(x2)--> one_minus
      one_minus --SQRT--> denom
      ATAN2(x_t, denom) or ATAN2(denom, x_t) --> y_t
      (optional) y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --MUL(X, X)--> x2
      c1 --SUB(x2)--> one_minus
      one_minus --SQRT--> denom
      ATAN2(X, denom) or ATAN2(denom, X) --> Y

    Safety:
    - Leading transpose input must be a model input tensor.
    - Strict local topology and single-consumer inner chain.
    - SUB constant side must be scalar-like.
    - Terminal must be graph output or one inverse post-transpose.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(str(v) for v in model_ir.inputs)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            if pre_input_name not in model_inputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = sorted({int(v) for v in consumers.get(pre_output_name, [])})
            if len(pre_users) != 2:
                continue

            mul_idx: Optional[int] = None
            atan2_idx: Optional[int] = None
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if user_type == "MUL":
                    mul_idx = int(user_idx)
                elif user_type == "ATAN2":
                    atan2_idx = int(user_idx)
            if mul_idx is None or atan2_idx is None:
                continue

            mul_op = model_ir.operators[int(mul_idx)]
            if len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            if not (mul_inputs[0] == pre_output_name and mul_inputs[1] == pre_output_name):
                continue
            mul_out_name = str(mul_op.outputs[0])
            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_users) != 1:
                continue

            sub_idx = int(mul_users[0])
            sub_op = model_ir.operators[int(sub_idx)]
            if str(sub_op.op_type) != "SUB" or len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue
            sub_inputs = [str(v) for v in list(sub_op.inputs)]
            if sub_inputs[1] != mul_out_name:
                continue
            sub_const_name = str(sub_inputs[0])
            if not _is_singleton_constant_tensor(model_ir, sub_const_name):
                continue
            sub_out_name = str(sub_op.outputs[0])
            sub_users = [int(v) for v in consumers.get(sub_out_name, [])]
            if len(sub_users) != 1:
                continue

            sqrt_idx = int(sub_users[0])
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if str(sqrt_op.op_type) != "SQRT" or len(sqrt_op.inputs) != 1 or len(sqrt_op.outputs) != 1:
                continue
            if str(sqrt_op.inputs[0]) != sub_out_name:
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])
            sqrt_users = [int(v) for v in consumers.get(sqrt_out_name, [])]
            if sqrt_users != [int(atan2_idx)]:
                continue

            atan2_op = model_ir.operators[int(atan2_idx)]
            if len(atan2_op.inputs) != 2 or len(atan2_op.outputs) != 1:
                continue
            atan2_inputs = [str(v) for v in list(atan2_op.inputs)]
            if set(atan2_inputs) != {str(pre_output_name), str(sqrt_out_name)}:
                continue
            atan2_out_name = str(atan2_op.outputs[0])
            if atan2_out_name in model_inputs:
                continue

            terminal_mode = ""
            post_idx: Optional[int] = None
            post_output_name = str(atan2_out_name)
            atan2_out_users = [int(v) for v in consumers.get(atan2_out_name, [])]
            if atan2_out_name in model_outputs and len(atan2_out_users) == 0:
                terminal_mode = "model_output"
            elif len(atan2_out_users) == 1:
                candidate_post_idx = int(atan2_out_users[0])
                candidate_post_op = model_ir.operators[int(candidate_post_idx)]
                if (
                    str(candidate_post_op.op_type) == "TRANSPOSE"
                    and len(candidate_post_op.inputs) >= 2
                    and len(candidate_post_op.outputs) == 1
                    and str(candidate_post_op.inputs[0]) == str(atan2_out_name)
                ):
                    perm_post = _read_transpose_perm(model_ir, candidate_post_op)
                    if perm_post is not None and list(perm_post) == list(perm_post_expected):
                        terminal_mode = "inverse_post_transpose"
                        post_idx = int(candidate_post_idx)
                        post_output_name = str(candidate_post_op.outputs[0])
            if terminal_mode == "":
                continue

            # Rewire chain heads from transposed input to external input.
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=[str(pre_input_name), str(pre_input_name)],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=atan2_op,
                new_inputs=[
                    str(pre_input_name) if str(v) == str(pre_output_name) else str(v)
                    for v in list(atan2_op.inputs)
                ],
            )

            if terminal_mode == "inverse_post_transpose":
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=atan2_op,
                    new_outputs=[str(post_output_name)],
                )

            # Update metadata to external layout.
            for name in [str(mul_out_name), str(sub_out_name), str(sqrt_out_name)]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(name, None),
                    perm_post_expected,
                )
            if terminal_mode == "inverse_post_transpose":
                atan2_out_tensor = model_ir.tensors.get(str(atan2_out_name), None)
                post_output_tensor = model_ir.tensors.get(str(post_output_name), None)
                if post_output_tensor is not None and atan2_out_tensor is not None:
                    post_output_tensor.dtype = str(atan2_out_tensor.dtype)
                    post_output_tensor.quantization = _clone_quantization(atan2_out_tensor.quantization)
                    _permute_tensor_metadata_if_rank_matches(
                        post_output_tensor,
                        perm_post_expected,
                    )
            else:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(post_output_name), None),
                    perm_post_expected,
                )

            remove_indices = {int(pre_idx)}
            if post_idx is not None:
                remove_indices.add(int(post_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_asin_transpose_passthrough_chains": int(rewritten)}


def _optimize_erf_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold leading input TRANSPOSE through Erf polynomial decomposition chains.

    Target (strict local topology):
      X --TRANSPOSE(P)--> x_t
      x_t --ABS--> a
      x_t --SIGN--> s
      a --MUL(p)--> px --ADD(1)--> one_plus --DIV(1)--> t
      a --MUL(a)--> a2 --MUL(-1)--> na2 --EXP--> e
      Horner chain with t:
        m1=MUL(a5,t); a1=ADD(m1,a4); m2=MUL(a1,t); a2=ADD(m2,a3);
        m3=MUL(a2,t); a3=ADD(m3,a2); m4=MUL(a3,t); a4=ADD(m4,a1);
        poly=MUL(a4,t)
      poly --MUL(e)--> pe --SUB(1)--> om
      MUL(s, om) --> y_t
      (optional) y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      Remove boundary transpose(s) and feed chain directly from X.
    """
    rewritten = 0

    def _single_user(
        *,
        tensor_name: str,
        consumers: Dict[str, List[int]],
    ) -> Optional[int]:
        users = sorted({int(v) for v in consumers.get(str(tensor_name), [])})
        if len(users) != 1:
            return None
        return int(users[0])

    def _is_scalar_const_tensor(name: str) -> bool:
        return _is_singleton_constant_tensor(model_ir, str(name))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(str(v) for v in model_ir.inputs)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            pre_input_name = str(pre_op.inputs[0])
            if pre_input_name not in model_inputs:
                continue
            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = sorted({int(v) for v in consumers.get(pre_output_name, [])})
            if len(pre_users) != 2:
                continue
            abs_idx: Optional[int] = None
            sign_idx: Optional[int] = None
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if user_type == "ABS":
                    abs_idx = int(user_idx)
                elif user_type == "SIGN":
                    sign_idx = int(user_idx)
            if abs_idx is None or sign_idx is None:
                continue
            abs_op = model_ir.operators[int(abs_idx)]
            sign_op = model_ir.operators[int(sign_idx)]
            if len(abs_op.inputs) != 1 or len(abs_op.outputs) != 1:
                continue
            if len(sign_op.inputs) != 1 or len(sign_op.outputs) != 1:
                continue
            if str(abs_op.inputs[0]) != pre_output_name or str(sign_op.inputs[0]) != pre_output_name:
                continue
            abs_out = str(abs_op.outputs[0])
            sign_out = str(sign_op.outputs[0])

            abs_users = sorted({int(v) for v in consumers.get(abs_out, [])})
            if len(abs_users) != 2:
                continue
            mul_px_idx: Optional[int] = None
            mul_sq_idx: Optional[int] = None
            for user_idx in abs_users:
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                if in0 == abs_out and in1 == abs_out:
                    mul_sq_idx = int(user_idx)
                    continue
                if abs_out in {in0, in1}:
                    side_name = in1 if in0 == abs_out else in0
                    if _is_scalar_const_tensor(side_name):
                        mul_px_idx = int(user_idx)
            if mul_px_idx is None or mul_sq_idx is None:
                continue

            mul_px_op = model_ir.operators[int(mul_px_idx)]
            mul_sq_op = model_ir.operators[int(mul_sq_idx)]
            mul_px_out = str(mul_px_op.outputs[0])
            mul_sq_out = str(mul_sq_op.outputs[0])

            add_px_idx = _single_user(tensor_name=mul_px_out, consumers=consumers)
            if add_px_idx is None:
                continue
            add_px_op = model_ir.operators[int(add_px_idx)]
            if str(add_px_op.op_type) != "ADD" or len(add_px_op.inputs) != 2 or len(add_px_op.outputs) != 1:
                continue
            if mul_px_out not in {str(v) for v in list(add_px_op.inputs)}:
                continue
            add_px_side = str(add_px_op.inputs[1]) if str(add_px_op.inputs[0]) == mul_px_out else str(add_px_op.inputs[0])
            if not _is_scalar_const_tensor(add_px_side):
                continue
            add_px_out = str(add_px_op.outputs[0])

            div_t_idx = _single_user(tensor_name=add_px_out, consumers=consumers)
            if div_t_idx is None:
                continue
            div_t_op = model_ir.operators[int(div_t_idx)]
            if str(div_t_op.op_type) != "DIV" or len(div_t_op.inputs) != 2 or len(div_t_op.outputs) != 1:
                continue
            div_t_inputs = [str(v) for v in list(div_t_op.inputs)]
            if add_px_out == div_t_inputs[0]:
                div_t_side = div_t_inputs[1]
            elif add_px_out == div_t_inputs[1]:
                div_t_side = div_t_inputs[0]
            else:
                continue
            if not _is_scalar_const_tensor(div_t_side):
                continue
            t_out = str(div_t_op.outputs[0])

            mul_neg_idx = _single_user(tensor_name=mul_sq_out, consumers=consumers)
            if mul_neg_idx is None:
                continue
            mul_neg_op = model_ir.operators[int(mul_neg_idx)]
            if str(mul_neg_op.op_type) != "MUL" or len(mul_neg_op.inputs) != 2 or len(mul_neg_op.outputs) != 1:
                continue
            if mul_sq_out not in {str(v) for v in list(mul_neg_op.inputs)}:
                continue
            mul_neg_side = (
                str(mul_neg_op.inputs[1])
                if str(mul_neg_op.inputs[0]) == mul_sq_out
                else str(mul_neg_op.inputs[0])
            )
            if not _is_scalar_const_tensor(mul_neg_side):
                continue
            mul_neg_out = str(mul_neg_op.outputs[0])

            exp_idx = _single_user(tensor_name=mul_neg_out, consumers=consumers)
            if exp_idx is None:
                continue
            exp_op = model_ir.operators[int(exp_idx)]
            if str(exp_op.op_type) != "EXP" or len(exp_op.inputs) != 1 or len(exp_op.outputs) != 1:
                continue
            if str(exp_op.inputs[0]) != mul_neg_out:
                continue
            exp_out = str(exp_op.outputs[0])

            t_users = sorted({int(v) for v in consumers.get(t_out, [])})
            if len(t_users) != 5:
                continue

            s1_mul_idx: Optional[int] = None
            for user_idx in t_users:
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                side_name: Optional[str] = None
                if in0 == t_out:
                    side_name = in1
                elif in1 == t_out:
                    side_name = in0
                if side_name is None:
                    continue
                if _is_scalar_const_tensor(side_name):
                    s1_mul_idx = int(user_idx)
                    break
            if s1_mul_idx is None:
                continue

            chain_indices: List[int] = []
            chain_ops: List[OperatorIR] = []
            s1_mul_op = model_ir.operators[int(s1_mul_idx)]
            s1_mul_out = str(s1_mul_op.outputs[0])
            chain_indices.append(int(s1_mul_idx))
            chain_ops.append(s1_mul_op)
            current = str(s1_mul_out)

            chain_ok = True
            for _ in range(4):
                add_idx = _single_user(tensor_name=current, consumers=consumers)
                if add_idx is None:
                    chain_ok = False
                    break
                add_op = model_ir.operators[int(add_idx)]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    chain_ok = False
                    break
                add_inputs = [str(v) for v in list(add_op.inputs)]
                if current == add_inputs[0]:
                    add_side = add_inputs[1]
                elif current == add_inputs[1]:
                    add_side = add_inputs[0]
                else:
                    chain_ok = False
                    break
                if not _is_scalar_const_tensor(add_side):
                    chain_ok = False
                    break
                chain_indices.append(int(add_idx))
                chain_ops.append(add_op)
                add_out = str(add_op.outputs[0])

                mul_idx = _single_user(tensor_name=add_out, consumers=consumers)
                if mul_idx is None:
                    chain_ok = False
                    break
                mul_op = model_ir.operators[int(mul_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    chain_ok = False
                    break
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                if add_out == mul_inputs[0]:
                    mul_side = mul_inputs[1]
                elif add_out == mul_inputs[1]:
                    mul_side = mul_inputs[0]
                else:
                    chain_ok = False
                    break
                if str(mul_side) != str(t_out):
                    chain_ok = False
                    break
                chain_indices.append(int(mul_idx))
                chain_ops.append(mul_op)
                current = str(mul_op.outputs[0])
            if not chain_ok:
                continue
            poly_out = str(current)

            mul_poly_exp_idx = _single_user(tensor_name=poly_out, consumers=consumers)
            if mul_poly_exp_idx is None:
                continue
            mul_poly_exp_op = model_ir.operators[int(mul_poly_exp_idx)]
            if str(mul_poly_exp_op.op_type) != "MUL" or len(mul_poly_exp_op.inputs) != 2 or len(mul_poly_exp_op.outputs) != 1:
                continue
            mpex_inputs = [str(v) for v in list(mul_poly_exp_op.inputs)]
            if set(mpex_inputs) != {str(poly_out), str(exp_out)}:
                continue
            mul_poly_exp_out = str(mul_poly_exp_op.outputs[0])

            if sorted({int(v) for v in consumers.get(exp_out, [])}) != [int(mul_poly_exp_idx)]:
                continue

            sub_idx = _single_user(tensor_name=mul_poly_exp_out, consumers=consumers)
            if sub_idx is None:
                continue
            sub_op = model_ir.operators[int(sub_idx)]
            if str(sub_op.op_type) != "SUB" or len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue
            sub_inputs = [str(v) for v in list(sub_op.inputs)]
            if mul_poly_exp_out == sub_inputs[0]:
                sub_side = sub_inputs[1]
            elif mul_poly_exp_out == sub_inputs[1]:
                sub_side = sub_inputs[0]
            else:
                continue
            if not _is_scalar_const_tensor(sub_side):
                continue
            sub_out = str(sub_op.outputs[0])

            final_mul_idx = _single_user(tensor_name=sub_out, consumers=consumers)
            if final_mul_idx is None:
                continue
            final_mul_op = model_ir.operators[int(final_mul_idx)]
            if str(final_mul_op.op_type) != "MUL" or len(final_mul_op.inputs) != 2 or len(final_mul_op.outputs) != 1:
                continue
            final_inputs = [str(v) for v in list(final_mul_op.inputs)]
            if set(final_inputs) != {str(sign_out), str(sub_out)}:
                continue
            if sorted({int(v) for v in consumers.get(sign_out, [])}) != [int(final_mul_idx)]:
                continue
            final_out = str(final_mul_op.outputs[0])

            terminal_mode = ""
            post_idx: Optional[int] = None
            post_output_name = str(final_out)
            final_users = sorted({int(v) for v in consumers.get(final_out, [])})
            if final_out in model_outputs and len(final_users) == 0:
                terminal_mode = "model_output"
            elif len(final_users) == 1:
                cand_idx = int(final_users[0])
                cand_op = model_ir.operators[int(cand_idx)]
                if (
                    str(cand_op.op_type) == "TRANSPOSE"
                    and len(cand_op.inputs) >= 2
                    and len(cand_op.outputs) == 1
                    and str(cand_op.inputs[0]) == str(final_out)
                ):
                    perm_post = _read_transpose_perm(model_ir, cand_op)
                    if perm_post is not None and list(perm_post) == list(perm_post_expected):
                        terminal_mode = "inverse_post_transpose"
                        post_idx = int(cand_idx)
                        post_output_name = str(cand_op.outputs[0])
            if terminal_mode == "":
                continue

            # Rewire chain head from transposed tensor to model input tensor.
            _set_operator_inputs(
                model_ir=model_ir,
                op=abs_op,
                new_inputs=[str(pre_input_name)],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=sign_op,
                new_inputs=[str(pre_input_name)],
            )
            if terminal_mode == "inverse_post_transpose":
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=final_mul_op,
                    new_outputs=[str(post_output_name)],
                )

            # Update tensor metadata from transposed layout to source layout.
            chain_all_ops: List[OperatorIR] = [
                abs_op,
                sign_op,
                mul_px_op,
                add_px_op,
                div_t_op,
                mul_sq_op,
                mul_neg_op,
                exp_op,
            ] + chain_ops + [mul_poly_exp_op, sub_op, final_mul_op]
            seen_names: set = set()
            for op in chain_all_ops:
                if len(op.outputs) != 1:
                    continue
                out_name = str(op.outputs[0])
                if out_name in seen_names:
                    continue
                seen_names.add(out_name)
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_post_expected,
                )

            remove_indices = {int(pre_idx)}
            if post_idx is not None:
                remove_indices.add(int(post_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_erf_transpose_passthrough_chains": int(rewritten)}


def _optimize_hardswish_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around pseudo-op-expanded HardSwish-like chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --ADD(c3)--> a --RELU6--> r --(DIV(c6)|MUL(c6_inv))--> d
        or
      x_t --ADD(c3)--> a --MUL(c6_inv)--> d
      MUL(x_t, d) --> y_t
      y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --ADD(c3)--> a --RELU6--> r --(DIV(c6)|MUL(c6_inv))--> d
        or
      X --ADD(c3)--> a --MUL(c6_inv)--> d
      MUL(X, d) --> Y

    Safety:
    - Strict residual topology (ADD -> optional RELU6 -> (DIV|MUL const) -> residual MUL).
    - ADD/(DIV|MUL) side inputs must be singleton constants.
    - Single-consumer chain on the main path.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 2:
                continue

            add_idx = None
            mul_idx = None
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if user_type == "ADD":
                    add_idx = int(user_idx)
                elif user_type == "MUL":
                    mul_idx = int(user_idx)
            if add_idx is None or mul_idx is None:
                continue

            add_op = model_ir.operators[int(add_idx)]
            if len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_inputs = [str(v) for v in add_op.inputs]
            if pre_output_name == add_inputs[0]:
                add_data_input_index = 0
                add_side_input_name = add_inputs[1]
            elif pre_output_name == add_inputs[1]:
                add_data_input_index = 1
                add_side_input_name = add_inputs[0]
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, add_side_input_name):
                continue
            add_out_name = str(add_op.outputs[0])

            add_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_users) != 1:
                continue
            stage_out_name: Optional[str] = None
            intermediates_to_permute: List[str] = [add_out_name]

            stage1_idx = int(add_users[0])
            stage1_op = model_ir.operators[int(stage1_idx)]
            stage1_type = str(stage1_op.op_type)

            if stage1_type == "RELU6":
                if len(stage1_op.inputs) != 1 or len(stage1_op.outputs) != 1:
                    continue
                if str(stage1_op.inputs[0]) != add_out_name:
                    continue
                relu_out_name = str(stage1_op.outputs[0])
                intermediates_to_permute.append(relu_out_name)

                relu_users = [int(v) for v in consumers.get(relu_out_name, [])]
                if len(relu_users) != 1:
                    continue
                stage2_idx = int(relu_users[0])
                stage2_op = model_ir.operators[int(stage2_idx)]
                stage2_type = str(stage2_op.op_type)
                if stage2_type not in {"DIV", "MUL"}:
                    continue
                if len(stage2_op.inputs) != 2 or len(stage2_op.outputs) != 1:
                    continue
                stage2_inputs = [str(v) for v in stage2_op.inputs]
                if relu_out_name == stage2_inputs[0]:
                    stage2_side_input_name = stage2_inputs[1]
                elif relu_out_name == stage2_inputs[1]:
                    stage2_side_input_name = stage2_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, stage2_side_input_name):
                    continue
                stage_out_name = str(stage2_op.outputs[0])
                intermediates_to_permute.append(stage_out_name)
            elif stage1_type in {"DIV", "MUL"}:
                if len(stage1_op.inputs) != 2 or len(stage1_op.outputs) != 1:
                    continue
                stage1_inputs = [str(v) for v in stage1_op.inputs]
                if add_out_name == stage1_inputs[0]:
                    stage1_side_input_name = stage1_inputs[1]
                elif add_out_name == stage1_inputs[1]:
                    stage1_side_input_name = stage1_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, stage1_side_input_name):
                    continue
                stage_out_name = str(stage1_op.outputs[0])
                intermediates_to_permute.append(stage_out_name)
            else:
                continue

            if stage_out_name is None:
                continue
            stage_users = [int(v) for v in consumers.get(stage_out_name, [])]
            if len(stage_users) != 1 or int(stage_users[0]) != int(mul_idx):
                continue

            mul_op = model_ir.operators[int(mul_idx)]
            if len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in mul_op.inputs]
            if pre_output_name == mul_inputs[0] and stage_out_name == mul_inputs[1]:
                mul_data_input_index = 0
            elif pre_output_name == mul_inputs[1] and stage_out_name == mul_inputs[0]:
                mul_data_input_index = 1
            else:
                continue
            mul_out_name = str(mul_op.outputs[0])

            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_users) != 1:
                continue
            post_idx = int(mul_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != mul_out_name:
                continue
            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or perm_post != perm_post_expected:
                continue
            post_output_name = str(post_op.outputs[0])

            # Rewire the pseudo-HardSwish chain to consume NHWC source directly.
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add_op,
                input_index=int(add_data_input_index),
                new_input_name=pre_input_name,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=int(mul_data_input_index),
                new_input_name=pre_input_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[post_output_name],
            )

            # Update intermediate metadata from transposed layout to source layout.
            for intermediate_name in intermediates_to_permute:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(intermediate_name, None),
                    perm_post_expected,
                )

            # Keep final tensor metadata stable on the post-transpose name.
            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            old_mul_tensor = model_ir.tensors.get(mul_out_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if post_output_tensor is not None:
                if pre_input_tensor is not None:
                    post_output_tensor.shape = [int(v) for v in list(pre_input_tensor.shape)]
                    post_output_tensor.shape_signature = (
                        [int(v) for v in list(pre_input_tensor.shape_signature)]
                        if pre_input_tensor.shape_signature is not None
                        else [int(v) for v in list(pre_input_tensor.shape)]
                    )
                if old_mul_tensor is not None:
                    post_output_tensor.dtype = str(old_mul_tensor.dtype)
                    post_output_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)

            remove_indices = sorted(list({int(pre_idx), int(post_idx)}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_hardswish_transpose_passthrough_chains": int(rewritten)}


def _optimize_hardsigmoid_mul_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around HardSigmoid-expanded residual MUL chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --MUL(alpha)--> m1 --ADD(beta)--> a1 --(MAXIMUM(0)->MINIMUM(1) | RELU_0_TO_1)--> h1
      MUL(x_t, h1) --> y_t
      y_t --TRANSPOSE(inv(P))--> Y
      (optional branch)
      y_t --MEAN(axes_nchw, keepDims=True)--> m_t --TRANSPOSE(inv(P))--> M

    Rewrite:
      X --MUL(alpha)--> m1 --ADD(beta)--> a1 --(MAXIMUM(0)->MINIMUM(1) | RELU_0_TO_1)--> h1
      MUL(X, h1) --> Y
      MEAN(Y, axes_nhwc, keepDims=True) --> M

    Safety:
    - Strict HardSigmoid decomposition topology on one branch.
    - HardSigmoid branch constants must be singleton constants.
    - Chain on the HardSigmoid branch must be linear.
    - Optional MEAN branch must be rank-4 constant-axes and only feed inverse post-transpose adapters.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 2:
                continue

            def _try_match_pair(hs_entry_idx: int, residual_mul_idx: int) -> Optional[Dict[str, Any]]:
                hs_entry_op = model_ir.operators[int(hs_entry_idx)]
                hs_entry_type = str(hs_entry_op.op_type)
                hs_data_idx: Optional[int] = None
                hs_data_input_index: Optional[int] = None
                hs_out_name: Optional[str] = None
                hs_intermediate_names: List[str] = []

                if hs_entry_type == "MUL" and len(hs_entry_op.inputs) == 2 and len(hs_entry_op.outputs) == 1:
                    hs_mul_inputs = [str(v) for v in list(hs_entry_op.inputs)]
                    if pre_output_name == hs_mul_inputs[0]:
                        hs_data_input_index = 0
                        hs_mul_side_name = hs_mul_inputs[1]
                    elif pre_output_name == hs_mul_inputs[1]:
                        hs_data_input_index = 1
                        hs_mul_side_name = hs_mul_inputs[0]
                    else:
                        return None
                    if not _is_singleton_constant_tensor(model_ir, hs_mul_side_name):
                        return None
                    hs_data_idx = int(hs_entry_idx)
                    hs_mul_out_name = str(hs_entry_op.outputs[0])

                    hs_mul_users = [int(v) for v in consumers.get(hs_mul_out_name, [])]
                    if len(hs_mul_users) != 1:
                        return None
                    hs_add_idx = int(hs_mul_users[0])
                    hs_add_op = model_ir.operators[int(hs_add_idx)]
                    if str(hs_add_op.op_type) != "ADD" or len(hs_add_op.inputs) != 2 or len(hs_add_op.outputs) != 1:
                        return None
                    hs_add_inputs = [str(v) for v in list(hs_add_op.inputs)]
                    if hs_mul_out_name == hs_add_inputs[0]:
                        hs_add_side_name = hs_add_inputs[1]
                    elif hs_mul_out_name == hs_add_inputs[1]:
                        hs_add_side_name = hs_add_inputs[0]
                    else:
                        return None
                    if not _is_singleton_constant_tensor(model_ir, hs_add_side_name):
                        return None
                    hs_add_out_name = str(hs_add_op.outputs[0])

                    hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
                    if len(hs_add_users) != 1:
                        return None
                    hs_clamp_idx = int(hs_add_users[0])
                    hs_clamp_op = model_ir.operators[int(hs_clamp_idx)]
                    hs_intermediate_names = [str(hs_mul_out_name), str(hs_add_out_name)]
                    if (
                        str(hs_clamp_op.op_type) == "RELU_0_TO_1"
                        and len(hs_clamp_op.inputs) == 1
                        and len(hs_clamp_op.outputs) == 1
                        and str(hs_clamp_op.inputs[0]) == hs_add_out_name
                    ):
                        hs_out_name = str(hs_clamp_op.outputs[0])
                    else:
                        if (
                            str(hs_clamp_op.op_type) != "MAXIMUM"
                            or len(hs_clamp_op.inputs) != 2
                            or len(hs_clamp_op.outputs) != 1
                        ):
                            return None
                        hs_max_inputs = [str(v) for v in list(hs_clamp_op.inputs)]
                        if hs_add_out_name == hs_max_inputs[0]:
                            hs_max_side_name = hs_max_inputs[1]
                        elif hs_add_out_name == hs_max_inputs[1]:
                            hs_max_side_name = hs_max_inputs[0]
                        else:
                            return None
                        if not _is_singleton_constant_tensor(model_ir, hs_max_side_name):
                            return None
                        hs_max_out_name = str(hs_clamp_op.outputs[0])
                        hs_intermediate_names.append(str(hs_max_out_name))

                        hs_max_users = [int(v) for v in consumers.get(hs_max_out_name, [])]
                        if len(hs_max_users) != 1:
                            return None
                        hs_min_idx = int(hs_max_users[0])
                        hs_min_op = model_ir.operators[int(hs_min_idx)]
                        if (
                            str(hs_min_op.op_type) != "MINIMUM"
                            or len(hs_min_op.inputs) != 2
                            or len(hs_min_op.outputs) != 1
                        ):
                            return None
                        hs_min_inputs = [str(v) for v in list(hs_min_op.inputs)]
                        if hs_max_out_name == hs_min_inputs[0]:
                            hs_min_side_name = hs_min_inputs[1]
                        elif hs_max_out_name == hs_min_inputs[1]:
                            hs_min_side_name = hs_min_inputs[0]
                        else:
                            return None
                        if not _is_singleton_constant_tensor(model_ir, hs_min_side_name):
                            return None
                        hs_out_name = str(hs_min_op.outputs[0])
                elif hs_entry_type == "ADD" and len(hs_entry_op.inputs) == 2 and len(hs_entry_op.outputs) == 1:
                    hs_add_inputs = [str(v) for v in list(hs_entry_op.inputs)]
                    if pre_output_name == hs_add_inputs[0]:
                        hs_data_input_index = 0
                        hs_add_side_name = hs_add_inputs[1]
                    elif pre_output_name == hs_add_inputs[1]:
                        hs_data_input_index = 1
                        hs_add_side_name = hs_add_inputs[0]
                    else:
                        return None
                    if not _is_singleton_constant_tensor(model_ir, hs_add_side_name):
                        return None
                    hs_data_idx = int(hs_entry_idx)
                    hs_add_out_name = str(hs_entry_op.outputs[0])
                    hs_intermediate_names.append(str(hs_add_out_name))

                    hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
                    if len(hs_add_users) != 1:
                        return None
                    hs_mul_idx = int(hs_add_users[0])
                    hs_mul_op = model_ir.operators[int(hs_mul_idx)]
                    if str(hs_mul_op.op_type) != "MUL" or len(hs_mul_op.inputs) != 2 or len(hs_mul_op.outputs) != 1:
                        return None
                    hs_mul_inputs = [str(v) for v in list(hs_mul_op.inputs)]
                    if hs_add_out_name == hs_mul_inputs[0]:
                        hs_mul_side_name = hs_mul_inputs[1]
                    elif hs_add_out_name == hs_mul_inputs[1]:
                        hs_mul_side_name = hs_mul_inputs[0]
                    else:
                        return None
                    if not _is_singleton_constant_tensor(model_ir, hs_mul_side_name):
                        return None
                    hs_out_name = str(hs_mul_op.outputs[0])
                    hs_intermediate_names.append(str(hs_out_name))
                else:
                    return None

                if hs_data_idx is None or hs_data_input_index is None or hs_out_name is None:
                    return None

                residual_mul_op = model_ir.operators[int(residual_mul_idx)]
                if (
                    str(residual_mul_op.op_type) != "MUL"
                    or len(residual_mul_op.inputs) != 2
                    or len(residual_mul_op.outputs) != 1
                ):
                    return None
                residual_mul_inputs = [str(v) for v in list(residual_mul_op.inputs)]
                if pre_output_name == residual_mul_inputs[0] and hs_out_name == residual_mul_inputs[1]:
                    residual_mul_data_input_index = 0
                elif pre_output_name == residual_mul_inputs[1] and hs_out_name == residual_mul_inputs[0]:
                    residual_mul_data_input_index = 1
                else:
                    return None
                residual_mul_out_name = str(residual_mul_op.outputs[0])
                residual_mul_users = [int(v) for v in consumers.get(residual_mul_out_name, []) if int(v) != int(residual_mul_idx)]
                if len(residual_mul_users) == 0:
                    return None

                post_indices: List[int] = []
                post_output_names: List[str] = []
                legacy_users: List[int] = []
                mean_branches: List[Dict[str, Any]] = []
                valid_users = True
                for user_idx in residual_mul_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == residual_mul_out_name
                    ):
                        perm_post = _read_transpose_perm(model_ir, user_op)
                        if perm_post is None or perm_post != perm_post_expected:
                            valid_users = False
                            break
                        post_indices.append(int(user_idx))
                        post_output_names.append(str(user_op.outputs[0]))
                    elif (
                        str(user_op.op_type) == "MEAN"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == residual_mul_out_name
                        and bool(user_op.options.get("keepDims", False))
                    ):
                        mean_idx = int(user_idx)
                        mean_axes_name = str(user_op.inputs[1])
                        mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                        mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
                        if mean_axes_vals is None or len(mean_axes_vals) == 0:
                            valid_users = False
                            break
                        mean_out_name = str(user_op.outputs[0])
                        if mean_out_name in model_outputs:
                            valid_users = False
                            break
                        mean_users = [int(v) for v in consumers.get(mean_out_name, [])]
                        if len(mean_users) == 0:
                            valid_users = False
                            break
                        # Keep legacy NCHW adapter path when MEAN branch terminates
                        # at graph output through a single inverse transpose.
                        if len(mean_users) == 1:
                            only_mean_user_op = model_ir.operators[int(mean_users[0])]
                            if (
                                str(only_mean_user_op.op_type) == "TRANSPOSE"
                                and len(only_mean_user_op.inputs) >= 2
                                and len(only_mean_user_op.outputs) == 1
                                and str(only_mean_user_op.inputs[0]) == mean_out_name
                                and _read_transpose_perm(model_ir, only_mean_user_op) == perm_post_expected
                                and str(only_mean_user_op.outputs[0]) in model_outputs
                            ):
                                legacy_users.append(int(mean_idx))
                                continue
                        mean_post_indices: List[int] = []
                        mean_post_output_names: List[str] = []
                        for mean_user_idx in mean_users:
                            mean_user_op = model_ir.operators[int(mean_user_idx)]
                            if (
                                str(mean_user_op.op_type) != "TRANSPOSE"
                                or len(mean_user_op.inputs) < 2
                                or len(mean_user_op.outputs) != 1
                                or str(mean_user_op.inputs[0]) != mean_out_name
                            ):
                                valid_users = False
                                break
                            mean_perm_post = _read_transpose_perm(model_ir, mean_user_op)
                            if mean_perm_post is None or mean_perm_post != perm_post_expected:
                                valid_users = False
                                break
                            mean_post_output_name = str(mean_user_op.outputs[0])
                            mean_post_indices.append(int(mean_user_idx))
                            mean_post_output_names.append(mean_post_output_name)
                        if not valid_users or len(mean_post_indices) == 0:
                            valid_users = False
                            break
                        mean_branches.append(
                            {
                                "mean_idx": int(mean_idx),
                                "mean_axes_name": str(mean_axes_name),
                                "mean_axes_vals": [int(v) for v in mean_axes_vals],
                                "mean_out_name": str(mean_out_name),
                                "mean_post_indices": [int(v) for v in mean_post_indices],
                                "mean_post_output_names": [str(v) for v in mean_post_output_names],
                            }
                        )
                    else:
                        legacy_users.append(int(user_idx))
                if not valid_users or len(post_indices) == 0:
                    return None

                return {
                    "hs_data_idx": int(hs_data_idx),
                    "hs_data_input_index": int(hs_data_input_index),
                    "hs_intermediate_names": [str(v) for v in hs_intermediate_names],
                    "residual_mul_idx": int(residual_mul_idx),
                    "residual_mul_data_input_index": int(residual_mul_data_input_index),
                    "residual_mul_out_name": str(residual_mul_out_name),
                    "post_indices": [int(v) for v in post_indices],
                    "post_output_names": [str(v) for v in post_output_names],
                    "legacy_users": [int(v) for v in legacy_users],
                    "mean_branches": list(mean_branches),
                }

            matched = _try_match_pair(int(pre_users[0]), int(pre_users[1]))
            if matched is None:
                matched = _try_match_pair(int(pre_users[1]), int(pre_users[0]))
            if matched is None:
                continue

            hs_data_op = model_ir.operators[int(matched["hs_data_idx"])]
            residual_mul_op = model_ir.operators[int(matched["residual_mul_idx"])]
            residual_mul_out_name = str(matched["residual_mul_out_name"])
            post_indices = [int(v) for v in list(matched.get("post_indices", []))]
            post_output_names = [str(v) for v in list(matched.get("post_output_names", []))]
            legacy_users = [int(v) for v in list(matched.get("legacy_users", []))]
            mean_branches = list(matched.get("mean_branches", []))
            if len(post_indices) == 0 or len(post_output_names) == 0:
                continue
            representative_output_name = str(post_output_names[0])

            _replace_operator_input_at(
                model_ir=model_ir,
                op=hs_data_op,
                input_index=int(matched["hs_data_input_index"]),
                new_input_name=pre_input_name,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=residual_mul_op,
                input_index=int(matched["residual_mul_data_input_index"]),
                new_input_name=pre_input_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=residual_mul_op,
                new_outputs=[representative_output_name],
            )
            for alias_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_output_name, representative_output_name)

            for intermediate_name in [str(v) for v in list(matched.get("hs_intermediate_names", []))]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(intermediate_name, None),
                    perm_post_expected,
                )

            old_mul_tensor = model_ir.tensors.get(residual_mul_out_name, None)
            representative_tensor = model_ir.tensors.get(representative_output_name, None)
            if representative_tensor is not None and old_mul_tensor is not None:
                representative_tensor.dtype = str(old_mul_tensor.dtype)
                representative_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                _permute_tensor_metadata_if_rank_matches(
                    representative_tensor,
                    perm_post_expected,
                )

            mean_post_remove_indices: List[int] = []
            mean_rewrite_ok = True
            for mean_branch in mean_branches:
                mean_idx = int(mean_branch["mean_idx"])
                mean_axes_name = str(mean_branch["mean_axes_name"])
                mean_axes_vals = [int(v) for v in list(mean_branch["mean_axes_vals"])]
                mean_out_name = str(mean_branch["mean_out_name"])
                mean_post_indices = [int(v) for v in list(mean_branch["mean_post_indices"])]
                mean_post_output_names = [str(v) for v in list(mean_branch["mean_post_output_names"])]
                if len(mean_post_indices) == 0 or len(mean_post_output_names) == 0:
                    mean_rewrite_ok = False
                    break

                mapped_axes: List[int] = []
                valid_axes = True
                for axis in mean_axes_vals:
                    a = int(axis)
                    if a < 0:
                        a += 4
                    if a < 0 or a >= 4:
                        valid_axes = False
                        break
                    mapped_axes.append(int(perm_pre[int(a)]))
                if not valid_axes:
                    mean_rewrite_ok = False
                    break

                mean_op = model_ir.operators[int(mean_idx)]
                mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                if mean_axes_tensor is None or not _write_const_ints_to_tensor(mean_axes_tensor, mapped_axes):
                    mean_rewrite_ok = False
                    break

                mean_inputs = [str(v) for v in list(mean_op.inputs)]
                if len(mean_inputs) == 0:
                    mean_rewrite_ok = False
                    break
                mean_inputs[0] = representative_output_name
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mean_op,
                    new_inputs=mean_inputs,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(mean_out_name, None),
                    perm_post_expected,
                )

                representative_mean_output_name = str(mean_post_output_names[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=mean_op,
                    new_outputs=[representative_mean_output_name],
                )
                for alias_mean_output_name in mean_post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, alias_mean_output_name, representative_mean_output_name)

                old_mean_tensor = model_ir.tensors.get(mean_out_name, None)
                representative_mean_tensor = model_ir.tensors.get(representative_mean_output_name, None)
                if representative_mean_tensor is not None and old_mean_tensor is not None:
                    representative_mean_tensor.dtype = str(old_mean_tensor.dtype)
                    representative_mean_tensor.quantization = _clone_quantization(old_mean_tensor.quantization)
                    _permute_tensor_metadata_if_rank_matches(
                        representative_mean_tensor,
                        perm_post_expected,
                    )

                mean_post_remove_indices.extend(int(v) for v in mean_post_indices)

            if not mean_rewrite_ok:
                continue

            preserve_nchw_adapter = len(legacy_users) > 0 or residual_mul_out_name in model_ir.outputs
            if preserve_nchw_adapter:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[keep_post_idx]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_pre, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[representative_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[residual_mul_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = sorted(
                list({int(pre_idx), *post_remove_indices, *[int(v) for v in mean_post_remove_indices]}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_hardsigmoid_mul_transpose_passthrough_chains": int(rewritten)}


def _optimize_hardsigmoid_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around standalone HardSigmoid-expanded chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --MUL(alpha)--> m1 --ADD(beta)--> a1 --(MAXIMUM(0)->MINIMUM(1) | RELU_0_TO_1)--> h1
      h1 --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --MUL(alpha)--> m1 --ADD(beta)--> a1 --(MAXIMUM(0)->MINIMUM(1) | RELU_0_TO_1)--> Y

    Safety:
    - Strict HardSigmoid decomposition topology.
    - All side inputs must be singleton constants.
    - Main branch must be strictly linear.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            hs_mul_idx = int(pre_users[0])
            hs_mul_op = model_ir.operators[int(hs_mul_idx)]
            if str(hs_mul_op.op_type) != "MUL" or len(hs_mul_op.inputs) != 2 or len(hs_mul_op.outputs) != 1:
                continue

            hs_mul_inputs = [str(v) for v in list(hs_mul_op.inputs)]
            if pre_output_name == hs_mul_inputs[0]:
                hs_mul_data_input_index = 0
                hs_mul_side_name = hs_mul_inputs[1]
            elif pre_output_name == hs_mul_inputs[1]:
                hs_mul_data_input_index = 1
                hs_mul_side_name = hs_mul_inputs[0]
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, hs_mul_side_name):
                continue
            hs_mul_out_name = str(hs_mul_op.outputs[0])

            hs_mul_users = [int(v) for v in consumers.get(hs_mul_out_name, [])]
            if len(hs_mul_users) != 1:
                continue
            hs_add_idx = int(hs_mul_users[0])
            hs_add_op = model_ir.operators[int(hs_add_idx)]
            if str(hs_add_op.op_type) != "ADD" or len(hs_add_op.inputs) != 2 or len(hs_add_op.outputs) != 1:
                continue
            hs_add_inputs = [str(v) for v in list(hs_add_op.inputs)]
            if hs_mul_out_name == hs_add_inputs[0]:
                hs_add_side_name = hs_add_inputs[1]
            elif hs_mul_out_name == hs_add_inputs[1]:
                hs_add_side_name = hs_add_inputs[0]
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, hs_add_side_name):
                continue
            hs_add_out_name = str(hs_add_op.outputs[0])

            hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
            if len(hs_add_users) != 1:
                continue
            hs_clamp_idx = int(hs_add_users[0])
            hs_clamp_op = model_ir.operators[int(hs_clamp_idx)]
            hs_intermediate_names: List[str] = [str(hs_mul_out_name), str(hs_add_out_name)]
            if (
                str(hs_clamp_op.op_type) == "RELU_0_TO_1"
                and len(hs_clamp_op.inputs) == 1
                and len(hs_clamp_op.outputs) == 1
                and str(hs_clamp_op.inputs[0]) == hs_add_out_name
            ):
                hs_terminal_op = hs_clamp_op
                hs_out_name = str(hs_clamp_op.outputs[0])
            else:
                if str(hs_clamp_op.op_type) != "MAXIMUM" or len(hs_clamp_op.inputs) != 2 or len(hs_clamp_op.outputs) != 1:
                    continue
                hs_max_inputs = [str(v) for v in list(hs_clamp_op.inputs)]
                if hs_add_out_name == hs_max_inputs[0]:
                    hs_max_side_name = hs_max_inputs[1]
                elif hs_add_out_name == hs_max_inputs[1]:
                    hs_max_side_name = hs_max_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_max_side_name):
                    continue
                hs_max_out_name = str(hs_clamp_op.outputs[0])
                hs_intermediate_names.append(str(hs_max_out_name))

                hs_max_users = [int(v) for v in consumers.get(hs_max_out_name, [])]
                if len(hs_max_users) != 1:
                    continue
                hs_min_idx = int(hs_max_users[0])
                hs_min_op = model_ir.operators[int(hs_min_idx)]
                if str(hs_min_op.op_type) != "MINIMUM" or len(hs_min_op.inputs) != 2 or len(hs_min_op.outputs) != 1:
                    continue
                hs_min_inputs = [str(v) for v in list(hs_min_op.inputs)]
                if hs_max_out_name == hs_min_inputs[0]:
                    hs_min_side_name = hs_min_inputs[1]
                elif hs_max_out_name == hs_min_inputs[1]:
                    hs_min_side_name = hs_min_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_min_side_name):
                    continue
                hs_terminal_op = hs_min_op
                hs_out_name = str(hs_min_op.outputs[0])

            hs_out_users = [int(v) for v in consumers.get(hs_out_name, [])]
            if len(hs_out_users) != 1:
                continue
            post_idx = int(hs_out_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != hs_out_name
            ):
                continue
            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or perm_post != perm_post_expected:
                continue
            post_output_name = str(post_op.outputs[0])

            _replace_operator_input_at(
                model_ir=model_ir,
                op=hs_mul_op,
                input_index=int(hs_mul_data_input_index),
                new_input_name=pre_input_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=hs_terminal_op,
                new_outputs=[post_output_name],
            )

            for intermediate_name in hs_intermediate_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(intermediate_name, None),
                    perm_post_expected,
                )

            old_hs_tensor = model_ir.tensors.get(hs_out_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if post_output_tensor is not None and old_hs_tensor is not None:
                post_output_tensor.dtype = str(old_hs_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_hs_tensor.quantization)
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_post_expected,
                )

            remove_indices = sorted(list({int(pre_idx), int(post_idx)}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_hardsigmoid_transpose_passthrough_chains": int(rewritten)}


def _optimize_transpose_hardswish_se_conv_hardsigmoid_mul_prepost_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NCHW bridge transposes in activation + SE-conv + HardSigmoid gating blocks.

    Supported activation roots:
    - HARD_SWISH
    - RELU
    - RELU6
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            hsw_out_name: Optional[str] = None
            hsw_input_rewrites: List[Tuple[OperatorIR, int]] = []
            hsw_root_intermediate_names: List[str] = []

            # Pattern 1: direct activation root
            #   pre_output -> {HARD_SWISH|RELU|RELU6} -> hsw_out
            if len(pre_users) == 1:
                hsw_idx = int(pre_users[0])
                hsw_op = model_ir.operators[int(hsw_idx)]
                if (
                    str(hsw_op.op_type) in {"HARD_SWISH", "RELU", "RELU6"}
                    and len(hsw_op.inputs) == 1
                    and len(hsw_op.outputs) == 1
                    and str(hsw_op.inputs[0]) == pre_output_name
                ):
                    hsw_out_name = str(hsw_op.outputs[0])
                    hsw_input_rewrites.append((hsw_op, 0))
                    hsw_root_intermediate_names.append(str(hsw_out_name))

            # Pattern 2: decomposed hard-swish root
            #   pre_output --ADD(+const, fused=RELU6)--> a
            #              --MUL(a, const)-------------> b
            #              --MUL(pre_output, b)--------> hsw_out
            if hsw_out_name is None and len(pre_users) == 2:
                add_idx: Optional[int] = None
                mul_data_idx: Optional[int] = None
                add_data_input_index: Optional[int] = None
                mul_data_input_index: Optional[int] = None
                for user_idx in pre_users:
                    user_op = model_ir.operators[int(user_idx)]
                    user_type = str(user_op.op_type)
                    if user_type == "ADD" and len(user_op.inputs) == 2 and len(user_op.outputs) == 1:
                        add_inputs = [str(v) for v in list(user_op.inputs)]
                        if pre_output_name == add_inputs[0]:
                            side_name = add_inputs[1]
                            add_data_input_index = 0
                        elif pre_output_name == add_inputs[1]:
                            side_name = add_inputs[0]
                            add_data_input_index = 1
                        else:
                            continue
                        if not _is_singleton_constant_tensor(model_ir, side_name):
                            continue
                        fused_act = str(user_op.options.get("fusedActivationFunction", "NONE")).upper()
                        if fused_act != "RELU6":
                            continue
                        add_idx = int(user_idx)
                    elif user_type == "MUL" and len(user_op.inputs) == 2 and len(user_op.outputs) == 1:
                        mul_inputs = [str(v) for v in list(user_op.inputs)]
                        if pre_output_name == mul_inputs[0]:
                            mul_data_input_index = 0
                            mul_data_idx = int(user_idx)
                        elif pre_output_name == mul_inputs[1]:
                            mul_data_input_index = 1
                            mul_data_idx = int(user_idx)
                if (
                    add_idx is not None
                    and mul_data_idx is not None
                    and add_data_input_index is not None
                    and mul_data_input_index is not None
                ):
                    add_op = model_ir.operators[int(add_idx)]
                    mul_data_op = model_ir.operators[int(mul_data_idx)]
                    add_out_name = str(add_op.outputs[0])
                    add_users = [int(v) for v in consumers.get(add_out_name, [])]
                    if len(add_users) == 1:
                        mul_scale_idx = int(add_users[0])
                        mul_scale_op = model_ir.operators[int(mul_scale_idx)]
                        if str(mul_scale_op.op_type) == "MUL" and len(mul_scale_op.inputs) == 2 and len(mul_scale_op.outputs) == 1:
                            mul_scale_inputs = [str(v) for v in list(mul_scale_op.inputs)]
                            if add_out_name == mul_scale_inputs[0]:
                                mul_scale_side_name = mul_scale_inputs[1]
                            elif add_out_name == mul_scale_inputs[1]:
                                mul_scale_side_name = mul_scale_inputs[0]
                            else:
                                mul_scale_side_name = ""
                            if _is_singleton_constant_tensor(model_ir, mul_scale_side_name):
                                mul_scale_out_name = str(mul_scale_op.outputs[0])
                                mul_data_inputs = [str(v) for v in list(mul_data_op.inputs)]
                                if (
                                    (mul_data_inputs[0] == pre_output_name and mul_data_inputs[1] == mul_scale_out_name)
                                    or (mul_data_inputs[1] == pre_output_name and mul_data_inputs[0] == mul_scale_out_name)
                                ):
                                    hsw_out_name = str(mul_data_op.outputs[0])
                                    hsw_input_rewrites.append((add_op, int(add_data_input_index)))
                                    hsw_input_rewrites.append((mul_data_op, int(mul_data_input_index)))
                                    hsw_root_intermediate_names.extend(
                                        [str(add_out_name), str(mul_scale_out_name), str(hsw_out_name)]
                                    )

            if hsw_out_name is None:
                continue
            if hsw_out_name in model_outputs:
                continue

            hsw_users = [int(v) for v in consumers.get(hsw_out_name, [])]
            if len(hsw_users) != 2:
                continue

            mean_idx: Optional[int] = None
            residual_mul_idx: Optional[int] = None
            for user_idx in hsw_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == hsw_out_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean_idx = int(user_idx)
                elif (
                    str(user_op.op_type) == "MUL"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and hsw_out_name in {str(v) for v in list(user_op.inputs)}
                ):
                    residual_mul_idx = int(user_idx)
            if mean_idx is None or residual_mul_idx is None:
                continue

            mean_op = model_ir.operators[int(mean_idx)]
            mean_axes_name = str(mean_op.inputs[1])
            mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_vals is None or len(mean_axes_vals) == 0:
                continue
            mean_out_name = str(mean_op.outputs[0])
            if mean_out_name in model_outputs:
                continue

            mean_users = [int(v) for v in consumers.get(mean_out_name, [])]
            if len(mean_users) != 1:
                continue
            post_mean_idx = int(mean_users[0])
            post_mean_op = model_ir.operators[int(post_mean_idx)]
            if (
                str(post_mean_op.op_type) != "TRANSPOSE"
                or len(post_mean_op.inputs) < 2
                or len(post_mean_op.outputs) != 1
                or str(post_mean_op.inputs[0]) != mean_out_name
                or _read_transpose_perm(model_ir, post_mean_op) != perm_nchw_to_nhwc
            ):
                continue
            post_mean_out_name = str(post_mean_op.outputs[0])
            if post_mean_out_name in model_outputs:
                continue

            post_mean_users = [int(v) for v in consumers.get(post_mean_out_name, [])]
            if len(post_mean_users) != 1:
                continue
            conv1_idx = int(post_mean_users[0])
            conv1_op = model_ir.operators[int(conv1_idx)]
            if (
                str(conv1_op.op_type) != "CONV_2D"
                or len(conv1_op.inputs) < 1
                or len(conv1_op.outputs) != 1
                or str(conv1_op.inputs[0]) != post_mean_out_name
            ):
                continue
            conv1_out_name = str(conv1_op.outputs[0])

            conv1_users = [int(v) for v in consumers.get(conv1_out_name, [])]
            if len(conv1_users) != 1:
                continue
            conv2_idx = int(conv1_users[0])
            conv2_op = model_ir.operators[int(conv2_idx)]
            if (
                str(conv2_op.op_type) != "CONV_2D"
                or len(conv2_op.inputs) < 1
                or len(conv2_op.outputs) != 1
                or str(conv2_op.inputs[0]) != conv1_out_name
            ):
                continue
            conv2_out_name = str(conv2_op.outputs[0])

            conv2_users = [int(v) for v in consumers.get(conv2_out_name, [])]
            if len(conv2_users) != 1:
                continue
            pre_gate_idx = int(conv2_users[0])
            pre_gate_op = model_ir.operators[int(pre_gate_idx)]
            if (
                str(pre_gate_op.op_type) != "TRANSPOSE"
                or len(pre_gate_op.inputs) < 2
                or len(pre_gate_op.outputs) != 1
                or str(pre_gate_op.inputs[0]) != conv2_out_name
                or _read_transpose_perm(model_ir, pre_gate_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_gate_out_name = str(pre_gate_op.outputs[0])

            pre_gate_users = [int(v) for v in consumers.get(pre_gate_out_name, [])]
            if len(pre_gate_users) != 1:
                continue
            hs_entry_idx = int(pre_gate_users[0])
            hs_entry_op = model_ir.operators[int(hs_entry_idx)]
            hs_data_op = hs_entry_op
            hs_data_input_index: Optional[int] = None
            hs_out_name: Optional[str] = None
            hs_gate_intermediate_names: List[str] = []

            # Pattern A:
            #   MUL(alpha) -> ADD(beta) -> RELU_0_TO_1
            if str(hs_entry_op.op_type) == "MUL":
                if len(hs_entry_op.inputs) != 2 or len(hs_entry_op.outputs) != 1:
                    continue
                hs_mul_inputs = [str(v) for v in list(hs_entry_op.inputs)]
                if pre_gate_out_name == hs_mul_inputs[0]:
                    hs_data_input_index = 0
                    hs_mul_side_name = hs_mul_inputs[1]
                elif pre_gate_out_name == hs_mul_inputs[1]:
                    hs_data_input_index = 1
                    hs_mul_side_name = hs_mul_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_mul_side_name):
                    continue
                hs_mul_out_name = str(hs_entry_op.outputs[0])
                hs_gate_intermediate_names.append(str(hs_mul_out_name))

                hs_mul_users = [int(v) for v in consumers.get(hs_mul_out_name, [])]
                if len(hs_mul_users) != 1:
                    continue
                hs_add_idx = int(hs_mul_users[0])
                hs_add_op = model_ir.operators[int(hs_add_idx)]
                if str(hs_add_op.op_type) != "ADD" or len(hs_add_op.inputs) != 2 or len(hs_add_op.outputs) != 1:
                    continue
                hs_add_inputs = [str(v) for v in list(hs_add_op.inputs)]
                if hs_mul_out_name == hs_add_inputs[0]:
                    hs_add_side_name = hs_add_inputs[1]
                elif hs_mul_out_name == hs_add_inputs[1]:
                    hs_add_side_name = hs_add_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_add_side_name):
                    continue
                hs_add_out_name = str(hs_add_op.outputs[0])
                hs_gate_intermediate_names.append(str(hs_add_out_name))

                hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
                if len(hs_add_users) != 1:
                    continue
                hs_relu_idx = int(hs_add_users[0])
                hs_relu_op = model_ir.operators[int(hs_relu_idx)]
                if (
                    str(hs_relu_op.op_type) != "RELU_0_TO_1"
                    or len(hs_relu_op.inputs) != 1
                    or len(hs_relu_op.outputs) != 1
                    or str(hs_relu_op.inputs[0]) != hs_add_out_name
                ):
                    continue
                hs_out_name = str(hs_relu_op.outputs[0])
                hs_gate_intermediate_names.append(str(hs_out_name))
            # Pattern B:
            #   ADD(beta, fusedActivation=RELU6) -> MUL(alpha)
            elif str(hs_entry_op.op_type) == "ADD":
                if len(hs_entry_op.inputs) != 2 or len(hs_entry_op.outputs) != 1:
                    continue
                hs_add_inputs = [str(v) for v in list(hs_entry_op.inputs)]
                if pre_gate_out_name == hs_add_inputs[0]:
                    hs_data_input_index = 0
                    hs_add_side_name = hs_add_inputs[1]
                elif pre_gate_out_name == hs_add_inputs[1]:
                    hs_data_input_index = 1
                    hs_add_side_name = hs_add_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_add_side_name):
                    continue
                fused_act = str(hs_entry_op.options.get("fusedActivationFunction", "NONE")).upper()
                if fused_act != "RELU6":
                    continue
                hs_add_out_name = str(hs_entry_op.outputs[0])
                hs_gate_intermediate_names.append(str(hs_add_out_name))

                hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
                if len(hs_add_users) != 1:
                    continue
                hs_mul_idx = int(hs_add_users[0])
                hs_mul_op = model_ir.operators[int(hs_mul_idx)]
                if str(hs_mul_op.op_type) != "MUL" or len(hs_mul_op.inputs) != 2 or len(hs_mul_op.outputs) != 1:
                    continue
                hs_mul_inputs = [str(v) for v in list(hs_mul_op.inputs)]
                if hs_add_out_name == hs_mul_inputs[0]:
                    hs_mul_side_name = hs_mul_inputs[1]
                elif hs_add_out_name == hs_mul_inputs[1]:
                    hs_mul_side_name = hs_mul_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_mul_side_name):
                    continue
                hs_out_name = str(hs_mul_op.outputs[0])
                hs_gate_intermediate_names.append(str(hs_out_name))
            else:
                continue

            if hs_data_input_index is None or hs_out_name is None:
                continue

            residual_mul_op = model_ir.operators[int(residual_mul_idx)]
            residual_inputs = [str(v) for v in list(residual_mul_op.inputs)]
            if not (
                (residual_inputs[0] == hsw_out_name and residual_inputs[1] == hs_out_name)
                or (residual_inputs[1] == hsw_out_name and residual_inputs[0] == hs_out_name)
            ):
                continue
            residual_out_name = str(residual_mul_op.outputs[0])
            if residual_out_name in model_outputs:
                continue

            residual_users = [int(v) for v in consumers.get(residual_out_name, [])]
            if len(residual_users) != 1:
                continue
            post_idx = int(residual_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != residual_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_output_name = str(post_op.outputs[0])

            mapped_axes: List[int] = []
            valid_axes = True
            for axis in mean_axes_vals:
                a = int(axis)
                if a < 0:
                    a += 4
                if a < 0 or a >= 4:
                    valid_axes = False
                    break
                mapped_axes.append(int(perm_nhwc_to_nchw[int(a)]))
            if not valid_axes:
                continue
            if mean_axes_tensor is None or not _write_const_ints_to_tensor(mean_axes_tensor, mapped_axes):
                continue

            for rewrite_op, rewrite_input_index in hsw_input_rewrites:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=rewrite_op,
                    input_index=int(rewrite_input_index),
                    new_input_name=pre_input_name,
                )
            for root_name in hsw_root_intermediate_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(root_name), None),
                    perm_nchw_to_nhwc,
                )

            mean_inputs = [str(v) for v in list(mean_op.inputs)]
            mean_inputs[0] = hsw_out_name
            _set_operator_inputs(model_ir=model_ir, op=mean_op, new_inputs=mean_inputs)
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mean_out_name, None),
                perm_nchw_to_nhwc,
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=mean_op,
                new_outputs=[post_mean_out_name],
            )
            old_mean_tensor = model_ir.tensors.get(mean_out_name, None)
            post_mean_tensor = model_ir.tensors.get(post_mean_out_name, None)
            if post_mean_tensor is not None and old_mean_tensor is not None:
                post_mean_tensor.dtype = str(old_mean_tensor.dtype)
                post_mean_tensor.quantization = _clone_quantization(old_mean_tensor.quantization)
                post_mean_tensor.shape = [int(v) for v in list(old_mean_tensor.shape)]
                post_mean_tensor.shape_signature = (
                    [int(v) for v in list(old_mean_tensor.shape_signature)]
                    if old_mean_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mean_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_mean_tensor,
                    perm_nchw_to_nhwc,
                )

            _replace_operator_input_at(
                model_ir=model_ir,
                op=hs_data_op,
                input_index=int(hs_data_input_index),
                new_input_name=conv2_out_name,
            )
            for name in hs_gate_intermediate_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(name, None),
                    perm_nchw_to_nhwc,
                )

            _set_operator_outputs(
                model_ir=model_ir,
                op=residual_mul_op,
                new_outputs=[post_output_name],
            )
            old_residual_tensor = model_ir.tensors.get(residual_out_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if post_output_tensor is not None and old_residual_tensor is not None:
                post_output_tensor.dtype = str(old_residual_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_residual_tensor.quantization)
                post_output_tensor.shape = [int(v) for v in list(old_residual_tensor.shape)]
                post_output_tensor.shape_signature = (
                    [int(v) for v in list(old_residual_tensor.shape_signature)]
                    if old_residual_tensor.shape_signature is not None
                    else [int(v) for v in list(old_residual_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = sorted(
                list({int(pre_idx), int(post_mean_idx), int(pre_gate_idx), int(post_idx)}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_hardswish_se_conv_hardsigmoid_mul_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_pad_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NHWC<->NCHW transpose wrappers around PAD/MIRROR_PAD when they can run in NHWC.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
      PAD|MIRROR_PAD(x_nchw, pads_nchw, [const]) -> p_nchw
      p_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      PAD|MIRROR_PAD(x_nhwc, pads_nhwc, [const]) -> y_nhwc

    Safety:
    - Exact NCHW/NHWC inverse transpose pair.
    - PAD output consumed only by the post-transpose.
    - Static 4x2 pads tensor available.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            pad_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if pad_output_name in model_outputs:
                continue

            pad_idx = producers.get(str(pad_output_name), None)
            if pad_idx is None:
                continue
            pad_op = model_ir.operators[int(pad_idx)]
            if (
                str(pad_op.op_type) not in {"PAD", "MIRROR_PAD"}
                or len(pad_op.inputs) < 2
                or len(pad_op.outputs) != 1
                or str(pad_op.outputs[0]) != str(pad_output_name)
            ):
                continue

            # Keep this rewrite strict to avoid changing mixed-layout fanout behavior.
            if set(int(v) for v in consumers.get(str(pad_output_name), [])) != {int(post_idx)}:
                continue

            pre_nchw_name = str(pad_op.inputs[0])
            pre_idx = producers.get(str(pre_nchw_name), None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != str(pre_nchw_name)
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            if pre_nchw_name in model_outputs:
                continue

            pads_tensor_name = str(pad_op.inputs[1])
            pads_tensor = model_ir.tensors.get(str(pads_tensor_name), None)
            if pads_tensor is None or pads_tensor.data is None:
                continue
            try:
                pads_array = np.asarray(pads_tensor.data)
                pads_pairs = np.asarray(pads_array).reshape(4, 2)
            except Exception:
                continue
            if int(pads_pairs.size) != 8:
                continue

            # Rewire PAD input to NHWC producer.
            new_pad_inputs = [str(pre_op.inputs[0])]
            new_pad_inputs.extend(str(v) for v in list(pad_op.inputs)[1:])
            _set_operator_inputs(
                model_ir=model_ir,
                op=pad_op,
                new_inputs=new_pad_inputs,
            )

            # Convert pad spec NCHW [N,C,H,W] -> NHWC [N,H,W,C].
            pads_nhwc = np.asarray(
                [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
                dtype=pads_pairs.dtype,
            )
            pads_tensor.data = np.asarray(pads_nhwc)
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

            # Bypass post-transpose by writing PAD directly to post output tensor.
            _set_operator_outputs(
                model_ir=model_ir,
                op=pad_op,
                new_outputs=[post_output_name],
            )

            old_pad_tensor = model_ir.tensors.get(str(pad_output_name), None)
            post_output_tensor = model_ir.tensors.get(str(post_output_name), None)
            if old_pad_tensor is not None and post_output_tensor is not None:
                post_output_tensor.dtype = str(old_pad_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_pad_tensor.quantization)
                post_output_tensor.shape = [int(v) for v in list(old_pad_tensor.shape)]
                post_output_tensor.shape_signature = (
                    [int(v) for v in list(old_pad_tensor.shape_signature)]
                    if old_pad_tensor.shape_signature is not None
                    else [int(v) for v in list(old_pad_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {int(post_idx)}
            pre_remaining_users = [int(v) for v in consumers.get(str(pre_nchw_name), []) if int(v) != int(pad_idx)]
            if len(pre_remaining_users) == 0:
                remove_indices.add(int(pre_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pad_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate NHWC<->NCHW transpose wrappers around UNARY->PAD/MIRROR_PAD tails.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw --UNARY--> u_nchw
      PAD|MIRROR_PAD(u_nchw, pads_nchw, [const]) -> p_nchw
      p_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      x_nhwc --UNARY--> u_nhwc
      PAD|MIRROR_PAD(u_nhwc, pads_nhwc, [const]) -> y_nhwc

    Notes:
    - Legacy side consumers of `u_nchw` are preserved through one local
      NHWC->NCHW transpose adapter.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    adapter_perm_name = "__unary_pad_tail_nhwc_to_nchw_perm_rank4__"
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "FLOOR",
        "CEIL",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            pad_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if pad_output_name in model_outputs:
                continue

            pad_idx = producers.get(str(pad_output_name), None)
            if pad_idx is None:
                continue
            pad_op = model_ir.operators[int(pad_idx)]
            if (
                str(pad_op.op_type) not in {"PAD", "MIRROR_PAD"}
                or len(pad_op.inputs) < 2
                or len(pad_op.outputs) != 1
                or str(pad_op.outputs[0]) != str(pad_output_name)
            ):
                continue
            if set(int(v) for v in consumers.get(str(pad_output_name), [])) != {int(post_idx)}:
                continue

            unary_output_name = str(pad_op.inputs[0])
            if unary_output_name in model_outputs:
                continue
            unary_idx = producers.get(str(unary_output_name), None)
            if unary_idx is None:
                continue
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_passthrough_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.outputs[0]) != str(unary_output_name)
            ):
                continue

            pre_nchw_name = str(unary_op.inputs[0])
            pre_idx = producers.get(str(pre_nchw_name), None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != str(pre_nchw_name)
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            if set(int(v) for v in consumers.get(str(pre_nchw_name), [])) != {int(unary_idx)}:
                continue

            pre_input_name = str(pre_op.inputs[0])

            pads_tensor_name = str(pad_op.inputs[1])
            pads_tensor = model_ir.tensors.get(str(pads_tensor_name), None)
            if pads_tensor is None or pads_tensor.data is None:
                continue
            try:
                pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
            except Exception:
                continue
            if int(pads_pairs.size) != 8:
                continue

            unary_users = [int(v) for v in consumers.get(str(unary_output_name), [])]
            legacy_consumer_slots: List[Tuple[int, int]] = []
            for user_idx in unary_users:
                if int(user_idx) == int(pad_idx):
                    continue
                user_op = model_ir.operators[int(user_idx)]
                for input_index, input_name in enumerate(list(user_op.inputs)):
                    if str(input_name) == str(unary_output_name):
                        legacy_consumer_slots.append((int(id(user_op)), int(input_index)))

            adapter_input_name = str(unary_output_name)
            old_unary_tensor = model_ir.tensors.get(str(unary_output_name), None)
            if len(legacy_consumer_slots) > 0:
                if old_unary_tensor is None:
                    continue
                unary_shape = [int(v) for v in list(old_unary_tensor.shape)] if old_unary_tensor.shape is not None else None
                unary_signature = (
                    [int(v) for v in list(old_unary_tensor.shape_signature)]
                    if old_unary_tensor.shape_signature is not None
                    else ([int(v) for v in list(old_unary_tensor.shape)] if old_unary_tensor.shape is not None else None)
                )
                unary_shape_nhwc = _permute_shape(unary_shape, perm_nchw_to_nhwc)
                unary_signature_nhwc = _permute_shape(unary_signature, perm_nchw_to_nhwc)
                if unary_shape_nhwc is None or unary_signature_nhwc is None:
                    continue
                adapter_input_name = _unique_tensor_name(f"{unary_output_name}_nhwc")
                model_ir.tensors[adapter_input_name] = TensorIR(
                    name=adapter_input_name,
                    dtype=str(old_unary_tensor.dtype),
                    shape=[int(v) for v in list(unary_shape_nhwc)],
                    shape_signature=[int(v) for v in list(unary_signature_nhwc)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(old_unary_tensor.quantization),
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_outputs=[str(adapter_input_name)],
                )
            else:
                _permute_tensor_metadata_if_rank_matches(
                    old_unary_tensor,
                    perm_nchw_to_nhwc,
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[str(pre_input_name)],
            )

            if len(legacy_consumer_slots) > 0:
                op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                rank = (
                    int(len(list(old_unary_tensor.shape)))
                    if old_unary_tensor is not None and old_unary_tensor.shape is not None
                    else 4
                )
                rewritten_legacy_slots: set[Tuple[int, int]] = set()
                for consumer_op_id, input_index in list(legacy_consumer_slots):
                    if int(input_index) != 0:
                        continue
                    consumer_index = op_index_by_id.get(int(consumer_op_id), None)
                    if consumer_index is None:
                        continue
                    if int(consumer_index) < 0 or int(consumer_index) >= len(model_ir.operators):
                        continue
                    consumer_op = model_ir.operators[int(consumer_index)]
                    if (
                        str(consumer_op.op_type) != "MEAN"
                        or len(consumer_op.inputs) < 2
                        or len(consumer_op.outputs) != 1
                        or str(consumer_op.inputs[0]) != str(unary_output_name)
                    ):
                        continue
                    if int(rank) != 4:
                        continue
                    mean_out_name = str(consumer_op.outputs[0])
                    if mean_out_name in model_outputs:
                        continue
                    mean_out_users = [int(v) for v in consumers.get(mean_out_name, [])]
                    if len(mean_out_users) != 1:
                        continue
                    tail_op = model_ir.operators[int(mean_out_users[0])]
                    if (
                        str(tail_op.op_type) != "RESHAPE"
                        or len(tail_op.inputs) < 2
                        or len(tail_op.outputs) != 1
                        or str(tail_op.inputs[0]) != mean_out_name
                        or _read_const_ints_from_tensor(model_ir.tensors.get(str(tail_op.inputs[1]), None)) is None
                    ):
                        continue
                    axes_name = str(consumer_op.inputs[1])
                    axes_tensor = model_ir.tensors.get(axes_name, None)
                    axes_values = _read_const_ints_from_tensor(axes_tensor)
                    if axes_values is None or len(axes_values) == 0:
                        continue
                    normalized_axes: List[int] = []
                    valid_axes = True
                    for axis in axes_values:
                        axis_value = int(axis)
                        if axis_value < 0:
                            axis_value += int(rank)
                        if axis_value < 0 or axis_value >= int(rank):
                            valid_axes = False
                            break
                        normalized_axes.append(int(axis_value))
                    if not valid_axes:
                        continue
                    mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]
                    _write_const_ints_to_tensor(axes_tensor, [int(v) for v in mapped_axes])
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=consumer_op,
                        input_index=0,
                        new_input_name=str(adapter_input_name),
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(mean_out_name, None),
                        perm_nchw_to_nhwc,
                    )
                    rewritten_legacy_slots.add((int(consumer_op_id), int(input_index)))
                if len(rewritten_legacy_slots) > 0:
                    legacy_consumer_slots = [
                        (int(consumer_op_id), int(input_index))
                        for consumer_op_id, input_index in legacy_consumer_slots
                        if (int(consumer_op_id), int(input_index)) not in rewritten_legacy_slots
                    ]

            pad_inputs = [str(v) for v in list(pad_op.inputs)]
            pad_inputs[0] = str(adapter_input_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=pad_op,
                new_inputs=pad_inputs,
            )

            pads_nhwc = np.asarray(
                [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
                dtype=pads_pairs.dtype,
            )
            pads_tensor.data = np.asarray(pads_nhwc)
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

            _set_operator_outputs(
                model_ir=model_ir,
                op=pad_op,
                new_outputs=[post_output_name],
            )

            old_pad_tensor = model_ir.tensors.get(str(pad_output_name), None)
            post_output_tensor = model_ir.tensors.get(str(post_output_name), None)
            if old_pad_tensor is not None and post_output_tensor is not None:
                post_output_tensor.dtype = str(old_pad_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_pad_tensor.quantization)
                post_output_tensor.shape = [int(v) for v in list(old_pad_tensor.shape)]
                post_output_tensor.shape_signature = (
                    [int(v) for v in list(old_pad_tensor.shape_signature)]
                    if old_pad_tensor.shape_signature is not None
                    else [int(v) for v in list(old_pad_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_nchw_to_nhwc,
                )

            for remove_idx in sorted([int(pre_idx), int(post_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            if len(legacy_consumer_slots) > 0:
                if adapter_perm_name not in model_ir.tensors:
                    model_ir.tensors[adapter_perm_name] = TensorIR(
                        name=adapter_perm_name,
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                slot_targets: List[Tuple[int, int]] = []
                for consumer_op_id, input_index in legacy_consumer_slots:
                    consumer_index = op_index_by_id.get(int(consumer_op_id), None)
                    if consumer_index is None:
                        continue
                    if int(consumer_index) < 0 or int(consumer_index) >= len(model_ir.operators):
                        continue
                    consumer_op = model_ir.operators[int(consumer_index)]
                    if int(input_index) < 0 or int(input_index) >= len(consumer_op.inputs):
                        continue
                    if str(consumer_op.inputs[int(input_index)]) != str(unary_output_name):
                        continue
                    slot_targets.append((int(consumer_index), int(input_index)))
                if len(slot_targets) > 0:
                    model_ir.operators.insert(
                        int(min(int(v[0]) for v in slot_targets)),
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(adapter_input_name), str(adapter_perm_name)],
                            outputs=[str(unary_output_name)],
                        ),
                    )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains": int(rewritten)}


def _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NHWC<->NCHW wrappers around closed rank-4 norm subgraphs ending in PAD.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw --(norm subgraph in NCHW)--> y_nchw
      PAD|MIRROR_PAD(y_nchw, pads_nchw, [const]) -> p_nchw
      p_nchw --T(0,2,3,1)--> z_nhwc

    Rewrite:
      x_nhwc --(same subgraph remapped to NHWC axes/broadcast)--> y_nhwc
      PAD|MIRROR_PAD(y_nhwc, pads_nhwc, [const]) -> z_nhwc

    Notes:
    - Accepts external non-rank4 inputs (e.g. rank-2 affine branches).
    - Accepts external rank-4 inputs only when they are NCHW adapters
      (TRANSPOSE NHWC->NCHW) consumed by this subgraph; those are bypassed.
    - Keeps rewrite strict with a limited op whitelist.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    allowed_op_types = {
        "MEAN",
        "SUB",
        "MUL",
        "ADD",
        "SQRT",
        "DIV",
        "RELU",
        "RELU6",
    }
    legacy_rank4_adapter_perm_name = "__norm_subgraph_shared_nhwc_to_nchw_perm_rank4__"
    external_rank4_adapter_perm_name = "__norm_subgraph_external_nchw_to_nhwc_perm_rank4__"

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _normalize_axes(axes: List[int], rank: int) -> Optional[List[int]]:
        normalized: List[int] = []
        for axis in axes:
            axis_value = int(axis)
            if axis_value < 0:
                axis_value += int(rank)
            if axis_value < 0 or axis_value >= int(rank):
                return None
            normalized.append(int(axis_value))
        return sorted(list(set(normalized)))

    def _is_channelwise_nchw_rank4_tensor(tensor: Optional[TensorIR]) -> bool:
        if tensor is None or tensor.shape is None:
            return False
        shape = [int(v) for v in list(tensor.shape)]
        return (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        )

    def _clone_const_if_shared_outside_region(
        *,
        input_name: str,
        op: OperatorIR,
        input_index: int,
        region_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> str:
        users = [int(v) for v in consumers.get(str(input_name), [])]
        shared_outside = any(int(user_idx) not in region_index_set for user_idx in users)
        if not shared_outside:
            return str(input_name)
        src_tensor = model_ir.tensors.get(str(input_name), None)
        if src_tensor is None or src_tensor.data is None:
            return str(input_name)
        clone_name = _unique_tensor_name(f"{input_name}_nhwc")
        clone_data = np.asarray(src_tensor.data)
        model_ir.tensors[clone_name] = TensorIR(
            name=clone_name,
            dtype=str(src_tensor.dtype),
            shape=[int(v) for v in list(clone_data.shape)],
            shape_signature=[int(v) for v in list(clone_data.shape)],
            data=np.asarray(clone_data),
            is_variable=False,
            quantization=_clone_quantization(src_tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=clone_name,
        )
        return str(clone_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)
        channel_last_hints = set(
            str(v) for v in model_ir.metadata.get("assume_channel_last_layout_tensor_names", [])
        )

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            pad_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if pad_out_name in model_outputs:
                continue

            pad_idx = producers.get(pad_out_name, None)
            if pad_idx is None:
                continue
            pad_op = model_ir.operators[int(pad_idx)]
            if (
                str(pad_op.op_type) not in {"PAD", "MIRROR_PAD"}
                or len(pad_op.inputs) < 2
                or len(pad_op.outputs) != 1
                or str(pad_op.outputs[0]) != pad_out_name
            ):
                continue
            if set(int(v) for v in consumers.get(pad_out_name, [])) != {int(post_idx)}:
                continue

            end_name = str(pad_op.inputs[0])
            if end_name in model_outputs:
                continue

            def _collect_region_indices(
                *,
                stop_at_nhwc_to_nchw_transpose: bool,
            ) -> set[int]:
                region_indices_local: set[int] = set()
                stack_local: List[str] = [str(end_name)]
                while len(stack_local) > 0:
                    current_name = str(stack_local.pop())
                    producer_idx = producers.get(current_name, None)
                    if producer_idx is None:
                        continue
                    producer_idx_i = int(producer_idx)
                    if producer_idx_i in region_indices_local:
                        continue
                    producer_op = model_ir.operators[int(producer_idx_i)]
                    producer_type = str(producer_op.op_type)
                    is_nhwc_to_nchw_transpose = (
                        producer_type == "TRANSPOSE"
                        and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                    )
                    if producer_type not in allowed_op_types and not is_nhwc_to_nchw_transpose:
                        # Stop slice expansion at non-target producers; they are treated
                        # as external inputs to the region.
                        continue
                    region_indices_local.add(int(producer_idx_i))
                    if stop_at_nhwc_to_nchw_transpose and is_nhwc_to_nchw_transpose:
                        continue
                    for input_name in list(producer_op.inputs):
                        stack_local.append(str(input_name))
                return set(int(v) for v in region_indices_local)

            # Collect minimal producer slice from end tensor backward.
            region_indices: set[int] = _collect_region_indices(
                stop_at_nhwc_to_nchw_transpose=False,
            )

            if len(region_indices) == 0:
                continue

            # Determine root pre-transpose in the slice.
            pre_candidates = [
                int(op_idx)
                for op_idx in sorted(list(region_indices))
                if str(model_ir.operators[int(op_idx)].op_type) == "TRANSPOSE"
                and _read_transpose_perm(model_ir, model_ir.operators[int(op_idx)]) == perm_nhwc_to_nchw
            ]
            if len(pre_candidates) == 0:
                continue

            if len(pre_candidates) > 1:
                side_pre_outputs = [
                    str(model_ir.operators[int(op_idx)].outputs[0])
                    for op_idx in list(pre_candidates)
                    if len(model_ir.operators[int(op_idx)].outputs) == 1
                ]
                has_layout_fix_side = any(
                    "layout_fix" in str(output_name)
                    for output_name in list(side_pre_outputs)
                )
                if has_layout_fix_side:
                    bounded_region_indices = _collect_region_indices(
                        stop_at_nhwc_to_nchw_transpose=True,
                    )
                    bounded_pre_candidates = [
                        int(op_idx)
                        for op_idx in sorted(list(bounded_region_indices))
                        if str(model_ir.operators[int(op_idx)].op_type) == "TRANSPOSE"
                        and _read_transpose_perm(model_ir, model_ir.operators[int(op_idx)])
                        == perm_nhwc_to_nchw
                    ]
                    if len(bounded_pre_candidates) > 0:
                        region_indices = set(int(v) for v in bounded_region_indices)
                        pre_candidates = [int(v) for v in list(bounded_pre_candidates)]

            # Prefer the dominant transpose (largest number of in-region users) as the
            # root; treat the rest as bypassable side adapters.
            candidate_scores: List[Tuple[int, int]] = []
            for candidate_idx in pre_candidates:
                candidate_out = str(model_ir.operators[int(candidate_idx)].outputs[0])
                score = sum(
                    1 for user_idx in consumers.get(candidate_out, [])
                    if int(user_idx) in set(int(v) for v in region_indices)
                )
                candidate_scores.append((int(score), int(candidate_idx)))
            candidate_scores = sorted(candidate_scores, key=lambda v: (v[0], -v[1]), reverse=True)
            if len(candidate_scores) >= 2 and int(candidate_scores[0][0]) == int(candidate_scores[1][0]):
                continue
            pre_idx = int(candidate_scores[0][1])
            pre_op = model_ir.operators[int(pre_idx)]
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            start_in_name = str(pre_op.inputs[0])
            start_out_name = str(pre_op.outputs[0])
            if start_out_name == end_name:
                continue
            if start_out_name in model_outputs:
                continue

            # Keep only ops reachable from the dominant pre-transpose output.
            # This avoids pulling unrelated upstream blocks into the candidate
            # region when residual branches feed the terminal ADD.
            reachable_region_indices: set[int] = set()
            reachable_stack: List[str] = [str(start_out_name)]
            while len(reachable_stack) > 0:
                current_tensor_name = str(reachable_stack.pop())
                for user_idx in list(consumers.get(current_tensor_name, [])):
                    user_idx_i = int(user_idx)
                    if user_idx_i not in region_indices:
                        continue
                    if user_idx_i in reachable_region_indices:
                        continue
                    reachable_region_indices.add(int(user_idx_i))
                    user_op = model_ir.operators[int(user_idx_i)]
                    for out_name in list(user_op.outputs):
                        reachable_stack.append(str(out_name))
            if len(reachable_region_indices) == 0:
                continue
            end_producer_idx = producers.get(str(end_name), None)
            if end_producer_idx is None or int(end_producer_idx) not in reachable_region_indices:
                continue
            region_indices = set(int(v) for v in reachable_region_indices)

            # Do not include root/side pre-transposes in the inner region.
            for transpose_idx in pre_candidates:
                if int(transpose_idx) in region_indices:
                    region_indices.remove(int(transpose_idx))
            if len(region_indices) == 0:
                continue

            # Strict op whitelist.
            if any(
                str(model_ir.operators[int(op_idx)].op_type) not in allowed_op_types
                for op_idx in region_indices
            ):
                continue

            # Pre-transpose output must be consumed only by this region.
            start_users = set(int(v) for v in consumers.get(start_out_name, []))
            if len(start_users) == 0 or not start_users.issubset(set(int(v) for v in region_indices)):
                continue

            # Region closure:
            # - Non-rank4 shared outputs are allowed as-is.
            # - Shared rank4 outputs are preserved for external users through local
            #   NHWC->NCHW transpose adapters.
            region_index_set = set(int(v) for v in region_indices)
            valid_closure = True
            shared_rank4_outputs: List[Tuple[str, int, List[int]]] = []
            for op_idx in list(region_index_set):
                op = model_ir.operators[int(op_idx)]
                for output_name in list(op.outputs):
                    output_name_s = str(output_name)
                    if output_name_s == str(end_name):
                        continue
                    if output_name_s in model_outputs:
                        valid_closure = False
                        break
                    outside_users = [
                        int(v) for v in consumers.get(output_name_s, [])
                        if int(v) not in region_index_set
                    ]
                    if len(outside_users) == 0:
                        continue
                    output_tensor = model_ir.tensors.get(output_name_s, None)
                    output_rank = (
                        int(len(list(output_tensor.shape)))
                        if output_tensor is not None and output_tensor.shape is not None
                        else 0
                    )
                    if int(output_rank) == 4:
                        if _is_channelwise_nchw_rank4_tensor(output_tensor):
                            shared_rank4_outputs.append((output_name_s, int(op_idx), outside_users))
                            continue
                        valid_closure = False
                        break
                    if int(output_rank) <= 3:
                        continue
                    valid_closure = False
                    break
                if not valid_closure:
                    break
            if not valid_closure:
                continue

            # Validate external inputs and collect bypassable NCHW adapters.
            bypass_transpose_indices: set[int] = set(
                int(v) for v in pre_candidates if int(v) != int(pre_idx)
            )
            external_channelwise_rank4_inputs: set[str] = set()
            external_rank4_ok = True
            for op_idx in list(region_index_set):
                op = model_ir.operators[int(op_idx)]
                for input_name in list(op.inputs):
                    in_name = str(input_name)
                    if in_name == start_out_name:
                        continue
                    producer_idx = producers.get(in_name, None)
                    if producer_idx is not None and int(producer_idx) in region_index_set:
                        continue

                    in_tensor = model_ir.tensors.get(in_name, None)
                    in_rank = (
                        int(len(list(in_tensor.shape)))
                        if in_tensor is not None and in_tensor.shape is not None
                        else 0
                    )

                    # Constants and rank<=3 external tensors are allowed as-is.
                    if in_tensor is not None and in_tensor.data is not None:
                        continue
                    if in_rank <= 3:
                        continue
                    if in_rank != 4:
                        external_rank4_ok = False
                        break

                    if producer_idx is None:
                        if in_name in channel_last_hints:
                            continue
                        if _is_channelwise_nchw_rank4_tensor(in_tensor):
                            external_channelwise_rank4_inputs.add(str(in_name))
                            continue
                        external_rank4_ok = False
                        break
                    producer_op = model_ir.operators[int(producer_idx)]
                    producer_type = str(producer_op.op_type)
                    if producer_type in {"CONV_2D", "DEPTHWISE_CONV_2D", "TRANSPOSE_CONV"}:
                        continue
                    if in_name in channel_last_hints:
                        continue
                    start_in_tensor = model_ir.tensors.get(str(start_in_name), None)
                    if (
                        start_in_tensor is not None
                        and start_in_tensor.shape is not None
                        and _is_fully_known_positive_shape(list(start_in_tensor.shape))
                        and _is_fully_known_positive_shape(list(in_tensor.shape))
                        and len(list(start_in_tensor.shape)) == 4
                    ):
                        start_shape = [int(v) for v in list(start_in_tensor.shape)]
                        in_shape = [int(v) for v in list(in_tensor.shape)]
                        if _broadcast_static_shapes(start_shape, in_shape) is not None:
                            # Accept external rank-4 NHWC-aligned residual inputs.
                            continue
                    if (
                        producer_type == "TRANSPOSE"
                        and len(producer_op.inputs) >= 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == in_name
                        and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                    ):
                        if not set(int(v) for v in consumers.get(in_name, [])).issubset(region_index_set):
                            external_rank4_ok = False
                            break
                        bypass_transpose_indices.add(int(producer_idx))
                        continue
                    if _is_channelwise_nchw_rank4_tensor(in_tensor):
                        external_channelwise_rank4_inputs.add(str(in_name))
                        continue
                    external_rank4_ok = False
                    break
                if not external_rank4_ok:
                    break
            if not external_rank4_ok:
                continue

            # Rewrite shared rank-4 outputs to private NHWC tensors, and keep
            # external NCHW consumers through local adapters inserted later.
            shared_output_adapter_specs: List[Tuple[str, str, List[int]]] = []
            rewrite_shared_outputs_ok = True
            for output_name, producer_idx_i, outside_users in shared_rank4_outputs:
                source_tensor = model_ir.tensors.get(str(output_name), None)
                if source_tensor is None or source_tensor.shape is None:
                    rewrite_shared_outputs_ok = False
                    break
                if len(list(source_tensor.shape)) != 4:
                    rewrite_shared_outputs_ok = False
                    break
                internal_output_name = _unique_tensor_name(f"{output_name}_nhwc")
                internal_shape = [int(v) for v in list(source_tensor.shape)]
                internal_signature = (
                    [int(v) for v in list(source_tensor.shape_signature)]
                    if source_tensor.shape_signature is not None
                    else [int(v) for v in list(source_tensor.shape)]
                )
                model_ir.tensors[internal_output_name] = TensorIR(
                    name=internal_output_name,
                    dtype=str(source_tensor.dtype),
                    shape=internal_shape,
                    shape_signature=internal_signature,
                    data=(
                        np.asarray(source_tensor.data).copy()
                        if isinstance(source_tensor.data, np.ndarray)
                        else None
                    ),
                    is_variable=bool(source_tensor.is_variable),
                    quantization=_clone_quantization(source_tensor.quantization),
                )
                producer_op = model_ir.operators[int(producer_idx_i)]
                producer_outputs = [str(v) for v in list(producer_op.outputs)]
                producer_outputs = [
                    str(internal_output_name) if str(v) == str(output_name) else str(v)
                    for v in producer_outputs
                ]
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=producer_op,
                    new_outputs=producer_outputs,
                )
                for region_user_idx in sorted(list(region_index_set)):
                    region_user_op = model_ir.operators[int(region_user_idx)]
                    for input_index, input_name in enumerate(list(region_user_op.inputs)):
                        if str(input_name) != str(output_name):
                            continue
                        _replace_operator_input_at(
                            model_ir=model_ir,
                            op=region_user_op,
                            input_index=int(input_index),
                            new_input_name=str(internal_output_name),
                        )
                outside_user_op_ids = [
                    int(id(model_ir.operators[int(user_idx)]))
                    for user_idx in list(outside_users)
                    if 0 <= int(user_idx) < len(model_ir.operators)
                ]
                shared_output_adapter_specs.append(
                    (str(output_name), str(internal_output_name), outside_user_op_ids)
                )
            if not rewrite_shared_outputs_ok:
                continue

            # Rewire pre-transpose root and bypassable external adapters.
            external_adapter_output_by_input: Dict[str, str] = {}
            external_adapter_consumer_op_ids: Dict[str, List[int]] = {}
            rewire_ok = True
            for op_idx in sorted(list(region_index_set)):
                op = model_ir.operators[int(op_idx)]
                new_inputs = [str(v) for v in list(op.inputs)]
                for input_index, input_name in enumerate(list(new_inputs)):
                    input_name_s = str(input_name)
                    if input_name_s == start_out_name:
                        new_inputs[int(input_index)] = str(start_in_name)
                        continue
                    producer_idx = producers.get(input_name_s, None)
                    if producer_idx is None or int(producer_idx) not in bypass_transpose_indices:
                        resolved_input_name = str(new_inputs[int(input_index)])
                    else:
                        producer_op = model_ir.operators[int(producer_idx)]
                        resolved_input_name = str(producer_op.inputs[0])
                        new_inputs[int(input_index)] = str(resolved_input_name)
                    if str(resolved_input_name) not in external_channelwise_rank4_inputs:
                        continue
                    if str(resolved_input_name) not in external_adapter_output_by_input:
                        source_tensor = model_ir.tensors.get(str(resolved_input_name), None)
                        if not _is_channelwise_nchw_rank4_tensor(source_tensor):
                            rewire_ok = False
                            break
                        source_shape = [int(v) for v in list(source_tensor.shape)]
                        source_signature = (
                            [int(v) for v in list(source_tensor.shape_signature)]
                            if source_tensor.shape_signature is not None
                            else [int(v) for v in list(source_tensor.shape)]
                        )
                        adapter_shape = _permute_shape(source_shape, perm_nchw_to_nhwc)
                        adapter_signature = _permute_shape(source_signature, perm_nchw_to_nhwc)
                        if adapter_shape is None or adapter_signature is None:
                            rewire_ok = False
                            break
                        adapter_output_name = _unique_tensor_name(f"{resolved_input_name}_nhwc")
                        model_ir.tensors[adapter_output_name] = TensorIR(
                            name=adapter_output_name,
                            dtype=str(source_tensor.dtype),
                            shape=[int(v) for v in list(adapter_shape)],
                            shape_signature=[int(v) for v in list(adapter_signature)],
                            data=None,
                            is_variable=False,
                            quantization=_clone_quantization(source_tensor.quantization),
                        )
                        external_adapter_output_by_input[str(resolved_input_name)] = str(adapter_output_name)
                        external_adapter_consumer_op_ids[str(resolved_input_name)] = []
                    adapter_output_name = str(external_adapter_output_by_input[str(resolved_input_name)])
                    new_inputs[int(input_index)] = str(adapter_output_name)
                    external_adapter_consumer_op_ids[str(resolved_input_name)].append(int(id(op)))
                if not rewire_ok:
                    break
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=op,
                    new_inputs=new_inputs,
                )
            if not rewire_ok:
                continue

            # Remap MEAN axes and rewrite rank-4 RESHAPE target shapes.
            rewritten_axes_names: set[str] = set()
            for op_idx in sorted(list(region_index_set)):
                op = model_ir.operators[int(op_idx)]
                if str(op.op_type) == "MEAN" and len(op.inputs) >= 2 and len(op.outputs) == 1:
                    mean_in_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
                    if mean_in_tensor is None or mean_in_tensor.shape is None or len(list(mean_in_tensor.shape)) != 4:
                        continue
                    axes_name = str(op.inputs[1])
                    axes_tensor = model_ir.tensors.get(axes_name, None)
                    axes_values = _read_const_ints_from_tensor(axes_tensor)
                    if axes_values is None or len(axes_values) == 0:
                        continue
                    normalized_axes = _normalize_axes([int(v) for v in list(axes_values)], 4)
                    if normalized_axes is None:
                        continue
                    mapped_axes = sorted(list(set(int(perm_nhwc_to_nchw[int(v)]) for v in normalized_axes)))
                    if set(int(v) for v in consumers.get(axes_name, [])) - region_index_set:
                        cloned_axes_name = _clone_const_if_shared_outside_region(
                            input_name=axes_name,
                            op=op,
                            input_index=1,
                            region_index_set=region_index_set,
                            consumers=consumers,
                        )
                        axes_name = str(cloned_axes_name)
                        axes_tensor = model_ir.tensors.get(axes_name, None)
                    if axes_name in rewritten_axes_names:
                        continue
                    _write_const_ints_to_tensor(axes_tensor, [int(v) for v in mapped_axes])
                    rewritten_axes_names.add(str(axes_name))
                elif str(op.op_type) == "RESHAPE" and len(op.inputs) >= 2 and len(op.outputs) == 1:
                    reshape_out_tensor = model_ir.tensors.get(str(op.outputs[0]), None)
                    if (
                        reshape_out_tensor is None
                        or reshape_out_tensor.shape is None
                        or len(list(reshape_out_tensor.shape)) != 4
                    ):
                        continue
                    target_shape = _permute_shape(
                        [int(v) for v in list(reshape_out_tensor.shape)],
                        perm_nchw_to_nhwc,
                    )
                    if target_shape is None:
                        continue
                    shape_name = str(op.inputs[1])
                    shape_tensor = model_ir.tensors.get(shape_name, None)
                    if _read_const_ints_from_tensor(shape_tensor) is None:
                        continue
                    if set(int(v) for v in consumers.get(shape_name, [])) - region_index_set:
                        cloned_shape_name = _clone_const_if_shared_outside_region(
                            input_name=shape_name,
                            op=op,
                            input_index=1,
                            region_index_set=region_index_set,
                            consumers=consumers,
                        )
                        shape_name = str(cloned_shape_name)
                        shape_tensor = model_ir.tensors.get(shape_name, None)
                    _write_const_ints_to_tensor(shape_tensor, [int(v) for v in list(target_shape)])

            # Rewrite rank-4 channelwise constants used inside region.
            for op_idx in sorted(list(region_index_set)):
                op = model_ir.operators[int(op_idx)]
                for input_index, input_name in enumerate(list(op.inputs)):
                    in_name = str(input_name)
                    producer_idx = producers.get(in_name, None)
                    if producer_idx is not None and int(producer_idx) in region_index_set:
                        continue
                    in_tensor = model_ir.tensors.get(in_name, None)
                    if in_tensor is None or in_tensor.data is None:
                        continue
                    data = np.asarray(in_tensor.data)
                    if data.ndim != 4:
                        continue
                    shape = [int(v) for v in list(data.shape)]
                    if (
                        len(shape) == 4
                        and int(shape[0]) == 1
                        and int(shape[2]) == 1
                        and int(shape[3]) == 1
                    ):
                        private_name = _clone_const_if_shared_outside_region(
                            input_name=in_name,
                            op=op,
                            input_index=int(input_index),
                            region_index_set=region_index_set,
                            consumers=consumers,
                        )
                        target_tensor = model_ir.tensors.get(str(private_name), None)
                        if target_tensor is None or target_tensor.data is None:
                            continue
                        rotated = np.transpose(np.asarray(target_tensor.data), perm_nchw_to_nhwc).astype(
                            np.asarray(target_tensor.data).dtype,
                            copy=False,
                        )
                        target_tensor.data = np.asarray(rotated)
                        target_tensor.shape = [int(v) for v in list(rotated.shape)]
                        target_tensor.shape_signature = [int(v) for v in list(rotated.shape)]

            # Convert pad spec NCHW [N,C,H,W] -> NHWC [N,H,W,C].
            pads_name = str(pad_op.inputs[1])
            pads_tensor = model_ir.tensors.get(str(pads_name), None)
            if pads_tensor is None or pads_tensor.data is None:
                continue
            try:
                pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
            except Exception:
                continue
            if int(pads_pairs.size) != 8:
                continue
            if set(int(v) for v in consumers.get(str(pads_name), [])) - {int(pad_idx)}:
                pads_name = _clone_const_if_shared_outside_region(
                    input_name=str(pads_name),
                    op=pad_op,
                    input_index=1,
                    region_index_set={int(pad_idx)},
                    consumers=consumers,
                )
                pads_tensor = model_ir.tensors.get(str(pads_name), None)
                if pads_tensor is None or pads_tensor.data is None:
                    continue
                try:
                    pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
                except Exception:
                    continue
            pads_nhwc = np.asarray(
                [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
                dtype=pads_pairs.dtype,
            )
            pads_tensor.data = np.asarray(pads_nhwc)
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

            # Permute rank-4 metadata in the rewritten region.
            for op_idx in sorted(list(region_index_set)):
                op = model_ir.operators[int(op_idx)]
                for output_name in list(op.outputs):
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(output_name), None),
                        perm_nchw_to_nhwc,
                    )

            # Bypass post transpose by writing PAD directly to its output tensor.
            _set_operator_outputs(
                model_ir=model_ir,
                op=pad_op,
                new_outputs=[post_out_name],
            )
            old_pad_tensor = model_ir.tensors.get(str(pad_out_name), None)
            post_out_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_pad_tensor is not None and post_out_tensor is not None:
                post_out_tensor.dtype = str(old_pad_tensor.dtype)
                post_out_tensor.quantization = _clone_quantization(old_pad_tensor.quantization)
                post_out_tensor.shape = [int(v) for v in list(old_pad_tensor.shape)]
                post_out_tensor.shape_signature = (
                    [int(v) for v in list(old_pad_tensor.shape_signature)]
                    if old_pad_tensor.shape_signature is not None
                    else [int(v) for v in list(old_pad_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_out_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {int(pre_idx), int(post_idx)}
            for bypass_idx in list(bypass_transpose_indices):
                bypass_out_name = str(model_ir.operators[int(bypass_idx)].outputs[0])
                remaining_users = [
                    int(v) for v in _build_tensor_consumer_map(model_ir).get(bypass_out_name, [])
                ]
                if len(remaining_users) == 0:
                    remove_indices.add(int(bypass_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            if len(shared_output_adapter_specs) > 0:
                if legacy_rank4_adapter_perm_name not in model_ir.tensors:
                    model_ir.tensors[legacy_rank4_adapter_perm_name] = TensorIR(
                        name=legacy_rank4_adapter_perm_name,
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                for output_name, internal_output_name, outside_user_op_ids in list(shared_output_adapter_specs):
                    op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                    insertion_candidates = [
                        int(op_index_by_id[int(op_id)])
                        for op_id in list(outside_user_op_ids)
                        if int(op_id) in op_index_by_id
                    ]
                    if len(insertion_candidates) > 0:
                        insertion_index = int(min(insertion_candidates))
                    elif str(output_name) in model_outputs:
                        insertion_index = int(len(model_ir.operators))
                    else:
                        continue
                    model_ir.operators.insert(
                        int(insertion_index),
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(internal_output_name), str(legacy_rank4_adapter_perm_name)],
                            outputs=[str(output_name)],
                        ),
                    )

            if len(external_adapter_output_by_input) > 0:
                if external_rank4_adapter_perm_name not in model_ir.tensors:
                    model_ir.tensors[external_rank4_adapter_perm_name] = TensorIR(
                        name=external_rank4_adapter_perm_name,
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nchw_to_nhwc, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                for source_input_name, adapter_output_name in list(external_adapter_output_by_input.items()):
                    consumer_op_ids = [int(v) for v in external_adapter_consumer_op_ids.get(str(source_input_name), [])]
                    if len(consumer_op_ids) == 0:
                        continue
                    op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                    insertion_candidates = [
                        int(op_index_by_id[int(op_id)])
                        for op_id in list(consumer_op_ids)
                        if int(op_id) in op_index_by_id
                    ]
                    if len(insertion_candidates) == 0:
                        continue
                    insertion_index = int(min(insertion_candidates))
                    model_ir.operators.insert(
                        int(insertion_index),
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(source_input_name), str(external_rank4_adapter_perm_name)],
                            outputs=[str(adapter_output_name)],
                        ),
                    )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_norm_subgraph_pad_prepost_nhwc_chains": int(rewritten)}


def _optimize_swish_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around pseudo-op-expanded Swish chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --LOGISTIC--> s_t
      MUL(x_t, s_t) --> y_t
      y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --LOGISTIC--> s
      MUL(X, s) --> Y

    Safety:
    - Strict Swish topology (LOGISTIC + residual MUL input).
    - LOGISTIC output must be consumed only by the MUL.
    - Requires at least one inverse post-transpose consumer.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 2:
                continue

            log_idx = None
            mul_idx = None
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if (
                    user_type == "LOGISTIC"
                    and len(user_op.inputs) == 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_output_name
                ):
                    log_idx = int(user_idx)
                elif user_type == "MUL":
                    mul_idx = int(user_idx)
            if log_idx is None or mul_idx is None:
                continue

            log_op = model_ir.operators[int(log_idx)]
            log_out_name = str(log_op.outputs[0])
            log_users = [int(v) for v in consumers.get(log_out_name, [])]
            if len(log_users) != 1 or int(log_users[0]) != int(mul_idx):
                continue

            mul_op = model_ir.operators[int(mul_idx)]
            if len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in mul_op.inputs]
            if mul_inputs[0] == pre_output_name and mul_inputs[1] == log_out_name:
                mul_data_input_index = 0
            elif mul_inputs[1] == pre_output_name and mul_inputs[0] == log_out_name:
                mul_data_input_index = 1
            else:
                continue

            mul_out_name = str(mul_op.outputs[0])
            mul_users = [int(v) for v in consumers.get(mul_out_name, []) if int(v) != int(mul_idx)]
            if len(mul_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in mul_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == mul_out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is None or perm_post != perm_post_expected:
                        valid_users = False
                        break
                    post_output_name = str(user_op.outputs[0])
                    if post_output_name in model_outputs:
                        valid_users = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                else:
                    legacy_users.append(int(user_idx))
            if not valid_users or len(post_indices) == 0:
                continue

            representative_output_name = str(post_output_names[0])

            _set_operator_inputs(
                model_ir=model_ir,
                op=log_op,
                new_inputs=[pre_input_name],
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=int(mul_data_input_index),
                new_input_name=pre_input_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[representative_output_name],
            )
            for alias_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_output_name, representative_output_name)

            # LOGISTIC is layout-agnostic; update metadata from transposed to source layout.
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(log_out_name, None),
                perm_post_expected,
            )

            old_mul_tensor = model_ir.tensors.get(mul_out_name, None)
            representative_tensor = model_ir.tensors.get(representative_output_name, None)
            if old_mul_tensor is not None and representative_tensor is not None:
                representative_tensor.dtype = str(old_mul_tensor.dtype)
                representative_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                representative_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                representative_tensor.shape_signature = (
                    [int(v) for v in list(old_mul_tensor.shape_signature)]
                    if old_mul_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    representative_tensor,
                    perm_post_expected,
                )

            preserve_nchw_adapter = len(legacy_users) > 0 or mul_out_name in model_outputs
            if preserve_nchw_adapter:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[keep_post_idx]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_pre, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[representative_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[mul_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = sorted(list({int(pre_idx), *post_remove_indices}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_swish_transpose_passthrough_chains": int(rewritten)}


def _optimize_gelu_tanh_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around tanh-approx GELU expansion chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      m2 = MUL(MUL(x_t, x_t), x_t)
      t0 = MUL(m2, c0)
      t1 = ADD(x_t, t0)
      t2 = MUL(t1, c1)
      t3 = TANH(t2)
      t4 = ADD(t3, c2)
      t5 = MUL(x_t, t4)
      y_t = MUL(t5, c3)
      y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      Same chain with x_t replaced by X and Y produced directly from final MUL.

    Safety:
    - Requires exact tanh-GELU topology with scalar-like constants c0..c3.
    - Final GELU output must be consumed by inverse post-transpose adapters only.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = sorted(list(set(int(v) for v in consumers.get(pre_output_name, []))))
            if len(pre_users) == 0:
                continue

            gelu_match_found = False
            for msq_idx in pre_users:
                msq_op = model_ir.operators[int(msq_idx)]
                if str(msq_op.op_type) != "MUL" or len(msq_op.inputs) != 2 or len(msq_op.outputs) != 1:
                    continue
                if not (str(msq_op.inputs[0]) == pre_output_name and str(msq_op.inputs[1]) == pre_output_name):
                    continue
                msq_out_name = str(msq_op.outputs[0])
                msq_users = [int(v) for v in consumers.get(msq_out_name, [])]
                if len(msq_users) != 1:
                    continue

                mcube_idx = int(msq_users[0])
                mcube_op = model_ir.operators[int(mcube_idx)]
                if str(mcube_op.op_type) != "MUL" or len(mcube_op.inputs) != 2 or len(mcube_op.outputs) != 1:
                    continue
                mcube_inputs = [str(v) for v in mcube_op.inputs]
                if pre_output_name == mcube_inputs[0] and msq_out_name == mcube_inputs[1]:
                    pass
                elif pre_output_name == mcube_inputs[1] and msq_out_name == mcube_inputs[0]:
                    pass
                else:
                    continue
                mcube_out_name = str(mcube_op.outputs[0])

                m1_users = [int(v) for v in consumers.get(mcube_out_name, [])]
                if len(m1_users) != 1:
                    continue
                m1_idx = int(m1_users[0])
                m1_op = model_ir.operators[int(m1_idx)]
                if str(m1_op.op_type) != "MUL" or len(m1_op.inputs) != 2 or len(m1_op.outputs) != 1:
                    continue
                m1_inputs = [str(v) for v in m1_op.inputs]
                if m1_inputs[0] == mcube_out_name:
                    m1_side_name = m1_inputs[1]
                elif m1_inputs[1] == mcube_out_name:
                    m1_side_name = m1_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, m1_side_name):
                    continue
                m1_out_name = str(m1_op.outputs[0])

                add1_users = [int(v) for v in consumers.get(m1_out_name, [])]
                if len(add1_users) != 1:
                    continue
                add1_idx = int(add1_users[0])
                add1_op = model_ir.operators[int(add1_idx)]
                if str(add1_op.op_type) != "ADD" or len(add1_op.inputs) != 2 or len(add1_op.outputs) != 1:
                    continue
                add1_inputs = [str(v) for v in add1_op.inputs]
                if pre_output_name == add1_inputs[0] and m1_out_name == add1_inputs[1]:
                    pass
                elif pre_output_name == add1_inputs[1] and m1_out_name == add1_inputs[0]:
                    pass
                else:
                    continue
                add1_out_name = str(add1_op.outputs[0])

                m2_users = [int(v) for v in consumers.get(add1_out_name, [])]
                if len(m2_users) != 1:
                    continue
                m2_idx = int(m2_users[0])
                m2_op = model_ir.operators[int(m2_idx)]
                if str(m2_op.op_type) != "MUL" or len(m2_op.inputs) != 2 or len(m2_op.outputs) != 1:
                    continue
                m2_inputs = [str(v) for v in m2_op.inputs]
                if m2_inputs[0] == add1_out_name:
                    m2_side_name = m2_inputs[1]
                elif m2_inputs[1] == add1_out_name:
                    m2_side_name = m2_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, m2_side_name):
                    continue
                m2_out_name = str(m2_op.outputs[0])

                tanh_users = [int(v) for v in consumers.get(m2_out_name, [])]
                if len(tanh_users) != 1:
                    continue
                tanh_idx = int(tanh_users[0])
                tanh_op = model_ir.operators[int(tanh_idx)]
                if (
                    str(tanh_op.op_type) != "TANH"
                    or len(tanh_op.inputs) != 1
                    or len(tanh_op.outputs) != 1
                    or str(tanh_op.inputs[0]) != m2_out_name
                ):
                    continue
                tanh_out_name = str(tanh_op.outputs[0])

                add2_users = [int(v) for v in consumers.get(tanh_out_name, [])]
                if len(add2_users) != 1:
                    continue
                add2_idx = int(add2_users[0])
                add2_op = model_ir.operators[int(add2_idx)]
                if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                    continue
                add2_inputs = [str(v) for v in add2_op.inputs]
                if add2_inputs[0] == tanh_out_name:
                    add2_side_name = add2_inputs[1]
                elif add2_inputs[1] == tanh_out_name:
                    add2_side_name = add2_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, add2_side_name):
                    continue
                add2_out_name = str(add2_op.outputs[0])

                m3_users = [int(v) for v in consumers.get(add2_out_name, [])]
                if len(m3_users) != 1:
                    continue
                m3_idx = int(m3_users[0])
                m3_op = model_ir.operators[int(m3_idx)]
                if str(m3_op.op_type) != "MUL" or len(m3_op.inputs) != 2 or len(m3_op.outputs) != 1:
                    continue
                m3_inputs = [str(v) for v in m3_op.inputs]
                if pre_output_name == m3_inputs[0] and add2_out_name == m3_inputs[1]:
                    pass
                elif pre_output_name == m3_inputs[1] and add2_out_name == m3_inputs[0]:
                    pass
                else:
                    continue
                m3_out_name = str(m3_op.outputs[0])

                m4_users = [int(v) for v in consumers.get(m3_out_name, [])]
                if len(m4_users) != 1:
                    continue
                m4_idx = int(m4_users[0])
                m4_op = model_ir.operators[int(m4_idx)]
                if str(m4_op.op_type) != "MUL" or len(m4_op.inputs) != 2 or len(m4_op.outputs) != 1:
                    continue
                m4_inputs = [str(v) for v in m4_op.inputs]
                if m4_inputs[0] == m3_out_name:
                    m4_side_name = m4_inputs[1]
                elif m4_inputs[1] == m3_out_name:
                    m4_side_name = m4_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, m4_side_name):
                    continue
                gelu_out_name = str(m4_op.outputs[0])

                gelu_users = [int(v) for v in consumers.get(gelu_out_name, []) if int(v) != int(m4_idx)]
                if len(gelu_users) == 0:
                    continue

                post_indices: List[int] = []
                post_output_names: List[str] = []
                valid_users = True
                for user_idx in gelu_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) != "TRANSPOSE"
                        or len(user_op.inputs) < 2
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != gelu_out_name
                    ):
                        valid_users = False
                        break
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is None or perm_post != perm_post_expected:
                        valid_users = False
                        break
                    post_output_name = str(user_op.outputs[0])
                    if post_output_name in model_outputs:
                        valid_users = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                if not valid_users or len(post_indices) == 0:
                    continue

                representative_output_name = str(post_output_names[0])
                for op_idx in [int(msq_idx), int(mcube_idx), int(add1_idx), int(m3_idx)]:
                    op = model_ir.operators[int(op_idx)]
                    updated_inputs = [
                        pre_input_name if str(v) == pre_output_name else str(v)
                        for v in list(op.inputs)
                    ]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=op,
                        new_inputs=updated_inputs,
                    )

                _set_operator_outputs(
                    model_ir=model_ir,
                    op=m4_op,
                    new_outputs=[representative_output_name],
                )
                for alias_name in post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, alias_name, representative_output_name)

                for tensor_name in [
                    msq_out_name,
                    mcube_out_name,
                    m1_out_name,
                    add1_out_name,
                    m2_out_name,
                    tanh_out_name,
                    add2_out_name,
                    m3_out_name,
                ]:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(tensor_name, None),
                        perm_post_expected,
                    )

                old_gelu_tensor = model_ir.tensors.get(gelu_out_name, None)
                representative_tensor = model_ir.tensors.get(representative_output_name, None)
                if old_gelu_tensor is not None and representative_tensor is not None:
                    representative_tensor.dtype = str(old_gelu_tensor.dtype)
                    representative_tensor.quantization = _clone_quantization(old_gelu_tensor.quantization)
                    representative_tensor.shape = [int(v) for v in list(old_gelu_tensor.shape)]
                    representative_tensor.shape_signature = (
                        [int(v) for v in list(old_gelu_tensor.shape_signature)]
                        if old_gelu_tensor.shape_signature is not None
                        else [int(v) for v in list(old_gelu_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        representative_tensor,
                        perm_post_expected,
                    )

                remove_indices = sorted(list({int(pre_idx), *post_indices}), reverse=True)
                for remove_idx in remove_indices:
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                gelu_match_found = True
                break

            if gelu_match_found:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_gelu_tanh_transpose_passthrough_chains": int(rewritten)}


def _optimize_center_size_offset_terminal_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant NHWC->NCHW adapters in center/size/offset terminal heads.

    Typical target structure:
      center_nhwc --T(0,3,1,2)--> center_nchw --LOGISTIC->MAXIMUM->MINIMUM->RESHAPE([N,HW])
      size_nhwc   --T(0,3,1,2)--> size_nchw   --LOGISTIC->MAXIMUM->MINIMUM->RESHAPE([N,C,HW])->GATHER_ND
      offset_nhwc --T(0,3,1,2)--> offset_nchw --RESHAPE([N,C,HW])->GATHER_ND

    Rewrite:
      - Bypass the three transposes.
      - Keep branch math in NHWC.
      - Convert size/offset reshape shapes from [N,C,HW] to [N,HW,C].
      - Reorder GATHER_ND coords concat from [batch, channel, axis] to [batch, axis, channel].
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _single_consumer_op(
        tensor_name: str,
        consumers: Dict[str, List[int]],
        expected_type: str,
    ) -> Optional[int]:
        user_indices = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(user_indices) != 1:
            return None
        op_idx = int(user_indices[0])
        op = model_ir.operators[int(op_idx)]
        if str(op.op_type) != str(expected_type):
            return None
        return int(op_idx)

    def _next_unary_chain(
        start_tensor_name: str,
        consumers: Dict[str, List[int]],
        op_type: str,
    ) -> Optional[Tuple[int, str]]:
        op_idx = _single_consumer_op(start_tensor_name, consumers, op_type)
        if op_idx is None:
            return None
        op = model_ir.operators[int(op_idx)]
        if len(op.inputs) == 0 or len(op.outputs) != 1:
            return None
        if str(op.inputs[0]) != str(start_tensor_name):
            return None
        return int(op_idx), str(op.outputs[0])

    def _next_binary_const_chain(
        start_tensor_name: str,
        consumers: Dict[str, List[int]],
        op_type: str,
    ) -> Optional[Tuple[int, str]]:
        op_idx = _single_consumer_op(start_tensor_name, consumers, op_type)
        if op_idx is None:
            return None
        op = model_ir.operators[int(op_idx)]
        if len(op.inputs) != 2 or len(op.outputs) != 1:
            return None
        in0 = str(op.inputs[0])
        in1 = str(op.inputs[1])
        if in0 == str(start_tensor_name):
            side_name = str(in1)
        elif in1 == str(start_tensor_name):
            side_name = str(in0)
        else:
            return None
        if not _is_singleton_constant_tensor(model_ir, side_name):
            return None
        return int(op_idx), str(op.outputs[0])

    def _resolve_gather_coords_reorder_plan(
        *,
        gather_idx: int,
        producers: Dict[str, int],
    ) -> Optional[Tuple[int, List[str]]]:
        gather_op = model_ir.operators[int(gather_idx)]
        if str(gather_op.op_type) != "GATHER_ND" or len(gather_op.inputs) != 2:
            return None
        coords_name = str(gather_op.inputs[1])
        concat_idx = producers.get(coords_name, None)
        if concat_idx is None:
            return None
        concat_op = model_ir.operators[int(concat_idx)]
        if (
            str(concat_op.op_type) != "CONCATENATION"
            or len(concat_op.inputs) != 3
            or len(concat_op.outputs) != 1
            or str(concat_op.outputs[0]) != coords_name
            or int(concat_op.options.get("axis", -1)) != 3
        ):
            return None

        concat_inputs = [str(v) for v in list(concat_op.inputs)]
        axis_pos: Optional[int] = None
        for pos, input_name in enumerate(concat_inputs):
            prod_idx = producers.get(str(input_name), None)
            if prod_idx is None:
                continue
            prod_op = model_ir.operators[int(prod_idx)]
            if str(prod_op.op_type) != "RESHAPE":
                continue
            axis_tensor = model_ir.tensors.get(str(input_name), None)
            if axis_tensor is None or len(list(axis_tensor.shape)) != 4:
                continue
            if int(axis_tensor.shape[-1]) != 1:
                continue
            axis_pos = int(pos)
            break
        if axis_pos is None:
            return None

        other_positions = [int(v) for v in [0, 1, 2] if int(v) != int(axis_pos)]
        if len(other_positions) != 2:
            return None

        channel_pos: Optional[int] = None
        for pos in other_positions:
            tensor = model_ir.tensors.get(concat_inputs[int(pos)], None)
            if tensor is None or tensor.data is None:
                continue
            values = np.asarray(tensor.data).reshape(-1)
            if values.size == 0:
                continue
            unique_vals = sorted(list({int(v) for v in values.tolist()}))
            if len(unique_vals) > 1:
                channel_pos = int(pos)
                break
        if channel_pos is None:
            channel_pos = int(other_positions[-1])
        batch_pos = int(other_positions[0] if int(other_positions[0]) != int(channel_pos) else other_positions[1])

        reordered = [
            concat_inputs[int(batch_pos)],
            concat_inputs[int(axis_pos)],
            concat_inputs[int(channel_pos)],
        ]
        return int(concat_idx), [str(v) for v in reordered]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for center_pre_idx, center_pre_op in enumerate(model_ir.operators):
            if str(center_pre_op.op_type) != "TRANSPOSE" or len(center_pre_op.inputs) < 2 or len(center_pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, center_pre_op) != perm_nhwc_to_nchw:
                continue
            center_nhwc_name = str(center_pre_op.inputs[0])
            center_nchw_name = str(center_pre_op.outputs[0])
            if center_nhwc_name in model_outputs or center_nchw_name in model_outputs:
                continue

            center_log_idx = _single_consumer_op(center_nchw_name, consumers, "LOGISTIC")
            if center_log_idx is None:
                continue
            center_log_op = model_ir.operators[int(center_log_idx)]
            if len(center_log_op.inputs) != 1 or len(center_log_op.outputs) != 1:
                continue
            if str(center_log_op.inputs[0]) != center_nchw_name:
                continue
            center_log_out = str(center_log_op.outputs[0])

            center_max_info = _next_binary_const_chain(center_log_out, consumers, "MAXIMUM")
            if center_max_info is None:
                continue
            center_max_idx, center_max_out = center_max_info
            center_min_info = _next_binary_const_chain(center_max_out, consumers, "MINIMUM")
            if center_min_info is None:
                continue
            center_min_idx, center_min_out = center_min_info

            center_flat_idx = _single_consumer_op(center_min_out, consumers, "RESHAPE")
            if center_flat_idx is None:
                continue
            center_flat_op = model_ir.operators[int(center_flat_idx)]
            if len(center_flat_op.inputs) < 2 or len(center_flat_op.outputs) != 1:
                continue
            if str(center_flat_op.inputs[0]) != center_min_out:
                continue
            center_flat_shape_tensor = model_ir.tensors.get(str(center_flat_op.inputs[1]), None)
            center_flat_shape = _read_const_ints_from_tensor(center_flat_shape_tensor)
            if center_flat_shape is None or len(center_flat_shape) != 2:
                continue

            # Find paired size/offset pre-transposes in the same region.
            size_pre_idx: Optional[int] = None
            size_pre_op: Optional[OperatorIR] = None
            offset_pre_idx: Optional[int] = None
            offset_pre_op: Optional[OperatorIR] = None
            for candidate_idx in range(int(center_pre_idx) + 1, min(len(model_ir.operators), int(center_pre_idx) + 16)):
                candidate_op = model_ir.operators[int(candidate_idx)]
                if str(candidate_op.op_type) != "TRANSPOSE" or len(candidate_op.inputs) < 2 or len(candidate_op.outputs) != 1:
                    continue
                if _read_transpose_perm(model_ir, candidate_op) != perm_nhwc_to_nchw:
                    continue
                candidate_out = str(candidate_op.outputs[0])
                candidate_users = [int(v) for v in consumers.get(candidate_out, [])]
                if len(candidate_users) != 1:
                    continue
                user_op = model_ir.operators[int(candidate_users[0])]
                if str(user_op.op_type) == "LOGISTIC" and size_pre_idx is None:
                    size_pre_idx = int(candidate_idx)
                    size_pre_op = candidate_op
                    continue
                if str(user_op.op_type) == "RESHAPE" and offset_pre_idx is None:
                    offset_pre_idx = int(candidate_idx)
                    offset_pre_op = candidate_op
                    continue
            if size_pre_idx is None or size_pre_op is None or offset_pre_idx is None or offset_pre_op is None:
                continue

            size_nhwc_name = str(size_pre_op.inputs[0])
            size_nchw_name = str(size_pre_op.outputs[0])
            offset_nhwc_name = str(offset_pre_op.inputs[0])
            offset_nchw_name = str(offset_pre_op.outputs[0])
            if (
                size_nhwc_name in model_outputs
                or size_nchw_name in model_outputs
                or offset_nhwc_name in model_outputs
                or offset_nchw_name in model_outputs
            ):
                continue

            size_log_idx = _single_consumer_op(size_nchw_name, consumers, "LOGISTIC")
            if size_log_idx is None:
                continue
            size_log_op = model_ir.operators[int(size_log_idx)]
            if len(size_log_op.inputs) != 1 or len(size_log_op.outputs) != 1:
                continue
            if str(size_log_op.inputs[0]) != size_nchw_name:
                continue
            size_log_out = str(size_log_op.outputs[0])

            size_max_info = _next_binary_const_chain(size_log_out, consumers, "MAXIMUM")
            if size_max_info is None:
                continue
            size_max_idx, size_max_out = size_max_info
            size_min_info = _next_binary_const_chain(size_max_out, consumers, "MINIMUM")
            if size_min_info is None:
                continue
            size_min_idx, size_min_out = size_min_info

            size_reshape_idx = _single_consumer_op(size_min_out, consumers, "RESHAPE")
            if size_reshape_idx is None:
                continue
            size_reshape_op = model_ir.operators[int(size_reshape_idx)]
            if len(size_reshape_op.inputs) < 2 or len(size_reshape_op.outputs) != 1:
                continue
            if str(size_reshape_op.inputs[0]) != size_min_out:
                continue
            size_reshape_shape_tensor = model_ir.tensors.get(str(size_reshape_op.inputs[1]), None)
            size_reshape_shape = _read_const_ints_from_tensor(size_reshape_shape_tensor)
            if size_reshape_shape is None or len(size_reshape_shape) != 3:
                continue
            size_reshape_out = str(size_reshape_op.outputs[0])

            offset_reshape_idx = _single_consumer_op(offset_nchw_name, consumers, "RESHAPE")
            if offset_reshape_idx is None:
                continue
            offset_reshape_op = model_ir.operators[int(offset_reshape_idx)]
            if len(offset_reshape_op.inputs) < 2 or len(offset_reshape_op.outputs) != 1:
                continue
            if str(offset_reshape_op.inputs[0]) != offset_nchw_name:
                continue
            offset_reshape_shape_tensor = model_ir.tensors.get(str(offset_reshape_op.inputs[1]), None)
            offset_reshape_shape = _read_const_ints_from_tensor(offset_reshape_shape_tensor)
            if offset_reshape_shape is None or len(offset_reshape_shape) != 3:
                continue
            if [int(v) for v in offset_reshape_shape] != [int(v) for v in size_reshape_shape]:
                continue
            offset_reshape_out = str(offset_reshape_op.outputs[0])

            size_gather_idx = _single_consumer_op(size_reshape_out, consumers, "GATHER_ND")
            if size_gather_idx is None:
                continue
            offset_gather_idx = _single_consumer_op(offset_reshape_out, consumers, "GATHER_ND")
            if offset_gather_idx is None:
                continue

            size_concat_plan = _resolve_gather_coords_reorder_plan(
                gather_idx=int(size_gather_idx),
                producers=producers,
            )
            if size_concat_plan is None:
                continue
            offset_concat_plan = _resolve_gather_coords_reorder_plan(
                gather_idx=int(offset_gather_idx),
                producers=producers,
            )
            if offset_concat_plan is None:
                continue

            # Rewire to NHWC sources.
            _set_operator_inputs(
                model_ir=model_ir,
                op=center_log_op,
                new_inputs=[center_nhwc_name],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=size_log_op,
                new_inputs=[size_nhwc_name],
            )
            offset_inputs = [str(v) for v in list(offset_reshape_op.inputs)]
            offset_inputs[0] = str(offset_nhwc_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=offset_reshape_op,
                new_inputs=offset_inputs,
            )

            for tensor_name in [
                center_log_out,
                center_max_out,
                center_min_out,
                size_log_out,
                size_max_out,
                size_min_out,
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(tensor_name, None),
                    perm_nchw_to_nhwc,
                )

            new_reshape_shape = [
                int(size_reshape_shape[0]),
                int(size_reshape_shape[2]),
                int(size_reshape_shape[1]),
            ]
            if not _write_const_ints_to_tensor(size_reshape_shape_tensor, [int(v) for v in new_reshape_shape]):
                continue
            if not _write_const_ints_to_tensor(offset_reshape_shape_tensor, [int(v) for v in new_reshape_shape]):
                continue

            size_opts = dict(size_reshape_op.options) if isinstance(size_reshape_op.options, dict) else {}
            if "newShape" in size_opts and isinstance(size_opts["newShape"], list):
                size_opts["newShape"] = [int(v) for v in list(new_reshape_shape)]
            if "onnxRawNewShape" in size_opts and isinstance(size_opts["onnxRawNewShape"], list) and len(size_opts["onnxRawNewShape"]) == 3:
                old_raw = [int(v) for v in list(size_opts["onnxRawNewShape"])]
                size_opts["onnxRawNewShape"] = [int(old_raw[0]), int(old_raw[2]), int(old_raw[1])]
            size_reshape_op.options = size_opts

            offset_opts = dict(offset_reshape_op.options) if isinstance(offset_reshape_op.options, dict) else {}
            if "newShape" in offset_opts and isinstance(offset_opts["newShape"], list):
                offset_opts["newShape"] = [int(v) for v in list(new_reshape_shape)]
            if "onnxRawNewShape" in offset_opts and isinstance(offset_opts["onnxRawNewShape"], list) and len(offset_opts["onnxRawNewShape"]) == 3:
                old_raw = [int(v) for v in list(offset_opts["onnxRawNewShape"])]
                offset_opts["onnxRawNewShape"] = [int(old_raw[0]), int(old_raw[2]), int(old_raw[1])]
            offset_reshape_op.options = offset_opts

            size_out_tensor = model_ir.tensors.get(size_reshape_out, None)
            if size_out_tensor is not None:
                size_out_tensor.shape = [int(v) for v in list(new_reshape_shape)]
                size_out_tensor.shape_signature = [int(v) for v in list(new_reshape_shape)]
            offset_out_tensor = model_ir.tensors.get(offset_reshape_out, None)
            if offset_out_tensor is not None:
                offset_out_tensor.shape = [int(v) for v in list(new_reshape_shape)]
                offset_out_tensor.shape_signature = [int(v) for v in list(new_reshape_shape)]

            size_concat_idx, size_concat_inputs = size_concat_plan
            _set_operator_inputs(
                model_ir=model_ir,
                op=model_ir.operators[int(size_concat_idx)],
                new_inputs=[str(v) for v in size_concat_inputs],
            )
            offset_concat_idx, offset_concat_inputs = offset_concat_plan
            _set_operator_inputs(
                model_ir=model_ir,
                op=model_ir.operators[int(offset_concat_idx)],
                new_inputs=[str(v) for v in offset_concat_inputs],
            )

            remove_indices = sorted(
                list({int(center_pre_idx), int(size_pre_idx), int(offset_pre_idx)}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_center_size_offset_terminal_transpose_chains": int(rewritten)}


def _optimize_leakyrelu_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around pseudo-op-expanded LeakyReLU chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --NEG--> n --RELU--> nr --MUL(alpha)--> m
      x_t --RELU--> p
      p --SUB(m)--> y_t
      y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --NEG--> n --RELU--> nr --MUL(alpha)--> m
      X --RELU--> p
      p --SUB(m)--> Y

    Safety:
    - Strict LeakyReLU topology (NEG/RELU fanout + RELU + MUL + SUB).
    - MUL side input must be singleton constant (alpha scalar-like).
    - Single-consumer chain on each branch and tail.
    """
    rewritten = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 2:
                continue

            neg_idx = None
            relu_pos_idx = None
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if (
                    user_type == "NEG"
                    and len(user_op.inputs) == 1
                    and str(user_op.inputs[0]) == pre_output_name
                    and len(user_op.outputs) == 1
                ):
                    neg_idx = int(user_idx)
                elif (
                    user_type == "RELU"
                    and len(user_op.inputs) == 1
                    and str(user_op.inputs[0]) == pre_output_name
                    and len(user_op.outputs) == 1
                ):
                    relu_pos_idx = int(user_idx)
            if neg_idx is None or relu_pos_idx is None:
                continue

            neg_op = model_ir.operators[int(neg_idx)]
            neg_out_name = str(neg_op.outputs[0])
            neg_users = [int(v) for v in consumers.get(neg_out_name, [])]
            if len(neg_users) != 1:
                continue
            relu_neg_idx = int(neg_users[0])
            relu_neg_op = model_ir.operators[int(relu_neg_idx)]
            if (
                str(relu_neg_op.op_type) != "RELU"
                or len(relu_neg_op.inputs) != 1
                or len(relu_neg_op.outputs) != 1
                or str(relu_neg_op.inputs[0]) != neg_out_name
            ):
                continue
            relu_neg_out_name = str(relu_neg_op.outputs[0])

            relu_neg_users = [int(v) for v in consumers.get(relu_neg_out_name, [])]
            if len(relu_neg_users) != 1:
                continue
            mul_idx = int(relu_neg_users[0])
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in mul_op.inputs]
            if relu_neg_out_name == mul_inputs[0]:
                mul_side_input_name = mul_inputs[1]
            elif relu_neg_out_name == mul_inputs[1]:
                mul_side_input_name = mul_inputs[0]
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, mul_side_input_name):
                continue
            mul_out_name = str(mul_op.outputs[0])

            relu_pos_op = model_ir.operators[int(relu_pos_idx)]
            relu_pos_out_name = str(relu_pos_op.outputs[0])
            relu_pos_users = [int(v) for v in consumers.get(relu_pos_out_name, [])]
            if len(relu_pos_users) != 1:
                continue

            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_users) != 1:
                continue
            sub_idx = int(mul_users[0])
            if relu_pos_users[0] != sub_idx:
                continue

            sub_op = model_ir.operators[int(sub_idx)]
            if str(sub_op.op_type) != "SUB" or len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue
            sub_inputs = [str(v) for v in sub_op.inputs]
            if sub_inputs != [relu_pos_out_name, mul_out_name]:
                continue
            sub_out_name = str(sub_op.outputs[0])

            sub_users = [int(v) for v in consumers.get(sub_out_name, []) if int(v) != int(sub_idx)]
            if len(sub_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in sub_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == sub_out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is None or perm_post != perm_post_expected:
                        valid_users = False
                        break
                    post_output_name = str(user_op.outputs[0])
                    if post_output_name in model_outputs:
                        valid_users = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                else:
                    legacy_users.append(int(user_idx))
            if not valid_users or len(post_indices) == 0:
                continue
            post_output_name = str(post_output_names[0])

            # Rewire the pseudo-LeakyReLU chain to consume source layout directly.
            _set_operator_inputs(
                model_ir=model_ir,
                op=neg_op,
                new_inputs=[pre_input_name],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=relu_pos_op,
                new_inputs=[pre_input_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=sub_op,
                new_outputs=[post_output_name],
            )
            for alias_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_output_name, post_output_name)

            # Update intermediate metadata from transposed layout to source layout.
            for intermediate_name in [neg_out_name, relu_pos_out_name, relu_neg_out_name, mul_out_name]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(intermediate_name, None),
                    perm_post_expected,
                )

            # Keep final tensor metadata stable on the post-transpose name.
            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            old_sub_tensor = model_ir.tensors.get(sub_out_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if post_output_tensor is not None:
                if pre_input_tensor is not None:
                    post_output_tensor.shape = [int(v) for v in list(pre_input_tensor.shape)]
                    post_output_tensor.shape_signature = (
                        [int(v) for v in list(pre_input_tensor.shape_signature)]
                        if pre_input_tensor.shape_signature is not None
                        else [int(v) for v in list(pre_input_tensor.shape)]
                    )
                if old_sub_tensor is not None:
                    post_output_tensor.dtype = str(old_sub_tensor.dtype)
                    post_output_tensor.quantization = _clone_quantization(old_sub_tensor.quantization)

            preserve_nchw_adapter = len(legacy_users) > 0 or sub_out_name in model_outputs
            if preserve_nchw_adapter:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[keep_post_idx]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_pre, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[sub_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = sorted(list({int(pre_idx), *post_remove_indices}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    fuse_stats = _optimize_fuse_pseudo_leakyrelu_chains(model_ir)
    _prune_unused_tensors(model_ir)
    return {
        "rewritten_leakyrelu_transpose_passthrough_chains": int(rewritten),
        "fused_pseudo_leakyrelu_chains": int(
            fuse_stats.get("fused_pseudo_leakyrelu_chains", 0)
        ),
    }


def _optimize_prelu_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around PRELU chains.

    Target:
      X --TRANSPOSE(P)--> x_t
      x_t --PRELU(alpha)--> y_t
      y_t --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --PRELU(alpha')--> Y

    Safety:
    - PRELU alpha must be constant.
    - Requires at least one inverse post-transpose consumer.
    - Keeps one adapter transpose when legacy NCHW consumers remain.
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue

            prelu_candidates: List[int] = []
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "PRELU"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_output_name
                ):
                    prelu_candidates.append(int(user_idx))
            if len(prelu_candidates) != 1:
                continue

            prelu_idx = int(prelu_candidates[0])
            prelu_op = model_ir.operators[int(prelu_idx)]
            prelu_out_name = str(prelu_op.outputs[0])
            if prelu_out_name in model_outputs:
                continue

            prelu_users = [int(v) for v in consumers.get(prelu_out_name, [])]
            if len(prelu_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in prelu_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == prelu_out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    post_output_name = str(user_op.outputs[0])
                    if (
                        perm_post is None
                        or perm_post != perm_post_expected
                        or post_output_name in model_outputs
                    ):
                        valid_users = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                else:
                    legacy_users.append(int(user_idx))
            if not valid_users or len(post_indices) == 0:
                continue

            alpha_name = str(prelu_op.inputs[1])
            alpha_tensor = model_ir.tensors.get(alpha_name, None)
            if alpha_tensor is None or not isinstance(alpha_tensor.data, np.ndarray):
                continue

            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            target_shape = (
                [int(v) for v in list(pre_input_tensor.shape)]
                if pre_input_tensor is not None and pre_input_tensor.shape is not None
                else None
            )

            alpha_data = np.asarray(alpha_tensor.data)
            alpha_candidates: List[np.ndarray] = []
            if int(alpha_data.ndim) == int(len(perm_post_expected)):
                # Prefer layout-remapped alpha first. For ambiguous static shapes
                # (e.g. [1,64,64,64]), both original/transposed alpha can be
                # broadcastable; original-first selection can silently keep
                # NCHW alpha on NHWC PRELU after transpose passthrough rewrite.
                alpha_candidates.append(
                    np.transpose(alpha_data, axes=perm_post_expected).astype(alpha_data.dtype, copy=False)
                )
            alpha_candidates.append(np.asarray(alpha_data))
            if perm_pre == [0, 3, 1, 2] and perm_post_expected == [0, 2, 3, 1] and int(alpha_data.ndim) == 3:
                alpha_candidates.append(
                    np.transpose(alpha_data, axes=[1, 2, 0]).astype(alpha_data.dtype, copy=False)
                )

            def _broadcastable_to_target(candidate: np.ndarray) -> bool:
                if target_shape is None:
                    return True
                return (
                    _broadcast_static_shapes(
                        [int(v) for v in list(target_shape)],
                        [int(v) for v in list(candidate.shape)],
                    )
                    is not None
                )

            selected_alpha: Optional[np.ndarray] = None
            for candidate in alpha_candidates:
                if _broadcastable_to_target(candidate):
                    selected_alpha = np.asarray(candidate)
                    break
            if selected_alpha is None:
                continue

            selected_alpha_name = str(alpha_name)
            alpha_users = [int(v) for v in consumers.get(alpha_name, [])]
            alpha_needs_rewrite = (
                selected_alpha.shape != alpha_data.shape
                or not np.array_equal(selected_alpha, alpha_data)
            )
            if alpha_needs_rewrite:
                if len(alpha_users) == 1 and int(alpha_users[0]) == int(prelu_idx):
                    shape, signature = normalize_onnx_shape(list(selected_alpha.shape))
                    alpha_tensor.data = np.asarray(selected_alpha)
                    alpha_tensor.shape = [int(v) for v in list(shape)]
                    alpha_tensor.shape_signature = [int(v) for v in list(signature)]
                else:
                    selected_alpha_name = _unique_tensor_name(f"{alpha_name}_nhwc")
                    shape, signature = normalize_onnx_shape(list(selected_alpha.shape))
                    model_ir.tensors[selected_alpha_name] = TensorIR(
                        name=selected_alpha_name,
                        dtype=str(alpha_tensor.dtype),
                        shape=[int(v) for v in list(shape)],
                        shape_signature=[int(v) for v in list(signature)],
                        data=np.asarray(selected_alpha),
                        is_variable=False,
                        quantization=_clone_quantization(alpha_tensor.quantization),
                    )

            representative_output_name = str(post_output_names[0])
            _set_operator_inputs(
                model_ir=model_ir,
                op=prelu_op,
                new_inputs=[pre_input_name, selected_alpha_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu_op,
                new_outputs=[representative_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_name, representative_output_name)

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(prelu_out_name, None),
                perm_post_expected,
            )

            old_prelu_tensor = model_ir.tensors.get(prelu_out_name, None)
            representative_tensor = model_ir.tensors.get(representative_output_name, None)
            if old_prelu_tensor is not None and representative_tensor is not None:
                representative_tensor.dtype = str(old_prelu_tensor.dtype)
                representative_tensor.quantization = _clone_quantization(old_prelu_tensor.quantization)
                representative_tensor.shape = [int(v) for v in list(old_prelu_tensor.shape)]
                representative_tensor.shape_signature = (
                    [int(v) for v in list(old_prelu_tensor.shape_signature)]
                    if old_prelu_tensor.shape_signature is not None
                    else [int(v) for v in list(old_prelu_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    representative_tensor,
                    perm_post_expected,
                )

            preserve_nchw_adapter = len(legacy_users) > 0 or prelu_out_name in model_outputs
            if preserve_nchw_adapter:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[keep_post_idx]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_pre, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[representative_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[prelu_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = set(int(v) for v in post_remove_indices)
            pre_remaining_users = [int(v) for v in pre_users if int(v) != int(prelu_idx)]
            if len(pre_remaining_users) == 0:
                remove_indices.add(int(pre_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_prelu_transpose_passthrough_chains": int(rewritten)}


def _optimize_transpose_elementwise_concat_conv_nhwc_groups(model_ir: ModelIR) -> Dict[str, int]:
    """
    Lift NCHW elementwise/concat groups back to NHWC when they are terminated by
    CONCAT(axis=1) -> TRANSPOSE(0,2,3,1) adapters into Conv inputs.

    This pass rewrites a whole connected concat group at once (not per-node),
    so shared upstream elementwise producers are converted consistently.
    """
    optimized_groups = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw_const_name = "__nhwc_to_nchw_perm_rank4__"
    unary_ops = {"RELU", "RELU6", "LEAKY_RELU", "LOGISTIC", "TANH", "NEG", "EXP", "ABS", "SQRT", "GELU", "ELU"}
    binary_ops = {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM"}
    legacy_consumer_safe_ops = set(unary_ops) | set(binary_ops)

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    if perm_nhwc_to_nchw_const_name not in model_ir.tensors:
        model_ir.tensors[perm_nhwc_to_nchw_const_name] = TensorIR(
            name=perm_nhwc_to_nchw_const_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        candidate_concats: List[int] = []
        post_by_concat: Dict[int, List[int]] = {}
        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            axis_old = int(concat_op.options.get("axis", 1))
            if axis_old < 0:
                axis_old += 4
            if axis_old != 1:
                continue
            concat_out = str(concat_op.outputs[0])
            concat_users = [int(v) for v in consumers.get(concat_out, [])]
            if len(concat_users) == 0:
                continue
            post_indices: List[int] = []
            valid_users = True
            for user_idx in concat_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != concat_out
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    valid_users = False
                    break
                post_indices.append(int(user_idx))
            if not valid_users or len(post_indices) == 0:
                continue
            candidate_concats.append(int(concat_idx))
            post_by_concat[int(concat_idx)] = [int(v) for v in post_indices]

        if len(candidate_concats) == 0:
            break

        for seed_concat_idx in candidate_concats:
            target_concats: set[int] = {int(seed_concat_idx)}
            closure_ops: set[int] = set()
            boundary_transpose_by_tensor: Dict[str, Tuple[str, int]] = {}
            const_tensor_names_to_permute: set[str] = set()

            def _trace_tensor_to_nhwc(
                tensor_name: str,
                visiting: set[int],
            ) -> bool:
                if str(tensor_name) in boundary_transpose_by_tensor:
                    return True
                producer_idx = producers.get(str(tensor_name), None)
                if producer_idx is None:
                    tensor = model_ir.tensors.get(str(tensor_name), None)
                    if (
                        str(tensor_name) in set(str(v) for v in model_ir.inputs)
                        and tensor is not None
                        and len(list(tensor.shape)) == 4
                    ):
                        boundary_transpose_by_tensor[str(tensor_name)] = (str(tensor_name), -1)
                        return True
                    return False
                producer_op = model_ir.operators[int(producer_idx)]
                producer_type = str(producer_op.op_type)

                if (
                    producer_type == "TRANSPOSE"
                    and len(producer_op.inputs) >= 2
                    and len(producer_op.outputs) == 1
                    and str(producer_op.outputs[0]) == str(tensor_name)
                    and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                ):
                    boundary_transpose_by_tensor[str(tensor_name)] = (
                        str(producer_op.inputs[0]),
                        int(producer_idx),
                    )
                    return True

                # Also accept an existing forward NCHW->NHWC adapter produced
                # from this tensor by a consumer TRANSPOSE.
                for consumer_idx in [int(v) for v in consumers.get(str(tensor_name), [])]:
                    consumer_op = model_ir.operators[int(consumer_idx)]
                    if (
                        str(consumer_op.op_type) == "TRANSPOSE"
                        and len(consumer_op.inputs) >= 2
                        and len(consumer_op.outputs) == 1
                        and str(consumer_op.inputs[0]) == str(tensor_name)
                        and _read_transpose_perm(model_ir, consumer_op) == perm_nchw_to_nhwc
                        and str(consumer_op.outputs[0]) not in model_outputs
                    ):
                        boundary_transpose_by_tensor[str(tensor_name)] = (
                            str(consumer_op.outputs[0]),
                            -1,
                        )
                        return True

                if int(producer_idx) in closure_ops:
                    return True
                if int(producer_idx) in visiting:
                    return False
                if str(tensor_name) in model_outputs:
                    return False

                if producer_type in unary_ops:
                    if len(producer_op.inputs) != 1 or len(producer_op.outputs) != 1:
                        return False
                    visiting.add(int(producer_idx))
                    closure_ops.add(int(producer_idx))
                    ok = _trace_tensor_to_nhwc(str(producer_op.inputs[0]), visiting)
                    visiting.remove(int(producer_idx))
                    return bool(ok)

                if producer_type in binary_ops:
                    if len(producer_op.inputs) != 2 or len(producer_op.outputs) != 1:
                        return False
                    visiting.add(int(producer_idx))
                    closure_ops.add(int(producer_idx))
                    ok = True
                    for input_name in [str(v) for v in list(producer_op.inputs)]:
                        input_tensor = model_ir.tensors.get(str(input_name), None)
                        if input_tensor is not None and input_tensor.data is not None:
                            try:
                                const_arr = np.asarray(input_tensor.data)
                            except Exception:
                                const_arr = None
                            if const_arr is not None and const_arr.ndim == 4 and int(const_arr.size) > 1:
                                const_tensor_names_to_permute.add(str(input_name))
                            continue
                        if not _trace_tensor_to_nhwc(str(input_name), visiting):
                            ok = False
                            break
                    visiting.remove(int(producer_idx))
                    return bool(ok)

                return False

            analysis_ok = True
            while True:
                closure_ops.clear()
                boundary_transpose_by_tensor.clear()
                const_tensor_names_to_permute.clear()

                for concat_idx in sorted(list(target_concats)):
                    concat_op = model_ir.operators[int(concat_idx)]
                    if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.inputs) == 0:
                        analysis_ok = False
                        break
                    for concat_input_name in [str(v) for v in list(concat_op.inputs)]:
                        if not _trace_tensor_to_nhwc(concat_input_name, set()):
                            analysis_ok = False
                            break
                    if not analysis_ok:
                        break
                if not analysis_ok:
                    break

                expanded = False
                for candidate_concat_idx in candidate_concats:
                    if int(candidate_concat_idx) in target_concats:
                        continue
                    candidate_concat_op = model_ir.operators[int(candidate_concat_idx)]
                    for candidate_input_name in [str(v) for v in list(candidate_concat_op.inputs)]:
                        candidate_input_prod = producers.get(str(candidate_input_name), None)
                        if candidate_input_prod is not None and int(candidate_input_prod) in closure_ops:
                            target_concats.add(int(candidate_concat_idx))
                            expanded = True
                            break
                if not expanded:
                    break

            if not analysis_ok:
                continue

            target_post_indices: set[int] = set()
            for concat_idx in target_concats:
                target_post_indices.update(set(int(v) for v in post_by_concat.get(int(concat_idx), [])))
            allowed_consumer_indices = set(int(v) for v in closure_ops) | set(int(v) for v in target_concats) | set(
                int(v) for v in target_post_indices
            )

            valid_group = True
            legacy_consumer_slots_by_output: Dict[str, List[Tuple[int, int]]] = {}
            for op_idx in sorted(list(closure_ops)):
                op = model_ir.operators[int(op_idx)]
                if any(str(output_name) in model_outputs for output_name in op.outputs):
                    valid_group = False
                    break
                for output_name in [str(v) for v in list(op.outputs)]:
                    external_slots: List[Tuple[int, int]] = []
                    for user_idx in [int(v) for v in consumers.get(str(output_name), [])]:
                        if int(user_idx) in allowed_consumer_indices:
                            continue
                        user_op = model_ir.operators[int(user_idx)]
                        if str(user_op.op_type) not in legacy_consumer_safe_ops:
                            valid_group = False
                            break
                        input_indices = [
                            int(input_idx)
                            for input_idx, input_name in enumerate(list(user_op.inputs))
                            if str(input_name) == str(output_name)
                        ]
                        if len(input_indices) == 0:
                            valid_group = False
                            break
                        for input_idx in input_indices:
                            external_slots.append((int(id(user_op)), int(input_idx)))
                    if not valid_group:
                        break
                    if len(external_slots) > 0:
                        legacy_consumer_slots_by_output[str(output_name)] = list(external_slots)
                if not valid_group:
                    break

            if not valid_group:
                continue

            for const_name in sorted(list(const_tensor_names_to_permute)):
                const_users = [int(v) for v in consumers.get(str(const_name), [])]
                if any(int(user_idx) not in closure_ops for user_idx in const_users):
                    valid_group = False
                    break
            if not valid_group:
                continue

            for concat_idx in sorted(list(target_concats)):
                concat_op = model_ir.operators[int(concat_idx)]
                concat_out_name = str(concat_op.outputs[0])
                concat_users = set(int(v) for v in consumers.get(str(concat_out_name), []))
                expected_users = set(int(v) for v in post_by_concat.get(int(concat_idx), []))
                if concat_users != expected_users:
                    valid_group = False
                    break
            if not valid_group:
                continue

            remove_indices: set[int] = set()
            legacy_nchw_tensor_meta_by_output: Dict[str, Tuple[str, Optional[Any], List[int], List[int]]] = {}
            for const_name in sorted(list(const_tensor_names_to_permute)):
                const_tensor = model_ir.tensors.get(str(const_name), None)
                if const_tensor is None or const_tensor.data is None:
                    continue
                const_arr = np.asarray(const_tensor.data)
                if const_arr.ndim != 4 or int(const_arr.size) <= 1:
                    continue
                transposed = np.transpose(const_arr, axes=perm_nchw_to_nhwc)
                const_tensor.data = np.asarray(transposed)
                const_tensor.shape = [int(v) for v in list(transposed.shape)]
                const_tensor.shape_signature = [int(v) for v in list(transposed.shape)]

            for op_idx in sorted(list(closure_ops)):
                op = model_ir.operators[int(op_idx)]
                old_inputs = [str(v) for v in list(op.inputs)]
                new_inputs: List[str] = []
                for input_name in old_inputs:
                    boundary = boundary_transpose_by_tensor.get(str(input_name), None)
                    if boundary is not None:
                        new_inputs.append(str(boundary[0]))
                        remove_indices.add(int(boundary[1]))
                    else:
                        new_inputs.append(str(input_name))
                if new_inputs != old_inputs:
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=op,
                        new_inputs=[str(v) for v in list(new_inputs)],
                    )
                for output_name in [str(v) for v in list(op.outputs)]:
                    if str(output_name) in legacy_consumer_slots_by_output:
                        output_tensor = model_ir.tensors.get(str(output_name), None)
                        if output_tensor is not None and len(list(output_tensor.shape)) == 4:
                            legacy_nchw_tensor_meta_by_output[str(output_name)] = (
                                str(output_tensor.dtype),
                                _clone_quantization(output_tensor.quantization),
                                [int(v) for v in list(output_tensor.shape)],
                                (
                                    [int(v) for v in list(output_tensor.shape_signature)]
                                    if output_tensor.shape_signature is not None
                                    else [int(v) for v in list(output_tensor.shape)]
                                ),
                            )
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(output_name), None),
                        perm_nchw_to_nhwc,
                    )

            post_indices_to_remove: set[int] = set()
            for concat_idx in sorted(list(target_concats)):
                concat_op = model_ir.operators[int(concat_idx)]
                old_concat_out_name = str(concat_op.outputs[0])
                old_concat_tensor = model_ir.tensors.get(old_concat_out_name, None)
                old_inputs = [str(v) for v in list(concat_op.inputs)]
                new_inputs: List[str] = []
                for input_name in old_inputs:
                    boundary = boundary_transpose_by_tensor.get(str(input_name), None)
                    if boundary is not None:
                        new_inputs.append(str(boundary[0]))
                        remove_indices.add(int(boundary[1]))
                    else:
                        new_inputs.append(str(input_name))
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=concat_op,
                    new_inputs=[str(v) for v in list(new_inputs)],
                )
                concat_op.options["axis"] = 3
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(old_concat_out_name, None),
                    perm_nchw_to_nhwc,
                )

                post_indices = sorted([int(v) for v in post_by_concat.get(int(concat_idx), [])])
                if len(post_indices) == 0:
                    valid_group = False
                    break
                canonical_post_idx = int(post_indices[0])
                canonical_post_op = model_ir.operators[int(canonical_post_idx)]
                canonical_post_out_name = str(canonical_post_op.outputs[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=concat_op,
                    new_outputs=[canonical_post_out_name],
                )
                for alias_post_idx in post_indices[1:]:
                    alias_post_op = model_ir.operators[int(alias_post_idx)]
                    alias_output_name = str(alias_post_op.outputs[0])
                    _replace_tensor_inputs(model_ir, alias_output_name, canonical_post_out_name)
                canonical_post_tensor = model_ir.tensors.get(canonical_post_out_name, None)
                if old_concat_tensor is not None and canonical_post_tensor is not None:
                    canonical_post_tensor.dtype = str(old_concat_tensor.dtype)
                    canonical_post_tensor.quantization = _clone_quantization(old_concat_tensor.quantization)
                    canonical_post_tensor.shape = [int(v) for v in list(old_concat_tensor.shape)]
                    canonical_post_tensor.shape_signature = (
                        [int(v) for v in list(old_concat_tensor.shape_signature)]
                        if old_concat_tensor.shape_signature is not None
                        else [int(v) for v in list(old_concat_tensor.shape)]
                    )
                for post_idx in post_indices:
                    post_indices_to_remove.add(int(post_idx))

            if not valid_group:
                continue

            for post_idx in post_indices_to_remove:
                remove_indices.add(int(post_idx))

            consumers_after = _build_tensor_consumer_map(model_ir)
            final_remove_indices: set[int] = set()
            for remove_idx in sorted(list(remove_indices)):
                if int(remove_idx) < 0 or int(remove_idx) >= len(model_ir.operators):
                    continue
                if int(remove_idx) in post_indices_to_remove:
                    final_remove_indices.add(int(remove_idx))
                    continue
                remove_op = model_ir.operators[int(remove_idx)]
                if str(remove_op.op_type) != "TRANSPOSE" or len(remove_op.outputs) != 1:
                    continue
                remove_out_name = str(remove_op.outputs[0])
                if remove_out_name in model_outputs:
                    continue
                if len(consumers_after.get(remove_out_name, [])) == 0:
                    final_remove_indices.add(int(remove_idx))

            for remove_idx in sorted(list(final_remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            adapter_insertions: List[Tuple[int, OperatorIR]] = []
            op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
            for output_name, legacy_slots in legacy_consumer_slots_by_output.items():
                output_tensor = model_ir.tensors.get(str(output_name), None)
                if output_tensor is None or len(list(output_tensor.shape)) != 4:
                    continue
                slot_targets: List[Tuple[int, int]] = []
                for consumer_op_id, input_index in legacy_slots:
                    consumer_index = op_index_by_id.get(int(consumer_op_id), None)
                    if consumer_index is None:
                        continue
                    if int(consumer_index) < 0 or int(consumer_index) >= len(model_ir.operators):
                        continue
                    consumer_op = model_ir.operators[int(consumer_index)]
                    if int(input_index) < 0 or int(input_index) >= len(consumer_op.inputs):
                        continue
                    if str(consumer_op.inputs[int(input_index)]) != str(output_name):
                        continue
                    slot_targets.append((int(consumer_index), int(input_index)))
                if len(slot_targets) == 0:
                    continue

                adapter_name = _unique_tensor_name(f"{output_name}_nchw_adapter")
                legacy_meta = legacy_nchw_tensor_meta_by_output.get(str(output_name), None)
                if legacy_meta is not None:
                    adapter_dtype, adapter_quant, adapter_shape, adapter_shape_signature = legacy_meta
                else:
                    adapter_dtype = str(output_tensor.dtype)
                    adapter_quant = _clone_quantization(output_tensor.quantization)
                    adapter_shape = [int(v) for v in list(output_tensor.shape)]
                    adapter_shape_signature = (
                        [int(v) for v in list(output_tensor.shape_signature)]
                        if output_tensor.shape_signature is not None
                        else [int(v) for v in list(output_tensor.shape)]
                    )
                model_ir.tensors[adapter_name] = TensorIR(
                    name=adapter_name,
                    dtype=str(adapter_dtype),
                    shape=[int(v) for v in list(adapter_shape)],
                    shape_signature=[int(v) for v in list(adapter_shape_signature)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(adapter_quant),
                )
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(output_name), perm_nhwc_to_nchw_const_name],
                    outputs=[adapter_name],
                )
                for consumer_index, input_index in slot_targets:
                    consumer_op = model_ir.operators[int(consumer_index)]
                    new_inputs = [str(v) for v in list(consumer_op.inputs)]
                    new_inputs[int(input_index)] = str(adapter_name)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=consumer_op,
                        new_inputs=new_inputs,
                    )
                insert_index = min(int(v[0]) for v in slot_targets)
                adapter_insertions.append((int(insert_index), adapter_op))

            if len(adapter_insertions) > 0:
                inserted = 0
                for insert_idx, adapter_op in sorted(adapter_insertions, key=lambda v: int(v[0])):
                    model_ir.operators.insert(int(insert_idx + inserted), adapter_op)
                    inserted += 1

            optimized_groups += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_elementwise_concat_conv_nhwc_groups": int(optimized_groups)}


def _optimize_transpose_pre_concat_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Convert NCHW concat blocks back to NHWC when they are wrapped by transpose adapters.

    Target:
      ... -> t_i_nchw (some inputs may be pseudo-LeakyRelu outputs from transpose-wrapped NHWC)
      CONCAT(axis=1, [t_0_nchw, ...]) -> y_nchw
      y_nchw -> TRANSPOSE(0,2,3,1) -> y_nhwc

    Rewrite:
      ... -> t_i_nhwc
      CONCAT(axis=3, [t_0_nhwc, ...]) -> y_nhwc
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    # NHWC <-> NHCW (self-inverse). Used to keep SOFTMAX "last-axis" semantics
    # while lifting pre-concat chains to NHWC.
    perm_nhwc_to_nhcw = [0, 1, 3, 2]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _find_or_create_perm_tensor(
        *,
        base_name: str,
        perm: List[int],
    ) -> str:
        perm_arr = np.asarray([int(v) for v in list(perm)], dtype=np.int32)
        for tensor_name, tensor in model_ir.tensors.items():
            if tensor is None or tensor.data is None:
                continue
            try:
                data = np.asarray(tensor.data)
            except Exception:
                continue
            if data.dtype != np.int32 or int(data.size) != int(len(perm)):
                continue
            if np.array_equal(data.reshape(-1), perm_arr):
                return str(tensor_name)
        new_name = _unique_tensor_name(str(base_name))
        model_ir.tensors[new_name] = TensorIR(
            name=new_name,
            dtype="INT32",
            shape=[int(len(perm))],
            shape_signature=[int(len(perm))],
            data=np.asarray(perm_arr),
            is_variable=False,
            quantization=None,
        )
        return str(new_name)

    def _select_prelu_alpha_for_nhwc(
        *,
        alpha_data: np.ndarray,
        target_nhwc_shape: Optional[List[int]],
    ) -> Optional[np.ndarray]:
        alpha_candidates: List[np.ndarray] = []
        if int(alpha_data.ndim) == int(len(perm_nchw_to_nhwc)):
            alpha_candidates.append(
                np.transpose(alpha_data, axes=perm_nchw_to_nhwc).astype(alpha_data.dtype, copy=False)
            )
        alpha_candidates.append(np.asarray(alpha_data))
        if int(alpha_data.ndim) == 3:
            alpha_candidates.append(
                np.transpose(alpha_data, axes=[1, 2, 0]).astype(alpha_data.dtype, copy=False)
            )

        def _broadcastable(candidate: np.ndarray) -> bool:
            if target_nhwc_shape is None or not _is_fully_known_positive_shape(target_nhwc_shape):
                return True
            return (
                _broadcast_static_shapes(
                    [int(v) for v in list(target_nhwc_shape)],
                    [int(v) for v in list(candidate.shape)],
                )
                is not None
            )

        for candidate in alpha_candidates:
            if _broadcastable(candidate):
                return np.asarray(candidate)
        return None

    def _analyze_swish_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
        allow_consumer_fanout_to_slice: bool = False,
    ) -> Optional[Dict[str, Any]]:
        mul_idx = producers.get(str(input_name), None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        if str(mul_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        input_users = set(int(v) for v in consumers.get(str(input_name), []))
        if not allow_consumer_fanout_to_slice:
            if input_users != {int(consumer_idx)}:
                return None
        else:
            if int(consumer_idx) not in input_users or len(input_users) <= 0:
                return None
            for user_idx in sorted(list(input_users)):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "SLICE" or len(user_op.inputs) < 3 or len(user_op.outputs) != 1:
                    return None
                if str(user_op.inputs[0]) != str(input_name):
                    return None
                begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(user_op.inputs[1]), None))
                size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(user_op.inputs[2]), None))
                if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
                    return None
                if int(size_vals[1]) <= 0:
                    return None
                if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
                    return None

        in0_name = str(mul_op.inputs[0])
        in1_name = str(mul_op.inputs[1])
        in0_prod_idx = producers.get(in0_name, None)
        in1_prod_idx = producers.get(in1_name, None)

        log_idx: Optional[int] = None
        log_out_name: Optional[str] = None
        pre_nchw_name: Optional[str] = None
        mul_data_input_index: Optional[int] = None

        if in0_prod_idx is not None:
            in0_prod_op = model_ir.operators[int(in0_prod_idx)]
            if (
                str(in0_prod_op.op_type) == "LOGISTIC"
                and len(in0_prod_op.inputs) == 1
                and len(in0_prod_op.outputs) == 1
                and str(in0_prod_op.outputs[0]) == in0_name
            ):
                log_idx = int(in0_prod_idx)
                log_out_name = in0_name
                pre_nchw_name = str(in0_prod_op.inputs[0])
                mul_data_input_index = 1

        if log_idx is None and in1_prod_idx is not None:
            in1_prod_op = model_ir.operators[int(in1_prod_idx)]
            if (
                str(in1_prod_op.op_type) == "LOGISTIC"
                and len(in1_prod_op.inputs) == 1
                and len(in1_prod_op.outputs) == 1
                and str(in1_prod_op.outputs[0]) == in1_name
            ):
                log_idx = int(in1_prod_idx)
                log_out_name = in1_name
                pre_nchw_name = str(in1_prod_op.inputs[0])
                mul_data_input_index = 0

        if log_idx is None or log_out_name is None or pre_nchw_name is None or mul_data_input_index is None:
            return None

        if str(mul_op.inputs[int(mul_data_input_index)]) != str(pre_nchw_name):
            return None

        pre_idx = producers.get(str(pre_nchw_name), None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != str(pre_nchw_name):
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None

        if set(int(v) for v in consumers.get(str(pre_nchw_name), [])) != {int(log_idx), int(mul_idx)}:
            return None
        if set(int(v) for v in consumers.get(str(log_out_name), [])) != {int(mul_idx)}:
            return None

        return {
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "log_idx": int(log_idx),
            "mul_idx": int(mul_idx),
            "mul_data_input_index": int(mul_data_input_index),
            "log_out_name": str(log_out_name),
            "mul_out_name": str(input_name),
        }

    def _apply_swish_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        log_idx = int(plan["log_idx"])
        mul_idx = int(plan["mul_idx"])
        mul_data_input_index = int(plan["mul_data_input_index"])
        log_out_name = str(plan["log_out_name"])
        mul_out_name = str(plan["mul_out_name"])

        mul_op = model_ir.operators[int(mul_idx)]
        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(log_idx)],
            new_inputs=[pre_input_name],
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=mul_op,
            input_index=int(mul_data_input_index),
            new_input_name=pre_input_name,
        )

        for tensor_name in [log_out_name, mul_out_name]:
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(tensor_name, None),
                perm_nchw_to_nhwc,
            )
        return [int(pre_idx)]

    def _try_rewrite_swish_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        return _analyze_swish_input_to_nhwc(
            input_name=input_name,
            consumer_idx=consumer_idx,
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )

    def _try_rewrite_leaky_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        sub_idx = producers.get(str(input_name), None)
        if sub_idx is None:
            return None
        sub_op = model_ir.operators[int(sub_idx)]
        if str(sub_op.op_type) != "SUB" or len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
            return None
        if str(sub_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        input_users = [int(v) for v in consumers.get(str(input_name), [])]
        if set(input_users) != {int(concat_idx)}:
            return None

        relu_pos_out_name = str(sub_op.inputs[0])
        mul_out_name = str(sub_op.inputs[1])

        relu_pos_idx = producers.get(relu_pos_out_name, None)
        mul_idx = producers.get(mul_out_name, None)
        if relu_pos_idx is None or mul_idx is None:
            return None
        relu_pos_op = model_ir.operators[int(relu_pos_idx)]
        mul_op = model_ir.operators[int(mul_idx)]
        if str(relu_pos_op.op_type) != "RELU" or len(relu_pos_op.inputs) != 1 or len(relu_pos_op.outputs) != 1:
            return None
        if str(relu_pos_op.outputs[0]) != relu_pos_out_name:
            return None
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        if str(mul_op.outputs[0]) != mul_out_name:
            return None
        if not _is_singleton_constant_tensor(model_ir, str(mul_op.inputs[0])) and not _is_singleton_constant_tensor(
            model_ir, str(mul_op.inputs[1])
        ):
            return None

        relu_pos_in_name = str(relu_pos_op.inputs[0])
        mul_in0 = str(mul_op.inputs[0])
        mul_in1 = str(mul_op.inputs[1])
        relu_neg_out_name = mul_in0 if mul_in0 != relu_pos_in_name else mul_in1
        if relu_neg_out_name == relu_pos_in_name:
            relu_neg_out_name = mul_in1 if mul_in0 == relu_pos_in_name else mul_in0
        relu_neg_idx = producers.get(relu_neg_out_name, None)
        if relu_neg_idx is None:
            return None
        relu_neg_op = model_ir.operators[int(relu_neg_idx)]
        if str(relu_neg_op.op_type) != "RELU" or len(relu_neg_op.inputs) != 1 or len(relu_neg_op.outputs) != 1:
            return None
        if str(relu_neg_op.outputs[0]) != relu_neg_out_name:
            return None

        neg_out_name = str(relu_neg_op.inputs[0])
        neg_idx = producers.get(neg_out_name, None)
        if neg_idx is None:
            return None
        neg_op = model_ir.operators[int(neg_idx)]
        if str(neg_op.op_type) != "NEG" or len(neg_op.inputs) != 1 or len(neg_op.outputs) != 1:
            return None
        if str(neg_op.outputs[0]) != neg_out_name:
            return None

        pre_nchw_name = str(neg_op.inputs[0])
        if relu_pos_in_name != pre_nchw_name:
            return None
        pre_idx = producers.get(pre_nchw_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pre_nchw_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None
        pre_users = [int(v) for v in consumers.get(pre_nchw_name, [])]
        if set(pre_users) != {int(neg_idx), int(relu_pos_idx)}:
            return None

        if set(int(v) for v in consumers.get(neg_out_name, [])) != {int(relu_neg_idx)}:
            return None
        if set(int(v) for v in consumers.get(relu_neg_out_name, [])) != {int(mul_idx)}:
            return None
        if set(int(v) for v in consumers.get(relu_pos_out_name, [])) != {int(sub_idx)}:
            return None
        if set(int(v) for v in consumers.get(mul_out_name, [])) != {int(sub_idx)}:
            return None

        return {
            "input_name": str(input_name),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "neg_idx": int(neg_idx),
            "relu_pos_idx": int(relu_pos_idx),
            "tensors_to_permute": [
                str(neg_out_name),
                str(relu_pos_out_name),
                str(relu_neg_out_name),
                str(mul_out_name),
                str(input_name),
            ],
        }

    def _apply_leaky_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        neg_idx = int(plan["neg_idx"])
        relu_pos_idx = int(plan["relu_pos_idx"])
        tensors_to_permute = [str(v) for v in list(plan.get("tensors_to_permute", []))]

        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(neg_idx)],
            new_inputs=[pre_input_name],
        )
        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(relu_pos_idx)],
            new_inputs=[pre_input_name],
        )
        for tensor_name in tensors_to_permute:
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(tensor_name), None),
                perm_nchw_to_nhwc,
            )
        return [int(pre_idx)]

    def _try_rewrite_split_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        split_idx = producers.get(str(input_name), None)
        if split_idx is None:
            return None
        split_op = model_ir.operators[int(split_idx)]
        if str(split_op.op_type) != "SPLIT" or len(split_op.inputs) < 2 or len(split_op.outputs) < 2:
            return None
        if str(input_name) not in set(str(v) for v in list(split_op.outputs)):
            return None
        if str(input_name) in model_outputs:
            return None

        split_axis_tensor_name = str(split_op.inputs[0])
        split_axis_vals = _read_const_ints_from_tensor(model_ir.tensors.get(split_axis_tensor_name, None))
        if split_axis_vals is None or len(split_axis_vals) != 1:
            return None

        split_input_name = str(split_op.inputs[1])
        split_input_tensor = model_ir.tensors.get(split_input_name, None)
        split_rank = int(len(list(split_input_tensor.shape))) if split_input_tensor is not None else 4
        if split_rank != 4:
            return None

        split_axis = int(split_axis_vals[0])
        if split_axis < 0:
            split_axis += int(split_rank)
        if split_axis != 1:
            return None

        source_plan: Dict[str, Any]
        swish_plan = _analyze_swish_input_to_nhwc(
            input_name=split_input_name,
            consumer_idx=int(split_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )
        if swish_plan is not None:
            source_plan = {
                "kind": "swish",
                "plan": dict(swish_plan),
            }
        else:
            pre_idx = producers.get(split_input_name, None)
            if pre_idx is None:
                return None
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != str(split_input_name)
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                or str(split_input_name) in model_outputs
            ):
                return None
            pre_users = set(int(v) for v in consumers.get(str(split_input_name), []))
            source_plan = {
                "kind": "direct",
                "pre_idx": int(pre_idx),
                "pre_input_name": str(pre_op.inputs[0]),
                "remove_pre": bool(pre_users == {int(split_idx)}),
            }

        post_transpose_indices: List[int] = []
        for split_out_name in [str(v) for v in list(split_op.outputs)]:
            if split_out_name in model_outputs:
                return None
            out_users = [int(v) for v in consumers.get(split_out_name, [])]
            for user_idx in out_users:
                if int(user_idx) in {int(concat_idx), int(consumer_idx)}:
                    continue
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)

                if user_type == "ADD":
                    if (
                        len(user_op.inputs) != 2
                        or len(user_op.outputs) != 1
                        or str(user_op.outputs[0]) in model_outputs
                    ):
                        return None
                    add_users = set(int(v) for v in consumers.get(str(user_op.outputs[0]), []))
                    if int(concat_idx) not in add_users:
                        return None
                    for add_user_idx in sorted(list(add_users)):
                        if int(add_user_idx) == int(concat_idx):
                            continue
                        add_user_op = model_ir.operators[int(add_user_idx)]
                        if (
                            str(add_user_op.op_type) == "TRANSPOSE"
                            and len(add_user_op.inputs) >= 2
                            and len(add_user_op.outputs) == 1
                            and str(add_user_op.inputs[0]) == str(user_op.outputs[0])
                            and _read_transpose_perm(model_ir, add_user_op) == perm_nchw_to_nhwc
                            and str(add_user_op.outputs[0]) not in model_outputs
                        ):
                            continue
                        if (
                            str(add_user_op.op_type) == "ADD"
                            and len(add_user_op.outputs) == 1
                            and str(add_user_op.outputs[0]) not in model_outputs
                            and int(concat_idx)
                            in set(
                                int(v)
                                for v in consumers.get(str(add_user_op.outputs[0]), [])
                            )
                        ):
                            continue
                        return None
                    continue

                if (
                    user_type == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(split_out_name)
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_transpose_indices.append(int(user_idx))
                    continue

                return None

        return {
            "input_name": str(input_name),
            "split_idx": int(split_idx),
            "source_plan": dict(source_plan),
            "post_transpose_indices": [int(v) for v in list(sorted(set(post_transpose_indices)))],
        }

    def _apply_split_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        split_idx = int(plan["split_idx"])
        if int(split_idx) < 0 or int(split_idx) >= len(model_ir.operators):
            return []
        split_op = model_ir.operators[int(split_idx)]
        if str(split_op.op_type) != "SPLIT" or len(split_op.inputs) < 2:
            return []

        remove_indices: List[int] = []
        source_plan = dict(plan.get("source_plan", {}))
        source_kind = str(source_plan.get("kind", ""))
        if source_kind == "swish":
            swish_plan = source_plan.get("plan", None)
            if swish_plan is not None:
                remove_indices.extend([int(v) for v in list(_apply_swish_nhwc_plan(plan=dict(swish_plan)))])
        elif source_kind == "direct":
            _replace_operator_input_at(
                model_ir=model_ir,
                op=split_op,
                input_index=1,
                new_input_name=str(source_plan["pre_input_name"]),
            )
            if bool(source_plan.get("remove_pre", False)):
                remove_indices.append(int(source_plan["pre_idx"]))
        else:
            return []

        axis_input_name = str(split_op.inputs[0])
        axis_tensor = model_ir.tensors.get(axis_input_name, None)
        axis_users = set(int(v) for v in _build_tensor_consumer_map(model_ir).get(axis_input_name, []))
        if axis_tensor is not None and axis_users == {int(split_idx)}:
            _write_const_ints_to_tensor(axis_tensor, [3])
        else:
            axis_name_new = _unique_tensor_name(f"{axis_input_name}_nhwc")
            model_ir.tensors[axis_name_new] = TensorIR(
                name=axis_name_new,
                dtype="INT32",
                shape=[1],
                shape_signature=[1],
                data=np.asarray([3], dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=split_op,
                input_index=0,
                new_input_name=str(axis_name_new),
            )

        for split_out_name in [str(v) for v in list(split_op.outputs)]:
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(split_out_name), None),
                perm_nchw_to_nhwc,
            )

        for post_idx in [int(v) for v in list(plan.get("post_transpose_indices", []))]:
            if int(post_idx) < 0 or int(post_idx) >= len(model_ir.operators):
                continue
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                or str(post_op.outputs[0]) in set(str(v) for v in model_ir.outputs)
            ):
                continue
            _replace_tensor_inputs(
                model_ir=model_ir,
                src_name=str(post_op.outputs[0]),
                dst_name=str(post_op.inputs[0]),
            )
            remove_indices.append(int(post_idx))

        return [int(v) for v in list(remove_indices)]

    def _try_rewrite_slice_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        slice_idx = producers.get(str(input_name), None)
        if slice_idx is None:
            return None
        slice_op = model_ir.operators[int(slice_idx)]
        if str(slice_op.op_type) != "SLICE" or len(slice_op.inputs) < 3 or len(slice_op.outputs) != 1:
            return None
        if str(slice_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None

        begin_tensor = model_ir.tensors.get(str(slice_op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(slice_op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return None
        if int(size_vals[1]) <= 0:
            return None
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return None

        slice_input_name = str(slice_op.inputs[0])
        slice_input_tensor = model_ir.tensors.get(slice_input_name, None)
        if slice_input_tensor is None or len(list(slice_input_tensor.shape)) != 4:
            return None

        source_plan: Dict[str, Any]
        slice_group_indices: List[int]
        swish_plan = _analyze_swish_input_to_nhwc(
            input_name=slice_input_name,
            consumer_idx=int(slice_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
            allow_consumer_fanout_to_slice=True,
        )
        if swish_plan is not None:
            source_plan = {
                "kind": "swish",
                "plan": dict(swish_plan),
            }
            slice_group_indices = sorted(
                set(int(v) for v in consumers.get(str(slice_input_name), []))
            )
        else:
            pre_idx = producers.get(slice_input_name, None)
            if pre_idx is None:
                return None
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != str(slice_input_name)
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                or str(slice_input_name) in model_outputs
            ):
                return None
            pre_users = set(int(v) for v in consumers.get(str(slice_input_name), []))
            source_plan = {
                "kind": "direct",
                "pre_idx": int(pre_idx),
                "pre_input_name": str(pre_op.inputs[0]),
                "remove_pre": bool(pre_users == {int(slice_idx)}),
            }
            slice_group_indices = [int(slice_idx)]

        post_transpose_indices: set[int] = set()
        for target_slice_idx in [int(v) for v in list(slice_group_indices)]:
            if int(target_slice_idx) < 0 or int(target_slice_idx) >= len(model_ir.operators):
                return None
            target_slice_op = model_ir.operators[int(target_slice_idx)]
            if (
                str(target_slice_op.op_type) != "SLICE"
                or len(target_slice_op.inputs) < 3
                or len(target_slice_op.outputs) != 1
                or str(target_slice_op.inputs[0]) != str(slice_input_name)
            ):
                return None
            target_out_name = str(target_slice_op.outputs[0])
            if target_out_name in model_outputs:
                return None
            for user_idx in [int(v) for v in consumers.get(target_out_name, [])]:
                if int(user_idx) == int(concat_idx):
                    continue
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == target_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_transpose_indices.add(int(user_idx))
                    continue
                return None

        return {
            "input_name": str(input_name),
            "slice_idx": int(slice_idx),
            "concat_idx": int(concat_idx),
            "slice_group_key": (
                str(slice_input_name)
                if str(source_plan.get("kind", "")) == "swish" and len(slice_group_indices) > 1
                else ""
            ),
            "slice_group_indices": [int(v) for v in list(slice_group_indices)],
            "source_plan": dict(source_plan),
            "post_transpose_indices": [int(v) for v in sorted(post_transpose_indices)],
        }

    def _apply_slice_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        slice_idx = int(plan["slice_idx"])
        if int(slice_idx) < 0 or int(slice_idx) >= len(model_ir.operators):
            return []
        slice_op = model_ir.operators[int(slice_idx)]
        if str(slice_op.op_type) != "SLICE" or len(slice_op.inputs) < 3 or len(slice_op.outputs) != 1:
            return []

        remove_indices: List[int] = []
        source_plan = dict(plan.get("source_plan", {}))
        source_kind = str(source_plan.get("kind", ""))
        if source_kind == "swish":
            swish_plan = source_plan.get("plan", None)
            if swish_plan is not None:
                remove_indices.extend([int(v) for v in list(_apply_swish_nhwc_plan(plan=dict(swish_plan)))])
        elif source_kind == "direct":
            _replace_operator_input_at(
                model_ir=model_ir,
                op=slice_op,
                input_index=0,
                new_input_name=str(source_plan["pre_input_name"]),
            )
            if bool(source_plan.get("remove_pre", False)):
                remove_indices.append(int(source_plan["pre_idx"]))
        else:
            return []

        target_slice_indices = [
            int(v) for v in list(plan.get("slice_group_indices", [int(slice_idx)]))
        ]
        for target_slice_idx in target_slice_indices:
            if int(target_slice_idx) < 0 or int(target_slice_idx) >= len(model_ir.operators):
                return []
            target_slice_op = model_ir.operators[int(target_slice_idx)]
            if (
                str(target_slice_op.op_type) != "SLICE"
                or len(target_slice_op.inputs) < 3
                or len(target_slice_op.outputs) != 1
            ):
                return []
            begin_tensor = model_ir.tensors.get(str(target_slice_op.inputs[1]), None)
            size_tensor = model_ir.tensors.get(str(target_slice_op.inputs[2]), None)
            begin_vals = _read_const_ints_from_tensor(begin_tensor)
            size_vals = _read_const_ints_from_tensor(size_tensor)
            if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
                return []
            _write_const_ints_to_tensor(
                begin_tensor,
                [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])],
            )
            _write_const_ints_to_tensor(
                size_tensor,
                [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(target_slice_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )

        for post_idx in [int(v) for v in list(plan.get("post_transpose_indices", []))]:
            if int(post_idx) < 0 or int(post_idx) >= len(model_ir.operators):
                continue
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                or str(post_op.outputs[0]) in set(str(v) for v in model_ir.outputs)
            ):
                continue
            _replace_tensor_inputs(
                model_ir=model_ir,
                src_name=str(post_op.outputs[0]),
                dst_name=str(post_op.inputs[0]),
            )
            remove_indices.append(int(post_idx))

        return [int(v) for v in list(remove_indices)]

    def _try_rewrite_add_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        root_concat_idx: Optional[int] = None,
        visited_add_outputs: Optional[set[str]] = None,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        if root_concat_idx is None:
            root_concat_idx = int(concat_idx)
        if visited_add_outputs is None:
            visited_add_outputs = set()
        if str(input_name) in visited_add_outputs:
            return None

        add_idx = producers.get(str(input_name), None)
        if add_idx is None:
            return None
        add_op = model_ir.operators[int(add_idx)]
        if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
            return None
        if str(add_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        input_users = set(int(v) for v in consumers.get(str(input_name), []))
        if int(concat_idx) not in input_users:
            return None

        removable_post_transpose_indices: List[int] = []
        for user_idx in sorted(list(input_users)):
            if int(user_idx) == int(concat_idx):
                continue
            if int(user_idx) == int(root_concat_idx):
                continue
            user_op = model_ir.operators[int(user_idx)]
            if (
                str(user_op.op_type) == "TRANSPOSE"
                and len(user_op.inputs) >= 2
                and len(user_op.outputs) == 1
                and str(user_op.inputs[0]) == str(input_name)
                and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                and str(user_op.outputs[0]) not in model_outputs
            ):
                removable_post_transpose_indices.append(int(user_idx))
                continue
            if (
                str(user_op.op_type) == "ADD"
                and len(user_op.outputs) == 1
                and str(user_op.outputs[0]) not in model_outputs
                and int(root_concat_idx) in set(int(v) for v in consumers.get(str(user_op.outputs[0]), []))
            ):
                continue
            return None

        actions: List[Dict[str, Any]] = []
        next_visited = set(str(v) for v in visited_add_outputs)
        next_visited.add(str(input_name))
        for add_input_name in [str(v) for v in list(add_op.inputs)]:
            input_producer_idx = producers.get(add_input_name, None)
            if input_producer_idx is not None:
                input_producer_op = model_ir.operators[int(input_producer_idx)]
                input_users = set(int(v) for v in consumers.get(add_input_name, []))
                if (
                    str(input_producer_op.op_type) == "TRANSPOSE"
                    and len(input_producer_op.inputs) >= 2
                    and len(input_producer_op.outputs) == 1
                    and str(input_producer_op.outputs[0]) == add_input_name
                    and _read_transpose_perm(model_ir, input_producer_op) == perm_nhwc_to_nchw
                    and int(add_idx) in input_users
                ):
                    extra_users = {int(v) for v in input_users if int(v) != int(add_idx)}
                    if extra_users.difference({int(root_concat_idx)}):
                        return None
                    actions.append(
                        {
                            "kind": "direct",
                            "new_input_name": str(input_producer_op.inputs[0]),
                            "remove_indices": (
                                [int(input_producer_idx)] if len(extra_users) == 0 else []
                            ),
                        }
                    )
                    continue

            unary_plan = _try_rewrite_unary_input_to_nhwc(
                input_name=add_input_name,
                concat_idx=int(add_idx),
                producers=producers,
                consumers=consumers,
                model_outputs=model_outputs,
            )
            if unary_plan is not None:
                actions.append(
                    {
                        "kind": "unary",
                        "new_input_name": str(add_input_name),
                        "plan": unary_plan,
                    }
                )
                continue

            swish_plan = _analyze_swish_input_to_nhwc(
                input_name=add_input_name,
                consumer_idx=int(add_idx),
                producers=producers,
                consumers=consumers,
                model_outputs=model_outputs,
            )
            if swish_plan is not None:
                actions.append(
                    {
                        "kind": "swish",
                        "new_input_name": str(add_input_name),
                        "plan": swish_plan,
                    }
                )
                continue

            split_plan = _try_rewrite_split_input_to_nhwc(
                input_name=add_input_name,
                consumer_idx=int(add_idx),
                concat_idx=int(root_concat_idx),
                producers=producers,
                consumers=consumers,
                model_outputs=model_outputs,
            )
            if split_plan is not None:
                actions.append(
                    {
                        "kind": "split",
                        "new_input_name": str(add_input_name),
                        "plan": split_plan,
                    }
                )
                continue

            nested_add_plan = _try_rewrite_add_input_to_nhwc(
                input_name=add_input_name,
                concat_idx=int(add_idx),
                root_concat_idx=int(root_concat_idx),
                visited_add_outputs=set(next_visited),
                producers=producers,
                consumers=consumers,
                model_outputs=model_outputs,
            )
            if nested_add_plan is not None:
                actions.append(
                    {
                        "kind": "add_chain",
                        "new_input_name": str(add_input_name),
                        "plan": dict(nested_add_plan),
                    }
                )
                continue

            return None

        new_add_inputs: List[str] = []
        remove_indices: List[int] = []
        for action in actions:
            action_kind = str(action.get("kind", ""))
            if action_kind == "direct":
                new_add_inputs.append(str(action["new_input_name"]))
                remove_indices.extend([int(v) for v in list(action.get("remove_indices", []))])
                continue
            if action_kind == "swish":
                new_add_inputs.append(str(action["new_input_name"]))
                remove_indices.extend([])
                continue
            if action_kind == "unary":
                new_add_inputs.append(str(action["new_input_name"]))
                remove_indices.extend([])
                continue
            if action_kind == "split":
                new_add_inputs.append(str(action["new_input_name"]))
                remove_indices.extend([])
                continue
            if action_kind == "add_chain":
                new_add_inputs.append(str(action["new_input_name"]))
                remove_indices.extend([])
                continue
            return None

        return {
            "input_name": str(input_name),
            "add_idx": int(add_idx),
            "root_concat_idx": int(root_concat_idx),
            "new_add_inputs": [str(v) for v in new_add_inputs],
            "remove_indices": [int(v) for v in remove_indices],
            "removable_post_transpose_indices": [int(v) for v in list(sorted(set(removable_post_transpose_indices)))],
            "actions": [dict(v) for v in actions],
        }

    def _apply_add_nhwc_plan(
        *,
        plan: Dict[str, Any],
        applied_split_indices: Optional[set[int]] = None,
        applied_add_indices: Optional[set[int]] = None,
    ) -> List[int]:
        input_name = str(plan["input_name"])
        add_idx = int(plan["add_idx"])
        new_add_inputs = [str(v) for v in list(plan.get("new_add_inputs", []))]
        remove_indices = [int(v) for v in list(plan.get("remove_indices", []))]
        removable_post_transpose_indices = [
            int(v) for v in list(plan.get("removable_post_transpose_indices", []))
        ]
        actions = [dict(v) for v in list(plan.get("actions", []))]
        if applied_split_indices is None:
            applied_split_indices = set()
        if applied_add_indices is None:
            applied_add_indices = set()
        if int(add_idx) in applied_add_indices:
            return []

        for action in actions:
            action_kind = str(action.get("kind", ""))
            if action_kind == "swish":
                swish_plan = action.get("plan", None)
                if swish_plan is None:
                    continue
                swish_pre_remove = _apply_swish_nhwc_plan(plan=dict(swish_plan))
                remove_indices.extend([int(v) for v in swish_pre_remove])
                continue
            if action_kind == "unary":
                unary_plan = action.get("plan", None)
                if unary_plan is None:
                    continue
                unary_pre_remove = _apply_unary_nhwc_plan(plan=dict(unary_plan))
                remove_indices.extend([int(v) for v in unary_pre_remove])
                continue
            if action_kind == "split":
                split_plan = action.get("plan", None)
                if split_plan is None:
                    continue
                split_idx = int(split_plan.get("split_idx", -1))
                if int(split_idx) in applied_split_indices:
                    continue
                split_pre_remove = _apply_split_nhwc_plan(plan=dict(split_plan))
                remove_indices.extend([int(v) for v in list(split_pre_remove)])
                applied_split_indices.add(int(split_idx))
                continue
            if action_kind == "add_chain":
                nested_plan = action.get("plan", None)
                if nested_plan is None:
                    continue
                nested_pre_remove = _apply_add_nhwc_plan(
                    plan=dict(nested_plan),
                    applied_split_indices=applied_split_indices,
                    applied_add_indices=applied_add_indices,
                )
                remove_indices.extend([int(v) for v in list(nested_pre_remove)])
                continue

        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(add_idx)],
            new_inputs=[str(v) for v in list(new_add_inputs)],
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )
        for post_idx in removable_post_transpose_indices:
            if int(post_idx) < 0 or int(post_idx) >= len(model_ir.operators):
                continue
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                or str(post_op.outputs[0]) in set(str(v) for v in model_ir.outputs)
            ):
                continue
            _replace_tensor_inputs(
                model_ir=model_ir,
                src_name=str(post_op.outputs[0]),
                dst_name=str(post_op.inputs[0]),
            )
            remove_indices.append(int(post_idx))
        applied_add_indices.add(int(add_idx))
        return [int(v) for v in remove_indices]

    def _try_rewrite_pad_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        pad_idx = producers.get(str(input_name), None)
        if pad_idx is None:
            return None
        pad_op = model_ir.operators[int(pad_idx)]
        if str(pad_op.op_type) != "PAD" or len(pad_op.inputs) < 2 or len(pad_op.outputs) != 1:
            return None
        if str(pad_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(concat_idx)}:
            return None

        pad_input_name = str(pad_op.inputs[0])
        pre_idx = producers.get(pad_input_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pad_input_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pad_input_name) in model_outputs:
            return None

        pads_tensor_name = str(pad_op.inputs[1])
        pads_tensor = model_ir.tensors.get(pads_tensor_name, None)
        if pads_tensor is None or pads_tensor.data is None:
            return None
        try:
            pads_array = np.asarray(pads_tensor.data)
            pads_pairs = pads_array.reshape(4, 2)
        except Exception:
            return None
        if int(pads_pairs.size) != 8:
            return None

        pre_users = set(int(v) for v in consumers.get(pad_input_name, []))
        remove_pre = pre_users == {int(pad_idx)}

        return {
            "input_name": str(input_name),
            "pad_idx": int(pad_idx),
            "pad_input_name": str(pad_input_name),
            "pad_input_rewritten_name": str(pre_op.inputs[0]),
            "pads_tensor_name": str(pads_tensor_name),
            "pre_idx": int(pre_idx),
            "remove_pre": bool(remove_pre),
        }

    def _apply_pad_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        input_name = str(plan["input_name"])
        pad_idx = int(plan["pad_idx"])
        pre_idx = int(plan["pre_idx"])
        remove_pre = bool(plan["remove_pre"])
        pad_rewritten_input_name = str(plan["pad_input_rewritten_name"])
        pads_tensor_name = str(plan["pads_tensor_name"])

        pad_op = model_ir.operators[int(pad_idx)]
        _set_operator_inputs(
            model_ir=model_ir,
            op=pad_op,
            new_inputs=[pad_rewritten_input_name, str(pad_op.inputs[1])],
        )

        pads_tensor = model_ir.tensors.get(pads_tensor_name, None)
        if pads_tensor is not None and pads_tensor.data is not None:
            pads_array = np.asarray(pads_tensor.data)
            pads_pairs = pads_array.reshape(4, 2)
            # NCHW [N,C,H,W] -> NHWC [N,H,W,C]
            pads_pairs = np.asarray(
                [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
                dtype=pads_array.dtype,
            )
            pads_tensor.data = pads_pairs
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )

        return [int(pre_idx)] if remove_pre else []

    def _try_rewrite_unary_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        unary_idx = producers.get(str(input_name), None)
        if unary_idx is None:
            return None
        unary_op = model_ir.operators[int(unary_idx)]
        if str(unary_op.op_type) not in {"RELU", "RELU6", "LOGISTIC", "TANH", "GELU"}:
            return None
        if len(unary_op.inputs) != 1 or len(unary_op.outputs) != 1:
            return None
        if str(unary_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(concat_idx)}:
            return None

        pre_nchw_name = str(unary_op.inputs[0])
        pre_idx = producers.get(pre_nchw_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pre_nchw_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(pre_nchw_name, [])) != {int(unary_idx)}:
            return None

        return {
            "input_name": str(input_name),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "unary_idx": int(unary_idx),
        }

    def _try_rewrite_softmax_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        softmax_idx = producers.get(str(input_name), None)
        if softmax_idx is None:
            return None
        softmax_op = model_ir.operators[int(softmax_idx)]
        if str(softmax_op.op_type) != "SOFTMAX":
            return None
        if len(softmax_op.inputs) != 1 or len(softmax_op.outputs) != 1:
            return None
        if str(softmax_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(concat_idx)}:
            return None

        pre_nchw_name = str(softmax_op.inputs[0])
        pre_idx = producers.get(pre_nchw_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pre_nchw_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None
        # Keep strict: the removed pre-transpose output must only feed SOFTMAX.
        if set(int(v) for v in consumers.get(pre_nchw_name, [])) != {int(softmax_idx)}:
            return None

        pre_input_tensor = model_ir.tensors.get(str(pre_op.inputs[0]), None)
        softmax_output_tensor = model_ir.tensors.get(str(input_name), None)
        if (
            pre_input_tensor is None
            or softmax_output_tensor is None
            or len(list(pre_input_tensor.shape)) != 4
            or len(list(softmax_output_tensor.shape)) != 4
        ):
            return None

        return {
            "input_name": str(input_name),
            "softmax_idx": int(softmax_idx),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "softmax_input_name": str(pre_nchw_name),
        }

    def _try_rewrite_dequant_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        dq_idx = producers.get(str(input_name), None)
        if dq_idx is None:
            return None
        dq_op = model_ir.operators[int(dq_idx)]
        if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
            return None
        if str(dq_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(concat_idx)}:
            return None

        pre_nchw_name = str(dq_op.inputs[0])
        pre_idx = producers.get(pre_nchw_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pre_nchw_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None

        pre_users = set(int(v) for v in consumers.get(pre_nchw_name, []))
        remove_pre = pre_users == {int(dq_idx)}
        return {
            "input_name": str(input_name),
            "dq_idx": int(dq_idx),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "remove_pre": bool(remove_pre),
        }

    def _apply_dequant_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        input_name = str(plan["input_name"])
        dq_idx = int(plan["dq_idx"])
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        remove_pre = bool(plan["remove_pre"])

        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(dq_idx)],
            new_inputs=[pre_input_name],
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(pre_idx)] if remove_pre else []

    def _try_rewrite_prelu_input_to_nhwc(
        *,
        input_name: str,
        concat_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        prelu_idx = producers.get(str(input_name), None)
        if prelu_idx is None:
            return None
        prelu_op = model_ir.operators[int(prelu_idx)]
        if str(prelu_op.op_type) != "PRELU" or len(prelu_op.inputs) != 2 or len(prelu_op.outputs) != 1:
            return None
        if str(prelu_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(concat_idx)}:
            return None

        pre_nchw_name = str(prelu_op.inputs[0])
        pre_idx = producers.get(pre_nchw_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pre_nchw_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(pre_nchw_name, [])) != {int(prelu_idx)}:
            return None

        alpha_name = str(prelu_op.inputs[1])
        alpha_tensor = model_ir.tensors.get(alpha_name, None)
        if alpha_tensor is None or alpha_tensor.data is None:
            return None
        pre_input_tensor = model_ir.tensors.get(str(pre_op.inputs[0]), None)
        target_nhwc_shape = (
            [int(v) for v in list(pre_input_tensor.shape)]
            if pre_input_tensor is not None and pre_input_tensor.shape is not None
            else None
        )
        alpha_data = np.asarray(alpha_tensor.data)
        if _select_prelu_alpha_for_nhwc(
            alpha_data=alpha_data,
            target_nhwc_shape=target_nhwc_shape,
        ) is None:
            return None

        return {
            "input_name": str(input_name),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "prelu_idx": int(prelu_idx),
            "target_nhwc_shape": (
                [int(v) for v in list(target_nhwc_shape)]
                if target_nhwc_shape is not None
                else None
            ),
        }

    def _apply_unary_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        input_name = str(plan["input_name"])
        unary_idx = int(plan["unary_idx"])
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])

        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(unary_idx)],
            new_inputs=[pre_input_name],
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(pre_idx)]

    def _apply_softmax_nhwc_plan(*, plan: Dict[str, Any]) -> Optional[List[int]]:
        input_name = str(plan["input_name"])
        softmax_idx = int(plan["softmax_idx"])
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        old_softmax_input_name = str(plan["softmax_input_name"])

        if softmax_idx < 0 or softmax_idx >= len(model_ir.operators):
            return None
        softmax_op = model_ir.operators[int(softmax_idx)]
        if str(softmax_op.op_type) != "SOFTMAX" or len(softmax_op.inputs) != 1 or len(softmax_op.outputs) != 1:
            return None
        if str(softmax_op.outputs[0]) != input_name:
            return None

        pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
        old_softmax_input_tensor = model_ir.tensors.get(old_softmax_input_name, None)
        softmax_output_tensor = model_ir.tensors.get(input_name, None)
        if (
            pre_input_tensor is None
            or old_softmax_input_tensor is None
            or softmax_output_tensor is None
            or len(list(pre_input_tensor.shape)) != 4
            or len(list(old_softmax_input_tensor.shape)) != 4
            or len(list(softmax_output_tensor.shape)) != 4
        ):
            return None

        pre_input_shape = [int(v) for v in list(pre_input_tensor.shape)]
        pre_input_signature = (
            [int(v) for v in list(pre_input_tensor.shape_signature)]
            if pre_input_tensor.shape_signature is not None
            else [int(v) for v in list(pre_input_tensor.shape)]
        )
        axis_last_shape = _permute_shape(pre_input_shape, perm_nhwc_to_nhcw)
        axis_last_signature = _permute_shape(pre_input_signature, perm_nhwc_to_nhcw)
        if axis_last_shape is None or axis_last_signature is None:
            return None

        axis_last_input_name = _unique_tensor_name(f"{old_softmax_input_name}_axis_last")
        axis_last_output_name = _unique_tensor_name(f"{input_name}_axis_last")
        perm_tensor_name = _find_or_create_perm_tensor(
            base_name=f"{input_name}_nhwc_to_nhcw_perm",
            perm=perm_nhwc_to_nhcw,
        )

        model_ir.tensors[axis_last_input_name] = TensorIR(
            name=axis_last_input_name,
            dtype=str(old_softmax_input_tensor.dtype),
            shape=[int(v) for v in list(axis_last_shape)],
            shape_signature=[int(v) for v in list(axis_last_signature)],
            data=None,
            is_variable=False,
            quantization=_clone_quantization(old_softmax_input_tensor.quantization),
        )
        model_ir.tensors[axis_last_output_name] = TensorIR(
            name=axis_last_output_name,
            dtype=str(softmax_output_tensor.dtype),
            shape=[int(v) for v in list(axis_last_shape)],
            shape_signature=[int(v) for v in list(axis_last_signature)],
            data=None,
            is_variable=False,
            quantization=_clone_quantization(softmax_output_tensor.quantization),
        )

        # Keep SOFTMAX last-axis semantics by locally rotating NHWC -> NHCW -> SOFTMAX -> NHWC.
        _set_operator_inputs(
            model_ir=model_ir,
            op=softmax_op,
            new_inputs=[axis_last_input_name],
        )
        _set_operator_outputs(
            model_ir=model_ir,
            op=softmax_op,
            new_outputs=[axis_last_output_name],
        )

        model_ir.operators.insert(
            int(softmax_idx),
            OperatorIR(
                op_type="TRANSPOSE",
                inputs=[pre_input_name, perm_tensor_name],
                outputs=[axis_last_input_name],
                options={},
            ),
        )
        model_ir.operators.insert(
            int(softmax_idx) + 2,
            OperatorIR(
                op_type="TRANSPOSE",
                inputs=[axis_last_output_name, perm_tensor_name],
                outputs=[input_name],
                options={},
            ),
        )

        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(pre_idx)]

    def _apply_prelu_nhwc_plan(*, plan: Dict[str, Any]) -> Optional[List[int]]:
        input_name = str(plan["input_name"])
        prelu_idx = int(plan["prelu_idx"])
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        target_nhwc_shape = (
            [int(v) for v in list(plan.get("target_nhwc_shape", []))]
            if plan.get("target_nhwc_shape", None) is not None
            else None
        )

        prelu_op = model_ir.operators[int(prelu_idx)]
        alpha_name = str(prelu_op.inputs[1])
        alpha_tensor = model_ir.tensors.get(alpha_name, None)
        if alpha_tensor is None or alpha_tensor.data is None:
            return None
        alpha_data = np.asarray(alpha_tensor.data)
        selected_alpha = _select_prelu_alpha_for_nhwc(
            alpha_data=alpha_data,
            target_nhwc_shape=target_nhwc_shape,
        )
        if selected_alpha is None:
            return None

        selected_alpha_name = str(alpha_name)
        alpha_users = [int(v) for v in _build_tensor_consumer_map(model_ir).get(alpha_name, [])]
        alpha_needs_rewrite = (
            selected_alpha.shape != alpha_data.shape
            or not np.array_equal(selected_alpha, alpha_data)
        )
        if alpha_needs_rewrite:
            shape, signature = normalize_onnx_shape(list(selected_alpha.shape))
            if len(alpha_users) == 1 and int(alpha_users[0]) == int(prelu_idx):
                alpha_tensor.data = np.asarray(selected_alpha)
                alpha_tensor.shape = [int(v) for v in list(shape)]
                alpha_tensor.shape_signature = [int(v) for v in list(signature)]
            else:
                selected_alpha_name = _unique_tensor_name(f"{alpha_name}_nhwc")
                model_ir.tensors[selected_alpha_name] = TensorIR(
                    name=selected_alpha_name,
                    dtype=str(alpha_tensor.dtype),
                    shape=[int(v) for v in list(shape)],
                    shape_signature=[int(v) for v in list(signature)],
                    data=np.asarray(selected_alpha),
                    is_variable=False,
                    quantization=_clone_quantization(alpha_tensor.quantization),
                )

        _set_operator_inputs(
            model_ir=model_ir,
            op=prelu_op,
            new_inputs=[pre_input_name, str(selected_alpha_name)],
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(pre_idx)]

    def _project_shape_after_nchw_to_nhwc(tensor_name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            return None
        shape = [int(v) for v in list(tensor.shape)]
        if len(shape) != 4:
            return None
        return [int(shape[0]), int(shape[2]), int(shape[3]), int(shape[1])]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_out_name = str(concat_op.outputs[0])
            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            if concat_out_tensor is None or len(list(concat_out_tensor.shape)) != 4:
                continue

            concat_axis_old = int(concat_op.options.get("axis", 1))
            if concat_axis_old < 0:
                concat_axis_old += 4
            if concat_axis_old != 1:
                continue

            concat_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(concat_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            post_quantize_idx: Optional[int] = None
            post_quantize_out_name: Optional[str] = None
            valid_post_only = True
            for user_idx in concat_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                elif (
                    len(concat_users) == 1
                    and str(user_op.op_type) == "QUANTIZE"
                    and len(user_op.inputs) == 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    q_out_name = str(user_op.outputs[0])
                    q_users = [int(v) for v in consumers.get(q_out_name, [])]
                    if len(q_users) == 0:
                        valid_post_only = False
                        break
                    valid_q_post = True
                    for q_user_idx in q_users:
                        q_user_op = model_ir.operators[int(q_user_idx)]
                        if (
                            str(q_user_op.op_type) == "TRANSPOSE"
                            and len(q_user_op.inputs) >= 2
                            and len(q_user_op.outputs) == 1
                            and str(q_user_op.inputs[0]) == q_out_name
                            and _read_transpose_perm(model_ir, q_user_op) == perm_nchw_to_nhwc
                            and str(q_user_op.outputs[0]) not in model_outputs
                        ):
                            post_indices.append(int(q_user_idx))
                            post_output_names.append(str(q_user_op.outputs[0]))
                        else:
                            valid_q_post = False
                            break
                    if not valid_q_post:
                        valid_post_only = False
                        break
                    post_quantize_idx = int(user_idx)
                    post_quantize_out_name = str(q_out_name)
                else:
                    valid_post_only = False
                    break
            if not valid_post_only or len(post_indices) == 0:
                continue
            post_op_ids: List[int] = [
                int(id(model_ir.operators[int(post_idx)]))
                for post_idx in post_indices
                if int(post_idx) >= 0 and int(post_idx) < len(model_ir.operators)
            ]

            concat_input_actions: List[Dict[str, Any]] = []
            rewritable = True
            for input_name in [str(v) for v in list(concat_op.inputs)]:
                input_producer_idx = producers.get(input_name, None)
                if input_producer_idx is not None:
                    input_producer_op = model_ir.operators[int(input_producer_idx)]
                    if (
                        str(input_producer_op.op_type) == "TRANSPOSE"
                        and len(input_producer_op.inputs) >= 2
                        and len(input_producer_op.outputs) == 1
                        and str(input_producer_op.outputs[0]) == input_name
                        and _read_transpose_perm(model_ir, input_producer_op) == perm_nhwc_to_nchw
                        and int(concat_idx) in set(int(v) for v in consumers.get(input_name, []))
                    ):
                        input_users = set(int(v) for v in consumers.get(input_name, []))
                        remove_indices = (
                            [int(input_producer_idx)]
                            if input_users == {int(concat_idx)}
                            else []
                        )
                        concat_input_actions.append(
                            {
                                "kind": "direct",
                                "input_name": str(input_name),
                                "new_input_name": str(input_producer_op.inputs[0]),
                                "remove_indices": remove_indices,
                            }
                        )
                        continue

                pad_plan = _try_rewrite_pad_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if pad_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "pad",
                            "input_name": str(input_name),
                            "plan": dict(pad_plan),
                        }
                    )
                    continue

                leaky_plan = _try_rewrite_leaky_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if leaky_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "leaky",
                            "input_name": str(input_name),
                            "plan": dict(leaky_plan),
                        }
                    )
                    continue

                dequant_plan = _try_rewrite_dequant_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if dequant_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "dequantize",
                            "input_name": str(input_name),
                            "plan": dict(dequant_plan),
                        }
                    )
                    continue

                unary_plan = _try_rewrite_unary_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if unary_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "unary",
                            "input_name": str(input_name),
                            "plan": dict(unary_plan),
                        }
                    )
                    continue

                softmax_plan = _try_rewrite_softmax_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if softmax_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "softmax",
                            "input_name": str(input_name),
                            "plan": dict(softmax_plan),
                        }
                    )
                    continue

                prelu_plan = _try_rewrite_prelu_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if prelu_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "prelu",
                            "input_name": str(input_name),
                            "plan": dict(prelu_plan),
                        }
                    )
                    continue

                swish_plan = _try_rewrite_swish_input_to_nhwc(
                    input_name=input_name,
                    consumer_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if swish_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "swish",
                            "input_name": str(input_name),
                            "plan": dict(swish_plan),
                        }
                    )
                    continue

                add_plan = _try_rewrite_add_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    root_concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if add_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "add",
                            "input_name": str(input_name),
                            "plan": dict(add_plan),
                        }
                    )
                    continue

                split_plan = _try_rewrite_split_input_to_nhwc(
                    input_name=input_name,
                    consumer_idx=int(concat_idx),
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if split_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "split",
                            "input_name": str(input_name),
                            "plan": dict(split_plan),
                        }
                    )
                    continue

                slice_plan = _try_rewrite_slice_input_to_nhwc(
                    input_name=input_name,
                    concat_idx=int(concat_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if slice_plan is not None:
                    concat_input_actions.append(
                        {
                            "kind": "slice",
                            "input_name": str(input_name),
                            "plan": dict(slice_plan),
                        }
                    )
                    continue

                rewritable = False
                break

            if not rewritable:
                continue

            softmax_action_count = int(
                sum(1 for action in concat_input_actions if str(action.get("kind", "")) == "softmax")
            )
            if softmax_action_count > 0:
                # Safety + benefit gate:
                # - support only float-path concat-post-transpose rewrite
                # - require at least one non-softmax input so total transpose count decreases
                if post_quantize_idx is not None:
                    continue
                if softmax_action_count != 1:
                    continue
                if int(len(concat_input_actions) - softmax_action_count) <= 0:
                    continue

            all_direct_input_actions = all(
                str(action.get("kind", "")) == "direct"
                for action in concat_input_actions
            )
            nhwc_inputs_ok = True
            nhwc_ref_shape: Optional[List[int]] = None
            for action in concat_input_actions:
                action_kind = str(action.get("kind", ""))
                if action_kind == "direct":
                    projected_input_name = str(action["new_input_name"])
                    input_tensor = model_ir.tensors.get(projected_input_name, None)
                    if input_tensor is None or len(list(input_tensor.shape)) != 4:
                        nhwc_inputs_ok = False
                        break
                    shape = [int(v) for v in list(input_tensor.shape)]
                elif action_kind in {"swish", "leaky", "add", "unary", "pad", "prelu", "dequantize", "softmax", "split", "slice"}:
                    projected_input_name = str(action["input_name"])
                    projected_shape = _project_shape_after_nchw_to_nhwc(projected_input_name)
                    if projected_shape is None:
                        nhwc_inputs_ok = False
                        break
                    shape = [int(v) for v in list(projected_shape)]
                else:
                    nhwc_inputs_ok = False
                    break
                if nhwc_ref_shape is None:
                    nhwc_ref_shape = list(shape)
                else:
                    for dim_idx in [0, 1, 2]:
                        if int(shape[dim_idx]) != int(nhwc_ref_shape[dim_idx]):
                            nhwc_inputs_ok = False
                            break
                if not nhwc_inputs_ok:
                    break
            if not nhwc_inputs_ok or nhwc_ref_shape is None:
                # Some paths keep stale rank-4 metadata even though runtime concat
                # is valid (e.g. resize branches). For strict direct transpose
                # wrappers, the algebraic rewrite is still safe.
                if not all_direct_input_actions:
                    continue

            new_concat_inputs: List[str] = []
            pre_remove_indices: List[int] = []
            applied_split_indices: set[int] = set()
            applied_add_indices: set[int] = set()
            applied_slice_group_keys: set[str] = set()
            for action in concat_input_actions:
                action_kind = str(action.get("kind", ""))
                if action_kind == "direct":
                    new_concat_inputs.append(str(action["new_input_name"]))
                    pre_remove_indices.extend([int(v) for v in list(action.get("remove_indices", []))])
                    continue
                if action_kind == "pad":
                    new_concat_inputs.append(str(action["input_name"]))
                    pre_remove_indices.extend(
                        _apply_pad_nhwc_plan(plan=dict(action["plan"]))
                    )
                    continue
                if action_kind == "swish":
                    new_concat_inputs.append(str(action["input_name"]))
                    pre_remove_indices.extend(
                        _apply_swish_nhwc_plan(plan=dict(action["plan"]))
                    )
                    continue
                if action_kind == "leaky":
                    new_concat_inputs.append(str(action["input_name"]))
                    pre_remove_indices.extend(
                        _apply_leaky_nhwc_plan(plan=dict(action["plan"]))
                    )
                    continue
                if action_kind == "dequantize":
                    new_concat_inputs.append(str(action["input_name"]))
                    pre_remove_indices.extend(
                        _apply_dequant_nhwc_plan(plan=dict(action["plan"]))
                    )
                    continue
                if action_kind == "add":
                    new_concat_inputs.append(str(action["input_name"]))
                    pre_remove_indices.extend(
                        _apply_add_nhwc_plan(
                            plan=dict(action["plan"]),
                            applied_split_indices=applied_split_indices,
                            applied_add_indices=applied_add_indices,
                        )
                    )
                    continue
                if action_kind == "split":
                    new_concat_inputs.append(str(action["input_name"]))
                    split_plan = dict(action["plan"])
                    split_idx = int(split_plan.get("split_idx", -1))
                    if int(split_idx) not in applied_split_indices:
                        pre_remove_indices.extend(
                            _apply_split_nhwc_plan(plan=split_plan)
                        )
                        applied_split_indices.add(int(split_idx))
                    continue
                if action_kind == "slice":
                    new_concat_inputs.append(str(action["input_name"]))
                    slice_plan = dict(action["plan"])
                    slice_group_key = str(slice_plan.get("slice_group_key", ""))
                    if slice_group_key != "" and slice_group_key in applied_slice_group_keys:
                        continue
                    pre_remove_indices.extend(_apply_slice_nhwc_plan(plan=slice_plan))
                    if slice_group_key != "":
                        applied_slice_group_keys.add(slice_group_key)
                    continue
                if action_kind == "unary":
                    new_concat_inputs.append(str(action["input_name"]))
                    pre_remove_indices.extend(
                        _apply_unary_nhwc_plan(plan=dict(action["plan"]))
                    )
                    continue
                if action_kind == "softmax":
                    new_concat_inputs.append(str(action["input_name"]))
                    softmax_remove = _apply_softmax_nhwc_plan(plan=dict(action["plan"]))
                    if softmax_remove is None:
                        rewritable = False
                        break
                    pre_remove_indices.extend([int(v) for v in list(softmax_remove)])
                    continue
                if action_kind == "prelu":
                    new_concat_inputs.append(str(action["input_name"]))
                    prelu_remove = _apply_prelu_nhwc_plan(plan=dict(action["plan"]))
                    if prelu_remove is None:
                        rewritable = False
                        break
                    pre_remove_indices.extend([int(v) for v in list(prelu_remove)])
                    continue
                rewritable = False
                break
            if not rewritable:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in new_concat_inputs],
            )
            concat_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(concat_out_name, None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            if post_quantize_idx is None:
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=concat_op,
                    new_outputs=[canonical_post_output_name],
                )
            else:
                post_quantize_op = model_ir.operators[int(post_quantize_idx)]
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=post_quantize_op,
                    new_outputs=[canonical_post_output_name],
                )
            for alias_post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_post_output_name, canonical_post_output_name)

            old_concat_tensor = model_ir.tensors.get(concat_out_name, None)
            canonical_post_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if canonical_post_tensor is not None:
                if post_quantize_idx is None:
                    if old_concat_tensor is not None:
                        canonical_post_tensor.dtype = str(old_concat_tensor.dtype)
                        canonical_post_tensor.quantization = _clone_quantization(old_concat_tensor.quantization)
                else:
                    old_q_out_tensor = (
                        model_ir.tensors.get(str(post_quantize_out_name), None)
                        if post_quantize_out_name is not None
                        else None
                    )
                    if old_q_out_tensor is not None:
                        canonical_post_tensor.dtype = str(old_q_out_tensor.dtype)
                        canonical_post_tensor.quantization = _clone_quantization(old_q_out_tensor.quantization)
                    if old_concat_tensor is not None:
                        canonical_post_tensor.shape = [int(v) for v in list(old_concat_tensor.shape)]
                        canonical_post_tensor.shape_signature = (
                            [int(v) for v in list(old_concat_tensor.shape_signature)]
                            if old_concat_tensor.shape_signature is not None
                            else [int(v) for v in list(old_concat_tensor.shape)]
                        )

            remove_op_ids: set[int] = set(int(v) for v in list(post_op_ids))
            for remove_idx in pre_remove_indices:
                if int(remove_idx) < 0 or int(remove_idx) >= len(model_ir.operators):
                    continue
                remove_op_ids.add(int(id(model_ir.operators[int(remove_idx)])))
            remove_indices = sorted(
                [
                    int(op_idx)
                    for op_idx, op in enumerate(model_ir.operators)
                    if int(id(op)) in remove_op_ids
                ],
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_concat_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_dequant_concat_quantize_post_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Convert strict NCHW concat quantize bridges back to NHWC and remove adapter transposes.

    Target:
      x_i_nhwc --TRANSPOSE(0,3,1,2)--> x_i_nchw --DEQUANTIZE--> f_i_nchw
      CONCAT(axis=1, [f_0_nchw, ...]) -> f_cat_nchw
      f_cat_nchw --QUANTIZE--> q_cat_nchw --TRANSPOSE(0,2,3,1)--> q_cat_nhwc

    Rewrite:
      x_i_nhwc --DEQUANTIZE--> f_i_nhwc
      CONCAT(axis=3, [f_0_nhwc, ...]) -> f_cat_nhwc
      f_cat_nhwc --QUANTIZE--> q_cat_nhwc
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue

            concat_axis = int(concat_op.options.get("axis", 1))
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue

            concat_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(concat_users) != 1:
                continue
            q_idx = int(concat_users[0])
            q_op = model_ir.operators[int(q_idx)]
            if (
                str(q_op.op_type) != "QUANTIZE"
                or len(q_op.inputs) != 1
                or len(q_op.outputs) != 1
                or str(q_op.inputs[0]) != concat_out_name
            ):
                continue

            q_out_name = str(q_op.outputs[0])
            if q_out_name in model_outputs:
                continue
            q_users = [int(v) for v in consumers.get(q_out_name, [])]
            if len(q_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            valid_posts = True
            for q_user_idx in q_users:
                q_user_op = model_ir.operators[int(q_user_idx)]
                if (
                    str(q_user_op.op_type) == "TRANSPOSE"
                    and len(q_user_op.inputs) >= 2
                    and len(q_user_op.outputs) == 1
                    and str(q_user_op.inputs[0]) == q_out_name
                    and _read_transpose_perm(model_ir, q_user_op) == perm_nchw_to_nhwc
                    and str(q_user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(q_user_idx))
                    post_output_names.append(str(q_user_op.outputs[0]))
                else:
                    valid_posts = False
                    break
            if not valid_posts or len(post_indices) == 0:
                continue

            dq_rewrite_plans: List[Dict[str, Any]] = []
            pre_indices_to_remove: List[int] = []
            rewrite_ok = True
            for concat_input_name in [str(v) for v in list(concat_op.inputs)]:
                dq_idx = producers.get(str(concat_input_name), None)
                if dq_idx is None:
                    rewrite_ok = False
                    break
                dq_op = model_ir.operators[int(dq_idx)]
                if (
                    str(dq_op.op_type) != "DEQUANTIZE"
                    or len(dq_op.inputs) != 1
                    or len(dq_op.outputs) != 1
                    or str(dq_op.outputs[0]) != str(concat_input_name)
                    or str(concat_input_name) in model_outputs
                ):
                    rewrite_ok = False
                    break
                dq_users = [int(v) for v in consumers.get(str(concat_input_name), [])]
                if set(dq_users) != {int(concat_idx)}:
                    rewrite_ok = False
                    break

                pre_name = str(dq_op.inputs[0])
                pre_idx = producers.get(pre_name, None)
                if pre_idx is None:
                    rewrite_ok = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != pre_name
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                ):
                    rewrite_ok = False
                    break
                if pre_name in model_outputs:
                    rewrite_ok = False
                    break

                dq_rewrite_plans.append(
                    {
                        "dq_idx": int(dq_idx),
                        "concat_input_name": str(concat_input_name),
                        "new_dq_input_name": str(pre_op.inputs[0]),
                    }
                )

                pre_users = [int(v) for v in consumers.get(pre_name, [])]
                if set(pre_users) == {int(dq_idx)}:
                    pre_indices_to_remove.append(int(pre_idx))

            if not rewrite_ok:
                continue

            for dq_plan in dq_rewrite_plans:
                dq_idx = int(dq_plan["dq_idx"])
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=model_ir.operators[int(dq_idx)],
                    new_inputs=[str(dq_plan["new_dq_input_name"])],
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(dq_plan["concat_input_name"]), None),
                    perm_nchw_to_nhwc,
                )

            concat_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(concat_out_name, None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_post_output_name, canonical_post_output_name)

            old_q_out_tensor = model_ir.tensors.get(q_out_name, None)
            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            canonical_post_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if canonical_post_tensor is not None:
                if old_q_out_tensor is not None:
                    canonical_post_tensor.dtype = str(old_q_out_tensor.dtype)
                    canonical_post_tensor.quantization = _clone_quantization(old_q_out_tensor.quantization)
                if concat_out_tensor is not None:
                    canonical_post_tensor.shape = [int(v) for v in list(concat_out_tensor.shape)]
                    canonical_post_tensor.shape_signature = (
                        [int(v) for v in list(concat_out_tensor.shape_signature)]
                        if concat_out_tensor.shape_signature is not None
                        else [int(v) for v in list(concat_out_tensor.shape)]
                    )

            remove_indices = sorted(
                list({*pre_indices_to_remove, *post_indices}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                if int(remove_idx) == int(concat_idx) or int(remove_idx) == int(q_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_dequant_concat_quantize_post_nhwc_chains": int(optimized)}


def _optimize_transpose_concat_unary_fanout_conv_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse NCHW concat blocks that are bracketed by transpose adapters.

    Target:
      x_i_nhwc --TRANSPOSE(0,3,1,2)--> x_i_nchw
      CONCAT(axis=1, [x_i_nchw...]) -> y_nchw
      (optional unary chain) -> z_nchw
      z_nchw --TRANSPOSE(0,2,3,1)--> z_i_nhwc -> CONV_2D/DEPTHWISE_CONV_2D

    Rewrite:
      CONCAT(axis=3, [x_i_nhwc...]) -> y_nhwc
      (same unary chain in NHWC semantics) -> z_nhwc
      z_nhwc -> CONV_2D/DEPTHWISE_CONV_2D
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {"RELU", "RELU6", "LEAKY_RELU", "LOGISTIC", "TANH", "HARD_SWISH"}
    conv_like_ops = {"CONV_2D", "DEPTHWISE_CONV_2D"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_axis = int(concat_op.options.get("axis", 1))
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue

            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue

            pre_transpose_indices: List[int] = []
            concat_inputs_nhwc: List[str] = []
            rewritable_inputs = True
            for concat_input_name in [str(v) for v in list(concat_op.inputs)]:
                pre_idx = producers.get(concat_input_name, None)
                if pre_idx is None:
                    rewritable_inputs = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(concat_input_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    or str(pre_op.outputs[0]) in model_outputs
                ):
                    rewritable_inputs = False
                    break
                if set(int(v) for v in consumers.get(str(concat_input_name), [])) != {int(concat_idx)}:
                    rewritable_inputs = False
                    break
                pre_transpose_indices.append(int(pre_idx))
                concat_inputs_nhwc.append(str(pre_op.inputs[0]))
            if not rewritable_inputs or len(pre_transpose_indices) == 0:
                continue

            unary_output_names: List[str] = []
            nchw_tail_name = str(concat_out_name)
            while True:
                tail_users = [int(v) for v in consumers.get(str(nchw_tail_name), [])]
                if len(tail_users) != 1:
                    break
                unary_idx = int(tail_users[0])
                unary_op = model_ir.operators[int(unary_idx)]
                if (
                    str(unary_op.op_type) not in unary_ops
                    or len(unary_op.inputs) != 1
                    or len(unary_op.outputs) != 1
                    or str(unary_op.inputs[0]) != str(nchw_tail_name)
                    or str(unary_op.outputs[0]) in model_outputs
                ):
                    break
                nchw_tail_name = str(unary_op.outputs[0])
                unary_output_names.append(str(nchw_tail_name))

            tail_users = [int(v) for v in consumers.get(str(nchw_tail_name), [])]
            if len(tail_users) == 0:
                continue

            post_transpose_indices: List[int] = []
            valid_tail = True
            for post_idx in tail_users:
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != str(nchw_tail_name)
                    or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                    or str(post_op.outputs[0]) in model_outputs
                ):
                    valid_tail = False
                    break
                post_out_name = str(post_op.outputs[0])
                post_out_users = [int(v) for v in consumers.get(post_out_name, [])]
                if len(post_out_users) == 0:
                    valid_tail = False
                    break
                if any(str(model_ir.operators[int(v)].op_type) not in conv_like_ops for v in post_out_users):
                    valid_tail = False
                    break
                post_transpose_indices.append(int(post_idx))
            if not valid_tail:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in list(concat_inputs_nhwc)],
            )
            concat_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat_out_name), None),
                perm_nchw_to_nhwc,
            )
            for unary_output_name in unary_output_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(unary_output_name), None),
                    perm_nchw_to_nhwc,
                )

            for post_idx in post_transpose_indices:
                post_op = model_ir.operators[int(post_idx)]
                post_out_name = str(post_op.outputs[0])
                _replace_tensor_inputs(model_ir, post_out_name, str(nchw_tail_name))

            remove_indices = sorted(
                list({int(v) for v in list(pre_transpose_indices) + list(post_transpose_indices)}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_concat_unary_fanout_conv_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_concat_ndhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Convert NCDHW concat blocks back to NDHWC when they are wrapped by transpose adapters.

    Target:
      ... -> t_i_ncdhw (some inputs may be unary outputs from transpose-wrapped NDHWC)
      CONCAT(axis=1, [t_0_ncdhw, ...]) -> y_ncdhw
      y_ncdhw -> TRANSPOSE(0,2,3,4,1) -> y_ndhwc

    Rewrite:
      ... -> t_i_ndhwc
      CONCAT(axis=4, [t_0_ndhwc, ...]) -> y_ndhwc
    """
    optimized = 0
    perm_ndhwc_to_ncdhw = [0, 4, 1, 2, 3]
    perm_ncdhw_to_ndhwc = [0, 2, 3, 4, 1]
    unary_ops = {
        "LOGISTIC",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "ELU",
        "LEAKY_RELU",
        "TANH",
        "GELU",
        "HARD_SWISH",
        "ABS",
        "EXP",
        "NEG",
        "SQRT",
    }

    def _project_shape_after_ncdhw_to_ndhwc(tensor_name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            return None
        shape = [int(v) for v in list(tensor.shape)]
        if len(shape) != 5:
            return None
        return _permute_shape(shape, perm_ncdhw_to_ndhwc)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue

            concat_axis_old = int(concat_op.options.get("axis", 1))
            if concat_axis_old < 0:
                concat_axis_old += 5
            if concat_axis_old != 1:
                continue

            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue

            concat_input_actions: List[Dict[str, Any]] = []
            rewritable = True
            for input_name in [str(v) for v in list(concat_op.inputs)]:
                input_producer_idx = producers.get(str(input_name), None)
                if input_producer_idx is None:
                    rewritable = False
                    break
                input_producer = model_ir.operators[int(input_producer_idx)]

                if (
                    str(input_producer.op_type) == "TRANSPOSE"
                    and len(input_producer.inputs) >= 2
                    and len(input_producer.outputs) == 1
                    and str(input_producer.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, input_producer) == perm_ndhwc_to_ncdhw
                    and set(int(v) for v in consumers.get(str(input_name), [])) == {int(concat_idx)}
                    and str(input_name) not in model_outputs
                ):
                    concat_input_actions.append(
                        {
                            "kind": "direct",
                            "new_input_name": str(input_producer.inputs[0]),
                            "remove_index": int(input_producer_idx),
                        }
                    )
                    continue

                if (
                    str(input_producer.op_type) in unary_ops
                    and len(input_producer.inputs) == 1
                    and len(input_producer.outputs) == 1
                    and str(input_producer.outputs[0]) == str(input_name)
                    and set(int(v) for v in consumers.get(str(input_name), [])) == {int(concat_idx)}
                    and str(input_name) not in model_outputs
                ):
                    unary_input_name = str(input_producer.inputs[0])
                    pre_idx = producers.get(unary_input_name, None)
                    if pre_idx is None:
                        rewritable = False
                        break
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != unary_input_name
                        or _read_transpose_perm(model_ir, pre_op) != perm_ndhwc_to_ncdhw
                        or unary_input_name in model_outputs
                        or set(int(v) for v in consumers.get(unary_input_name, [])) != {int(input_producer_idx)}
                    ):
                        rewritable = False
                        break
                    concat_input_actions.append(
                        {
                            "kind": "unary",
                            "input_name": str(input_name),
                            "unary_index": int(input_producer_idx),
                            "new_unary_input_name": str(pre_op.inputs[0]),
                            "remove_index": int(pre_idx),
                        }
                    )
                    continue

                rewritable = False
                break

            if not rewritable or len(concat_input_actions) == 0:
                continue

            concat_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(concat_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            for user_idx in concat_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_ncdhw_to_ndhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                    continue
                rewritable = False
                break
            if not rewritable or len(post_indices) == 0:
                continue

            nhwc_ref_shape: Optional[List[int]] = None
            nhwc_inputs_ok = True
            for action in concat_input_actions:
                action_kind = str(action.get("kind", ""))
                if action_kind == "direct":
                    projected_input_name = str(action["new_input_name"])
                    input_tensor = model_ir.tensors.get(projected_input_name, None)
                    if input_tensor is None or len(list(input_tensor.shape)) != 5:
                        nhwc_inputs_ok = False
                        break
                    shape = [int(v) for v in list(input_tensor.shape)]
                elif action_kind == "unary":
                    projected_shape = _project_shape_after_ncdhw_to_ndhwc(str(action["input_name"]))
                    if projected_shape is None:
                        nhwc_inputs_ok = False
                        break
                    shape = [int(v) for v in list(projected_shape)]
                else:
                    nhwc_inputs_ok = False
                    break
                if nhwc_ref_shape is None:
                    nhwc_ref_shape = list(shape)
                else:
                    for dim_idx in [0, 1, 2, 3]:
                        if int(shape[dim_idx]) != int(nhwc_ref_shape[dim_idx]):
                            nhwc_inputs_ok = False
                            break
                if not nhwc_inputs_ok:
                    break
            if not nhwc_inputs_ok:
                continue

            new_concat_inputs: List[str] = []
            remove_indices: List[int] = []
            for action in concat_input_actions:
                action_kind = str(action.get("kind", ""))
                remove_indices.append(int(action["remove_index"]))
                if action_kind == "direct":
                    new_concat_inputs.append(str(action["new_input_name"]))
                    continue
                if action_kind == "unary":
                    unary_idx = int(action["unary_index"])
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=model_ir.operators[int(unary_idx)],
                        new_inputs=[str(action["new_unary_input_name"])],
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(action["input_name"]), None),
                        perm_ncdhw_to_ndhwc,
                    )
                    new_concat_inputs.append(str(action["input_name"]))
                    continue
                rewritable = False
                break
            if not rewritable:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in new_concat_inputs],
            )
            concat_op.options["axis"] = 4

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_post_output_name, canonical_post_output_name)

            old_concat_tensor = model_ir.tensors.get(concat_out_name, None)
            canonical_post_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_concat_tensor is not None and canonical_post_tensor is not None:
                canonical_post_tensor.dtype = str(old_concat_tensor.dtype)
                canonical_post_tensor.quantization = _clone_quantization(old_concat_tensor.quantization)
                canonical_post_tensor.shape = [int(v) for v in list(old_concat_tensor.shape)]
                canonical_post_tensor.shape_signature = (
                    [int(v) for v in list(old_concat_tensor.shape_signature)]
                    if old_concat_tensor.shape_signature is not None
                    else [int(v) for v in list(old_concat_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_post_tensor,
                    perm_ncdhw_to_ndhwc,
                )

            remove_indices.extend(post_indices)
            for remove_idx in sorted(list({int(v) for v in remove_indices}), reverse=True):
                if int(remove_idx) == int(concat_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_concat_ndhwc_chains": int(optimized)}


def _optimize_transpose_slice_prepost_nhwc_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate strict NHWC<->NCHW transpose wrappers around rank-4 SLICE.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --SLICE(begin,size)--> y_nchw
      y_nchw --TRANSPOSE(0,2,3,1)--> z_nhwc

    Rewrite:
      x_nhwc --SLICE(begin',size')--> z_nhwc

    Safety:
    - Rank-4 only.
    - SLICE begin/size must be constant tensors.
    - Intermediate transpose/slice tensors must be strict single-consumer links.
    - begin/size tensors must be used only by the target SLICE.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            slice_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if slice_out_name in model_outputs:
                continue

            slice_idx = producers.get(slice_out_name, None)
            if slice_idx is None:
                continue
            slice_op = model_ir.operators[int(slice_idx)]
            if str(slice_op.op_type) != "SLICE" or len(slice_op.inputs) < 3 or len(slice_op.outputs) != 1:
                continue
            if str(slice_op.outputs[0]) != slice_out_name:
                continue
            if set(int(v) for v in consumers.get(slice_out_name, [])) != {int(post_idx)}:
                continue

            begin_name = str(slice_op.inputs[1])
            size_name = str(slice_op.inputs[2])
            begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(begin_name, None))
            size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(size_name, None))
            if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
                continue
            if set(int(v) for v in consumers.get(begin_name, [])) != {int(slice_idx)}:
                continue
            if set(int(v) for v in consumers.get(size_name, [])) != {int(slice_idx)}:
                continue

            pre_out_name = str(slice_op.inputs[0])
            if pre_out_name in model_outputs:
                continue
            pre_idx = producers.get(pre_out_name, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != pre_out_name
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            if set(int(v) for v in consumers.get(pre_out_name, [])) != {int(slice_idx)}:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            post_out_tensor = model_ir.tensors.get(post_out_name, None)
            if (
                pre_input_tensor is None
                or post_out_tensor is None
                or len(list(pre_input_tensor.shape)) != 4
                or len(list(post_out_tensor.shape)) != 4
            ):
                continue

            pre_input_shape = [int(v) for v in list(pre_input_tensor.shape)]
            post_out_shape = [int(v) for v in list(post_out_tensor.shape)]
            begin_vals_i = [int(v) for v in list(begin_vals)]
            size_vals_i = [int(v) for v in list(size_vals)]

            # Select remapped params only when they reproduce the expected NHWC post shape.
            # This prevents double-remap corruption when constants were already converted.
            as_is_shape, _, _ = _infer_slice_output_shape_and_resolved_params(
                input_shape=pre_input_shape,
                begin_vals=begin_vals_i,
                size_vals=size_vals_i,
            )
            remapped_begin = [int(begin_vals_i[int(axis)]) for axis in perm_nchw_to_nhwc]
            remapped_size = [int(size_vals_i[int(axis)]) for axis in perm_nchw_to_nhwc]
            remapped_shape, _, _ = _infer_slice_output_shape_and_resolved_params(
                input_shape=pre_input_shape,
                begin_vals=remapped_begin,
                size_vals=remapped_size,
            )

            selected_begin: Optional[List[int]] = None
            selected_size: Optional[List[int]] = None
            if as_is_shape is not None and [int(v) for v in list(as_is_shape)] == post_out_shape:
                selected_begin = begin_vals_i
                selected_size = size_vals_i
            elif remapped_shape is not None and [int(v) for v in list(remapped_shape)] == post_out_shape:
                selected_begin = remapped_begin
                selected_size = remapped_size
            else:
                continue

            _write_const_ints_to_tensor(model_ir.tensors.get(begin_name, None), selected_begin)
            _write_const_ints_to_tensor(model_ir.tensors.get(size_name, None), selected_size)
            _replace_operator_input_at(
                model_ir=model_ir,
                op=slice_op,
                input_index=0,
                new_input_name=pre_input_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=slice_op,
                new_outputs=[post_out_name],
            )

            for remove_idx in sorted([int(post_idx), int(pre_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_slice_prepost_nhwc_passthrough_chains": int(optimized)}


def _optimize_transpose_shape_extract_nhwc_to_nchw_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove NHWC->NCHW transposes used only for SHAPE extraction.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw --SHAPE--> s_nchw
      s_nchw --GATHER/SLICE(...)--> ...

    Rewrite:
      x_nhwc --SHAPE--> s_nchw'
      - Remap GATHER indices by transpose permutation.
      - For SLICE on rank-1 shape vectors, remap selected indices.
        If remapped indices are non-contiguous, convert SLICE to GATHER.
      - Remove the now-redundant TRANSPOSE.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    def _unique_tensor_name(base: str) -> str:
        if base not in model_ir.tensors:
            return base
        suffix = 1
        while f"{base}_{suffix}" in model_ir.tensors:
            suffix += 1
        return f"{base}_{suffix}"

    def _materialize_const_tensor(
        *,
        base_name: str,
        template_tensor: Optional[TensorIR],
        values: List[int],
    ) -> str:
        new_name = _unique_tensor_name(base_name)
        np_dtype = np.int32
        dtype_name = "INT32"
        quant = None
        if template_tensor is not None:
            dtype_name = str(template_tensor.dtype)
            quant = _clone_quantization(template_tensor.quantization)
            if template_tensor.data is not None:
                try:
                    np_dtype = np.asarray(template_tensor.data).dtype
                except Exception:
                    np_dtype = np.int32
        const_values = [int(v) for v in list(values)]
        model_ir.tensors[new_name] = TensorIR(
            name=new_name,
            dtype=dtype_name,
            shape=[int(len(const_values))],
            shape_signature=[int(len(const_values))],
            data=np.asarray(const_values, dtype=np_dtype),
            is_variable=False,
            quantization=quant,
        )
        return new_name

    def _assign_const_input(
        *,
        op_idx: int,
        input_index: int,
        values: List[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        op = model_ir.operators[int(op_idx)]
        if int(input_index) < 0 or int(input_index) >= len(op.inputs):
            return False
        input_name = str(op.inputs[int(input_index)])
        input_tensor = model_ir.tensors.get(input_name, None)
        if input_tensor is None:
            return False
        user_set = set(int(v) for v in consumers.get(input_name, []))
        if user_set == {int(op_idx)}:
            return _write_const_ints_to_tensor(input_tensor, [int(v) for v in list(values)])
        cloned_name = _materialize_const_tensor(
            base_name=f"{input_name}_shape_remap",
            template_tensor=input_tensor,
            values=[int(v) for v in list(values)],
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=cloned_name,
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for transpose_idx, transpose_op in enumerate(model_ir.operators):
            if (
                str(transpose_op.op_type) != "TRANSPOSE"
                or len(transpose_op.inputs) < 2
                or len(transpose_op.outputs) != 1
            ):
                continue
            perm = _read_transpose_perm(model_ir, transpose_op)
            if perm != perm_nhwc_to_nchw:
                continue

            transpose_input_name = str(transpose_op.inputs[0])
            transpose_output_name = str(transpose_op.outputs[0])
            if transpose_output_name in model_outputs:
                continue

            transpose_users = [int(v) for v in consumers.get(transpose_output_name, [])]
            if len(transpose_users) != 1:
                continue
            shape_idx = int(transpose_users[0])
            shape_op = model_ir.operators[int(shape_idx)]
            if (
                str(shape_op.op_type) != "SHAPE"
                or len(shape_op.inputs) != 1
                or len(shape_op.outputs) != 1
                or str(shape_op.inputs[0]) != transpose_output_name
            ):
                continue

            shape_output_name = str(shape_op.outputs[0])
            if shape_output_name in model_outputs:
                continue
            shape_users = [int(v) for v in consumers.get(shape_output_name, [])]
            if len(shape_users) == 0:
                continue

            gather_rewrites: List[Tuple[int, List[int]]] = []
            slice_contiguous_rewrites: List[Tuple[int, List[int], List[int]]] = []
            slice_to_gather_rewrites: List[Tuple[int, List[int]]] = []
            rewrite_supported = True
            rank = int(len(perm))

            for user_idx in shape_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if user_type == "GATHER":
                    if (
                        len(user_op.inputs) < 2
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != shape_output_name
                    ):
                        rewrite_supported = False
                        break
                    axis = int(user_op.options.get("axis", 0))
                    if axis < 0:
                        axis += 1
                    if axis != 0:
                        rewrite_supported = False
                        break
                    indices_tensor = model_ir.tensors.get(str(user_op.inputs[1]), None)
                    indices_vals = _read_const_ints_from_tensor(indices_tensor)
                    if indices_vals is None or len(indices_vals) == 0:
                        rewrite_supported = False
                        break
                    remapped_indices: List[int] = []
                    for index_value in indices_vals:
                        index = int(index_value)
                        if index < 0:
                            index += int(rank)
                        if index < 0 or index >= int(rank):
                            rewrite_supported = False
                            break
                        remapped_indices.append(int(perm[int(index)]))
                    if not rewrite_supported:
                        break
                    gather_rewrites.append((int(user_idx), remapped_indices))
                    continue

                if user_type == "SLICE":
                    if (
                        len(user_op.inputs) < 3
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != shape_output_name
                    ):
                        rewrite_supported = False
                        break
                    begin_tensor = model_ir.tensors.get(str(user_op.inputs[1]), None)
                    size_tensor = model_ir.tensors.get(str(user_op.inputs[2]), None)
                    begin_vals = _read_const_ints_from_tensor(begin_tensor)
                    size_vals = _read_const_ints_from_tensor(size_tensor)
                    if begin_vals is None or size_vals is None:
                        rewrite_supported = False
                        break
                    if len(begin_vals) != 1 or len(size_vals) != 1:
                        rewrite_supported = False
                        break
                    begin = int(begin_vals[0])
                    size = int(size_vals[0])
                    if begin < 0:
                        begin += int(rank)
                    begin = max(0, min(int(begin), int(rank)))
                    if size == -1:
                        end = int(rank)
                    elif size < -1:
                        rewrite_supported = False
                        break
                    else:
                        end = min(int(rank), int(begin + size))
                    if int(end) <= int(begin):
                        rewrite_supported = False
                        break
                    selected = [int(v) for v in range(int(begin), int(end))]
                    remapped_selected = [int(perm[int(v)]) for v in selected]
                    is_contiguous = all(
                        int(remapped_selected[int(i)]) == int(remapped_selected[0] + i)
                        for i in range(len(remapped_selected))
                    )
                    if is_contiguous:
                        new_begin = [int(remapped_selected[0])]
                        new_size = [int(len(remapped_selected))]
                        slice_contiguous_rewrites.append((int(user_idx), new_begin, new_size))
                    else:
                        slice_to_gather_rewrites.append((int(user_idx), remapped_selected))
                    continue

                rewrite_supported = False
                break

            if not rewrite_supported:
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=shape_op,
                input_index=0,
                new_input_name=transpose_input_name,
            )

            for gather_idx, remapped_indices in gather_rewrites:
                _assign_const_input(
                    op_idx=int(gather_idx),
                    input_index=1,
                    values=[int(v) for v in remapped_indices],
                    consumers=consumers,
                )

            for slice_idx, new_begin, new_size in slice_contiguous_rewrites:
                _assign_const_input(
                    op_idx=int(slice_idx),
                    input_index=1,
                    values=[int(v) for v in list(new_begin)],
                    consumers=consumers,
                )
                _assign_const_input(
                    op_idx=int(slice_idx),
                    input_index=2,
                    values=[int(v) for v in list(new_size)],
                    consumers=consumers,
                )

            for slice_idx, remapped_indices in slice_to_gather_rewrites:
                slice_op = model_ir.operators[int(slice_idx)]
                begin_input_name = str(slice_op.inputs[1])
                begin_tensor = model_ir.tensors.get(begin_input_name, None)
                gather_indices_name = _materialize_const_tensor(
                    base_name=f"{begin_input_name}_gather",
                    template_tensor=begin_tensor,
                    values=[int(v) for v in remapped_indices],
                )
                slice_op.op_type = "GATHER"
                slice_op.version = 1
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=slice_op,
                    new_inputs=[shape_output_name, gather_indices_name],
                )
                slice_op.options = {"axis": 0, "batchDims": 0}
                slice_out_name = str(slice_op.outputs[0])
                slice_out_tensor = model_ir.tensors.get(slice_out_name, None)
                if slice_out_tensor is not None:
                    slice_out_tensor.shape = [int(len(remapped_indices))]
                    slice_out_tensor.shape_signature = [int(len(remapped_indices))]

            del model_ir.operators[int(transpose_idx)]
            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_shape_extract_nhwc_to_nchw_chains": int(optimized)}


def _optimize_transpose_stridedslice_pre_concat_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around stride-slice fan-in concat blocks.

    Target:
      X_nhwc --TRANSPOSE(0,3,1,2)--> X_nchw
      X_nchw --STRIDED_SLICE(...)--> s_i_nchw   (all consumed by one CONCAT axis=1)
      CONCAT(axis=1, [s_i_nchw...]) --> y_nchw
      y_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      X_nhwc --STRIDED_SLICE'(NHWC params)--> s_i_nhwc
      CONCAT(axis=3, [s_i_nhwc...]) --> y_nhwc

    If legacy NCHW consumers of `y_nchw` exist, keep one adapter:
      y_nhwc --TRANSPOSE(0,3,1,2)--> y_nchw
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _is_supported_stridedslice(op: OperatorIR) -> bool:
        if str(op.op_type) != "STRIDED_SLICE":
            return False
        if len(op.inputs) != 4 or len(op.outputs) != 1:
            return False
        opts = op.options if isinstance(op.options, dict) else {}
        if int(opts.get("beginMask", 0)) != 0:
            return False
        if int(opts.get("endMask", 0)) != 0:
            return False
        if int(opts.get("ellipsisMask", 0)) != 0:
            return False
        if int(opts.get("newAxisMask", 0)) != 0:
            return False
        if int(opts.get("shrinkAxisMask", 0)) != 0:
            return False
        if bool(opts.get("offset", False)):
            return False
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) < 2:
                continue

            slice_indices: List[int] = []
            slice_output_names: List[str] = []
            stridedslice_consts: List[Tuple[str, str, str]] = []
            valid_slices = True
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                if not _is_supported_stridedslice(user_op):
                    valid_slices = False
                    break
                if str(user_op.inputs[0]) != pre_output_name:
                    valid_slices = False
                    break
                out_name = str(user_op.outputs[0])
                if out_name in model_outputs:
                    valid_slices = False
                    break

                begin_name = str(user_op.inputs[1])
                end_name = str(user_op.inputs[2])
                strides_name = str(user_op.inputs[3])
                begin_tensor = model_ir.tensors.get(begin_name, None)
                end_tensor = model_ir.tensors.get(end_name, None)
                strides_tensor = model_ir.tensors.get(strides_name, None)
                begin_vals = _read_const_ints_from_tensor(begin_tensor)
                end_vals = _read_const_ints_from_tensor(end_tensor)
                strides_vals = _read_const_ints_from_tensor(strides_tensor)
                if begin_vals is None or end_vals is None or strides_vals is None:
                    valid_slices = False
                    break
                if len(begin_vals) != 4 or len(end_vals) != 4 or len(strides_vals) != 4:
                    valid_slices = False
                    break
                if any(int(v) == 0 for v in strides_vals):
                    valid_slices = False
                    break
                # Avoid mutating shared constants.
                if set(int(v) for v in consumers.get(begin_name, [])) != {int(user_idx)}:
                    valid_slices = False
                    break
                if set(int(v) for v in consumers.get(end_name, [])) != {int(user_idx)}:
                    valid_slices = False
                    break
                if set(int(v) for v in consumers.get(strides_name, [])) != {int(user_idx)}:
                    valid_slices = False
                    break

                slice_indices.append(int(user_idx))
                slice_output_names.append(out_name)
                stridedslice_consts.append((begin_name, end_name, strides_name))

            if not valid_slices:
                continue

            concat_idx: Optional[int] = None
            for slice_out_name in slice_output_names:
                users = [int(v) for v in consumers.get(slice_out_name, [])]
                if len(users) != 1:
                    concat_idx = None
                    break
                if concat_idx is None:
                    concat_idx = int(users[0])
                elif int(concat_idx) != int(users[0]):
                    concat_idx = None
                    break
            if concat_idx is None:
                continue

            concat_op = model_ir.operators[int(concat_idx)]
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_axis_old = int(concat_op.options.get("axis", 1))
            if concat_axis_old < 0:
                concat_axis_old += 4
            if concat_axis_old != 1:
                continue

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) != len(slice_output_names):
                continue
            if set(concat_inputs) != set(slice_output_names):
                continue

            concat_out_name = str(concat_op.outputs[0])
            concat_users = [int(v) for v in consumers.get(concat_out_name, []) if int(v) != int(concat_idx)]
            if len(concat_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_posts = True
            for user_idx in concat_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))
            if not valid_posts or len(post_indices) == 0:
                continue

            # Rewrite all STRIDED_SLICE params/layout from NCHW to NHWC.
            rewrite_ok = True
            for slice_idx, (begin_name, end_name, strides_name), slice_out_name in zip(
                slice_indices,
                stridedslice_consts,
                slice_output_names,
            ):
                slice_op = model_ir.operators[int(slice_idx)]
                begin_tensor = model_ir.tensors.get(begin_name, None)
                end_tensor = model_ir.tensors.get(end_name, None)
                strides_tensor = model_ir.tensors.get(strides_name, None)
                begin_vals = _read_const_ints_from_tensor(begin_tensor)
                end_vals = _read_const_ints_from_tensor(end_tensor)
                strides_vals = _read_const_ints_from_tensor(strides_tensor)
                if begin_vals is None or end_vals is None or strides_vals is None:
                    rewrite_ok = False
                    break
                begin_nhwc = _permute_shape(begin_vals, perm_nchw_to_nhwc)
                end_nhwc = _permute_shape(end_vals, perm_nchw_to_nhwc)
                strides_nhwc = _permute_shape(strides_vals, perm_nchw_to_nhwc)
                if begin_nhwc is None or end_nhwc is None or strides_nhwc is None:
                    rewrite_ok = False
                    break
                _write_const_ints_to_tensor(begin_tensor, [int(v) for v in begin_nhwc])
                _write_const_ints_to_tensor(end_tensor, [int(v) for v in end_nhwc])
                _write_const_ints_to_tensor(strides_tensor, [int(v) for v in strides_nhwc])
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=slice_op,
                    input_index=0,
                    new_input_name=pre_input_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(slice_out_name, None),
                    perm_nchw_to_nhwc,
                )

            if not rewrite_ok:
                continue

            concat_op.options["axis"] = 3
            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_post_output_name, canonical_post_output_name)

            old_concat_tensor = model_ir.tensors.get(concat_out_name, None)
            canonical_post_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_concat_tensor is not None and canonical_post_tensor is not None:
                canonical_post_tensor.dtype = str(old_concat_tensor.dtype)
                canonical_post_tensor.quantization = _clone_quantization(old_concat_tensor.quantization)
                canonical_post_tensor.shape = [int(v) for v in list(old_concat_tensor.shape)]
                canonical_post_tensor.shape_signature = (
                    [int(v) for v in list(old_concat_tensor.shape_signature)]
                    if old_concat_tensor.shape_signature is not None
                    else [int(v) for v in list(old_concat_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_post_tensor,
                    perm_nchw_to_nhwc,
                )

            preserve_nchw_adapter = len(legacy_users) > 0 or concat_out_name in model_outputs
            if preserve_nchw_adapter:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[concat_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = sorted(
                list({int(pre_idx), *post_remove_indices}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_stridedslice_pre_concat_nhwc_chains": int(optimized)}


def _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir: ModelIR) -> Dict[str, int]:
    """
    Collapse per-input NHWC->NCHW adapters before CONCAT into a single post-CONCAT adapter.

    Target:
      x0_nhwc --TRANSPOSE(0,3,1,2)--> x0_nchw --(optional unary)--> i0_nchw
      x1_nhwc --TRANSPOSE(0,3,1,2)--> x1_nchw --(optional unary)--> i1_nchw
      ...
      CONCAT(axis=1, [i0_nchw, i1_nchw, ...]) -> y_nchw
      y_nchw -> downstream (e.g. RESHAPE)

    Rewrite:
      x0_nhwc --(optional unary)--> i0_nhwc
      x1_nhwc --(optional unary)--> i1_nhwc
      ...
      CONCAT(axis=3, [i0_nhwc, i1_nhwc, ...]) -> y_nhwc
      y_nhwc --TRANSPOSE(0,3,1,2)--> y_nchw

    This keeps downstream NCHW semantics unchanged while reducing redundant pre-concat
    transpose chains.

    Supported layout adapters:
      - TRANSPOSE(0,3,1,2)
      - RESHAPE [N,H,W,1] -> [N,1,H,W] (singleton-channel layout adapter)
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {
        "LOGISTIC",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "ELU",
        "LEAKY_RELU",
        "TANH",
        "GELU",
        "HARD_SWISH",
        "ABS",
        "EXP",
        "NEG",
        "SQRT",
    }

    def _unique_tensor_name(base: str) -> str:
        if base not in model_ir.tensors:
            return base
        suffix = 1
        while True:
            candidate = f"{base}_{suffix}"
            if candidate not in model_ir.tensors:
                return candidate
            suffix += 1

    def _is_singleton_channel_nhwc_to_nchw_reshape(
        *,
        input_name: str,
        output_name: str,
    ) -> bool:
        in_tensor = model_ir.tensors.get(str(input_name), None)
        out_tensor = model_ir.tensors.get(str(output_name), None)
        if in_tensor is None or out_tensor is None:
            return False
        in_shape = [int(v) for v in list(in_tensor.shape)]
        out_shape = [int(v) for v in list(out_tensor.shape)]
        if len(in_shape) != 4 or len(out_shape) != 4:
            return False
        return (
            int(in_shape[0]) == int(out_shape[0])
            and int(in_shape[1]) == int(out_shape[2])
            and int(in_shape[2]) == int(out_shape[3])
            and int(in_shape[3]) == 1
            and int(out_shape[1]) == 1
        )

    def _find_or_create_nhwc_to_nchw_perm_tensor() -> str:
        target = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
        for tensor_name, tensor in model_ir.tensors.items():
            if tensor is None or tensor.data is None:
                continue
            data = np.asarray(tensor.data)
            if data.dtype != np.int32 or data.size != 4:
                continue
            if np.array_equal(data.reshape(-1), target):
                return str(tensor_name)
        perm_name = _unique_tensor_name("transpose_input_chains_nhwc_to_nchw_perm")
        model_ir.tensors[perm_name] = TensorIR(
            name=perm_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=target,
            is_variable=False,
        )
        return perm_name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue

            concat_axis_old = int(concat_op.options.get("axis", 1))
            if concat_axis_old < 0:
                concat_axis_old += 4
            if concat_axis_old != 1:
                continue

            concat_out_name = str(concat_op.outputs[0])
            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            if concat_out_tensor is None or len(list(concat_out_tensor.shape)) != 4:
                continue

            new_concat_inputs: List[str] = []
            pre_remove_indices: List[int] = []
            unary_input_rewrites: Dict[int, str] = {}
            unary_output_names_to_permute: List[str] = []
            adapter_perm_tensor_name: Optional[str] = None
            rewritable = True

            for input_name in [str(v) for v in list(concat_op.inputs)]:
                input_producer_idx = producers.get(str(input_name), None)
                if input_producer_idx is None:
                    rewritable = False
                    break
                input_producer = model_ir.operators[int(input_producer_idx)]

                if (
                    str(input_producer.op_type) == "TRANSPOSE"
                    and len(input_producer.inputs) >= 2
                    and len(input_producer.outputs) == 1
                    and str(input_producer.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, input_producer) == perm_nhwc_to_nchw
                    and set(int(v) for v in consumers.get(str(input_name), [])) == {int(concat_idx)}
                    and str(input_name) not in model_outputs
                ):
                    new_concat_inputs.append(str(input_producer.inputs[0]))
                    pre_remove_indices.append(int(input_producer_idx))
                    if adapter_perm_tensor_name is None:
                        adapter_perm_tensor_name = str(input_producer.inputs[1])
                    continue

                if (
                    str(input_producer.op_type) in unary_ops
                    and len(input_producer.inputs) == 1
                    and len(input_producer.outputs) == 1
                    and str(input_producer.outputs[0]) == str(input_name)
                    and set(int(v) for v in consumers.get(str(input_name), [])) == {int(concat_idx)}
                    and str(input_name) not in model_outputs
                ):
                    unary_in_name = str(input_producer.inputs[0])
                    pre_idx = producers.get(unary_in_name, None)
                    if pre_idx is None:
                        rewritable = False
                        break
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        set(int(v) for v in consumers.get(unary_in_name, [])) != {int(input_producer_idx)}
                        or unary_in_name in model_outputs
                    ):
                        rewritable = False
                        break
                    if str(pre_op.op_type) == "TRANSPOSE":
                        if (
                            len(pre_op.inputs) < 2
                            or len(pre_op.outputs) != 1
                            or str(pre_op.outputs[0]) != unary_in_name
                            or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                        ):
                            rewritable = False
                            break
                        if adapter_perm_tensor_name is None:
                            adapter_perm_tensor_name = str(pre_op.inputs[1])
                    elif str(pre_op.op_type) == "RESHAPE":
                        if (
                            len(pre_op.inputs) < 1
                            or len(pre_op.outputs) != 1
                            or str(pre_op.outputs[0]) != unary_in_name
                            or not _is_singleton_channel_nhwc_to_nchw_reshape(
                                input_name=str(pre_op.inputs[0]),
                                output_name=unary_in_name,
                            )
                        ):
                            rewritable = False
                            break
                    else:
                        rewritable = False
                        break
                    unary_input_rewrites[int(input_producer_idx)] = str(pre_op.inputs[0])
                    unary_output_names_to_permute.append(str(input_name))
                    new_concat_inputs.append(str(input_name))
                    pre_remove_indices.append(int(pre_idx))
                    continue

                rewritable = False
                break

            if not rewritable:
                continue
            if adapter_perm_tensor_name is None:
                adapter_perm_tensor_name = _find_or_create_nhwc_to_nchw_perm_tensor()

            unary_outputs_to_permute = {str(v) for v in unary_output_names_to_permute}
            nhwc_inputs_ok = True
            nhwc_ref_shape: Optional[List[int]] = None
            for input_name in new_concat_inputs:
                input_tensor = model_ir.tensors.get(str(input_name), None)
                if input_tensor is None or len(list(input_tensor.shape)) != 4:
                    nhwc_inputs_ok = False
                    break
                shape = [int(v) for v in list(input_tensor.shape)]
                if str(input_name) in unary_outputs_to_permute:
                    permuted_shape = _permute_shape(shape, perm_nchw_to_nhwc)
                    if permuted_shape is None:
                        nhwc_inputs_ok = False
                        break
                    shape = [int(v) for v in list(permuted_shape)]
                if nhwc_ref_shape is None:
                    nhwc_ref_shape = list(shape)
                else:
                    for dim_idx in [0, 1, 2]:
                        if int(shape[dim_idx]) != int(nhwc_ref_shape[dim_idx]):
                            nhwc_inputs_ok = False
                            break
                if not nhwc_inputs_ok:
                    break
            if not nhwc_inputs_ok:
                continue

            new_concat_shape = _permute_shape(list(concat_out_tensor.shape), perm_nchw_to_nhwc)
            concat_sig_source = (
                list(concat_out_tensor.shape_signature)
                if concat_out_tensor.shape_signature is not None
                else list(concat_out_tensor.shape)
            )
            new_concat_signature = _permute_shape(concat_sig_source, perm_nchw_to_nhwc)
            if new_concat_shape is None or new_concat_signature is None:
                continue

            for unary_idx, rewrite_input_name in unary_input_rewrites.items():
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=model_ir.operators[int(unary_idx)],
                    new_inputs=[str(rewrite_input_name)],
                )
            for tensor_name in unary_output_names_to_permute:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            new_concat_out_name = _unique_tensor_name(f"{concat_out_name}_nhwc")
            model_ir.tensors[new_concat_out_name] = TensorIR(
                name=new_concat_out_name,
                dtype=str(concat_out_tensor.dtype),
                shape=[int(v) for v in list(new_concat_shape)],
                data=None,
                quantization=_clone_quantization(concat_out_tensor.quantization),
                shape_signature=[int(v) for v in list(new_concat_signature)],
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in new_concat_inputs],
            )
            concat_op.options["axis"] = 3
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[str(new_concat_out_name)],
            )

            adapter_op = OperatorIR(
                op_type="TRANSPOSE",
                inputs=[str(new_concat_out_name), str(adapter_perm_tensor_name)],
                outputs=[str(concat_out_name)],
                options={},
            )
            model_ir.operators.insert(int(concat_idx) + 1, adapter_op)

            remove_indices = sorted(list({int(v) for v in pre_remove_indices}), reverse=True)
            for remove_idx in remove_indices:
                if int(remove_idx) == int(concat_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_input_chains_pre_concat_to_single_post_adapter": int(optimized)}


def _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse mixed pre-CONCAT NCHW adapters with SPLIT branches into NHWC + single post adapter.

    Target:
      split_in_nchw --SPLIT(axis=1)--> s0_nchw, s1_nchw, ...
      s1_nchw --TRANSPOSE(0,2,3,1)--> s1_nhwc --...
      b0_nhwc --TRANSPOSE(0,3,1,2)--> b0_nchw
      b1_nhwc --TRANSPOSE(0,3,1,2)--> b1_nchw
      CONCAT(axis=1, [s0_nchw, b0_nchw, b1_nchw, ...]) -> cat_nchw

    Rewrite:
      split_in_nchw --TRANSPOSE(0,2,3,1)--> split_in_nhwc
      split_in_nhwc --SPLIT(axis=3)--> s0_nhwc, s1_nhwc, ...
      bypass/remove s_i_nchw -> TRANSPOSE(0,2,3,1)
      CONCAT(axis=3, [s0_nhwc, b0_nhwc, b1_nhwc, ...]) -> cat_nhwc
      cat_nhwc --TRANSPOSE(0,3,1,2)--> cat_nchw

    This keeps downstream NCHW semantics unchanged while reducing redundant local
    transpose chains around split/concat bridges.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _find_op_index_by_id(op_id: int) -> Optional[int]:
        for idx, op in enumerate(model_ir.operators):
            if int(id(op)) == int(op_id):
                return int(idx)
        return None

    def _find_or_create_perm_tensor(
        *,
        base_name: str,
        perm: List[int],
    ) -> str:
        target = np.asarray([int(v) for v in list(perm)], dtype=np.int32)
        for tensor_name, tensor in model_ir.tensors.items():
            if tensor is None or tensor.data is None:
                continue
            try:
                data = np.asarray(tensor.data)
            except Exception:
                continue
            if data.dtype != np.int32 or int(data.size) != int(len(perm)):
                continue
            if np.array_equal(data.reshape(-1), target):
                return str(tensor_name)
        perm_name = _unique_tensor_name(str(base_name))
        model_ir.tensors[perm_name] = TensorIR(
            name=perm_name,
            dtype="INT32",
            shape=[int(len(perm))],
            shape_signature=[int(len(perm))],
            data=np.asarray(target, dtype=np.int32),
            is_variable=False,
        )
        return perm_name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue

            concat_axis_old = int(concat_op.options.get("axis", 1))
            if concat_axis_old < 0:
                concat_axis_old += 4
            if concat_axis_old != 1:
                continue

            concat_out_name = str(concat_op.outputs[0])
            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            if concat_out_tensor is None or len(list(concat_out_tensor.shape)) != 4:
                continue
            if concat_out_name in model_outputs:
                continue

            concat_input_names = [str(v) for v in list(concat_op.inputs)]
            if len(concat_input_names) < 2:
                continue

            remove_op_ids: set[int] = set()
            new_concat_inputs: List[str] = []
            split_plans: Dict[int, Dict[str, Any]] = {}
            adapter_perm_tensor_name: Optional[str] = None
            nhwc_ref_shape: Optional[List[int]] = None
            rewritable = True

            for input_name in concat_input_names:
                input_producer_idx = producers.get(str(input_name), None)
                if input_producer_idx is None:
                    rewritable = False
                    break
                input_producer = model_ir.operators[int(input_producer_idx)]

                if (
                    str(input_producer.op_type) == "TRANSPOSE"
                    and len(input_producer.inputs) >= 2
                    and len(input_producer.outputs) == 1
                    and str(input_producer.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, input_producer) == perm_nhwc_to_nchw
                    and set(int(v) for v in consumers.get(str(input_name), [])) == {int(concat_idx)}
                    and str(input_name) not in model_outputs
                ):
                    source_name = str(input_producer.inputs[0])
                    source_tensor = model_ir.tensors.get(source_name, None)
                    if source_tensor is None or len(list(source_tensor.shape)) != 4:
                        rewritable = False
                        break
                    source_shape = [int(v) for v in list(source_tensor.shape)]
                    if nhwc_ref_shape is None:
                        nhwc_ref_shape = list(source_shape)
                    else:
                        for dim_idx in [0, 1, 2]:
                            if int(source_shape[dim_idx]) != int(nhwc_ref_shape[dim_idx]):
                                rewritable = False
                                break
                        if not rewritable:
                            break
                    new_concat_inputs.append(str(source_name))
                    remove_op_ids.add(int(id(input_producer)))
                    if adapter_perm_tensor_name is None:
                        adapter_perm_tensor_name = str(input_producer.inputs[1])
                    continue

                if str(input_producer.op_type) != "SPLIT" or len(input_producer.inputs) < 2:
                    rewritable = False
                    break
                split_idx = int(input_producer_idx)
                split_op_id = int(id(input_producer))

                split_plan = split_plans.get(split_idx, None)
                if split_plan is None:
                    split_axis_name = str(input_producer.inputs[0])
                    split_axis_tensor = model_ir.tensors.get(split_axis_name, None)
                    split_axis_vals = _read_const_ints_from_tensor(split_axis_tensor)
                    if split_axis_vals is None or len(split_axis_vals) != 1:
                        rewritable = False
                        break
                    split_axis = int(split_axis_vals[0])
                    if split_axis < 0:
                        split_axis += 4
                    if split_axis != 1:
                        rewritable = False
                        break

                    split_input_name = str(input_producer.inputs[1])
                    split_input_tensor = model_ir.tensors.get(split_input_name, None)
                    if split_input_tensor is None or len(list(split_input_tensor.shape)) != 4:
                        rewritable = False
                        break

                    split_output_names = [str(v) for v in list(input_producer.outputs)]
                    if str(input_name) not in set(split_output_names):
                        rewritable = False
                        break

                    split_output_shapes_nhwc: Dict[str, List[int]] = {}
                    split_post_rewires: List[Tuple[str, str]] = []
                    split_post_remove_ids: set[int] = set()
                    valid_split = True
                    for split_output_name in split_output_names:
                        split_output_tensor = model_ir.tensors.get(str(split_output_name), None)
                        if split_output_tensor is None or len(list(split_output_tensor.shape)) != 4:
                            valid_split = False
                            break
                        projected_shape = _permute_shape(
                            [int(v) for v in list(split_output_tensor.shape)],
                            perm_nchw_to_nhwc,
                        )
                        if projected_shape is None:
                            valid_split = False
                            break
                        split_output_shapes_nhwc[str(split_output_name)] = [int(v) for v in list(projected_shape)]

                        output_users = [int(v) for v in consumers.get(str(split_output_name), [])]
                        for user_idx in output_users:
                            if int(user_idx) == int(concat_idx):
                                continue
                            user_op = model_ir.operators[int(user_idx)]
                            if (
                                str(user_op.op_type) == "TRANSPOSE"
                                and len(user_op.inputs) >= 2
                                and len(user_op.outputs) == 1
                                and str(user_op.inputs[0]) == str(split_output_name)
                                and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                                and str(user_op.outputs[0]) not in model_outputs
                            ):
                                split_post_rewires.append((str(user_op.outputs[0]), str(split_output_name)))
                                split_post_remove_ids.add(int(id(user_op)))
                                continue
                            valid_split = False
                            break
                        if not valid_split:
                            break

                    if not valid_split:
                        rewritable = False
                        break

                    split_plan = {
                        "split_idx": int(split_idx),
                        "split_op_id": int(split_op_id),
                        "split_axis_name": str(split_axis_name),
                        "split_input_name": str(split_input_name),
                        "split_output_names": [str(v) for v in split_output_names],
                        "split_output_shapes_nhwc": dict(split_output_shapes_nhwc),
                        "split_post_rewires": list(split_post_rewires),
                        "split_post_remove_ids": set(int(v) for v in split_post_remove_ids),
                    }
                    split_plans[int(split_idx)] = dict(split_plan)

                projected_input_shape = split_plan["split_output_shapes_nhwc"].get(str(input_name), None)
                if projected_input_shape is None:
                    rewritable = False
                    break
                if nhwc_ref_shape is None:
                    nhwc_ref_shape = [int(v) for v in list(projected_input_shape)]
                else:
                    for dim_idx in [0, 1, 2]:
                        if int(projected_input_shape[dim_idx]) != int(nhwc_ref_shape[dim_idx]):
                            rewritable = False
                            break
                    if not rewritable:
                        break
                new_concat_inputs.append(str(input_name))

            if not rewritable or nhwc_ref_shape is None:
                continue
            if len(split_plans) == 0:
                continue

            if adapter_perm_tensor_name is None:
                adapter_perm_tensor_name = _find_or_create_perm_tensor(
                    base_name="mixed_pre_concat_nhwc_to_nchw_perm",
                    perm=perm_nhwc_to_nchw,
                )
            split_adapter_perm_name = _find_or_create_perm_tensor(
                base_name="mixed_pre_concat_nchw_to_nhwc_perm",
                perm=perm_nchw_to_nhwc,
            )

            pre_split_insertions: List[Tuple[int, OperatorIR]] = []
            for split_plan in split_plans.values():
                split_op_idx = _find_op_index_by_id(int(split_plan["split_op_id"]))
                if split_op_idx is None:
                    rewritable = False
                    break
                split_op = model_ir.operators[int(split_op_idx)]

                split_axis_name = str(split_plan["split_axis_name"])
                split_axis_tensor = model_ir.tensors.get(split_axis_name, None)
                if split_axis_tensor is None:
                    rewritable = False
                    break
                split_axis_users = [int(v) for v in consumers.get(split_axis_name, [])]
                if set(split_axis_users) == {int(split_op_idx)}:
                    _write_const_ints_to_tensor(split_axis_tensor, [3])
                else:
                    axis_dtype = np.int32
                    try:
                        if split_axis_tensor.data is not None:
                            axis_dtype = np.asarray(split_axis_tensor.data).dtype
                    except Exception:
                        axis_dtype = np.int32
                    split_axis_nhwc_name = _unique_tensor_name(f"{split_axis_name}_nhwc")
                    model_ir.tensors[split_axis_nhwc_name] = TensorIR(
                        name=split_axis_nhwc_name,
                        dtype=str(split_axis_tensor.dtype),
                        shape=[1],
                        shape_signature=[1],
                        data=np.asarray([3], dtype=axis_dtype),
                        is_variable=False,
                        quantization=_clone_quantization(split_axis_tensor.quantization),
                    )
                    split_inputs = [str(v) for v in list(split_op.inputs)]
                    split_inputs[0] = str(split_axis_nhwc_name)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=split_op,
                        new_inputs=split_inputs,
                    )

                split_input_name = str(split_plan["split_input_name"])
                split_input_tensor = model_ir.tensors.get(split_input_name, None)
                if split_input_tensor is None or len(list(split_input_tensor.shape)) != 4:
                    rewritable = False
                    break
                split_input_shape_nhwc = _permute_shape(
                    [int(v) for v in list(split_input_tensor.shape)],
                    perm_nchw_to_nhwc,
                )
                split_input_sig = (
                    [int(v) for v in list(split_input_tensor.shape_signature)]
                    if split_input_tensor.shape_signature is not None
                    else [int(v) for v in list(split_input_tensor.shape)]
                )
                split_input_sig_nhwc = _permute_shape(split_input_sig, perm_nchw_to_nhwc)
                if split_input_shape_nhwc is None or split_input_sig_nhwc is None:
                    rewritable = False
                    break
                split_input_nhwc_name = _unique_tensor_name(f"{split_input_name}_nhwc")
                model_ir.tensors[split_input_nhwc_name] = TensorIR(
                    name=split_input_nhwc_name,
                    dtype=str(split_input_tensor.dtype),
                    shape=[int(v) for v in list(split_input_shape_nhwc)],
                    shape_signature=[int(v) for v in list(split_input_sig_nhwc)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(split_input_tensor.quantization),
                )
                split_inputs = [str(v) for v in list(split_op.inputs)]
                split_inputs[1] = str(split_input_nhwc_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=split_op,
                    new_inputs=split_inputs,
                )
                pre_split_insertions.append(
                    (
                        int(split_op_idx),
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(split_input_name), str(split_adapter_perm_name)],
                            outputs=[str(split_input_nhwc_name)],
                            options={},
                        ),
                    )
                )

                for split_output_name in [str(v) for v in list(split_plan["split_output_names"])]:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(split_output_name), None),
                        perm_nchw_to_nhwc,
                    )

                for post_output_name, replacement_name in list(split_plan["split_post_rewires"]):
                    _replace_tensor_inputs(model_ir, str(post_output_name), str(replacement_name))
                remove_op_ids.update(set(int(v) for v in list(split_plan["split_post_remove_ids"])))

            if not rewritable:
                continue

            new_concat_shape = _permute_shape(list(concat_out_tensor.shape), perm_nchw_to_nhwc)
            concat_sig_src = (
                list(concat_out_tensor.shape_signature)
                if concat_out_tensor.shape_signature is not None
                else list(concat_out_tensor.shape)
            )
            new_concat_sig = _permute_shape(concat_sig_src, perm_nchw_to_nhwc)
            if new_concat_shape is None or new_concat_sig is None:
                continue

            new_concat_out_name = _unique_tensor_name(f"{concat_out_name}_nhwc")
            model_ir.tensors[new_concat_out_name] = TensorIR(
                name=new_concat_out_name,
                dtype=str(concat_out_tensor.dtype),
                shape=[int(v) for v in list(new_concat_shape)],
                shape_signature=[int(v) for v in list(new_concat_sig)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(concat_out_tensor.quantization),
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in list(new_concat_inputs)],
            )
            concat_op.options["axis"] = 3
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[str(new_concat_out_name)],
            )

            concat_op_idx = _find_op_index_by_id(int(id(concat_op)))
            if concat_op_idx is None:
                continue
            model_ir.operators.insert(
                int(concat_op_idx) + 1,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(new_concat_out_name), str(adapter_perm_tensor_name)],
                    outputs=[str(concat_out_name)],
                    options={},
                ),
            )

            inserted = 0
            for split_op_idx, split_adapter_op in sorted(pre_split_insertions, key=lambda v: int(v[0])):
                split_plan_target_idx = int(split_op_idx) + int(inserted)
                model_ir.operators.insert(int(split_plan_target_idx), split_adapter_op)
                inserted += 1

            remove_indices = sorted(
                [
                    int(op_idx)
                    for op_idx, op in enumerate(model_ir.operators)
                    if int(id(op)) in remove_op_ids
                ],
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains": int(optimized)}


def _optimize_transpose_slice_logistic_concat_reshape_tail_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Collapse repeated NCHW adapters in detection heads into a single 3D post-CONCAT transpose.

    Target (per branch):
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --SLICE--> a_nchw --(optional LOGISTIC)--> a'_nchw
      x_nchw --SLICE--> b_nchw
      CONCAT(axis=1, [a'_nchw, b_nchw]) -> c_nchw
      RESHAPE([N,C,S]) -> r_nchw

    Tail:
      CONCAT(axis=2, [r0_nchw, r1_nchw, ...]) -> y_nchw

    Rewrite:
      - Per branch: keep everything in NHWC
          x_nhwc --SLICE--> a_nhwc --(optional LOGISTIC)--> a'_nhwc
          x_nhwc --SLICE--> b_nhwc
          CONCAT(axis=3) -> c_nhwc
          RESHAPE([N,S,C]) -> r_nhwc
      - Tail:
          CONCAT(axis=1, [r0_nhwc, r1_nhwc, ...]) -> y_nhwc
          y_nhwc --TRANSPOSE(0,2,1)--> y_nchw
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_3d_nchw_to_nhwc = [0, 2, 1]
    unary_ops = {"LOGISTIC"}

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _op_index(op_ref: OperatorIR) -> Optional[int]:
        for idx, op in enumerate(model_ir.operators):
            if int(id(op)) == int(id(op_ref)):
                return int(idx)
        return None

    def _find_or_create_perm_3d_nhwc_to_nchw() -> str:
        target = np.asarray(perm_3d_nchw_to_nhwc, dtype=np.int32)
        for tensor_name, tensor in model_ir.tensors.items():
            if tensor is None or tensor.data is None:
                continue
            data = np.asarray(tensor.data)
            if data.dtype != np.int32 or int(data.size) != 3:
                continue
            if np.array_equal(data.reshape(-1), target):
                return str(tensor_name)
        perm_name = _unique_tensor_name("transpose_tail_3d_nhwc_to_nchw_perm")
        model_ir.tensors[perm_name] = TensorIR(
            name=perm_name,
            dtype="INT32",
            shape=[3],
            shape_signature=[3],
            data=np.asarray(target, dtype=np.int32),
            is_variable=False,
        )
        return perm_name

    def _swap_reshape_shape_tensor(
        *,
        reshape_op: OperatorIR,
        reshape_idx: int,
        consumers: Dict[str, List[int]],
    ) -> bool:
        if len(reshape_op.inputs) < 2:
            return False
        shape_name = str(reshape_op.inputs[1])
        shape_tensor = model_ir.tensors.get(shape_name, None)
        shape_vals = _read_const_ints_from_tensor(shape_tensor)
        if shape_vals is None or len(shape_vals) != 3:
            return False
        swapped = [int(shape_vals[0]), int(shape_vals[2]), int(shape_vals[1])]
        shape_consumers = [int(v) for v in consumers.get(shape_name, [])]
        if any(int(v) != int(reshape_idx) for v in shape_consumers):
            if shape_tensor is None:
                return False
            cloned_name = _unique_tensor_name(f"{shape_name}_nhwc")
            data_dtype = np.int32
            if shape_tensor.data is not None:
                try:
                    data_dtype = np.asarray(shape_tensor.data).dtype
                except Exception:
                    data_dtype = np.int32
            swapped_arr = np.asarray(swapped, dtype=data_dtype)
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(shape_tensor.dtype),
                shape=[int(v) for v in list(swapped_arr.shape)],
                shape_signature=[int(v) for v in list(swapped_arr.shape)],
                data=np.asarray(swapped_arr),
                is_variable=False,
                quantization=_clone_quantization(shape_tensor.quantization),
            )
            reshape_inputs = [str(v) for v in list(reshape_op.inputs)]
            reshape_inputs[1] = str(cloned_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=reshape_inputs,
            )
            return True
        _write_const_ints_to_tensor(shape_tensor, swapped)
        return True

    def _swap_reshape_option_shape(reshape_op: OperatorIR) -> None:
        if not isinstance(reshape_op.options, dict):
            return
        opts = dict(reshape_op.options)
        changed = False
        for key in ["newShape", "onnxRawNewShape"]:
            value = opts.get(key, None)
            if not isinstance(value, list) or len(value) != 3:
                continue
            swapped = [int(value[0]), int(value[2]), int(value[1])]
            if swapped != [int(v) for v in value]:
                opts[key] = swapped
                changed = True
        if changed:
            reshape_op.options = opts

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for tail_concat_idx, tail_concat_op in enumerate(model_ir.operators):
            if str(tail_concat_op.op_type) != "CONCATENATION" or len(tail_concat_op.outputs) != 1:
                continue
            tail_axis = int(tail_concat_op.options.get("axis", 2))
            if tail_axis < 0:
                tail_axis += 3
            if tail_axis != 2:
                continue

            tail_out_name = str(tail_concat_op.outputs[0])
            tail_out_tensor = model_ir.tensors.get(tail_out_name, None)
            if tail_out_tensor is None or len(list(tail_out_tensor.shape)) != 3:
                continue

            tail_inputs = [str(v) for v in list(tail_concat_op.inputs)]
            if len(tail_inputs) < 2:
                continue

            branch_plans: List[Dict[str, Any]] = []
            removable_pre_transpose_ids: set[int] = set()
            rewritable = True

            for reshape_out_name in tail_inputs:
                reshape_idx = producers.get(str(reshape_out_name), None)
                if reshape_idx is None:
                    rewritable = False
                    break
                reshape_op = model_ir.operators[int(reshape_idx)]
                if (
                    str(reshape_op.op_type) != "RESHAPE"
                    or len(reshape_op.inputs) < 2
                    or len(reshape_op.outputs) != 1
                    or str(reshape_op.outputs[0]) != str(reshape_out_name)
                ):
                    rewritable = False
                    break
                if set(int(v) for v in consumers.get(str(reshape_out_name), [])) != {int(tail_concat_idx)}:
                    rewritable = False
                    break

                branch_concat_out_name = str(reshape_op.inputs[0])
                branch_concat_idx = producers.get(branch_concat_out_name, None)
                if branch_concat_idx is None:
                    rewritable = False
                    break
                branch_concat_op = model_ir.operators[int(branch_concat_idx)]
                if str(branch_concat_op.op_type) != "CONCATENATION" or len(branch_concat_op.outputs) != 1:
                    rewritable = False
                    break
                branch_axis = int(branch_concat_op.options.get("axis", 1))
                if branch_axis < 0:
                    branch_axis += 4
                if branch_axis != 1:
                    rewritable = False
                    break
                if set(int(v) for v in consumers.get(branch_concat_out_name, [])) != {int(reshape_idx)}:
                    rewritable = False
                    break

                branch_concat_inputs = [str(v) for v in list(branch_concat_op.inputs)]
                if len(branch_concat_inputs) != 2:
                    rewritable = False
                    break

                pre_transpose_idx: Optional[int] = None
                unary_indices: List[int] = []
                slice_plans: List[Dict[str, Any]] = []
                expected_pre_users: set[int] = set()

                for branch_input_name in branch_concat_inputs:
                    slice_out_name = str(branch_input_name)
                    unary_idx: Optional[int] = None

                    maybe_unary_idx = producers.get(slice_out_name, None)
                    if maybe_unary_idx is not None:
                        maybe_unary_op = model_ir.operators[int(maybe_unary_idx)]
                        if (
                            str(maybe_unary_op.op_type) in unary_ops
                            and len(maybe_unary_op.inputs) == 1
                            and len(maybe_unary_op.outputs) == 1
                            and str(maybe_unary_op.outputs[0]) == slice_out_name
                            and set(int(v) for v in consumers.get(slice_out_name, [])) == {int(branch_concat_idx)}
                            and slice_out_name not in model_outputs
                        ):
                            unary_idx = int(maybe_unary_idx)
                            unary_indices.append(int(maybe_unary_idx))
                            slice_out_name = str(maybe_unary_op.inputs[0])

                    slice_idx = producers.get(slice_out_name, None)
                    if slice_idx is None:
                        rewritable = False
                        break
                    slice_op = model_ir.operators[int(slice_idx)]
                    if (
                        str(slice_op.op_type) != "SLICE"
                        or len(slice_op.inputs) < 3
                        or len(slice_op.outputs) != 1
                        or str(slice_op.outputs[0]) != slice_out_name
                    ):
                        rewritable = False
                        break
                    expected_slice_users = {int(unary_idx)} if unary_idx is not None else {int(branch_concat_idx)}
                    if set(int(v) for v in consumers.get(slice_out_name, [])) != expected_slice_users:
                        rewritable = False
                        break

                    begin_name = str(slice_op.inputs[1])
                    size_name = str(slice_op.inputs[2])
                    begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(begin_name, None))
                    size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(size_name, None))
                    if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
                        rewritable = False
                        break
                    if int(begin_vals[0]) != 0 or int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
                        rewritable = False
                        break
                    if int(begin_vals[1]) < 0 or int(size_vals[1]) <= 0:
                        rewritable = False
                        break

                    pre_idx = producers.get(str(slice_op.inputs[0]), None)
                    if pre_idx is None:
                        rewritable = False
                        break
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != str(slice_op.inputs[0])
                        or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                        or str(pre_op.outputs[0]) in model_outputs
                    ):
                        rewritable = False
                        break
                    if pre_transpose_idx is None:
                        pre_transpose_idx = int(pre_idx)
                    elif int(pre_transpose_idx) != int(pre_idx):
                        rewritable = False
                        break

                    expected_pre_users.add(int(slice_idx))
                    slice_plans.append(
                        {
                            "slice_idx": int(slice_idx),
                            "slice_out_name": str(slice_out_name),
                            "begin_name": str(begin_name),
                            "size_name": str(size_name),
                        }
                    )

                if not rewritable or pre_transpose_idx is None:
                    rewritable = False
                    break

                pre_op = model_ir.operators[int(pre_transpose_idx)]
                pre_out_name = str(pre_op.outputs[0])
                if set(int(v) for v in consumers.get(pre_out_name, [])) != set(expected_pre_users):
                    rewritable = False
                    break

                pre_input_name = str(pre_op.inputs[0])
                pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
                pre_out_tensor = model_ir.tensors.get(pre_out_name, None)
                if (
                    pre_input_tensor is None
                    or pre_out_tensor is None
                    or len(list(pre_input_tensor.shape)) != 4
                    or len(list(pre_out_tensor.shape)) != 4
                ):
                    rewritable = False
                    break

                branch_plans.append(
                    {
                        "reshape_idx": int(reshape_idx),
                        "reshape_out_name": str(reshape_out_name),
                        "branch_concat_idx": int(branch_concat_idx),
                        "branch_concat_out_name": str(branch_concat_out_name),
                        "pre_transpose_idx": int(pre_transpose_idx),
                        "pre_input_name": str(pre_input_name),
                        "unary_indices": [int(v) for v in unary_indices],
                        "slice_plans": list(slice_plans),
                    }
                )
                removable_pre_transpose_ids.add(int(id(pre_op)))

            if not rewritable or len(branch_plans) < 2:
                continue

            # Keep this rewrite strict: all branch reshapes must share the same channel dimension.
            channel_dim: Optional[int] = None
            for plan in branch_plans:
                reshape_tensor = model_ir.tensors.get(str(plan["reshape_out_name"]), None)
                if reshape_tensor is None or len(list(reshape_tensor.shape)) != 3:
                    rewritable = False
                    break
                cdim = int(reshape_tensor.shape[1]) if len(list(reshape_tensor.shape)) > 1 else -1
                if cdim <= 0:
                    rewritable = False
                    break
                if channel_dim is None:
                    channel_dim = int(cdim)
                elif int(channel_dim) != int(cdim):
                    rewritable = False
                    break
            if not rewritable:
                continue

            apply_ok = True
            for plan in branch_plans:
                pre_input_name = str(plan["pre_input_name"])

                for slice_plan in list(plan["slice_plans"]):
                    slice_idx = int(slice_plan["slice_idx"])
                    begin_name = str(slice_plan["begin_name"])
                    size_name = str(slice_plan["size_name"])
                    if set(int(v) for v in consumers.get(begin_name, [])) != {int(slice_idx)}:
                        apply_ok = False
                        break
                    if set(int(v) for v in consumers.get(size_name, [])) != {int(slice_idx)}:
                        apply_ok = False
                        break

                    slice_op = model_ir.operators[int(slice_idx)]
                    begin_tensor = model_ir.tensors.get(begin_name, None)
                    size_tensor = model_ir.tensors.get(size_name, None)
                    begin_vals = _read_const_ints_from_tensor(begin_tensor)
                    size_vals = _read_const_ints_from_tensor(size_tensor)
                    if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
                        apply_ok = False
                        break

                    new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
                    new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
                    _write_const_ints_to_tensor(begin_tensor, new_begin)
                    _write_const_ints_to_tensor(size_tensor, new_size)
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=slice_op,
                        input_index=0,
                        new_input_name=pre_input_name,
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(slice_plan["slice_out_name"]), None),
                        perm_nchw_to_nhwc,
                    )

                if not apply_ok:
                    break

                for unary_idx in list(plan["unary_indices"]):
                    unary_out_name = str(model_ir.operators[int(unary_idx)].outputs[0])
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(unary_out_name, None),
                        perm_nchw_to_nhwc,
                    )

                branch_concat_op = model_ir.operators[int(plan["branch_concat_idx"])]
                branch_concat_op.options["axis"] = 3
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(plan["branch_concat_out_name"]), None),
                    perm_nchw_to_nhwc,
                )

                reshape_op = model_ir.operators[int(plan["reshape_idx"])]
                if not _swap_reshape_shape_tensor(
                    reshape_op=reshape_op,
                    reshape_idx=int(plan["reshape_idx"]),
                    consumers=consumers,
                ):
                    apply_ok = False
                    break
                _swap_reshape_option_shape(reshape_op)
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(plan["reshape_out_name"]), None),
                    perm_3d_nchw_to_nhwc,
                )

            if not apply_ok:
                continue

            tail_shape = [int(v) for v in list(tail_out_tensor.shape)]
            tail_signature = (
                [int(v) for v in list(tail_out_tensor.shape_signature)]
                if tail_out_tensor.shape_signature is not None
                else [int(v) for v in list(tail_out_tensor.shape)]
            )
            tail_nhwc_shape = _permute_shape(tail_shape, perm_3d_nchw_to_nhwc)
            tail_nhwc_signature = _permute_shape(tail_signature, perm_3d_nchw_to_nhwc)
            if tail_nhwc_shape is None or tail_nhwc_signature is None:
                continue

            tail_concat_nhwc_name = _unique_tensor_name(f"{tail_out_name}_nhwc")
            model_ir.tensors[tail_concat_nhwc_name] = TensorIR(
                name=tail_concat_nhwc_name,
                dtype=str(tail_out_tensor.dtype),
                shape=[int(v) for v in list(tail_nhwc_shape)],
                shape_signature=[int(v) for v in list(tail_nhwc_signature)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(tail_out_tensor.quantization),
            )
            tail_concat_op.options["axis"] = 1
            _set_operator_outputs(
                model_ir=model_ir,
                op=tail_concat_op,
                new_outputs=[tail_concat_nhwc_name],
            )

            perm_3d_name = _find_or_create_perm_3d_nhwc_to_nchw()
            tail_concat_current_idx = _op_index(tail_concat_op)
            if tail_concat_current_idx is None:
                continue
            model_ir.operators.insert(
                int(tail_concat_current_idx) + 1,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[tail_concat_nhwc_name, perm_3d_name],
                    outputs=[tail_out_name],
                    options={},
                ),
            )

            remove_indices = [
                int(op_idx)
                for op_idx, op in enumerate(model_ir.operators)
                if int(id(op)) in removable_pre_transpose_ids
            ]
            for remove_idx in sorted(list(set(remove_indices)), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_slice_logistic_concat_reshape_tail_nhwc_chains": int(optimized)}


def _optimize_shufflenet_transpose_shuffle_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Optimize ShuffleNet-style channel-shuffle blocks that rely on NCHW transpose chains.

    Target (representative):
      x_nhwc
        -> TRANSPOSE(0,3,1,2)
        -> RESHAPE -> TRANSPOSE(1,0,2) -> RESHAPE
        -> GATHER(idx=0/1)
      one gather -> TRANSPOSE(0,2,3,1) -> conv branch -> TRANSPOSE(0,3,1,2) -> (optional unary)
      other gather -> skip branch
      CONCAT(axis=1, [skip_nchw, branch_nchw]) -> y_nchw

    Rewrite:
      x_nhwc -> GATHER(axis=3, even/odd channels) -> skip_nhwc / branch_nhwc
      conv branch runs in NHWC directly
      CONCAT(axis=3, [skip_nhwc, branch_nhwc]) -> y_nhwc
      y_nhwc -> TRANSPOSE(0,3,1,2) -> y_nchw

    This preserves downstream NCHW contracts while reducing redundant transpose chains.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {
        "LOGISTIC",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "TANH",
        "GELU",
        "HARD_SWISH",
        "ABS",
        "EXP",
        "NEG",
        "SQRT",
    }
    branch_passthrough_ops = {
        "CONV_2D",
        "DEPTHWISE_CONV_2D",
        "ADD",
        "SUB",
        "MUL",
        "DIV",
        "PAD",
        "PADV2",
        "MAX_POOL_2D",
        "AVERAGE_POOL_2D",
        "CAST",
    }.union(unary_ops)

    def _unique_tensor_name(base: str) -> str:
        if base not in model_ir.tensors:
            return base
        suffix = 1
        while True:
            candidate = f"{base}_{suffix}"
            if candidate not in model_ir.tensors:
                return candidate
            suffix += 1

    def _op_index(op_ref: OperatorIR) -> Optional[int]:
        for idx, op in enumerate(model_ir.operators):
            if id(op) == id(op_ref):
                return int(idx)
        return None

    def _is_singleton_channel_nhwc_to_nchw_reshape(
        *,
        input_name: str,
        output_name: str,
    ) -> bool:
        input_tensor = model_ir.tensors.get(str(input_name), None)
        output_tensor = model_ir.tensors.get(str(output_name), None)
        if input_tensor is None or output_tensor is None:
            return False
        in_shape = [int(v) for v in list(input_tensor.shape)]
        out_shape = [int(v) for v in list(output_tensor.shape)]
        if len(in_shape) != 4 or len(out_shape) != 4:
            return False
        return (
            int(in_shape[0]) == int(out_shape[0])
            and int(in_shape[1]) == int(out_shape[2])
            and int(in_shape[2]) == int(out_shape[3])
            and int(in_shape[3]) == 1
            and int(out_shape[1]) == 1
        )

    def _parse_shuffle_split_candidate(
        *,
        op: OperatorIR,
        split_base_name: str,
        half_channels: int,
    ) -> Optional[Tuple[int, str]]:
        """
        Parse one shuffle split branch selector and return:
          (selector_idx_value, output_name), where selector_idx_value in {0, 1}.

        Supported forms:
        - GATHER(axis=0, indices scalar in {0,1}) on rank-5 split base
        - SLICE(begin/size const) selecting channel-halves on NCHW rank-4 split base
        """
        if len(op.outputs) != 1:
            return None
        output_name = str(op.outputs[0])

        if str(op.op_type) == "GATHER":
            if len(op.inputs) < 2:
                return None
            if int(op.options.get("axis", 0)) != 0:
                return None
            idx_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
            if idx_vals is None or len(idx_vals) != 1:
                return None
            idx_value = int(idx_vals[0])
            if idx_value not in [0, 1]:
                return None
            return int(idx_value), output_name

        if str(op.op_type) == "SLICE":
            if len(op.inputs) < 3:
                return None
            if str(op.inputs[0]) != str(split_base_name):
                return None
            begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
            size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
            if begin_vals is None or size_vals is None:
                return None
            if len(begin_vals) != len(size_vals) or len(begin_vals) != 4:
                return None

            axis = 1
            begin = [int(v) for v in begin_vals]
            size = [int(v) for v in size_vals]
            start = int(begin[axis])
            if start not in [0, int(half_channels)]:
                return None
            idx_value = 0 if int(start) == 0 else 1
            if int(size[axis]) != int(half_channels):
                return None
            for dim_idx in range(4):
                if dim_idx == axis:
                    continue
                if int(begin[dim_idx]) != 0:
                    return None
                # Non-channel dims can be full extent (-1) or explicit full size.
                # Anything else means this is not a simple channel-halving split.
                if int(size[dim_idx]) not in [-1]:
                    split_base_tensor = model_ir.tensors.get(str(split_base_name), None)
                    if split_base_tensor is None or len(list(split_base_tensor.shape)) != 4:
                        return None
                    split_dim = int(split_base_tensor.shape[dim_idx])
                    if int(split_dim) <= 0 or int(size[dim_idx]) != int(split_dim):
                        return None
            return int(idx_value), output_name

        return None

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat0_idx, concat0_op in enumerate(model_ir.operators):
            if str(concat0_op.op_type) != "CONCATENATION" or len(concat0_op.outputs) != 1:
                continue
            if int(concat0_op.options.get("axis", -1)) != 3:
                continue
            x_nhwc_name = str(concat0_op.outputs[0])
            x_nhwc_tensor = model_ir.tensors.get(x_nhwc_name, None)
            if x_nhwc_tensor is None or len(list(x_nhwc_tensor.shape)) != 4:
                continue
            x_shape = [int(v) for v in list(x_nhwc_tensor.shape)]
            channels = int(x_shape[3])
            if int(channels) <= 0 or int(channels % 2) != 0:
                continue
            half_channels = int(channels // 2)

            t0_users = [int(v) for v in consumers.get(x_nhwc_name, [])]
            if len(t0_users) != 1:
                continue
            t0_idx = int(t0_users[0])
            t0_op = model_ir.operators[t0_idx]
            if str(t0_op.op_type) != "TRANSPOSE" or len(t0_op.inputs) < 2 or len(t0_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, t0_op) != perm_nhwc_to_nchw:
                continue
            x_nchw_name = str(t0_op.outputs[0])
            if x_nchw_name in model_outputs:
                continue

            r1_users = [int(v) for v in consumers.get(x_nchw_name, [])]
            if len(r1_users) != 1:
                continue
            r1_op = model_ir.operators[int(r1_users[0])]
            if str(r1_op.op_type) != "RESHAPE" or len(r1_op.inputs) < 1 or len(r1_op.outputs) != 1:
                continue
            r1_out = str(r1_op.outputs[0])
            if r1_out in model_outputs:
                continue

            t1_users = [int(v) for v in consumers.get(r1_out, [])]
            if len(t1_users) != 1:
                continue
            t1_op = model_ir.operators[int(t1_users[0])]
            if str(t1_op.op_type) != "TRANSPOSE" or len(t1_op.inputs) < 2 or len(t1_op.outputs) != 1:
                continue
            t1_perm = _read_transpose_perm(model_ir, t1_op)
            if t1_perm not in ([1, 0, 2], [0, 2, 1, 3, 4]):
                continue
            is_rank5_shuffle_swap = bool(t1_perm == [0, 2, 1, 3, 4])
            t1_out = str(t1_op.outputs[0])
            if t1_out in model_outputs:
                continue

            r2_users = [int(v) for v in consumers.get(t1_out, [])]
            if len(r2_users) != 1:
                continue
            r2_op = model_ir.operators[int(r2_users[0])]
            if str(r2_op.op_type) != "RESHAPE" or len(r2_op.inputs) < 1 or len(r2_op.outputs) != 1:
                continue
            split_base = str(r2_op.outputs[0])
            if split_base in model_outputs:
                continue

            split_user_indices = [int(v) for v in consumers.get(split_base, [])]
            if len(split_user_indices) != 2:
                continue
            split_candidates: Dict[int, Tuple[int, OperatorIR, str]] = {}
            split_ids: set[int] = set()
            valid_split_candidates = True
            for split_idx in split_user_indices:
                split_op = model_ir.operators[int(split_idx)]
                parsed = _parse_shuffle_split_candidate(
                    op=split_op,
                    split_base_name=split_base,
                    half_channels=int(half_channels),
                )
                if parsed is None:
                    valid_split_candidates = False
                    break
                idx_value, out_name = parsed
                if int(idx_value) in split_ids:
                    valid_split_candidates = False
                    break
                split_ids.add(int(idx_value))
                split_candidates[int(idx_value)] = (
                    int(split_idx),
                    split_op,
                    str(out_name),
                )
            if not valid_split_candidates or set(split_candidates.keys()) != {0, 1}:
                continue

            # Determine which split candidate feeds NHWC conv-branch through NCHW->NHWC transpose.
            conv_split_idx_value: Optional[int] = None
            pre_conv_t_idx: Optional[int] = None
            pre_conv_t_op: Optional[OperatorIR] = None
            pre_conv_out_name: Optional[str] = None
            for idx_value in [0, 1]:
                _, _, split_out = split_candidates[idx_value]
                users = [int(v) for v in consumers.get(split_out, [])]
                if len(users) != 1:
                    continue
                maybe_t_idx = int(users[0])
                maybe_t_op = model_ir.operators[maybe_t_idx]
                if (
                    str(maybe_t_op.op_type) == "TRANSPOSE"
                    and len(maybe_t_op.inputs) >= 2
                    and len(maybe_t_op.outputs) == 1
                    and _read_transpose_perm(model_ir, maybe_t_op) == perm_nchw_to_nhwc
                    and str(maybe_t_op.outputs[0]) not in model_outputs
                ):
                    conv_split_idx_value = int(idx_value)
                    pre_conv_t_idx = int(maybe_t_idx)
                    pre_conv_t_op = maybe_t_op
                    pre_conv_out_name = str(maybe_t_op.outputs[0])
                    break
            if (
                conv_split_idx_value is None
                or pre_conv_t_idx is None
                or pre_conv_t_op is None
                or pre_conv_out_name is None
            ):
                continue
            skip_split_idx_value = 1 - int(conv_split_idx_value)
            _, _, skip_split_out_name = split_candidates[int(skip_split_idx_value)]
            if skip_split_out_name in model_outputs:
                continue

            pre_conv_out_users = [int(v) for v in consumers.get(pre_conv_out_name, [])]
            if len(pre_conv_out_users) != 1:
                continue
            first_branch_op = model_ir.operators[int(pre_conv_out_users[0])]

            # Trace branch NHWC flow until NHWC->NCHW transpose.
            transpose_post_idx: Optional[int] = None
            transpose_post_op: Optional[OperatorIR] = None
            branch_nhwc_cursor = str(pre_conv_out_name)
            visited_tensors: set[str] = set()
            branch_trace_valid = True
            while True:
                if branch_nhwc_cursor in visited_tensors:
                    branch_trace_valid = False
                    break
                visited_tensors.add(branch_nhwc_cursor)
                users = [int(v) for v in consumers.get(branch_nhwc_cursor, [])]
                if len(users) != 1:
                    branch_trace_valid = False
                    break
                op_idx = int(users[0])
                op = model_ir.operators[op_idx]
                if (
                    str(op.op_type) == "TRANSPOSE"
                    and len(op.inputs) >= 2
                    and len(op.outputs) == 1
                    and _read_transpose_perm(model_ir, op) == perm_nhwc_to_nchw
                ):
                    transpose_post_idx = int(op_idx)
                    transpose_post_op = op
                    break
                if str(op.op_type) not in branch_passthrough_ops or len(op.outputs) != 1:
                    branch_trace_valid = False
                    break
                branch_nhwc_cursor = str(op.outputs[0])
            if (
                not branch_trace_valid
                or transpose_post_idx is None
                or transpose_post_op is None
            ):
                continue

            branch_nhwc_tail_name = str(transpose_post_op.inputs[0])
            branch_nchw_start_name = str(transpose_post_op.outputs[0])
            if branch_nchw_start_name in model_outputs:
                continue

            skip_users = [int(v) for v in consumers.get(skip_split_out_name, [])]
            if len(skip_users) != 1:
                continue
            concat1_idx = int(skip_users[0])
            concat1_op = model_ir.operators[concat1_idx]
            if str(concat1_op.op_type) != "CONCATENATION" or len(concat1_op.outputs) != 1:
                continue
            if int(concat1_op.options.get("axis", -1)) != 1:
                continue
            concat1_inputs = [str(v) for v in list(concat1_op.inputs)]
            if len(concat1_inputs) != 2 or skip_split_out_name not in concat1_inputs:
                continue

            branch_nchw_cursor = branch_nchw_start_name
            unary_chain_indices: List[int] = []
            branch_trace_to_concat_ok = True
            while True:
                users = [int(v) for v in consumers.get(branch_nchw_cursor, [])]
                if len(users) != 1:
                    branch_trace_to_concat_ok = False
                    break
                user_idx = int(users[0])
                if user_idx == int(concat1_idx):
                    break
                user_op = model_ir.operators[user_idx]
                if (
                    str(user_op.op_type) in unary_ops
                    and len(user_op.inputs) == 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == branch_nchw_cursor
                    and branch_nchw_cursor not in model_outputs
                ):
                    unary_chain_indices.append(int(user_idx))
                    branch_nchw_cursor = str(user_op.outputs[0])
                    continue
                branch_trace_to_concat_ok = False
                break
            if not branch_trace_to_concat_ok:
                continue
            branch_concat_input_nchw = str(branch_nchw_cursor)
            if branch_concat_input_nchw not in concat1_inputs:
                continue

            # Build NHWC branch gather tensors from concat0 output.
            x_signature = (
                list(x_nhwc_tensor.shape_signature)
                if x_nhwc_tensor.shape_signature is not None
                else list(x_shape)
            )
            even_odd_shape = [int(x_shape[0]), int(x_shape[1]), int(x_shape[2]), int(half_channels)]
            even_odd_signature = [int(v) for v in list(x_signature)]
            if len(even_odd_signature) == 4 and int(even_odd_signature[3]) > 0:
                even_odd_signature[3] = int(half_channels)
            even_name = _unique_tensor_name(f"{x_nhwc_name}_shuffle_even_nhwc")
            odd_name = _unique_tensor_name(f"{x_nhwc_name}_shuffle_odd_nhwc")
            model_ir.tensors[even_name] = TensorIR(
                name=even_name,
                dtype=str(x_nhwc_tensor.dtype),
                shape=list(even_odd_shape),
                shape_signature=[int(v) for v in list(even_odd_signature)],
                quantization=_clone_quantization(x_nhwc_tensor.quantization),
            )
            model_ir.tensors[odd_name] = TensorIR(
                name=odd_name,
                dtype=str(x_nhwc_tensor.dtype),
                shape=list(even_odd_shape),
                shape_signature=[int(v) for v in list(even_odd_signature)],
                quantization=_clone_quantization(x_nhwc_tensor.quantization),
            )

            if is_rank5_shuffle_swap:
                r1_tensor = model_ir.tensors.get(str(r1_out), None)
                if r1_tensor is None:
                    continue
                r1_shape = [int(v) for v in list(r1_tensor.shape)]
                if len(r1_shape) != 5 or not _is_fully_known_positive_shape(r1_shape):
                    continue
                groups = int(r1_shape[1])
                cpg = int(r1_shape[2])
                if (
                    int(groups) <= 1
                    or int(cpg) <= 0
                    or int(groups) != 2
                    or int(groups * cpg) != int(channels)
                ):
                    continue
                shuffle_indices = np.asarray(
                    [int((k % groups) * cpg + (k // groups)) for k in range(int(channels))],
                    dtype=np.int32,
                )
                first_split_indices = np.asarray(shuffle_indices[:int(half_channels)], dtype=np.int32)
                second_split_indices = np.asarray(shuffle_indices[int(half_channels):], dtype=np.int32)
            else:
                first_split_indices = np.asarray(
                    [int(v) for v in range(0, int(channels), 2)],
                    dtype=np.int32,
                )
                second_split_indices = np.asarray(
                    [int(v) for v in range(1, int(channels), 2)],
                    dtype=np.int32,
                )
            if (
                len(first_split_indices) != int(half_channels)
                or len(second_split_indices) != int(half_channels)
            ):
                continue

            even_indices_name = _unique_tensor_name(f"{x_nhwc_name}_shuffle_even_indices")
            odd_indices_name = _unique_tensor_name(f"{x_nhwc_name}_shuffle_odd_indices")
            model_ir.tensors[even_indices_name] = TensorIR(
                name=even_indices_name,
                dtype="INT32",
                shape=[int(half_channels)],
                shape_signature=[int(half_channels)],
                data=np.asarray(first_split_indices, dtype=np.int32),
                is_variable=False,
            )
            model_ir.tensors[odd_indices_name] = TensorIR(
                name=odd_indices_name,
                dtype="INT32",
                shape=[int(half_channels)],
                shape_signature=[int(half_channels)],
                data=np.asarray(second_split_indices, dtype=np.int32),
                is_variable=False,
            )

            even_gather_op = OperatorIR(
                op_type="GATHER",
                inputs=[x_nhwc_name, even_indices_name],
                outputs=[even_name],
                options={"axis": 3, "batchDims": 0},
            )
            odd_gather_op = OperatorIR(
                op_type="GATHER",
                inputs=[x_nhwc_name, odd_indices_name],
                outputs=[odd_name],
                options={"axis": 3, "batchDims": 0},
            )

            conv_branch_input_nhwc = even_name if int(conv_split_idx_value) == 0 else odd_name
            skip_branch_input_nhwc = odd_name if int(conv_split_idx_value) == 0 else even_name

            # Rewrite branch start input.
            _replace_tensor_inputs(model_ir, pre_conv_out_name, conv_branch_input_nhwc)

            # Rewire unary tail (if present) to NHWC by bypassing post transpose.
            if len(unary_chain_indices) > 0:
                first_unary_idx = int(unary_chain_indices[0])
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=model_ir.operators[first_unary_idx],
                    new_inputs=[branch_nhwc_tail_name],
                )
                for unary_idx in unary_chain_indices:
                    unary_out_name = str(model_ir.operators[int(unary_idx)].outputs[0])
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(unary_out_name, None),
                        perm_nchw_to_nhwc,
                    )
                branch_concat_input_nhwc = str(model_ir.operators[int(unary_chain_indices[-1])].outputs[0])
            else:
                branch_concat_input_nhwc = branch_nhwc_tail_name

            concat1_out_name = str(concat1_op.outputs[0])
            concat1_out_tensor = model_ir.tensors.get(concat1_out_name, None)
            if concat1_out_tensor is None:
                continue
            concat1_out_signature = (
                list(concat1_out_tensor.shape_signature)
                if concat1_out_tensor.shape_signature is not None
                else list(concat1_out_tensor.shape)
            )
            concat1_out_nhwc_shape = _permute_shape(list(concat1_out_tensor.shape), perm_nchw_to_nhwc)
            concat1_out_nhwc_signature = _permute_shape(list(concat1_out_signature), perm_nchw_to_nhwc)
            if concat1_out_nhwc_shape is None or concat1_out_nhwc_signature is None:
                continue

            concat1_out_nhwc_name = _unique_tensor_name(f"{concat1_out_name}_nhwc")
            model_ir.tensors[concat1_out_nhwc_name] = TensorIR(
                name=concat1_out_nhwc_name,
                dtype=str(concat1_out_tensor.dtype),
                shape=[int(v) for v in list(concat1_out_nhwc_shape)],
                shape_signature=[int(v) for v in list(concat1_out_nhwc_signature)],
                quantization=_clone_quantization(concat1_out_tensor.quantization),
            )

            new_concat1_inputs: List[str] = []
            for input_name in concat1_inputs:
                if str(input_name) == str(skip_split_out_name):
                    new_concat1_inputs.append(str(skip_branch_input_nhwc))
                elif str(input_name) == str(branch_concat_input_nchw):
                    new_concat1_inputs.append(str(branch_concat_input_nhwc))
                else:
                    new_concat1_inputs.append(str(input_name))
            if str(skip_branch_input_nhwc) not in set(new_concat1_inputs):
                continue
            if str(branch_concat_input_nhwc) not in set(new_concat1_inputs):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat1_op,
                new_inputs=[str(v) for v in list(new_concat1_inputs)],
            )
            concat1_op.options["axis"] = 3
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat1_op,
                new_outputs=[concat1_out_nhwc_name],
            )

            post_perm_name = str(transpose_post_op.inputs[1]) if len(transpose_post_op.inputs) >= 2 else ""
            post_perm_tensor = model_ir.tensors.get(post_perm_name, None)
            post_perm_vals = _read_const_ints_from_tensor(post_perm_tensor)
            if post_perm_vals != perm_nhwc_to_nchw:
                post_perm_name = _unique_tensor_name("shufflenet_concat_post_perm")
                model_ir.tensors[post_perm_name] = TensorIR(
                    name=post_perm_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                )

            concat1_current_idx = _op_index(concat1_op)
            if concat1_current_idx is None:
                continue
            model_ir.operators.insert(
                int(concat1_current_idx) + 1,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[concat1_out_nhwc_name, post_perm_name],
                    outputs=[concat1_out_name],
                ),
            )

            concat0_current_idx = _op_index(concat0_op)
            if concat0_current_idx is None:
                continue
            model_ir.operators.insert(int(concat0_current_idx) + 1, even_gather_op)
            model_ir.operators.insert(int(concat0_current_idx) + 2, odd_gather_op)

            remove_ids = {
                int(id(t0_op)),
                int(id(r1_op)),
                int(id(t1_op)),
                int(id(r2_op)),
                int(id(split_candidates[0][1])),
                int(id(split_candidates[1][1])),
                int(id(pre_conv_t_op)),
                int(id(transpose_post_op)),
            }
            remove_indices = [
                int(idx)
                for idx, op in enumerate(model_ir.operators)
                if int(id(op)) in remove_ids
            ]
            for remove_idx in sorted(list(set(remove_indices)), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_shufflenet_transpose_shuffle_chains": int(optimized)}


def _optimize_shufflenet_reshape_transpose_shuffle_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Collapse ShuffleNet channel-shuffle chains expressed as:

      x_nhwc
        -> TRANSPOSE(0,3,1,2)
        -> RESHAPE([N,g,cpg,H,W])
        -> TRANSPOSE(0,2,1,3,4)
        -> RESHAPE([N,C,H,W])
        -> TRANSPOSE(0,2,3,1)
        -> y_nhwc

    into:

      x_nhwc -> GATHER(axis=3, shuffle_indices) -> y_nhwc

    where C=g*cpg and shuffle_indices[k] = (k % g) * cpg + (k // g).
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_shuffle_swap = [0, 2, 1, 3, 4]
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "LOGISTIC",
        "TANH",
        "LEAKY_RELU",
        "HARD_SWISH",
        "ABS",
        "EXP",
        "NEG",
        "SQRT",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for t0_idx, t0_op in enumerate(model_ir.operators):
            if str(t0_op.op_type) != "TRANSPOSE" or len(t0_op.inputs) < 2 or len(t0_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, t0_op) != perm_nhwc_to_nchw:
                continue

            x_nhwc_name = str(t0_op.inputs[0])
            x_nchw_name = str(t0_op.outputs[0])
            if x_nchw_name in model_outputs:
                continue

            t0_users = [int(v) for v in consumers.get(x_nchw_name, [])]
            if len(t0_users) == 0:
                continue

            for t0_user_idx in t0_users:
                gather_src_name = str(x_nhwc_name)
                unary_bridge_op: Optional[OperatorIR] = None
                r1_idx: Optional[int] = None
                r1_op: Optional[OperatorIR] = None

                t0_user_op = model_ir.operators[int(t0_user_idx)]
                if (
                    str(t0_user_op.op_type) == "RESHAPE"
                    and len(t0_user_op.inputs) >= 1
                    and len(t0_user_op.outputs) == 1
                    and str(t0_user_op.inputs[0]) == x_nchw_name
                ):
                    r1_idx = int(t0_user_idx)
                    r1_op = t0_user_op
                elif (
                    str(t0_user_op.op_type) in unary_passthrough_ops
                    and len(t0_user_op.inputs) == 1
                    and len(t0_user_op.outputs) == 1
                    and str(t0_user_op.inputs[0]) == x_nchw_name
                    and str(t0_user_op.outputs[0]) not in model_outputs
                ):
                    unary_out_name = str(t0_user_op.outputs[0])
                    unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
                    if len(unary_users) != 1:
                        continue
                    unary_r1_idx = int(unary_users[0])
                    unary_r1_op = model_ir.operators[int(unary_r1_idx)]
                    if (
                        str(unary_r1_op.op_type) != "RESHAPE"
                        or len(unary_r1_op.inputs) < 1
                        or len(unary_r1_op.outputs) != 1
                        or str(unary_r1_op.inputs[0]) != unary_out_name
                    ):
                        continue
                    r1_idx = int(unary_r1_idx)
                    r1_op = unary_r1_op
                    unary_bridge_op = t0_user_op
                    gather_src_name = str(unary_out_name)
                else:
                    continue

                if r1_idx is None or r1_op is None:
                    continue
                r1_out_name = str(r1_op.outputs[0])
                if r1_out_name in model_outputs:
                    continue

                r1_users = [int(v) for v in consumers.get(r1_out_name, [])]
                if len(r1_users) != 1:
                    continue
                t1_idx = int(r1_users[0])
                t1_op = model_ir.operators[int(t1_idx)]
                if (
                    str(t1_op.op_type) != "TRANSPOSE"
                    or len(t1_op.inputs) < 2
                    or len(t1_op.outputs) != 1
                    or str(t1_op.inputs[0]) != r1_out_name
                    or _read_transpose_perm(model_ir, t1_op) != perm_shuffle_swap
                ):
                    continue
                t1_out_name = str(t1_op.outputs[0])
                if t1_out_name in model_outputs:
                    continue

                t1_users = [int(v) for v in consumers.get(t1_out_name, [])]
                if len(t1_users) != 1:
                    continue
                r2_idx = int(t1_users[0])
                r2_op = model_ir.operators[int(r2_idx)]
                if (
                    str(r2_op.op_type) != "RESHAPE"
                    or len(r2_op.inputs) < 1
                    or len(r2_op.outputs) != 1
                    or str(r2_op.inputs[0]) != t1_out_name
                ):
                    continue
                r2_out_name = str(r2_op.outputs[0])
                if r2_out_name in model_outputs:
                    continue

                r2_users = [int(v) for v in consumers.get(r2_out_name, [])]
                if len(r2_users) != 1:
                    continue
                t2_idx = int(r2_users[0])
                t2_op = model_ir.operators[int(t2_idx)]
                if (
                    str(t2_op.op_type) != "TRANSPOSE"
                    or len(t2_op.inputs) < 2
                    or len(t2_op.outputs) != 1
                    or str(t2_op.inputs[0]) != r2_out_name
                    or _read_transpose_perm(model_ir, t2_op) != perm_nchw_to_nhwc
                ):
                    continue

                x_tensor = model_ir.tensors.get(x_nhwc_name, None)
                r1_tensor = model_ir.tensors.get(r1_out_name, None)
                t1_tensor = model_ir.tensors.get(t1_out_name, None)
                y_tensor = model_ir.tensors.get(str(t2_op.outputs[0]), None)
                if x_tensor is None or r1_tensor is None or t1_tensor is None:
                    continue

                x_shape = [int(v) for v in list(x_tensor.shape)]
                r1_shape = [int(v) for v in list(r1_tensor.shape)]
                t1_shape = [int(v) for v in list(t1_tensor.shape)]
                if (
                    not _is_fully_known_positive_shape(x_shape)
                    or not _is_fully_known_positive_shape(r1_shape)
                    or not _is_fully_known_positive_shape(t1_shape)
                ):
                    continue
                if len(x_shape) != 4 or len(r1_shape) != 5 or len(t1_shape) != 5:
                    continue

                n, h, w, c = [int(v) for v in x_shape]
                if int(c) <= 1:
                    continue

                # r1:[N,g,cpg,H,W], t1:[N,cpg,g,H,W]
                if (
                    int(r1_shape[0]) != int(n)
                    or int(r1_shape[3]) != int(h)
                    or int(r1_shape[4]) != int(w)
                    or int(t1_shape[0]) != int(n)
                    or int(t1_shape[3]) != int(h)
                    or int(t1_shape[4]) != int(w)
                    or int(t1_shape[1]) != int(r1_shape[2])
                    or int(t1_shape[2]) != int(r1_shape[1])
                ):
                    continue
                groups = int(r1_shape[1])
                cpg = int(r1_shape[2])
                if int(groups) <= 1 or int(cpg) <= 0 or int(groups * cpg) != int(c):
                    continue

                shuffle_indices = np.asarray(
                    [int((k % groups) * cpg + (k // groups)) for k in range(int(c))],
                    dtype=np.int32,
                )
                if np.array_equal(shuffle_indices, np.arange(int(c), dtype=np.int32)):
                    continue

                gather_idx_name = _unique_tensor_name(f"{x_nhwc_name}_shuffle_indices")
                model_ir.tensors[gather_idx_name] = TensorIR(
                    name=gather_idx_name,
                    dtype="INT32",
                    shape=[int(c)],
                    shape_signature=[int(c)],
                    data=shuffle_indices,
                    is_variable=False,
                )

                t2_op.op_type = "GATHER"
                t2_op.version = 1
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=t2_op,
                    new_inputs=[gather_src_name, gather_idx_name],
                )
                t2_op.options = {"axis": 3, "batchDims": 0}

                if unary_bridge_op is not None:
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=unary_bridge_op,
                        new_inputs=[x_nhwc_name],
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(unary_bridge_op.outputs[0]), None),
                        perm_nchw_to_nhwc,
                    )

                if y_tensor is not None:
                    y_tensor.shape = [int(v) for v in list(x_shape)]
                    y_signature = (
                        [int(v) for v in list(x_tensor.shape_signature)]
                        if x_tensor.shape_signature is not None
                        else [int(v) for v in list(x_shape)]
                    )
                    y_tensor.shape_signature = [int(v) for v in list(y_signature)]

                remove_indices = {int(r1_idx), int(t1_idx), int(r2_idx)}
                t0_remaining_users = [
                    int(v)
                    for v in t0_users
                    if int(v) != int(t0_user_idx)
                ]
                if len(t0_remaining_users) == 0:
                    remove_indices.add(int(t0_idx))

                for remove_idx in sorted(list(remove_indices), reverse=True):
                    del model_ir.operators[int(remove_idx)]

                optimized += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_shufflenet_reshape_transpose_shuffle_nhwc_chains": int(optimized)}


def _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse NCHW channel-shuffle blocks into a single GATHER(axis=1).

    Target:
      x_nchw
        -> RESHAPE([N,g,cpg,H,W])
        -> TRANSPOSE([0,2,1,3,4])
        -> RESHAPE([N,C,H,W]) -> y_nchw

    Rewrite:
      x_nchw -> GATHER(axis=1, shuffle_indices) -> y_nchw

    where C=g*cpg and shuffle_indices[k] = (k % g) * cpg + (k // g).
    """
    optimized = 0
    perm_shuffle_swap = [0, 2, 1, 3, 4]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for r1_idx, r1_op in enumerate(model_ir.operators):
            if str(r1_op.op_type) != "RESHAPE" or len(r1_op.inputs) < 1 or len(r1_op.outputs) != 1:
                continue
            x_nchw_name = str(r1_op.inputs[0])
            r1_out_name = str(r1_op.outputs[0])

            r1_users = [int(v) for v in consumers.get(r1_out_name, [])]
            if len(r1_users) != 1:
                continue
            t1_idx = int(r1_users[0])
            t1_op = model_ir.operators[int(t1_idx)]
            if (
                str(t1_op.op_type) != "TRANSPOSE"
                or len(t1_op.inputs) < 2
                or len(t1_op.outputs) != 1
                or str(t1_op.inputs[0]) != r1_out_name
                or _read_transpose_perm(model_ir, t1_op) != perm_shuffle_swap
            ):
                continue
            t1_out_name = str(t1_op.outputs[0])

            t1_users = [int(v) for v in consumers.get(t1_out_name, [])]
            if len(t1_users) != 1:
                continue
            r2_idx = int(t1_users[0])
            r2_op = model_ir.operators[int(r2_idx)]
            if (
                str(r2_op.op_type) != "RESHAPE"
                or len(r2_op.inputs) < 1
                or len(r2_op.outputs) != 1
                or str(r2_op.inputs[0]) != t1_out_name
            ):
                continue
            y_nchw_name = str(r2_op.outputs[0])

            x_tensor = model_ir.tensors.get(x_nchw_name, None)
            r1_tensor = model_ir.tensors.get(r1_out_name, None)
            t1_tensor = model_ir.tensors.get(t1_out_name, None)
            y_tensor = model_ir.tensors.get(y_nchw_name, None)
            if x_tensor is None or r1_tensor is None or t1_tensor is None or y_tensor is None:
                continue

            x_shape = [int(v) for v in list(x_tensor.shape)]
            r1_shape = [int(v) for v in list(r1_tensor.shape)]
            t1_shape = [int(v) for v in list(t1_tensor.shape)]
            y_shape = [int(v) for v in list(y_tensor.shape)]
            if (
                not _is_fully_known_positive_shape(x_shape)
                or not _is_fully_known_positive_shape(r1_shape)
                or not _is_fully_known_positive_shape(t1_shape)
                or not _is_fully_known_positive_shape(y_shape)
            ):
                continue
            if len(x_shape) != 4 or len(r1_shape) != 5 or len(t1_shape) != 5 or len(y_shape) != 4:
                continue

            n, c, h, w = [int(v) for v in x_shape]
            groups = int(r1_shape[1])
            cpg = int(r1_shape[2])
            if (
                int(groups) <= 1
                or int(cpg) <= 0
                or int(groups * cpg) != int(c)
                or int(r1_shape[0]) != int(n)
                or int(r1_shape[3]) != int(h)
                or int(r1_shape[4]) != int(w)
                or int(t1_shape[0]) != int(n)
                or int(t1_shape[1]) != int(cpg)
                or int(t1_shape[2]) != int(groups)
                or int(t1_shape[3]) != int(h)
                or int(t1_shape[4]) != int(w)
                or [int(v) for v in list(y_shape)] != [int(n), int(c), int(h), int(w)]
            ):
                continue

            shuffle_indices = np.asarray(
                [int((k % groups) * cpg + (k // groups)) for k in range(int(c))],
                dtype=np.int32,
            )
            if np.array_equal(shuffle_indices, np.arange(int(c), dtype=np.int32)):
                continue

            gather_idx_name = _unique_tensor_name(f"{x_nchw_name}_shuffle_indices_nchw")
            model_ir.tensors[gather_idx_name] = TensorIR(
                name=gather_idx_name,
                dtype="INT32",
                shape=[int(c)],
                shape_signature=[int(c)],
                data=np.asarray(shuffle_indices, dtype=np.int32),
                is_variable=False,
            )

            r2_op.op_type = "GATHER"
            r2_op.version = 1
            _set_operator_inputs(
                model_ir=model_ir,
                op=r2_op,
                new_inputs=[x_nchw_name, gather_idx_name],
            )
            r2_op.options = {"axis": 1, "batchDims": 0}

            for remove_idx in sorted([int(r1_idx), int(t1_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_nchw_channel_shuffle_reshape_transpose_reshape_to_gather": int(optimized)}


def _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC<->NCHW transpose wrappers around GATHER by remapping axis.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --GATHER(axis=a_nchw, batchDims=0)--> g_nchw
      g_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      x_nhwc --GATHER(axis=a_nhwc, batchDims=0)--> y_nhwc

    where a_nhwc = perm_nhwc_to_nchw[a_nchw].
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for gather_idx, gather_op in enumerate(model_ir.operators):
            if str(gather_op.op_type) != "GATHER" or len(gather_op.inputs) < 2 or len(gather_op.outputs) != 1:
                continue
            gather_in_name = str(gather_op.inputs[0])
            gather_out_name = str(gather_op.outputs[0])
            if gather_out_name in model_outputs:
                continue

            gather_options = dict(gather_op.options) if isinstance(gather_op.options, dict) else {}
            batch_dims = int(gather_options.get("batchDims", 0))
            if int(batch_dims) != 0:
                continue

            pre_idx = producers.get(gather_in_name, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != gather_in_name
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_out_name = str(pre_op.outputs[0])

            gather_users = [int(v) for v in consumers.get(gather_out_name, [])]
            if len(gather_users) != 1:
                continue
            post_idx = int(gather_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != gather_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue

            gather_in_tensor = model_ir.tensors.get(gather_in_name, None)
            if gather_in_tensor is None or len(list(gather_in_tensor.shape)) != 4:
                continue

            axis = int(gather_options.get("axis", 0))
            if axis < 0:
                axis += 4
            if int(axis) < 0 or int(axis) >= 4:
                continue
            remapped_axis = int(perm_nhwc_to_nchw[int(axis)])

            pre_input_name = str(pre_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if post_output_name not in model_ir.tensors:
                gather_out_tensor = model_ir.tensors.get(gather_out_name, None)
                if gather_out_tensor is not None:
                    model_ir.tensors[post_output_name] = TensorIR(
                        name=post_output_name,
                        dtype=str(gather_out_tensor.dtype),
                        shape=[int(v) for v in list(gather_out_tensor.shape)],
                        shape_signature=(
                            [int(v) for v in list(gather_out_tensor.shape_signature)]
                            if gather_out_tensor.shape_signature is not None
                            else [int(v) for v in list(gather_out_tensor.shape)]
                        ),
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(gather_out_tensor.quantization),
                    )

            gather_inputs = [str(v) for v in list(gather_op.inputs)]
            gather_inputs[0] = pre_input_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=gather_op,
                new_inputs=gather_inputs,
            )
            gather_options["axis"] = int(remapped_axis)
            gather_options["batchDims"] = 0
            gather_op.options = gather_options
            _set_operator_outputs(
                model_ir=model_ir,
                op=gather_op,
                new_outputs=[str(post_output_name)],
            )

            remove_indices: List[int] = [int(post_idx)]
            remaining_pre_users = [
                int(v)
                for v in consumers.get(pre_out_name, [])
                if int(v) != int(gather_idx)
            ]
            if len(remaining_pre_users) == 0 and pre_out_name not in model_outputs:
                remove_indices.append(int(pre_idx))
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_gather_transpose_axis_remap_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_add_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Convert NCHW Add blocks back to NHWC when wrapped by transpose adapters.

    Target:
      a_nhwc -> (optional swish-wrapped NCHW adapters) -> a_nchw
      b_nhwc -> (optional swish-wrapped NCHW adapters) -> b_nchw
      ADD(a_nchw, b_nchw) -> y_nchw
      (optional unary RELU/RELU6/LOGISTIC/TANH) -> y2_nchw
      y_or_y2_nchw -> TRANSPOSE(0,2,3,1) -> y_nhwc

    Rewrite:
      ADD(a_nhwc, b_nhwc) -> y_nhwc
      (if unary exists, keep unary and bridge its output instead of ADD output directly)

    Notes:
    - If y_nchw has legacy consumers, keep one adapter TRANSPOSE(0,3,1,2): y_nhwc -> y_nchw.
    - Swish-wrapped NCHW inputs are recognized as:
        pre_nhwc --T(0,3,1,2)--> x_nchw --LOGISTIC--> s --MUL(x_nchw,s)--> x_swish_nchw
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_passthrough_ops = {"RELU", "RELU6", "LOGISTIC", "TANH", "GELU", "HARD_SWISH", "LEAKY_RELU"}
    skip_add_activation_fuse_marker = "__skip_add_activation_fuse__"

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_prelu_alpha_to_nhwc(
        *,
        prelu_idx: int,
        prelu_op: OperatorIR,
        target_nhwc_shape: Optional[List[int]],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        alpha_name = str(prelu_op.inputs[1]) if len(prelu_op.inputs) >= 2 else ""
        alpha_tensor = model_ir.tensors.get(alpha_name, None)
        if alpha_tensor is None or alpha_tensor.data is None:
            return None

        alpha_data = np.asarray(alpha_tensor.data)
        alpha_candidates: List[np.ndarray] = []
        if int(alpha_data.ndim) == int(len(perm_nchw_to_nhwc)):
            # Prefer layout-remapped alpha first when both can broadcast.
            alpha_candidates.append(
                np.transpose(alpha_data, axes=perm_nchw_to_nhwc).astype(alpha_data.dtype, copy=False)
            )
        alpha_candidates.append(np.asarray(alpha_data))
        if int(alpha_data.ndim) == 3:
            alpha_candidates.append(
                np.transpose(alpha_data, axes=[1, 2, 0]).astype(alpha_data.dtype, copy=False)
            )

        def _broadcastable_to_target(candidate: np.ndarray) -> bool:
            if target_nhwc_shape is None or not _is_fully_known_positive_shape(target_nhwc_shape):
                return True
            return (
                _broadcast_static_shapes(
                    [int(v) for v in list(target_nhwc_shape)],
                    [int(v) for v in list(candidate.shape)],
                )
                is not None
            )

        selected_alpha: Optional[np.ndarray] = None
        for candidate in alpha_candidates:
            if _broadcastable_to_target(candidate):
                selected_alpha = np.asarray(candidate)
                break
        if selected_alpha is None:
            return None

        selected_alpha_name = str(alpha_name)
        alpha_users = [int(v) for v in consumers.get(alpha_name, [])]
        alpha_needs_rewrite = (
            selected_alpha.shape != alpha_data.shape
            or not np.array_equal(selected_alpha, alpha_data)
        )
        if alpha_needs_rewrite:
            shape, signature = normalize_onnx_shape(list(selected_alpha.shape))
            if len(alpha_users) == 1 and int(alpha_users[0]) == int(prelu_idx):
                alpha_tensor.data = np.asarray(selected_alpha)
                alpha_tensor.shape = [int(v) for v in list(shape)]
                alpha_tensor.shape_signature = [int(v) for v in list(signature)]
            else:
                selected_alpha_name = _unique_tensor_name(f"{alpha_name}_nhwc")
                model_ir.tensors[selected_alpha_name] = TensorIR(
                    name=selected_alpha_name,
                    dtype=str(alpha_tensor.dtype),
                    shape=[int(v) for v in list(shape)],
                    shape_signature=[int(v) for v in list(signature)],
                    data=np.asarray(selected_alpha),
                    is_variable=False,
                    quantization=_clone_quantization(alpha_tensor.quantization),
                )
        return str(selected_alpha_name)

    def _analyze_swish_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        mul_idx = producers.get(str(input_name), None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        if str(mul_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(consumer_idx)}:
            return None

        in0_name = str(mul_op.inputs[0])
        in1_name = str(mul_op.inputs[1])
        in0_prod_idx = producers.get(in0_name, None)
        in1_prod_idx = producers.get(in1_name, None)

        log_idx: Optional[int] = None
        log_out_name: Optional[str] = None
        pre_nchw_name: Optional[str] = None
        mul_data_input_index: Optional[int] = None

        if in0_prod_idx is not None:
            in0_prod_op = model_ir.operators[int(in0_prod_idx)]
            if (
                str(in0_prod_op.op_type) == "LOGISTIC"
                and len(in0_prod_op.inputs) == 1
                and len(in0_prod_op.outputs) == 1
                and str(in0_prod_op.outputs[0]) == in0_name
            ):
                log_idx = int(in0_prod_idx)
                log_out_name = in0_name
                pre_nchw_name = str(in0_prod_op.inputs[0])
                mul_data_input_index = 1

        if log_idx is None and in1_prod_idx is not None:
            in1_prod_op = model_ir.operators[int(in1_prod_idx)]
            if (
                str(in1_prod_op.op_type) == "LOGISTIC"
                and len(in1_prod_op.inputs) == 1
                and len(in1_prod_op.outputs) == 1
                and str(in1_prod_op.outputs[0]) == in1_name
            ):
                log_idx = int(in1_prod_idx)
                log_out_name = in1_name
                pre_nchw_name = str(in1_prod_op.inputs[0])
                mul_data_input_index = 0

        if log_idx is None or log_out_name is None or pre_nchw_name is None or mul_data_input_index is None:
            return None
        if str(mul_op.inputs[int(mul_data_input_index)]) != str(pre_nchw_name):
            return None

        pre_idx = producers.get(str(pre_nchw_name), None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != str(pre_nchw_name):
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None

        if set(int(v) for v in consumers.get(str(pre_nchw_name), [])) != {int(log_idx), int(mul_idx)}:
            return None
        if set(int(v) for v in consumers.get(str(log_out_name), [])) != {int(mul_idx)}:
            return None

        return {
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "log_idx": int(log_idx),
            "mul_idx": int(mul_idx),
            "mul_data_input_index": int(mul_data_input_index),
            "log_out_name": str(log_out_name),
            "mul_out_name": str(input_name),
        }

    def _apply_swish_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        log_idx = int(plan["log_idx"])
        mul_idx = int(plan["mul_idx"])
        mul_data_input_index = int(plan["mul_data_input_index"])
        log_out_name = str(plan["log_out_name"])
        mul_out_name = str(plan["mul_out_name"])

        mul_op = model_ir.operators[int(mul_idx)]
        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(log_idx)],
            new_inputs=[pre_input_name],
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=mul_op,
            input_index=int(mul_data_input_index),
            new_input_name=pre_input_name,
        )

        for tensor_name in [log_out_name, mul_out_name]:
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(tensor_name, None),
                perm_nchw_to_nhwc,
            )
        return [int(pre_idx)]

    def _analyze_unary_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        unary_idx = producers.get(str(input_name), None)
        if unary_idx is None:
            return None
        unary_op = model_ir.operators[int(unary_idx)]
        if str(unary_op.op_type) not in unary_passthrough_ops:
            return None
        if len(unary_op.inputs) != 1 or len(unary_op.outputs) != 1:
            return None
        if str(unary_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(consumer_idx)}:
            return None

        pre_nchw_name = str(unary_op.inputs[0])
        pre_idx = producers.get(pre_nchw_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
            return None
        if str(pre_op.outputs[0]) != pre_nchw_name:
            return None
        if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
            return None
        if str(pre_nchw_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(pre_nchw_name), [])) != {int(unary_idx)}:
            return None

        return {
            "input_name": str(input_name),
            "unary_idx": int(unary_idx),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
        }

    def _apply_unary_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        input_name = str(plan["input_name"])
        unary_idx = int(plan["unary_idx"])
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])

        _set_operator_inputs(
            model_ir=model_ir,
            op=model_ir.operators[int(unary_idx)],
            new_inputs=[pre_input_name],
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(input_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(pre_idx)]

    def _analyze_mul_const_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        mul_idx = producers.get(str(input_name), None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        if str(mul_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(consumer_idx)}:
            return None

        data_input_index: Optional[int] = None
        side_input_index: Optional[int] = None
        side_input_name: Optional[str] = None
        pre_idx: Optional[int] = None
        pre_input_name: Optional[str] = None
        pre_output_name: Optional[str] = None

        for candidate_data_index, candidate_side_index in [(0, 1), (1, 0)]:
            data_name = str(mul_op.inputs[int(candidate_data_index)])
            side_name = str(mul_op.inputs[int(candidate_side_index)])
            side_tensor = model_ir.tensors.get(side_name, None)
            if side_tensor is None or side_tensor.data is None:
                continue

            data_prod_idx = producers.get(data_name, None)
            if data_prod_idx is None:
                continue
            data_prod_op = model_ir.operators[int(data_prod_idx)]
            if (
                str(data_prod_op.op_type) != "TRANSPOSE"
                or len(data_prod_op.inputs) < 2
                or len(data_prod_op.outputs) != 1
                or str(data_prod_op.outputs[0]) != data_name
                or _read_transpose_perm(model_ir, data_prod_op) != perm_nhwc_to_nchw
                or str(data_name) in model_outputs
            ):
                continue
            data_users = set(int(v) for v in consumers.get(data_name, []))
            if int(mul_idx) not in data_users:
                continue

            data_tensor = model_ir.tensors.get(data_name, None)
            pre_in_name = str(data_prod_op.inputs[0])
            pre_in_tensor = model_ir.tensors.get(pre_in_name, None)
            if data_tensor is None or pre_in_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(list(data_tensor.shape))
                or not _is_fully_known_positive_shape(list(pre_in_tensor.shape))
                or len(list(data_tensor.shape)) != 4
                or len(list(pre_in_tensor.shape)) != 4
            ):
                continue

            side_data = np.asarray(side_tensor.data)
            if int(side_data.size) != 1 and side_data.ndim not in {3, 4}:
                continue
            if int(side_data.size) != 1:
                if _broadcast_static_shapes(list(data_tensor.shape), [int(v) for v in list(side_data.shape)]) is None:
                    continue

            data_input_index = int(candidate_data_index)
            side_input_index = int(candidate_side_index)
            side_input_name = str(side_name)
            pre_idx = int(data_prod_idx)
            pre_input_name = pre_in_name
            pre_output_name = data_name
            break

        if (
            data_input_index is None
            or side_input_index is None
            or side_input_name is None
            or pre_idx is None
            or pre_input_name is None
            or pre_output_name is None
        ):
            return None

        side_tensor = model_ir.tensors.get(side_input_name, None)
        if side_tensor is None or side_tensor.data is None:
            return None
        side_data = np.asarray(side_tensor.data)
        nhwc_side_data: Optional[np.ndarray] = None
        side_needs_update = False
        if int(side_data.size) != 1:
            target_nhwc_shape = [int(v) for v in list(model_ir.tensors[str(pre_input_name)].shape)]
            side_shape = [int(v) for v in list(side_data.shape)]
            # Prefer semantic channel remap for NCHW channel-wise constants even
            # when static shapes are broadcast-ambiguous (e.g. [1,64,64,64]).
            is_nchw_channelwise = (
                int(side_data.ndim) == 4
                and len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            )
            if is_nchw_channelwise:
                rotated = np.transpose(side_data, axes=perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_nhwc_shape, rotated_shape) is None:
                    return None
                nhwc_side_data = np.asarray(rotated)
                side_needs_update = not np.array_equal(nhwc_side_data, side_data)
            elif _broadcast_static_shapes(target_nhwc_shape, side_shape) is not None:
                nhwc_side_data = np.asarray(side_data)
                side_needs_update = False
            else:
                rotated = np.asarray(side_data)
                found = False
                transpose_perm = (
                    perm_nchw_to_nhwc
                    if int(rotated.ndim) == 4
                    else [1, 2, 0]
                )
                max_rotate = 3 if int(rotated.ndim) == 3 else 1
                for _ in range(int(max_rotate)):
                    rotated = np.transpose(rotated, transpose_perm).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if _broadcast_static_shapes(target_nhwc_shape, rotated_shape) is not None:
                        nhwc_side_data = np.asarray(rotated)
                        side_needs_update = True
                        found = True
                        break
                if not found:
                    return None

        shared_outside_mul = any(
            int(v) != int(mul_idx)
            for v in consumers.get(str(side_input_name), [])
        )

        return {
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_input_name),
            "mul_idx": int(mul_idx),
            "mul_data_input_index": int(data_input_index),
            "mul_side_input_index": int(side_input_index),
            "mul_side_input_name": str(side_input_name),
            "mul_out_name": str(input_name),
            "pre_removable": bool(
                set(int(v) for v in consumers.get(str(pre_output_name), [])) == {int(mul_idx)}
            ),
            "side_needs_update": bool(side_needs_update),
            "side_shared_outside_mul": bool(shared_outside_mul),
            "side_nhwc_data": (None if nhwc_side_data is None else np.asarray(nhwc_side_data)),
        }

    def _apply_mul_const_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        pre_idx = int(plan["pre_idx"])
        pre_input_name = str(plan["pre_input_name"])
        mul_idx = int(plan["mul_idx"])
        mul_data_input_index = int(plan["mul_data_input_index"])
        mul_side_input_index = int(plan["mul_side_input_index"])
        mul_side_input_name = str(plan["mul_side_input_name"])
        mul_out_name = str(plan["mul_out_name"])

        side_input_name_for_mul = str(mul_side_input_name)
        if bool(plan.get("side_needs_update", False)):
            nhwc_data = np.asarray(plan.get("side_nhwc_data"))
            side_tensor = model_ir.tensors.get(str(mul_side_input_name), None)
            if side_tensor is None:
                return []
            if bool(plan.get("side_shared_outside_mul", False)):
                side_input_name_for_mul = _unique_tensor_name(f"{mul_side_input_name}_nhwc")
                model_ir.tensors[side_input_name_for_mul] = TensorIR(
                    name=side_input_name_for_mul,
                    dtype=str(side_tensor.dtype),
                    shape=[int(v) for v in list(nhwc_data.shape)],
                    shape_signature=[int(v) for v in list(nhwc_data.shape)],
                    data=np.asarray(nhwc_data),
                    is_variable=False,
                    quantization=_clone_quantization(side_tensor.quantization),
                )
            else:
                side_tensor.data = np.asarray(nhwc_data)
                side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
                side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        mul_op = model_ir.operators[int(mul_idx)]
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        if int(mul_data_input_index) < 0 or int(mul_data_input_index) >= len(mul_inputs):
            return []
        if int(mul_side_input_index) < 0 or int(mul_side_input_index) >= len(mul_inputs):
            return []
        mul_inputs[int(mul_data_input_index)] = str(pre_input_name)
        mul_inputs[int(mul_side_input_index)] = str(side_input_name_for_mul)
        _set_operator_inputs(
            model_ir=model_ir,
            op=mul_op,
            new_inputs=mul_inputs,
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(mul_out_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(pre_idx)] if bool(plan.get("pre_removable", False)) else []

    def _analyze_gather_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        gather_idx = producers.get(str(input_name), None)
        if gather_idx is None:
            return None
        gather_op = model_ir.operators[int(gather_idx)]
        if str(gather_op.op_type) != "GATHER" or len(gather_op.inputs) < 2 or len(gather_op.outputs) != 1:
            return None
        if str(gather_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(consumer_idx)}:
            return None

        gather_options = dict(gather_op.options) if isinstance(gather_op.options, dict) else {}
        if int(gather_options.get("batchDims", 0)) != 0:
            return None

        gather_data_name = str(gather_op.inputs[0])
        pre_idx = producers.get(gather_data_name, None)
        if pre_idx is None:
            return None
        pre_op = model_ir.operators[int(pre_idx)]
        if (
            str(pre_op.op_type) != "TRANSPOSE"
            or len(pre_op.inputs) < 2
            or len(pre_op.outputs) != 1
            or str(pre_op.outputs[0]) != gather_data_name
            or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            or str(gather_data_name) in model_outputs
        ):
            return None

        gather_data_tensor = model_ir.tensors.get(gather_data_name, None)
        if gather_data_tensor is None or len(list(gather_data_tensor.shape)) != 4:
            return None
        axis = int(gather_options.get("axis", 0))
        if axis < 0:
            axis += 4
        if axis < 0 or axis >= 4:
            return None
        remapped_axis = int(perm_nhwc_to_nchw[axis])

        pre_users = set(int(v) for v in consumers.get(gather_data_name, []))
        return {
            "gather_idx": int(gather_idx),
            "gather_in_name": str(gather_data_name),
            "gather_out_name": str(input_name),
            "pre_idx": int(pre_idx),
            "pre_input_name": str(pre_op.inputs[0]),
            "remapped_axis": int(remapped_axis),
            "pre_removable": bool(pre_users == {int(gather_idx)}),
        }

    def _apply_gather_input_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        gather_idx = int(plan["gather_idx"])
        gather_op = model_ir.operators[int(gather_idx)]
        gather_inputs = [str(v) for v in list(gather_op.inputs)]
        if len(gather_inputs) < 2:
            return []
        gather_inputs[0] = str(plan["pre_input_name"])
        _set_operator_inputs(
            model_ir=model_ir,
            op=gather_op,
            new_inputs=gather_inputs,
        )
        gather_options = dict(gather_op.options) if isinstance(gather_op.options, dict) else {}
        gather_options["axis"] = int(plan["remapped_axis"])
        gather_options["batchDims"] = 0
        gather_op.options = gather_options
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(str(plan["gather_out_name"]), None),
            perm_nchw_to_nhwc,
        )
        return [int(plan["pre_idx"])] if bool(plan.get("pre_removable", False)) else []

    def _analyze_mul_sub_const_input_to_nhwc(
        *,
        input_name: str,
        consumer_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        mul_idx = producers.get(str(input_name), None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        if str(mul_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        if set(int(v) for v in consumers.get(str(input_name), [])) != {int(consumer_idx)}:
            return None

        mul_data_input_index: Optional[int] = None
        mul_side_input_index: Optional[int] = None
        mul_side_input_name: Optional[str] = None
        sub_idx: Optional[int] = None
        sub_out_name: Optional[str] = None
        sub_input_rewrites: List[Dict[str, Any]] = []
        removable_pre_indices: List[int] = []
        target_nhwc_shape: Optional[List[int]] = None

        for candidate_data_index, candidate_side_index in [(0, 1), (1, 0)]:
            data_name = str(mul_op.inputs[int(candidate_data_index)])
            side_name = str(mul_op.inputs[int(candidate_side_index)])
            side_tensor = model_ir.tensors.get(side_name, None)
            if side_tensor is None or side_tensor.data is None:
                continue

            candidate_sub_idx = producers.get(data_name, None)
            if candidate_sub_idx is None:
                continue
            candidate_sub_op = model_ir.operators[int(candidate_sub_idx)]
            if (
                str(candidate_sub_op.op_type) != "SUB"
                or len(candidate_sub_op.inputs) != 2
                or len(candidate_sub_op.outputs) != 1
                or str(candidate_sub_op.outputs[0]) != data_name
                or str(data_name) in model_outputs
            ):
                continue
            if set(int(v) for v in consumers.get(data_name, [])) != {int(mul_idx)}:
                continue

            candidate_rewrites: List[Dict[str, Any]] = []
            candidate_removable: List[int] = []
            candidate_target_shape: Optional[List[int]] = None
            valid_sub_inputs = True
            for sub_input_index, sub_input_name in enumerate([str(v) for v in list(candidate_sub_op.inputs)]):
                sub_pre_idx = producers.get(str(sub_input_name), None)
                if sub_pre_idx is None:
                    valid_sub_inputs = False
                    break
                sub_pre_op = model_ir.operators[int(sub_pre_idx)]
                if (
                    str(sub_pre_op.op_type) != "TRANSPOSE"
                    or len(sub_pre_op.inputs) < 2
                    or len(sub_pre_op.outputs) != 1
                    or str(sub_pre_op.outputs[0]) != str(sub_input_name)
                    or _read_transpose_perm(model_ir, sub_pre_op) != perm_nhwc_to_nchw
                    or str(sub_input_name) in model_outputs
                ):
                    valid_sub_inputs = False
                    break
                sub_input_users = set(int(v) for v in consumers.get(str(sub_input_name), []))
                if not sub_input_users.issubset({int(candidate_sub_idx), int(consumer_idx)}):
                    valid_sub_inputs = False
                    break
                nhwc_name = str(sub_pre_op.inputs[0])
                candidate_rewrites.append(
                    {
                        "sub_input_index": int(sub_input_index),
                        "nhwc_name": nhwc_name,
                        "pre_idx": int(sub_pre_idx),
                    }
                )
                candidate_removable.append(int(sub_pre_idx))
                nhwc_tensor = model_ir.tensors.get(nhwc_name, None)
                if (
                    candidate_target_shape is None
                    and nhwc_tensor is not None
                    and _is_fully_known_positive_shape(list(nhwc_tensor.shape))
                ):
                    candidate_target_shape = [int(v) for v in list(nhwc_tensor.shape)]

            if not valid_sub_inputs or len(candidate_rewrites) != 2:
                continue

            mul_data_input_index = int(candidate_data_index)
            mul_side_input_index = int(candidate_side_index)
            mul_side_input_name = str(side_name)
            sub_idx = int(candidate_sub_idx)
            sub_out_name = str(data_name)
            sub_input_rewrites = [dict(v) for v in candidate_rewrites]
            removable_pre_indices = [int(v) for v in candidate_removable]
            target_nhwc_shape = (
                [int(v) for v in list(candidate_target_shape)]
                if candidate_target_shape is not None
                else None
            )
            break

        if (
            mul_data_input_index is None
            or mul_side_input_index is None
            or mul_side_input_name is None
            or sub_idx is None
            or sub_out_name is None
            or len(sub_input_rewrites) != 2
        ):
            return None

        side_tensor = model_ir.tensors.get(str(mul_side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return None

        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) != 1 and side_data.ndim not in {1, 3, 4}:
            return None

        nhwc_side_data: Optional[np.ndarray] = None
        side_needs_update = False
        if int(side_data.size) != 1 and side_data.ndim in {3, 4}:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                target_nhwc_shape is not None
                and _broadcast_static_shapes(target_nhwc_shape, side_shape) is not None
            ):
                nhwc_side_data = np.asarray(side_data)
            else:
                rotated = np.asarray(side_data)
                transpose_perm = perm_nchw_to_nhwc if int(rotated.ndim) == 4 else [1, 2, 0]
                max_rotate = 1 if int(rotated.ndim) == 4 else 3
                found = False
                for _ in range(int(max_rotate)):
                    rotated = np.transpose(rotated, transpose_perm).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if (
                        target_nhwc_shape is None
                        or _broadcast_static_shapes(target_nhwc_shape, rotated_shape) is not None
                    ):
                        nhwc_side_data = np.asarray(rotated)
                        side_needs_update = True
                        found = True
                        break
                if not found:
                    return None

        side_shared_outside_mul = any(
            int(v) != int(mul_idx)
            for v in consumers.get(str(mul_side_input_name), [])
        )

        return {
            "sub_idx": int(sub_idx),
            "sub_input_rewrites": [dict(v) for v in sub_input_rewrites],
            "sub_out_name": str(sub_out_name),
            "mul_idx": int(mul_idx),
            "mul_data_input_index": int(mul_data_input_index),
            "mul_side_input_index": int(mul_side_input_index),
            "mul_side_input_name": str(mul_side_input_name),
            "mul_out_name": str(input_name),
            "removable_pre_indices": sorted(list({int(v) for v in removable_pre_indices})),
            "side_needs_update": bool(side_needs_update),
            "side_shared_outside_mul": bool(side_shared_outside_mul),
            "side_nhwc_data": (None if nhwc_side_data is None else np.asarray(nhwc_side_data)),
        }

    def _apply_mul_sub_const_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
        sub_idx = int(plan["sub_idx"])
        sub_input_rewrites = [dict(v) for v in list(plan["sub_input_rewrites"])]
        sub_out_name = str(plan["sub_out_name"])
        mul_idx = int(plan["mul_idx"])
        mul_data_input_index = int(plan["mul_data_input_index"])
        mul_side_input_index = int(plan["mul_side_input_index"])
        mul_side_input_name = str(plan["mul_side_input_name"])
        mul_out_name = str(plan["mul_out_name"])
        removable_pre_indices = [int(v) for v in list(plan.get("removable_pre_indices", []))]

        side_input_name_for_mul = str(mul_side_input_name)
        if bool(plan.get("side_needs_update", False)):
            nhwc_data = np.asarray(plan.get("side_nhwc_data"))
            side_tensor = model_ir.tensors.get(str(mul_side_input_name), None)
            if side_tensor is None:
                return []
            if bool(plan.get("side_shared_outside_mul", False)):
                side_input_name_for_mul = _unique_tensor_name(f"{mul_side_input_name}_nhwc")
                model_ir.tensors[side_input_name_for_mul] = TensorIR(
                    name=side_input_name_for_mul,
                    dtype=str(side_tensor.dtype),
                    shape=[int(v) for v in list(nhwc_data.shape)],
                    shape_signature=[int(v) for v in list(nhwc_data.shape)],
                    data=np.asarray(nhwc_data),
                    is_variable=False,
                    quantization=_clone_quantization(side_tensor.quantization),
                )
            else:
                side_tensor.data = np.asarray(nhwc_data)
                side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
                side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        sub_op = model_ir.operators[int(sub_idx)]
        sub_inputs = [str(v) for v in list(sub_op.inputs)]
        for rewrite in sub_input_rewrites:
            input_index = int(rewrite["sub_input_index"])
            if int(input_index) < 0 or int(input_index) >= len(sub_inputs):
                return []
            sub_inputs[int(input_index)] = str(rewrite["nhwc_name"])
        _set_operator_inputs(
            model_ir=model_ir,
            op=sub_op,
            new_inputs=sub_inputs,
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(sub_out_name, None),
            perm_nchw_to_nhwc,
        )

        mul_op = model_ir.operators[int(mul_idx)]
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        if int(mul_data_input_index) < 0 or int(mul_data_input_index) >= len(mul_inputs):
            return []
        if int(mul_side_input_index) < 0 or int(mul_side_input_index) >= len(mul_inputs):
            return []
        mul_inputs[int(mul_data_input_index)] = str(sub_out_name)
        mul_inputs[int(mul_side_input_index)] = str(side_input_name_for_mul)
        _set_operator_inputs(
            model_ir=model_ir,
            op=mul_op,
            new_inputs=mul_inputs,
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(mul_out_name, None),
            perm_nchw_to_nhwc,
        )

        return [int(v) for v in list(removable_pre_indices)]

    def _analyze_const_add_input_to_nhwc(
        *,
        input_name: str,
        add_idx: int,
        consumers: Dict[str, List[int]],
    ) -> Optional[Dict[str, Any]]:
        const_tensor = model_ir.tensors.get(str(input_name), None)
        if const_tensor is None or const_tensor.data is None:
            return None

        add_op = model_ir.operators[int(add_idx)]
        if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
            return None

        const_data = np.asarray(const_tensor.data)
        const_shape = [int(v) for v in list(const_data.shape)]
        target_nhwc_shape: Optional[List[int]] = None
        add_out_tensor = model_ir.tensors.get(str(add_op.outputs[0]), None)
        if (
            add_out_tensor is not None
            and _is_fully_known_positive_shape(list(add_out_tensor.shape))
            and len(list(add_out_tensor.shape)) == 4
        ):
            target_nhwc_shape = _permute_shape(list(add_out_tensor.shape), perm_nchw_to_nhwc)

        nhwc_data: Optional[np.ndarray] = None
        const_needs_update = False
        if int(const_data.size) == 1:
            pass
        elif const_data.ndim == 4:
            is_nchw_channelwise = (
                len(const_shape) == 4
                and int(const_shape[0]) == 1
                and int(const_shape[1]) > 0
                and int(const_shape[2]) == 1
                and int(const_shape[3]) == 1
            )
            if is_nchw_channelwise:
                rotated = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if target_nhwc_shape is not None and _broadcast_static_shapes(target_nhwc_shape, rotated_shape) is None:
                    return None
                nhwc_data = np.asarray(rotated)
                const_needs_update = not np.array_equal(nhwc_data, const_data)
            elif target_nhwc_shape is not None and _broadcast_static_shapes(target_nhwc_shape, const_shape) is not None:
                pass
            else:
                rotated: Optional[np.ndarray] = None
                candidate = np.asarray(const_data)
                for _ in range(3):
                    candidate = np.transpose(candidate, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
                    candidate_shape = [int(v) for v in list(candidate.shape)]
                    if target_nhwc_shape is None or _broadcast_static_shapes(target_nhwc_shape, candidate_shape) is not None:
                        rotated = np.asarray(candidate)
                        break
                if rotated is None:
                    return None
                nhwc_data = np.asarray(rotated)
                const_needs_update = True
        elif const_data.ndim == 3:
            if target_nhwc_shape is not None and _broadcast_static_shapes(target_nhwc_shape, const_shape) is not None:
                pass
            else:
                candidate = np.asarray(const_data)
                rotated: Optional[np.ndarray] = None
                for _ in range(3):
                    candidate = np.transpose(candidate, [1, 2, 0]).astype(const_data.dtype, copy=False)
                    candidate_shape = [int(v) for v in list(candidate.shape)]
                    if target_nhwc_shape is None or _broadcast_static_shapes(target_nhwc_shape, candidate_shape) is not None:
                        rotated = np.asarray(candidate)
                        break
                if rotated is None:
                    return None
                nhwc_data = np.asarray(rotated)
                const_needs_update = True
        else:
            if target_nhwc_shape is not None and _broadcast_static_shapes(target_nhwc_shape, const_shape) is None:
                return None

        shared_outside_add = any(
            int(v) != int(add_idx)
            for v in consumers.get(str(input_name), [])
        )
        return {
            "const_input_name": str(input_name),
            "const_needs_update": bool(const_needs_update),
            "const_shared_outside_add": bool(shared_outside_add),
            "const_nhwc_data": (None if nhwc_data is None else np.asarray(nhwc_data)),
        }

    def _apply_const_add_input_nhwc_plan(*, plan: Dict[str, Any]) -> str:
        const_input_name = str(plan["const_input_name"])
        const_input_name_for_add = str(const_input_name)
        if not bool(plan.get("const_needs_update", False)):
            return const_input_name_for_add

        nhwc_data = np.asarray(plan.get("const_nhwc_data"))
        const_tensor = model_ir.tensors.get(const_input_name, None)
        if const_tensor is None:
            return const_input_name_for_add
        if bool(plan.get("const_shared_outside_add", False)):
            const_input_name_for_add = _unique_tensor_name(f"{const_input_name}_nhwc")
            model_ir.tensors[const_input_name_for_add] = TensorIR(
                name=const_input_name_for_add,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
        else:
            const_tensor.data = np.asarray(nhwc_data)
            const_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            const_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
        return str(const_input_name_for_add)

    def _build_direct_nchw_fallback_plan(
        *,
        input_name: str,
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        if str(input_name) in model_outputs:
            return None
        input_tensor = model_ir.tensors.get(str(input_name), None)
        if input_tensor is None:
            return None
        # Restrict fallback to dynamic rank-4 tensors. Constants are handled by
        # const-specific remap logic; scalar/rank!=4 paths are left unchanged.
        if input_tensor.data is not None:
            return None
        input_shape = [int(v) for v in list(input_tensor.shape)]
        if len(input_shape) != 4:
            return None

        bridge_name = _unique_tensor_name(f"{input_name}_nhwc_bridge")
        return {
            "nhwc_input_name": str(bridge_name),
            "pre_remove_indices": [],
            "direct_nchw_plan": {
                "input_name": str(input_name),
                "bridge_name": str(bridge_name),
            },
        }

    def _materialize_direct_nchw_bridge(
        *,
        plan: Dict[str, Any],
    ) -> Optional[Dict[str, str]]:
        input_name = str(plan.get("input_name", ""))
        bridge_name = str(plan.get("bridge_name", ""))
        if input_name == "" or bridge_name == "":
            return None
        input_tensor = model_ir.tensors.get(input_name, None)
        if input_tensor is None:
            return None
        if len(list(input_tensor.shape)) != 4:
            return None

        perm_arr = np.asarray([0, 2, 3, 1], dtype=np.int32)
        # Keep per-bridge dedicated perm tensor to avoid cross-rewrite mutation.
        perm_const_name = _unique_tensor_name(f"{bridge_name}_perm")
        model_ir.tensors[perm_const_name] = TensorIR(
            name=perm_const_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=perm_arr,
            is_variable=False,
            quantization=None,
        )

        bridge_shape = _permute_shape(
            [int(v) for v in list(input_tensor.shape)],
            [0, 2, 3, 1],
        )
        bridge_signature_src = (
            [int(v) for v in list(input_tensor.shape_signature)]
            if input_tensor.shape_signature is not None
            else [int(v) for v in list(input_tensor.shape)]
        )
        bridge_signature = _permute_shape(bridge_signature_src, [0, 2, 3, 1])
        if bridge_shape is None or bridge_signature is None:
            return None

        if bridge_name not in model_ir.tensors:
            model_ir.tensors[bridge_name] = TensorIR(
                name=bridge_name,
                dtype=str(input_tensor.dtype),
                shape=[int(v) for v in list(bridge_shape)],
                shape_signature=[int(v) for v in list(bridge_signature)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(input_tensor.quantization),
            )

        return {
            "input_name": str(input_name),
            "bridge_name": str(bridge_name),
            "perm_name": str(perm_const_name),
        }

    def _resolve_add_input_to_nhwc(
        *,
        input_name: str,
        add_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
        allow_nested_add: bool = True,
    ) -> Optional[Dict[str, Any]]:
        input_producer_idx = producers.get(str(input_name), None)
        if input_producer_idx is not None:
            input_producer_op = model_ir.operators[int(input_producer_idx)]
            input_users = set(int(v) for v in consumers.get(str(input_name), []))

            if (
                str(input_producer_op.op_type) == "TRANSPOSE"
                and len(input_producer_op.inputs) >= 2
                and len(input_producer_op.outputs) == 1
                and str(input_producer_op.outputs[0]) == str(input_name)
                and _read_transpose_perm(model_ir, input_producer_op) == perm_nhwc_to_nchw
                and int(add_idx) in input_users
                and str(input_name) not in model_outputs
            ):
                extra_user_indices = [int(v) for v in sorted(input_users) if int(v) != int(add_idx)]
                return {
                    "nhwc_input_name": str(input_producer_op.inputs[0]),
                    "pre_remove_indices": (
                        [int(input_producer_idx)] if len(extra_user_indices) == 0 else []
                    ),
                    "swish_plan": None,
                    "unary_plan": None,
                    "mul_const_plan": None,
                    "nested_add_plan": None,
                }

        swish_plan = _analyze_swish_input_to_nhwc(
            input_name=input_name,
            consumer_idx=int(add_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )
        if swish_plan is not None:
            return {
                "nhwc_input_name": str(input_name),
                "pre_remove_indices": [],
                "swish_plan": swish_plan,
                "unary_plan": None,
                "mul_const_plan": None,
                "mul_sub_const_plan": None,
                "gather_plan": None,
                "nested_add_plan": None,
            }

        unary_plan = _analyze_unary_input_to_nhwc(
            input_name=input_name,
            consumer_idx=int(add_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )
        if unary_plan is not None:
            return {
                "nhwc_input_name": str(input_name),
                "pre_remove_indices": [],
                "swish_plan": None,
                "unary_plan": unary_plan,
                "mul_const_plan": None,
                "mul_sub_const_plan": None,
                "gather_plan": None,
                "nested_add_plan": None,
            }

        mul_const_plan = _analyze_mul_const_input_to_nhwc(
            input_name=input_name,
            consumer_idx=int(add_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )
        if mul_const_plan is not None:
            return {
                "nhwc_input_name": str(input_name),
                "pre_remove_indices": [],
                "swish_plan": None,
                "unary_plan": None,
                "mul_const_plan": mul_const_plan,
                "mul_sub_const_plan": None,
                "gather_plan": None,
                "nested_add_plan": None,
            }

        mul_sub_const_plan = _analyze_mul_sub_const_input_to_nhwc(
            input_name=input_name,
            consumer_idx=int(add_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )
        if mul_sub_const_plan is not None:
            return {
                "nhwc_input_name": str(input_name),
                "pre_remove_indices": [],
                "swish_plan": None,
                "unary_plan": None,
                "mul_const_plan": None,
                "mul_sub_const_plan": mul_sub_const_plan,
                "gather_plan": None,
                "nested_add_plan": None,
            }

        gather_plan = _analyze_gather_input_to_nhwc(
            input_name=input_name,
            consumer_idx=int(add_idx),
            producers=producers,
            consumers=consumers,
            model_outputs=model_outputs,
        )
        if gather_plan is not None:
            return {
                "nhwc_input_name": str(input_name),
                "pre_remove_indices": [],
                "swish_plan": None,
                "unary_plan": None,
                "mul_const_plan": None,
                "mul_sub_const_plan": None,
                "gather_plan": gather_plan,
                "nested_add_plan": None,
            }

        const_add_plan = _analyze_const_add_input_to_nhwc(
            input_name=input_name,
            add_idx=int(add_idx),
            consumers=consumers,
        )
        if const_add_plan is not None:
            return {
                "nhwc_input_name": str(input_name),
                "pre_remove_indices": [],
                "swish_plan": None,
                "unary_plan": None,
                "mul_const_plan": None,
                "mul_sub_const_plan": None,
                "gather_plan": None,
                "nested_add_plan": None,
                "const_add_plan": const_add_plan,
            }

        if input_producer_idx is not None and allow_nested_add:
            input_producer_op = model_ir.operators[int(input_producer_idx)]
            if (
                str(input_producer_op.op_type) == "ADD"
                and len(input_producer_op.inputs) == 2
                and len(input_producer_op.outputs) == 1
                and str(input_producer_op.outputs[0]) == str(input_name)
                and str(input_name) not in model_outputs
            ):
                input_users = set(int(v) for v in consumers.get(str(input_name), []))
                if input_users == {int(add_idx)}:
                    nested_input_plans: List[Dict[str, Any]] = []
                    nested_rewritable = True
                    for nested_input_name in [str(v) for v in list(input_producer_op.inputs)]:
                        nested_plan = _resolve_add_input_to_nhwc(
                            input_name=str(nested_input_name),
                            add_idx=int(input_producer_idx),
                            producers=producers,
                            consumers=consumers,
                            model_outputs=model_outputs,
                            allow_nested_add=False,
                        )
                        if nested_plan is None:
                            nested_rewritable = False
                            break
                        nested_input_plans.append(dict(nested_plan))
                    if nested_rewritable:
                        return {
                            "nhwc_input_name": str(input_name),
                            "pre_remove_indices": [],
                            "swish_plan": None,
                            "unary_plan": None,
                            "mul_const_plan": None,
                            "mul_sub_const_plan": None,
                            "gather_plan": None,
                            "nested_add_plan": {
                                "input_name": str(input_name),
                                "add_idx": int(input_producer_idx),
                                "input_plans": [dict(v) for v in nested_input_plans],
                            },
                        }

        if input_producer_idx is not None:
            input_producer_op = model_ir.operators[int(input_producer_idx)]
            if (
                str(input_producer_op.op_type) == "ADD"
                and len(input_producer_op.outputs) == 1
                and str(input_producer_op.outputs[0]) == str(input_name)
            ):
                input_users = [int(v) for v in consumers.get(str(input_name), []) if int(v) != int(input_producer_idx)]
                post_candidates: List[str] = []
                for user_idx in input_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == str(input_name)
                        and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        post_candidates.append(str(user_op.outputs[0]))
                if len(post_candidates) > 0:
                    return {
                        "nhwc_input_name": str(post_candidates[0]),
                        "pre_remove_indices": [],
                        "swish_plan": None,
                        "unary_plan": None,
                        "mul_const_plan": None,
                        "mul_sub_const_plan": None,
                        "gather_plan": None,
                        "nested_add_plan": None,
                    }

        return None

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_out_name = str(add_op.outputs[0])

            bridge_output_name = add_out_name
            bridge_output_producer_idx = int(add_idx)
            bridge_output_producer_op = add_op
            bridge_prelu_idx: Optional[int] = None

            # Also support: ADD -> unary -> post-transpose fanout.
            add_out_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
            if len(add_out_users) == 1:
                unary_idx = int(add_out_users[0])
                unary_op = model_ir.operators[int(unary_idx)]
                if (
                    str(unary_op.op_type) in unary_passthrough_ops
                    and len(unary_op.inputs) == 1
                    and len(unary_op.outputs) == 1
                    and str(unary_op.inputs[0]) == add_out_name
                ):
                    unary_out_name = str(unary_op.outputs[0])
                    if add_out_name not in model_outputs and unary_out_name not in model_outputs:
                        bridge_output_name = unary_out_name
                        bridge_output_producer_idx = int(unary_idx)
                        bridge_output_producer_op = unary_op
                elif (
                    str(unary_op.op_type) == "PRELU"
                    and len(unary_op.inputs) == 2
                    and len(unary_op.outputs) == 1
                    and str(unary_op.inputs[0]) == add_out_name
                ):
                    prelu_out_name = str(unary_op.outputs[0])
                    if add_out_name not in model_outputs and prelu_out_name not in model_outputs:
                        bridge_output_name = prelu_out_name
                        bridge_output_producer_idx = int(unary_idx)
                        bridge_output_producer_op = unary_op
                        bridge_prelu_idx = int(unary_idx)

            out_users = [
                int(v)
                for v in consumers.get(bridge_output_name, [])
                if int(v) != int(bridge_output_producer_idx)
            ]
            if len(out_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == bridge_output_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))
            if not valid_users or len(post_indices) == 0:
                continue

            input_plans: List[Dict[str, Any]] = []
            rewritable = True
            for input_name in [str(v) for v in list(add_op.inputs)]:
                resolved = _resolve_add_input_to_nhwc(
                    input_name=input_name,
                    add_idx=int(add_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if resolved is None:
                    resolved = _build_direct_nchw_fallback_plan(
                        input_name=str(input_name),
                        model_outputs=model_outputs,
                    )
                    if resolved is None:
                        rewritable = False
                        break
                input_plans.append(dict(resolved))
            if not rewritable:
                continue

            planned_direct_bridge_count = int(
                sum(1 for plan in input_plans if plan.get("direct_nchw_plan", None) is not None)
            )
            if planned_direct_bridge_count > 0:
                planned_pre_remove_count = int(
                    sum(
                        len([int(v) for v in list(plan.get("pre_remove_indices", []))])
                        for plan in input_plans
                    )
                )
                planned_post_remove_count = int(len(post_indices) - 1) if len(legacy_users) > 0 else int(len(post_indices))
                # Require strict local transpose reduction before introducing
                # direct NCHW->NHWC bridge inserts.
                if int(planned_post_remove_count + planned_pre_remove_count) <= int(planned_direct_bridge_count):
                    if int(planned_post_remove_count + planned_pre_remove_count) < int(planned_direct_bridge_count):
                        continue

            def _apply_nested_add_nhwc_plan(*, plan: Dict[str, Any]) -> List[int]:
                nested_input_name = str(plan["input_name"])
                nested_add_idx = int(plan["add_idx"])
                nested_input_plans = [dict(v) for v in list(plan.get("input_plans", []))]

                nested_remove_indices: List[int] = []
                nested_new_inputs: List[str] = []
                for nested_input_plan in nested_input_plans:
                    nested_nhwc_input_name = str(nested_input_plan["nhwc_input_name"])
                    nested_swish_plan = nested_input_plan.get("swish_plan", None)
                    if nested_swish_plan is not None:
                        nested_remove_indices.extend(_apply_swish_nhwc_plan(plan=dict(nested_swish_plan)))
                    nested_unary_plan = nested_input_plan.get("unary_plan", None)
                    if nested_unary_plan is not None:
                        nested_remove_indices.extend(_apply_unary_nhwc_plan(plan=dict(nested_unary_plan)))
                    nested_mul_const_plan = nested_input_plan.get("mul_const_plan", None)
                    if nested_mul_const_plan is not None:
                        nested_remove_indices.extend(_apply_mul_const_nhwc_plan(plan=dict(nested_mul_const_plan)))
                    nested_mul_sub_const_plan = nested_input_plan.get("mul_sub_const_plan", None)
                    if nested_mul_sub_const_plan is not None:
                        nested_remove_indices.extend(
                            _apply_mul_sub_const_nhwc_plan(plan=dict(nested_mul_sub_const_plan))
                        )
                    nested_gather_plan = nested_input_plan.get("gather_plan", None)
                    if nested_gather_plan is not None:
                        nested_remove_indices.extend(
                            _apply_gather_input_nhwc_plan(plan=dict(nested_gather_plan))
                        )
                    nested_nested_add_plan = nested_input_plan.get("nested_add_plan", None)
                    if nested_nested_add_plan is not None:
                        nested_remove_indices.extend(
                            _apply_nested_add_nhwc_plan(plan=dict(nested_nested_add_plan))
                        )
                    nested_const_add_plan = nested_input_plan.get("const_add_plan", None)
                    if nested_const_add_plan is not None:
                        nested_nhwc_input_name = _apply_const_add_input_nhwc_plan(
                            plan=dict(nested_const_add_plan)
                        )
                    nested_new_inputs.append(str(nested_nhwc_input_name))
                    nested_remove_indices.extend(
                        [int(v) for v in list(nested_input_plan.get("pre_remove_indices", []))]
                    )

                nested_add_op = model_ir.operators[int(nested_add_idx)]
                if (
                    str(nested_add_op.op_type) != "ADD"
                    or len(nested_add_op.inputs) != 2
                    or len(nested_add_op.outputs) != 1
                ):
                    return nested_remove_indices
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=nested_add_op,
                    new_inputs=[str(v) for v in list(nested_new_inputs)],
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(nested_input_name, None),
                    perm_nchw_to_nhwc,
                )
                return nested_remove_indices

            new_add_inputs: List[str] = []
            pre_remove_indices: List[int] = []
            direct_bridge_specs: List[Dict[str, str]] = []
            for plan in input_plans:
                nhwc_input_name = str(plan["nhwc_input_name"])
                swish_plan = plan.get("swish_plan", None)
                if swish_plan is not None:
                    pre_remove_indices.extend(_apply_swish_nhwc_plan(plan=dict(swish_plan)))
                unary_plan = plan.get("unary_plan", None)
                if unary_plan is not None:
                    pre_remove_indices.extend(_apply_unary_nhwc_plan(plan=dict(unary_plan)))
                mul_const_plan = plan.get("mul_const_plan", None)
                if mul_const_plan is not None:
                    pre_remove_indices.extend(_apply_mul_const_nhwc_plan(plan=dict(mul_const_plan)))
                mul_sub_const_plan = plan.get("mul_sub_const_plan", None)
                if mul_sub_const_plan is not None:
                    pre_remove_indices.extend(_apply_mul_sub_const_nhwc_plan(plan=dict(mul_sub_const_plan)))
                gather_plan = plan.get("gather_plan", None)
                if gather_plan is not None:
                    pre_remove_indices.extend(_apply_gather_input_nhwc_plan(plan=dict(gather_plan)))
                nested_add_plan = plan.get("nested_add_plan", None)
                if nested_add_plan is not None:
                    pre_remove_indices.extend(_apply_nested_add_nhwc_plan(plan=dict(nested_add_plan)))
                const_add_plan = plan.get("const_add_plan", None)
                if const_add_plan is not None:
                    nhwc_input_name = _apply_const_add_input_nhwc_plan(plan=dict(const_add_plan))
                direct_nchw_plan = plan.get("direct_nchw_plan", None)
                if direct_nchw_plan is not None:
                    bridge_spec = _materialize_direct_nchw_bridge(plan=dict(direct_nchw_plan))
                    if bridge_spec is None:
                        rewritable = False
                        break
                    nhwc_input_name = str(bridge_spec["bridge_name"])
                    direct_bridge_specs.append(dict(bridge_spec))
                new_add_inputs.append(str(nhwc_input_name))
                pre_remove_indices.extend([int(v) for v in list(plan.get("pre_remove_indices", []))])
            if not rewritable:
                continue

            if bridge_prelu_idx is not None:
                target_nhwc_shape: Optional[List[int]] = None
                if len(new_add_inputs) > 0:
                    target_tensor = model_ir.tensors.get(str(new_add_inputs[0]), None)
                    if target_tensor is not None and target_tensor.shape is not None:
                        target_nhwc_shape = [int(v) for v in list(target_tensor.shape)]
                bridge_prelu_op = model_ir.operators[int(bridge_prelu_idx)]
                selected_alpha_name = _rewrite_prelu_alpha_to_nhwc(
                    prelu_idx=int(bridge_prelu_idx),
                    prelu_op=bridge_prelu_op,
                    target_nhwc_shape=target_nhwc_shape,
                    consumers=consumers,
                )
                if selected_alpha_name is None:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=bridge_prelu_op,
                    new_inputs=[add_out_name, str(selected_alpha_name)],
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=[str(v) for v in new_add_inputs],
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=bridge_output_producer_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_post_output_name, canonical_post_output_name)

            old_bridge_tensor = model_ir.tensors.get(bridge_output_name, None)
            canonical_post_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_bridge_tensor is not None and canonical_post_tensor is not None:
                canonical_post_tensor.dtype = str(old_bridge_tensor.dtype)
                canonical_post_tensor.quantization = _clone_quantization(old_bridge_tensor.quantization)
                canonical_post_tensor.shape = [int(v) for v in list(old_bridge_tensor.shape)]
                canonical_post_tensor.shape_signature = (
                    [int(v) for v in list(old_bridge_tensor.shape_signature)]
                    if old_bridge_tensor.shape_signature is not None
                    else [int(v) for v in list(old_bridge_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_post_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[bridge_output_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = sorted(
                list({int(v) for v in pre_remove_indices + post_remove_indices}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                if int(remove_idx) == int(add_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            if len(direct_bridge_specs) > 0:
                try:
                    add_current_idx = int(model_ir.operators.index(add_op))
                except ValueError:
                    add_current_idx = int(max(0, min(add_idx, len(model_ir.operators))))
                insert_offset = 0
                for bridge_spec in direct_bridge_specs:
                    model_ir.operators.insert(
                        int(add_current_idx + insert_offset),
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(bridge_spec["input_name"]), str(bridge_spec["perm_name"])],
                            outputs=[str(bridge_spec["bridge_name"])],
                        ),
                    )
                    insert_offset += 1

            # Keep explicit unary ops after transpose-bridge rewrites for stability.
            add_opts = dict(add_op.options) if isinstance(add_op.options, dict) else {}
            add_opts[skip_add_activation_fuse_marker] = True
            add_op.options = add_opts

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_add_nhwc_chains": int(optimized)}


def _optimize_transpose_dual_pre_add_to_single_post_adapter_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Reduce dual NHWC->NCHW adapters feeding one ADD into a single post-ADD adapter.

    Target:
      a_nhwc --T(0,3,1,2)--> a_nchw
      b_nhwc --T(0,3,1,2)--> b_nchw
      ADD(a_nchw, b_nchw) -> y_nchw

    Rewrite:
      ADD(a_nhwc, b_nhwc) -> y_nhwc
      y_nhwc --T(0,3,1,2)--> y_nchw

    Safety:
    - Both pre-transpose outputs are consumed only by the target ADD.
    - ADD output is rank-4 and is not a graph output.
    - Skip when ADD already fans out to NCHW->NHWC post-transpose consumers.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_const_name = "__nhwc_to_nchw_perm_rank4__"

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _op_index(op_ref: OperatorIR) -> Optional[int]:
        for idx, op in enumerate(model_ir.operators):
            if id(op) == id(op_ref):
                return int(idx)
        return None

    if perm_const_name not in model_ir.tensors:
        model_ir.tensors[perm_const_name] = TensorIR(
            name=perm_const_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            add_out_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
            if any(
                str(model_ir.operators[int(user_idx)].op_type) == "TRANSPOSE"
                and _read_transpose_perm(model_ir, model_ir.operators[int(user_idx)]) == perm_nchw_to_nhwc
                for user_idx in add_out_users
            ):
                continue

            pre_indices: List[int] = []
            new_add_inputs: List[str] = []
            rewritable = True
            for add_input_name in [str(v) for v in list(add_op.inputs)]:
                pre_idx = producers.get(add_input_name, None)
                if pre_idx is None:
                    rewritable = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != add_input_name
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    or str(add_input_name) in model_outputs
                ):
                    rewritable = False
                    break
                if set(int(v) for v in consumers.get(add_input_name, [])) != {int(add_idx)}:
                    rewritable = False
                    break
                pre_input_name = str(pre_op.inputs[0])
                pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
                if pre_input_tensor is None or len(list(pre_input_tensor.shape)) != 4:
                    rewritable = False
                    break
                pre_indices.append(int(pre_idx))
                new_add_inputs.append(pre_input_name)
            if not rewritable:
                continue

            add_out_tensor = model_ir.tensors.get(add_out_name, None)
            if add_out_tensor is None or len(list(add_out_tensor.shape)) != 4:
                continue
            add_out_shape = [int(v) for v in list(add_out_tensor.shape)]
            add_out_signature = (
                [int(v) for v in list(add_out_tensor.shape_signature)]
                if add_out_tensor.shape_signature is not None
                else [int(v) for v in list(add_out_tensor.shape)]
            )
            add_out_nhwc_shape = _permute_shape(add_out_shape, perm_nchw_to_nhwc)
            add_out_nhwc_signature = _permute_shape(add_out_signature, perm_nchw_to_nhwc)
            if add_out_nhwc_shape is None or add_out_nhwc_signature is None:
                continue

            add_out_nhwc_name = _unique_tensor_name(f"{add_out_name}_nhwc")
            model_ir.tensors[add_out_nhwc_name] = TensorIR(
                name=add_out_nhwc_name,
                dtype=str(add_out_tensor.dtype),
                shape=[int(v) for v in list(add_out_nhwc_shape)],
                shape_signature=[int(v) for v in list(add_out_nhwc_signature)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(add_out_tensor.quantization),
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=[str(v) for v in list(new_add_inputs)],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[str(add_out_nhwc_name)],
            )

            for remove_idx in sorted(list(set(int(v) for v in pre_indices)), reverse=True):
                del model_ir.operators[int(remove_idx)]

            add_new_idx = _op_index(add_op)
            if add_new_idx is None:
                changed = False
                break
            model_ir.operators.insert(
                int(add_new_idx) + 1,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(add_out_nhwc_name), str(perm_const_name)],
                    outputs=[str(add_out_name)],
                    options={},
                ),
            )

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_dual_pre_add_to_single_post_adapter_nhwc_chains": int(optimized)}


def _optimize_terminal_transpose_mul_add_reshape_fc_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove terminal NHWC->NCHW transpose before MUL/ADD/RESHAPE/FULLY_CONNECTED.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw --MUL(c1)--> m_nchw
      m_nchw --ADD(c2)--> y_nchw
      y_nchw --RESHAPE([N,-1]|[N,K])--> r
      r --FULLY_CONNECTED(W,b)--> o

    Rewrite:
      x_nhwc --MUL(c1_nhwc)--> m_nhwc
      m_nhwc --ADD(c2_nhwc)--> y_nhwc
      y_nhwc --RESHAPE([N,-1]|[N,K])--> r
      r --FULLY_CONNECTED(W',b)--> o

    where W' is input-axis permuted for NHWC flatten ordering.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _materialize_nhwc_const(
        *,
        tensor_name: str,
        target_nhwc_shape: List[int],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> Optional[str]:
        side_tensor = model_ir.tensors.get(str(tensor_name), None)
        if side_tensor is None or side_tensor.data is None:
            return None
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return str(tensor_name)

        if side_data.ndim != 4:
            return None
        side_shape = [int(v) for v in list(side_data.shape)]
        if _broadcast_static_shapes(target_nhwc_shape, side_shape) is not None:
            return str(tensor_name)

        rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
        rotated_shape = [int(v) for v in list(rotated.shape)]
        if _broadcast_static_shapes(target_nhwc_shape, rotated_shape) is None:
            return None

        side_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        shared_outside_chain = any(int(u) not in chain_indices for u in side_users)
        if shared_outside_chain:
            replacement_name = _unique_tensor_name(f"{tensor_name}_nhwc")
            model_ir.tensors[replacement_name] = TensorIR(
                name=replacement_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(rotated_shape)],
                shape_signature=[int(v) for v in list(rotated_shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
            return str(replacement_name)

        side_tensor.data = np.asarray(rotated)
        side_tensor.shape = [int(v) for v in list(rotated_shape)]
        side_tensor.shape_signature = [int(v) for v in list(rotated_shape)]
        return str(tensor_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) != 1:
                continue
            mul_idx = int(pre_users[0])
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue

            mul_side_index: Optional[int] = None
            mul_side_name: Optional[str] = None
            if str(mul_op.inputs[0]) == pre_out_name:
                mul_side_index = 1
                mul_side_name = str(mul_op.inputs[1])
            elif str(mul_op.inputs[1]) == pre_out_name:
                mul_side_index = 0
                mul_side_name = str(mul_op.inputs[0])
            if mul_side_index is None or mul_side_name is None:
                continue

            mul_out_name = str(mul_op.outputs[0])
            if mul_out_name in model_outputs:
                continue
            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_users) != 1:
                continue
            add_idx = int(mul_users[0])
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_side_index: Optional[int] = None
            add_side_name: Optional[str] = None
            if str(add_op.inputs[0]) == mul_out_name:
                add_side_index = 1
                add_side_name = str(add_op.inputs[1])
            elif str(add_op.inputs[1]) == mul_out_name:
                add_side_index = 0
                add_side_name = str(add_op.inputs[0])
            if add_side_index is None or add_side_name is None:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue
            add_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_users) != 1:
                continue
            reshape_idx = int(add_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != add_out_name
            ):
                continue

            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue
            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) != 1:
                continue
            fc_idx = int(reshape_users[0])
            fc_op = model_ir.operators[int(fc_idx)]
            if (
                str(fc_op.op_type) != "FULLY_CONNECTED"
                or len(fc_op.inputs) < 2
                or len(fc_op.outputs) != 1
                or str(fc_op.inputs[0]) != reshape_out_name
            ):
                continue

            pre_out_tensor = model_ir.tensors.get(pre_out_name, None)
            pre_in_tensor = model_ir.tensors.get(pre_in_name, None)
            add_out_tensor = model_ir.tensors.get(add_out_name, None)
            if (
                pre_out_tensor is None
                or pre_in_tensor is None
                or add_out_tensor is None
                or len(list(pre_out_tensor.shape)) != 4
                or len(list(pre_in_tensor.shape)) != 4
                or len(list(add_out_tensor.shape)) != 4
            ):
                continue

            nchw_shape = [int(v) for v in list(pre_out_tensor.shape)]
            nhwc_shape = [int(v) for v in list(pre_in_tensor.shape)]
            if _permute_shape(nchw_shape, perm_nchw_to_nhwc) != nhwc_shape:
                continue
            if not _is_fully_known_positive_shape(nchw_shape):
                continue

            n, c, h, w = [int(v) for v in nchw_shape]
            input_size = int(c * h * w)

            weight_name = str(fc_op.inputs[1])
            weight_tensor = model_ir.tensors.get(weight_name, None)
            if weight_tensor is None or weight_tensor.data is None:
                continue
            weight_data = np.asarray(weight_tensor.data)
            if weight_data.ndim != 2:
                continue

            perm_index = np.transpose(
                np.arange(input_size, dtype=np.int64).reshape(c, h, w),
                (1, 2, 0),
            ).reshape(-1)
            if weight_data.shape[1] == input_size:
                weight_nhwc = np.asarray(weight_data[:, perm_index])
            elif weight_data.shape[0] == input_size:
                weight_nhwc = np.asarray(weight_data[perm_index, :])
            else:
                continue

            chain_indices = {int(pre_idx), int(mul_idx), int(add_idx), int(reshape_idx), int(fc_idx)}
            mul_side_replacement = _materialize_nhwc_const(
                tensor_name=str(mul_side_name),
                target_nhwc_shape=[int(v) for v in list(nhwc_shape)],
                consumers=consumers,
                chain_indices=chain_indices,
            )
            if mul_side_replacement is None:
                continue
            add_side_replacement = _materialize_nhwc_const(
                tensor_name=str(add_side_name),
                target_nhwc_shape=[int(v) for v in list(nhwc_shape)],
                consumers=consumers,
                chain_indices=chain_indices,
            )
            if add_side_replacement is None:
                continue

            weight_users = [int(v) for v in consumers.get(weight_name, [])]
            selected_weight_name = str(weight_name)
            if len(weight_users) == 1 and int(weight_users[0]) == int(fc_idx):
                weight_tensor.data = np.asarray(weight_nhwc)
                weight_tensor.shape = [int(v) for v in list(weight_nhwc.shape)]
                weight_tensor.shape_signature = [int(v) for v in list(weight_nhwc.shape)]
            else:
                selected_weight_name = _unique_tensor_name(f"{weight_name}_nhwc")
                model_ir.tensors[selected_weight_name] = TensorIR(
                    name=selected_weight_name,
                    dtype=str(weight_tensor.dtype),
                    shape=[int(v) for v in list(weight_nhwc.shape)],
                    shape_signature=[int(v) for v in list(weight_nhwc.shape)],
                    data=np.asarray(weight_nhwc),
                    is_variable=False,
                    quantization=_clone_quantization(weight_tensor.quantization),
                )

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=0 if int(mul_side_index) == 1 else 1,
                new_input_name=str(pre_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=int(mul_side_index),
                new_input_name=str(mul_side_replacement),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add_op,
                input_index=int(add_side_index),
                new_input_name=str(add_side_replacement),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=fc_op,
                input_index=1,
                new_input_name=str(selected_weight_name),
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mul_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_out_name, None),
                perm_nchw_to_nhwc,
            )

            del model_ir.operators[int(pre_idx)]
            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_terminal_transpose_mul_add_reshape_fc_nhwc_chains": int(optimized)}


def _optimize_terminal_transpose_prelu_reshape_batchmatmul_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove terminal NHWC->NCHW transpose before PRELU/RESHAPE/BATCH_MATMUL.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw --PRELU(alpha_nchw)--> p_nchw
      p_nchw --RESHAPE([N,-1]|[N,K])--> r
      r --BATCH_MATMUL(W)--> o

    Rewrite:
      x_nhwc --PRELU(alpha_nhwc)--> p_nhwc
      p_nhwc --RESHAPE([N,-1]|[N,K])--> r
      r --BATCH_MATMUL(W')--> o

    where W' is input-axis permuted for NHWC flatten ordering.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _materialize_prelu_alpha_nhwc(
        *,
        alpha_name: str,
        channel_dim: int,
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> Optional[str]:
        alpha_tensor = model_ir.tensors.get(str(alpha_name), None)
        if alpha_tensor is None or alpha_tensor.data is None:
            return None
        alpha_data = np.asarray(alpha_tensor.data)
        if int(alpha_data.size) == 1:
            return str(alpha_name)

        alpha_nhwc: Optional[np.ndarray] = None
        alpha_shape = [int(v) for v in list(alpha_data.shape)]
        if alpha_data.ndim == 4:
            if (
                int(alpha_shape[0]) == 1
                and int(alpha_shape[1]) == int(channel_dim)
                and int(alpha_shape[2]) == 1
                and int(alpha_shape[3]) == 1
            ):
                alpha_nhwc = np.transpose(alpha_data, perm_nchw_to_nhwc).astype(alpha_data.dtype, copy=False)
            elif (
                int(alpha_shape[0]) == 1
                and int(alpha_shape[1]) == 1
                and int(alpha_shape[2]) == 1
                and int(alpha_shape[3]) == int(channel_dim)
            ):
                return str(alpha_name)
        elif alpha_data.ndim == 3:
            if (
                int(alpha_shape[0]) == int(channel_dim)
                and int(alpha_shape[1]) == 1
                and int(alpha_shape[2]) == 1
            ):
                alpha_nhwc = np.transpose(alpha_data, (1, 2, 0)).astype(alpha_data.dtype, copy=False)
        elif alpha_data.ndim == 1:
            # 1D slope is ambiguous for rank-4 layout in this converter; keep strict.
            return None

        if alpha_nhwc is None:
            return None

        alpha_users = [int(v) for v in consumers.get(str(alpha_name), [])]
        shared_outside_chain = any(int(u) not in chain_indices for u in alpha_users)
        if shared_outside_chain:
            replacement_name = _unique_tensor_name(f"{alpha_name}_nhwc")
            replacement_shape = [int(v) for v in list(alpha_nhwc.shape)]
            model_ir.tensors[replacement_name] = TensorIR(
                name=replacement_name,
                dtype=str(alpha_tensor.dtype),
                shape=replacement_shape,
                shape_signature=replacement_shape,
                data=np.asarray(alpha_nhwc),
                is_variable=False,
                quantization=_clone_quantization(alpha_tensor.quantization),
            )
            return str(replacement_name)

        alpha_tensor.data = np.asarray(alpha_nhwc)
        alpha_tensor.shape = [int(v) for v in list(alpha_nhwc.shape)]
        alpha_tensor.shape_signature = [int(v) for v in list(alpha_nhwc.shape)]
        return str(alpha_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) != 1:
                continue
            prelu_idx = int(pre_users[0])
            prelu_op = model_ir.operators[int(prelu_idx)]
            if (
                str(prelu_op.op_type) != "PRELU"
                or len(prelu_op.inputs) != 2
                or len(prelu_op.outputs) != 1
                or str(prelu_op.inputs[0]) != pre_out_name
            ):
                continue
            prelu_out_name = str(prelu_op.outputs[0])
            if prelu_out_name in model_outputs:
                continue

            prelu_users = [int(v) for v in consumers.get(prelu_out_name, [])]
            if len(prelu_users) != 1:
                continue
            reshape_idx = int(prelu_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != prelu_out_name
            ):
                continue
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue

            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) != 1:
                continue
            bmm_idx = int(reshape_users[0])
            bmm_op = model_ir.operators[int(bmm_idx)]
            if (
                str(bmm_op.op_type) != "BATCH_MATMUL"
                or len(bmm_op.inputs) != 2
                or len(bmm_op.outputs) != 1
                or str(bmm_op.inputs[0]) != reshape_out_name
            ):
                continue
            bmm_opts = dict(bmm_op.options) if isinstance(bmm_op.options, dict) else {}
            if bool(bmm_opts.get("adjX", False)) or bool(bmm_opts.get("adjY", False)):
                continue

            pre_out_tensor = model_ir.tensors.get(pre_out_name, None)
            pre_in_tensor = model_ir.tensors.get(pre_in_name, None)
            prelu_out_tensor = model_ir.tensors.get(prelu_out_name, None)
            if (
                pre_out_tensor is None
                or pre_in_tensor is None
                or prelu_out_tensor is None
                or len(list(pre_out_tensor.shape)) != 4
                or len(list(pre_in_tensor.shape)) != 4
                or len(list(prelu_out_tensor.shape)) != 4
            ):
                continue
            nchw_shape = [int(v) for v in list(pre_out_tensor.shape)]
            nhwc_shape = [int(v) for v in list(pre_in_tensor.shape)]
            if not _is_fully_known_positive_shape(nchw_shape):
                continue
            if _permute_shape(nchw_shape, perm_nchw_to_nhwc) != nhwc_shape:
                continue
            _, c, h, w = [int(v) for v in nchw_shape]
            input_size = int(c * h * w)

            rhs_name = str(bmm_op.inputs[1])
            rhs_tensor = model_ir.tensors.get(rhs_name, None)
            if rhs_tensor is None or rhs_tensor.data is None:
                continue
            rhs_data = np.asarray(rhs_tensor.data)
            if rhs_data.ndim != 2 or int(rhs_data.shape[0]) != int(input_size):
                continue

            perm_index = np.transpose(
                np.arange(input_size, dtype=np.int64).reshape(c, h, w),
                (1, 2, 0),
            ).reshape(-1)
            rhs_nhwc = np.asarray(rhs_data[perm_index, :])

            chain_indices = {int(pre_idx), int(prelu_idx), int(reshape_idx), int(bmm_idx)}
            alpha_name = str(prelu_op.inputs[1])
            alpha_replacement = _materialize_prelu_alpha_nhwc(
                alpha_name=str(alpha_name),
                channel_dim=int(c),
                consumers=consumers,
                chain_indices=chain_indices,
            )
            if alpha_replacement is None:
                continue

            rhs_users = [int(v) for v in consumers.get(rhs_name, [])]
            selected_rhs_name = str(rhs_name)
            if len(rhs_users) == 1 and int(rhs_users[0]) == int(bmm_idx):
                rhs_tensor.data = np.asarray(rhs_nhwc)
                rhs_tensor.shape = [int(v) for v in list(rhs_nhwc.shape)]
                rhs_tensor.shape_signature = [int(v) for v in list(rhs_nhwc.shape)]
            else:
                selected_rhs_name = _unique_tensor_name(f"{rhs_name}_nhwc")
                model_ir.tensors[selected_rhs_name] = TensorIR(
                    name=selected_rhs_name,
                    dtype=str(rhs_tensor.dtype),
                    shape=[int(v) for v in list(rhs_nhwc.shape)],
                    shape_signature=[int(v) for v in list(rhs_nhwc.shape)],
                    data=np.asarray(rhs_nhwc),
                    is_variable=False,
                    quantization=_clone_quantization(rhs_tensor.quantization),
                )

            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu_op,
                input_index=0,
                new_input_name=str(pre_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu_op,
                input_index=1,
                new_input_name=str(alpha_replacement),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=bmm_op,
                input_index=1,
                new_input_name=str(selected_rhs_name),
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(prelu_out_name, None),
                perm_nchw_to_nhwc,
            )

            del model_ir.operators[int(pre_idx)]
            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_terminal_transpose_prelu_reshape_batchmatmul_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Lift NCHW residual blocks back to NHWC when they follow:
      ADD -> MUL(const) -> ADD(const) -> PRELU -> TRANSPOSE(0,2,3,1)

    Inputs to the first ADD must be wrapped by TRANSPOSE(0,3,1,2) adapters.
    If PRELU keeps legacy NCHW consumers, preserve one adapter transpose.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_binary_const_to_nhwc(
        *,
        op_idx: int,
        side_input_index: int,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False

        nhwc_data: Optional[np.ndarray] = None
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is not None:
            side_shape = [int(v) for v in list(side_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is not None:
                return True
            # NCHW per-channel const [1,C,1,1] is mapped once to NHWC [1,1,1,C].
            # Avoid multi-step rotations, which can accidentally produce [1,1,C,1]
            # when target metadata is stale/non-canonical.
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                    nhwc_data = np.asarray(rotated)
            else:
                rotated = np.asarray(side_data)
                for _ in range(3):
                    rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                        nhwc_data = np.asarray(rotated)
                        break
            if nhwc_data is None:
                return False
        else:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 1
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                nhwc_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            else:
                # Unknown target rank-4 shape: avoid speculative re-rotation to keep
                # repeated optimization sweeps idempotent.
                return True

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        target_name = str(side_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{side_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
        else:
            side_tensor.data = np.asarray(nhwc_data)
            side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        _replace_operator_input_at(
            model_ir=model_ir,
            op=model_ir.operators[int(op_idx)],
            input_index=int(side_input_index),
            new_input_name=str(target_name),
        )
        return True

    def _can_rewrite_binary_const_to_nhwc(
        *,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is None:
            return True
        side_shape = [int(v) for v in list(side_data.shape)]
        if _broadcast_static_shapes(target_shape, side_shape) is not None:
            return True
        if (
            len(side_shape) == 4
            and int(side_shape[0]) == 1
            and int(side_shape[1]) > 0
            and int(side_shape[2]) == 1
            and int(side_shape[3]) == 1
        ):
            rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            return _broadcast_static_shapes(target_shape, rotated_shape) is not None
        rotated = np.asarray(side_data)
        for _ in range(3):
            rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue
            add_out_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
            if len(add_out_users) != 1:
                continue
            mul_idx = int(add_out_users[0])
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_in0 = str(mul_op.inputs[0])
            mul_in1 = str(mul_op.inputs[1])
            if mul_in0 == add_out_name:
                mul_side_input_index = 1
                mul_side_name = mul_in1
            elif mul_in1 == add_out_name:
                mul_side_input_index = 0
                mul_side_name = mul_in0
            else:
                continue

            mul_out_name = str(mul_op.outputs[0])
            if mul_out_name in model_outputs:
                continue
            mul_out_users = [int(v) for v in consumers.get(mul_out_name, []) if int(v) != int(mul_idx)]
            if len(mul_out_users) != 1:
                continue
            add2_idx = int(mul_out_users[0])
            add2_op = model_ir.operators[int(add2_idx)]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])
            if add2_in0 == mul_out_name:
                add2_side_input_index = 1
                add2_side_name = add2_in1
            elif add2_in1 == mul_out_name:
                add2_side_input_index = 0
                add2_side_name = add2_in0
            else:
                continue

            add2_out_name = str(add2_op.outputs[0])
            if add2_out_name in model_outputs:
                continue
            add2_out_users = [int(v) for v in consumers.get(add2_out_name, []) if int(v) != int(add2_idx)]
            if len(add2_out_users) != 1:
                continue
            prelu_idx = int(add2_out_users[0])
            prelu_op = model_ir.operators[int(prelu_idx)]
            if (
                str(prelu_op.op_type) != "PRELU"
                or len(prelu_op.inputs) != 2
                or len(prelu_op.outputs) != 1
                or str(prelu_op.inputs[0]) != add2_out_name
            ):
                continue
            prelu_side_input_index = 1
            prelu_side_name = str(prelu_op.inputs[1])
            prelu_out_name = str(prelu_op.outputs[0])
            if prelu_out_name in model_outputs:
                continue

            prelu_users = [int(v) for v in consumers.get(prelu_out_name, []) if int(v) != int(prelu_idx)]
            if len(prelu_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            for user_idx in prelu_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == prelu_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))
            if len(post_indices) == 0:
                continue

            input_plans: List[Dict[str, Any]] = []
            rewritable_inputs = True
            for add_input_name in [str(v) for v in list(add_op.inputs)]:
                pre_idx = producers.get(str(add_input_name), None)
                if pre_idx is None:
                    rewritable_inputs = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(add_input_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    or str(add_input_name) in model_outputs
                ):
                    rewritable_inputs = False
                    break
                users = [int(v) for v in consumers.get(str(add_input_name), [])]
                if int(add_idx) not in users:
                    rewritable_inputs = False
                    break
                input_plans.append(
                    {
                        "nhwc_input_name": str(pre_op.inputs[0]),
                        "pre_remove_indices": [int(pre_idx)] if set(users) == {int(add_idx)} else [],
                    }
                )
            if not rewritable_inputs or len(input_plans) != 2:
                continue

            target_nhwc_shape = (
                list(model_ir.tensors[str(input_plans[0]["nhwc_input_name"])].shape)
                if str(input_plans[0]["nhwc_input_name"]) in model_ir.tensors
                else None
            )
            chain_index_set = {int(mul_idx), int(add2_idx), int(prelu_idx)}
            if not _can_rewrite_binary_const_to_nhwc(
                side_input_name=str(mul_side_name),
                target_broadcast_shape=target_nhwc_shape,
            ):
                continue
            if not _can_rewrite_binary_const_to_nhwc(
                side_input_name=str(add2_side_name),
                target_broadcast_shape=target_nhwc_shape,
            ):
                continue
            if not _can_rewrite_binary_const_to_nhwc(
                side_input_name=str(prelu_side_name),
                target_broadcast_shape=target_nhwc_shape,
            ):
                continue
            if not _rewrite_binary_const_to_nhwc(
                op_idx=int(mul_idx),
                side_input_index=int(mul_side_input_index),
                side_input_name=str(mul_side_name),
                target_broadcast_shape=target_nhwc_shape,
                chain_index_set=chain_index_set,
                consumers=consumers,
            ):
                continue
            if not _rewrite_binary_const_to_nhwc(
                op_idx=int(add2_idx),
                side_input_index=int(add2_side_input_index),
                side_input_name=str(add2_side_name),
                target_broadcast_shape=target_nhwc_shape,
                chain_index_set=chain_index_set,
                consumers=consumers,
            ):
                continue
            if not _rewrite_binary_const_to_nhwc(
                op_idx=int(prelu_idx),
                side_input_index=int(prelu_side_input_index),
                side_input_name=str(prelu_side_name),
                target_broadcast_shape=target_nhwc_shape,
                chain_index_set=chain_index_set,
                consumers=consumers,
            ):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=[str(plan["nhwc_input_name"]) for plan in input_plans],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mul_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add2_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(prelu_out_name, None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            old_prelu_out_tensor = model_ir.tensors.get(prelu_out_name, None)
            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_prelu_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_prelu_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_prelu_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_prelu_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_prelu_out_tensor.shape_signature)]
                    if old_prelu_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_prelu_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[prelu_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = set(int(v) for v in post_remove_indices)
            for plan in input_plans:
                for pre_remove_idx in list(plan.get("pre_remove_indices", [])):
                    remove_indices.add(int(pre_remove_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_add_mul_add_prelu_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Lift NCHW residual add fanouts back to NHWC when they follow:
      [x_nhwc --T(0,3,1,2)--> x_nchw] + [y_nhwc --T(0,3,1,2)--> y_nchw]
      ADD(x_nchw, y_nchw) -> a_nchw
      branch_i: a_nchw -> MUL(const) -> ADD(const) -> T(0,2,3,1) -> out_i_nhwc

    The ADD output may also keep legacy NCHW consumers. In that case, keep one
    adapter transpose: a_nhwc --T(0,3,1,2)--> a_nchw(legacy).
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_binary_const_to_nhwc(
        *,
        op_idx: int,
        side_input_index: int,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False

        nhwc_data: Optional[np.ndarray] = None
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is not None:
            side_shape = [int(v) for v in list(side_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is not None:
                return True
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                    nhwc_data = np.asarray(rotated)
            else:
                rotated = np.asarray(side_data)
                for _ in range(3):
                    rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                        nhwc_data = np.asarray(rotated)
                        break
            if nhwc_data is None:
                return False
        else:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 1
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                nhwc_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            else:
                return True

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        target_name = str(side_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{side_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
        else:
            side_tensor.data = np.asarray(nhwc_data)
            side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        _replace_operator_input_at(
            model_ir=model_ir,
            op=model_ir.operators[int(op_idx)],
            input_index=int(side_input_index),
            new_input_name=str(target_name),
        )
        return True

    def _can_rewrite_binary_const_to_nhwc(
        *,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is None:
            return True
        side_shape = [int(v) for v in list(side_data.shape)]
        if _broadcast_static_shapes(target_shape, side_shape) is not None:
            return True
        if (
            len(side_shape) == 4
            and int(side_shape[0]) == 1
            and int(side_shape[1]) > 0
            and int(side_shape[2]) == 1
            and int(side_shape[3]) == 1
        ):
            rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            return _broadcast_static_shapes(target_shape, rotated_shape) is not None
        rotated = np.asarray(side_data)
        for _ in range(3):
            rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            input_plans: List[Dict[str, Any]] = []
            rewritable_inputs = True
            for add_input_name in [str(v) for v in list(add_op.inputs)]:
                pre_idx = producers.get(str(add_input_name), None)
                if pre_idx is None:
                    rewritable_inputs = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(add_input_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    or str(add_input_name) in model_outputs
                ):
                    rewritable_inputs = False
                    break
                input_users = [int(v) for v in consumers.get(str(add_input_name), [])]
                if int(add_idx) not in input_users:
                    rewritable_inputs = False
                    break
                input_plans.append(
                    {
                        "nhwc_input_name": str(pre_op.inputs[0]),
                        "pre_remove_indices": [int(pre_idx)] if set(input_users) == {int(add_idx)} else [],
                    }
                )
            if not rewritable_inputs or len(input_plans) != 2:
                continue

            add_out_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
            if len(add_out_users) == 0:
                continue

            branch_plans: List[Dict[str, Any]] = []
            legacy_users: List[int] = []
            for user_idx in add_out_users:
                mul_op = model_ir.operators[int(user_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    legacy_users.append(int(user_idx))
                    continue
                mul_in0 = str(mul_op.inputs[0])
                mul_in1 = str(mul_op.inputs[1])
                if mul_in0 == add_out_name:
                    mul_data_input_index = 0
                    mul_side_input_index = 1
                    mul_side_name = mul_in1
                elif mul_in1 == add_out_name:
                    mul_data_input_index = 1
                    mul_side_input_index = 0
                    mul_side_name = mul_in0
                else:
                    legacy_users.append(int(user_idx))
                    continue

                mul_out_name = str(mul_op.outputs[0])
                if mul_out_name in model_outputs:
                    legacy_users.append(int(user_idx))
                    continue
                mul_out_users = [int(v) for v in consumers.get(mul_out_name, []) if int(v) != int(user_idx)]
                if len(mul_out_users) != 1:
                    legacy_users.append(int(user_idx))
                    continue
                add2_idx = int(mul_out_users[0])
                add2_op = model_ir.operators[int(add2_idx)]
                if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                    legacy_users.append(int(user_idx))
                    continue
                add2_in0 = str(add2_op.inputs[0])
                add2_in1 = str(add2_op.inputs[1])
                if add2_in0 == mul_out_name:
                    add2_side_input_index = 1
                    add2_side_name = add2_in1
                elif add2_in1 == mul_out_name:
                    add2_side_input_index = 0
                    add2_side_name = add2_in0
                else:
                    legacy_users.append(int(user_idx))
                    continue

                add2_out_name = str(add2_op.outputs[0])
                if add2_out_name in model_outputs:
                    legacy_users.append(int(user_idx))
                    continue
                add2_out_users = [int(v) for v in consumers.get(add2_out_name, []) if int(v) != int(add2_idx)]
                if len(add2_out_users) != 1:
                    legacy_users.append(int(user_idx))
                    continue
                post_idx = int(add2_out_users[0])
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != add2_out_name
                    or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                    or str(post_op.outputs[0]) in model_outputs
                ):
                    legacy_users.append(int(user_idx))
                    continue

                branch_plans.append(
                    {
                        "mul_idx": int(user_idx),
                        "mul_data_input_index": int(mul_data_input_index),
                        "mul_side_input_index": int(mul_side_input_index),
                        "mul_side_name": str(mul_side_name),
                        "mul_out_name": str(mul_out_name),
                        "add2_idx": int(add2_idx),
                        "add2_side_input_index": int(add2_side_input_index),
                        "add2_side_name": str(add2_side_name),
                        "add2_out_name": str(add2_out_name),
                        "post_idx": int(post_idx),
                        "post_output_name": str(post_op.outputs[0]),
                    }
                )

            if len(branch_plans) == 0:
                continue

            planned_pre_remove_count = int(
                sum(len([int(v) for v in list(plan.get("pre_remove_indices", []))]) for plan in input_plans)
            )
            planned_post_remove_count = int(len(branch_plans))
            planned_adapter_count = int(1 if len(legacy_users) > 0 else 0)
            if int(planned_pre_remove_count + planned_post_remove_count) <= int(planned_adapter_count):
                continue

            target_nhwc_shape = (
                list(model_ir.tensors[str(input_plans[0]["nhwc_input_name"])].shape)
                if str(input_plans[0]["nhwc_input_name"]) in model_ir.tensors
                else None
            )
            chain_index_set: set[int] = set()
            for branch in branch_plans:
                chain_index_set.add(int(branch["mul_idx"]))
                chain_index_set.add(int(branch["add2_idx"]))
            constants_rewritable = True
            for branch in branch_plans:
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(branch["mul_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    constants_rewritable = False
                    break
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(branch["add2_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    constants_rewritable = False
                    break
            if not constants_rewritable:
                continue
            constants_applied = True
            for branch in branch_plans:
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(branch["mul_idx"]),
                    side_input_index=int(branch["mul_side_input_index"]),
                    side_input_name=str(branch["mul_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    constants_applied = False
                    break
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(branch["add2_idx"]),
                    side_input_index=int(branch["add2_side_input_index"]),
                    side_input_name=str(branch["add2_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    constants_applied = False
                    break
            if not constants_applied:
                continue

            add_nhwc_name = (
                _unique_tensor_name(f"{add_out_name}_nhwc")
                if len(legacy_users) > 0
                else str(add_out_name)
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=[str(plan["nhwc_input_name"]) for plan in input_plans],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[str(add_nhwc_name)],
            )

            old_add_tensor = model_ir.tensors.get(add_out_name, None)
            if add_nhwc_name != add_out_name and old_add_tensor is not None and add_nhwc_name not in model_ir.tensors:
                model_ir.tensors[add_nhwc_name] = TensorIR(
                    name=add_nhwc_name,
                    dtype=str(old_add_tensor.dtype),
                    shape=[int(v) for v in list(old_add_tensor.shape)],
                    shape_signature=(
                        [int(v) for v in list(old_add_tensor.shape_signature)]
                        if old_add_tensor.shape_signature is not None
                        else [int(v) for v in list(old_add_tensor.shape)]
                    ),
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(old_add_tensor.quantization),
                )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_nhwc_name, None),
                perm_nchw_to_nhwc,
            )

            for branch in branch_plans:
                mul_op = model_ir.operators[int(branch["mul_idx"])]
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                data_input_index = int(branch["mul_data_input_index"])
                if int(data_input_index) >= len(mul_inputs):
                    constants_applied = False
                    break
                if str(mul_inputs[int(data_input_index)]) != add_out_name:
                    constants_applied = False
                    break
                mul_inputs[int(data_input_index)] = str(add_nhwc_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )

                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(branch["mul_out_name"]), None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(branch["add2_out_name"]), None),
                    perm_nchw_to_nhwc,
                )

                add2_op = model_ir.operators[int(branch["add2_idx"])]
                canonical_post_output_name = str(branch["post_output_name"])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=add2_op,
                    new_outputs=[canonical_post_output_name],
                )

                old_add2_out_tensor = model_ir.tensors.get(str(branch["add2_out_name"]), None)
                canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
                if old_add2_out_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_add2_out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_add2_out_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_add2_out_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_add2_out_tensor.shape_signature)]
                        if old_add2_out_tensor.shape_signature is not None
                        else [int(v) for v in list(old_add2_out_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        canonical_tensor,
                        perm_nchw_to_nhwc,
                    )
            if not constants_applied:
                continue

            remove_indices: set[int] = set()
            for plan in input_plans:
                for pre_remove_idx in list(plan.get("pre_remove_indices", [])):
                    remove_indices.add(int(pre_remove_idx))
            for branch in branch_plans:
                remove_indices.add(int(branch["post_idx"]))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                if int(remove_idx) == int(add_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            if len(legacy_users) > 0:
                adapter_perm_name = _unique_tensor_name(f"{add_out_name}_adapter_perm")
                model_ir.tensors[adapter_perm_name] = TensorIR(
                    name=adapter_perm_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                try:
                    add_current_idx = int(model_ir.operators.index(add_op))
                except ValueError:
                    add_current_idx = int(max(0, min(add_idx, len(model_ir.operators))))
                model_ir.operators.insert(
                    int(add_current_idx + 1),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(add_nhwc_name), str(adapter_perm_name)],
                        outputs=[str(add_out_name)],
                    ),
                )

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Lift strict fanout chains:
      x_nhwc --T(0,3,1,2)--> x_nchw --UNARY-->
      branch_i: MUL(const)->ADD(const)->T(0,2,3,1)

    into NHWC without adapter transposes.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {"RELU", "RELU6", "LOGISTIC", "TANH", "HARD_SWISH", "LEAKY_RELU", "GELU"}

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_binary_const_to_nhwc(
        *,
        op_idx: int,
        side_input_index: int,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False

        nhwc_data: Optional[np.ndarray] = None
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is not None:
            side_shape = [int(v) for v in list(side_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is not None:
                return True
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                    nhwc_data = np.asarray(rotated)
            else:
                rotated = np.asarray(side_data)
                for _ in range(3):
                    rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                        nhwc_data = np.asarray(rotated)
                        break
            if nhwc_data is None:
                return False
        else:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 1
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                nhwc_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            else:
                return True

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        target_name = str(side_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{side_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
        else:
            side_tensor.data = np.asarray(nhwc_data)
            side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        _replace_operator_input_at(
            model_ir=model_ir,
            op=model_ir.operators[int(op_idx)],
            input_index=int(side_input_index),
            new_input_name=str(target_name),
        )
        return True

    def _can_rewrite_binary_const_to_nhwc(
        *,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is None:
            return True
        side_shape = [int(v) for v in list(side_data.shape)]
        if _broadcast_static_shapes(target_shape, side_shape) is not None:
            return True
        if (
            len(side_shape) == 4
            and int(side_shape[0]) == 1
            and int(side_shape[1]) > 0
            and int(side_shape[2]) == 1
            and int(side_shape[3]) == 1
        ):
            rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            return _broadcast_static_shapes(target_shape, rotated_shape) is not None
        rotated = np.asarray(side_data)
        for _ in range(3):
            rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_outputs or pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            unary_idx = int(pre_users[0])
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.inputs[0]) != pre_output_name
            ):
                continue

            unary_out_name = str(unary_op.outputs[0])
            if unary_out_name in model_outputs:
                continue
            unary_users = [int(v) for v in consumers.get(unary_out_name, []) if int(v) != int(unary_idx)]
            if len(unary_users) == 0:
                continue

            branch_plans: List[Dict[str, Any]] = []
            valid = True
            for user_idx in unary_users:
                mul_op = model_ir.operators[int(user_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    valid = False
                    break
                mul_in0 = str(mul_op.inputs[0])
                mul_in1 = str(mul_op.inputs[1])
                if mul_in0 == unary_out_name:
                    mul_data_input_index = 0
                    mul_side_input_index = 1
                    mul_side_name = mul_in1
                elif mul_in1 == unary_out_name:
                    mul_data_input_index = 1
                    mul_side_input_index = 0
                    mul_side_name = mul_in0
                else:
                    valid = False
                    break

                mul_out_name = str(mul_op.outputs[0])
                if mul_out_name in model_outputs:
                    valid = False
                    break
                mul_out_users = [int(v) for v in consumers.get(mul_out_name, []) if int(v) != int(user_idx)]
                if len(mul_out_users) != 1:
                    valid = False
                    break
                add_idx = int(mul_out_users[0])
                add_op = model_ir.operators[int(add_idx)]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    valid = False
                    break
                add_in0 = str(add_op.inputs[0])
                add_in1 = str(add_op.inputs[1])
                if add_in0 == mul_out_name:
                    add_side_input_index = 1
                    add_side_name = add_in1
                elif add_in1 == mul_out_name:
                    add_side_input_index = 0
                    add_side_name = add_in0
                else:
                    valid = False
                    break

                add_out_name = str(add_op.outputs[0])
                if add_out_name in model_outputs:
                    valid = False
                    break
                add_out_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
                if len(add_out_users) != 1:
                    valid = False
                    break
                post_idx = int(add_out_users[0])
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != add_out_name
                    or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                    or str(post_op.outputs[0]) in model_outputs
                ):
                    valid = False
                    break

                branch_plans.append(
                    {
                        "mul_idx": int(user_idx),
                        "mul_data_input_index": int(mul_data_input_index),
                        "mul_side_input_index": int(mul_side_input_index),
                        "mul_side_name": str(mul_side_name),
                        "mul_out_name": str(mul_out_name),
                        "add_idx": int(add_idx),
                        "add_side_input_index": int(add_side_input_index),
                        "add_side_name": str(add_side_name),
                        "add_out_name": str(add_out_name),
                        "post_idx": int(post_idx),
                        "post_output_name": str(post_op.outputs[0]),
                    }
                )
            if not valid or len(branch_plans) == 0:
                continue

            target_nhwc_shape = (
                [int(v) for v in list(model_ir.tensors[pre_input_name].shape)]
                if pre_input_name in model_ir.tensors
                else None
            )
            chain_index_set: set[int] = {int(unary_idx)}
            for branch in branch_plans:
                chain_index_set.add(int(branch["mul_idx"]))
                chain_index_set.add(int(branch["add_idx"]))

            constants_rewritable = True
            for branch in branch_plans:
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(branch["mul_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    constants_rewritable = False
                    break
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(branch["add_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    constants_rewritable = False
                    break
            if not constants_rewritable:
                continue

            constants_applied = True
            for branch in branch_plans:
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(branch["mul_idx"]),
                    side_input_index=int(branch["mul_side_input_index"]),
                    side_input_name=str(branch["mul_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    constants_applied = False
                    break
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(branch["add_idx"]),
                    side_input_index=int(branch["add_side_input_index"]),
                    side_input_name=str(branch["add_side_name"]),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    constants_applied = False
                    break
            if not constants_applied:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[str(pre_input_name)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(unary_out_name, None),
                perm_nchw_to_nhwc,
            )

            for branch in branch_plans:
                mul_op = model_ir.operators[int(branch["mul_idx"])]
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                data_input_index = int(branch["mul_data_input_index"])
                if int(data_input_index) >= len(mul_inputs):
                    constants_applied = False
                    break
                if str(mul_inputs[int(data_input_index)]) != unary_out_name:
                    constants_applied = False
                    break
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )

                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(branch["mul_out_name"]), None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(branch["add_out_name"]), None),
                    perm_nchw_to_nhwc,
                )

                add_op = model_ir.operators[int(branch["add_idx"])]
                canonical_post_output_name = str(branch["post_output_name"])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_outputs=[canonical_post_output_name],
                )

                old_add_out_tensor = model_ir.tensors.get(str(branch["add_out_name"]), None)
                canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
                if old_add_out_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_add_out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_add_out_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_add_out_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_add_out_tensor.shape_signature)]
                        if old_add_out_tensor.shape_signature is not None
                        else [int(v) for v in list(old_add_out_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        canonical_tensor,
                        perm_nchw_to_nhwc,
                    )
            if not constants_applied:
                continue

            remove_indices: set[int] = {int(pre_idx)}
            for branch in branch_plans:
                remove_indices.add(int(branch["post_idx"]))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains": int(optimized)}


def _optimize_transpose_pre_add_reshape_transpose_suffix_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW add bridges that immediately feed [N,C,HW] reshape + [0,2,1] transpose.

    Target:
      a_nhwc --TRANSPOSE(0,3,1,2)--> a_nchw
      b_nhwc --TRANSPOSE(0,3,1,2)--> b_nchw
      ADD(a_nchw, b_nchw) -> y_nchw
      y_nchw --RESHAPE([N,C,HW])--> r
      r --TRANSPOSE([0,2,1])--> z
      (optional legacy users consume y_nchw)

    Rewrite:
      ADD(a_nhwc, b_nhwc) -> y_nhwc
      y_nhwc --RESHAPE([N,HW,C])--> z
      (if legacy users exist, keep one adapter TRANSPOSE(0,3,1,2): y_nhwc -> y_nchw)
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_3d_nchw_to_nhwc = [0, 2, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _maybe_swap_reshape_shape_tensor(
        *,
        reshape_op: OperatorIR,
        reshape_idx: int,
        consumers: Dict[str, List[int]],
    ) -> None:
        if len(reshape_op.inputs) < 2:
            return
        shape_name = str(reshape_op.inputs[1])
        shape_tensor = model_ir.tensors.get(shape_name, None)
        if shape_tensor is None or shape_tensor.data is None:
            return
        shape_data = np.asarray(shape_tensor.data)
        if int(shape_data.size) != 3:
            return
        flat = shape_data.reshape(-1).copy()
        flat[[1, 2]] = flat[[2, 1]]
        swapped = flat.reshape(shape_data.shape).astype(shape_data.dtype, copy=False)

        shape_consumers = [int(v) for v in consumers.get(shape_name, [])]
        if any(int(v) != int(reshape_idx) for v in shape_consumers):
            cloned_name = _unique_tensor_name(f"{shape_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(shape_tensor.dtype),
                shape=[int(v) for v in list(swapped.shape)],
                shape_signature=[int(v) for v in list(swapped.shape)],
                data=np.asarray(swapped),
                is_variable=False,
                quantization=_clone_quantization(shape_tensor.quantization),
            )
            reshape_inputs = [str(v) for v in list(reshape_op.inputs)]
            reshape_inputs[1] = str(cloned_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=reshape_inputs,
            )
        else:
            shape_tensor.data = np.asarray(swapped)
            shape_tensor.shape = [int(v) for v in list(swapped.shape)]
            shape_tensor.shape_signature = [int(v) for v in list(swapped.shape)]

    def _maybe_swap_reshape_option_shape(reshape_op: OperatorIR) -> None:
        if not isinstance(reshape_op.options, dict):
            return
        opts = dict(reshape_op.options)
        changed = False
        for key in ["newShape", "onnxRawNewShape"]:
            value = opts.get(key, None)
            if not isinstance(value, list) or len(value) != 3:
                continue
            swapped = [value[0], value[2], value[1]]
            if swapped != value:
                opts[key] = [int(v) for v in swapped]
                changed = True
        if changed:
            reshape_op.options = opts

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            add_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
            if len(add_users) == 0:
                continue

            suffix_reshape_idx: Optional[int] = None
            suffix_transpose_idx: Optional[int] = None
            suffix_output_name: Optional[str] = None
            legacy_users: List[int] = []

            valid = True
            for user_idx in add_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "RESHAPE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == add_out_name
                ):
                    reshape_out_name = str(user_op.outputs[0])
                    reshape_users = [
                        int(v)
                        for v in consumers.get(reshape_out_name, [])
                        if int(v) != int(user_idx)
                    ]
                    if len(reshape_users) != 1:
                        legacy_users.append(int(user_idx))
                        continue
                    post_idx = int(reshape_users[0])
                    post_op = model_ir.operators[int(post_idx)]
                    if (
                        str(post_op.op_type) != "TRANSPOSE"
                        or len(post_op.inputs) < 2
                        or len(post_op.outputs) != 1
                        or str(post_op.inputs[0]) != reshape_out_name
                        or _read_transpose_perm(model_ir, post_op) != perm_3d_nchw_to_nhwc
                    ):
                        legacy_users.append(int(user_idx))
                        continue
                    post_out_name = str(post_op.outputs[0])
                    if post_out_name in model_outputs:
                        valid = False
                        break
                    if suffix_reshape_idx is not None:
                        valid = False
                        break
                    suffix_reshape_idx = int(user_idx)
                    suffix_transpose_idx = int(post_idx)
                    suffix_output_name = post_out_name
                else:
                    legacy_users.append(int(user_idx))
            if not valid or suffix_reshape_idx is None or suffix_transpose_idx is None or suffix_output_name is None:
                continue

            input_plans: List[Dict[str, Any]] = []
            can_rewrite_inputs = True
            for input_name in [str(v) for v in list(add_op.inputs)]:
                pre_idx = producers.get(input_name, None)
                if pre_idx is None:
                    can_rewrite_inputs = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != input_name
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    or input_name in model_outputs
                ):
                    can_rewrite_inputs = False
                    break

                input_users = [int(v) for v in consumers.get(input_name, [])]
                if int(add_idx) not in input_users:
                    can_rewrite_inputs = False
                    break
                pre_remove = [int(pre_idx)] if set(input_users) == {int(add_idx)} else []
                input_plans.append(
                    {
                        "nhwc_input_name": str(pre_op.inputs[0]),
                        "pre_remove_indices": pre_remove,
                    }
                )
            if not can_rewrite_inputs:
                continue

            add_nhwc_name = (
                _unique_tensor_name(f"{add_out_name}_nhwc")
                if len(legacy_users) > 0
                else str(add_out_name)
            )
            new_add_inputs = [str(plan["nhwc_input_name"]) for plan in input_plans]
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=new_add_inputs,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[add_nhwc_name],
            )

            old_add_tensor = model_ir.tensors.get(add_out_name, None)
            if add_nhwc_name != add_out_name:
                if old_add_tensor is not None and add_nhwc_name not in model_ir.tensors:
                    model_ir.tensors[add_nhwc_name] = TensorIR(
                        name=add_nhwc_name,
                        dtype=str(old_add_tensor.dtype),
                        shape=[int(v) for v in list(old_add_tensor.shape)],
                        shape_signature=(
                            [int(v) for v in list(old_add_tensor.shape_signature)]
                            if old_add_tensor.shape_signature is not None
                            else [int(v) for v in list(old_add_tensor.shape)]
                        ),
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(old_add_tensor.quantization),
                    )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_nhwc_name, None),
                perm_nchw_to_nhwc,
            )

            reshape_op = model_ir.operators[int(suffix_reshape_idx)]
            reshape_inputs = [str(v) for v in list(reshape_op.inputs)]
            reshape_inputs[0] = add_nhwc_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=reshape_inputs,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=reshape_op,
                new_outputs=[str(suffix_output_name)],
            )
            _maybe_swap_reshape_shape_tensor(
                reshape_op=reshape_op,
                reshape_idx=int(suffix_reshape_idx),
                consumers=consumers,
            )
            _maybe_swap_reshape_option_shape(reshape_op)

            remove_indices = [int(suffix_transpose_idx)]
            for plan in input_plans:
                remove_indices.extend([int(v) for v in list(plan.get("pre_remove_indices", []))])
            remove_indices = sorted(list(set(remove_indices)), reverse=True)
            removed_before_add = sum(1 for idx in remove_indices if int(idx) < int(add_idx))
            new_add_idx = int(add_idx) - int(removed_before_add)
            for remove_idx in remove_indices:
                if int(remove_idx) == int(add_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            if len(legacy_users) > 0:
                adapter_perm_name = _unique_tensor_name(f"{add_out_name}_adapter_perm")
                model_ir.tensors[adapter_perm_name] = TensorIR(
                    name=adapter_perm_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[add_nhwc_name, adapter_perm_name],
                    outputs=[add_out_name],
                    options={},
                )
                model_ir.operators.insert(int(new_add_idx) + 1, adapter_op)

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_add_reshape_transpose_suffix_nhwc_chains": int(rewritten)}


def _optimize_transpose_pre_add_mulconst_reshape_transpose_suffix_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW add bridges with mul-const inputs that feed [N,C,HW] reshape + [0,2,1] transpose.

    Target:
      a_nhwc --TRANSPOSE(0,3,1,2)--> a_nchw
      b_nhwc --TRANSPOSE(0,3,1,2)--> b_nchw --MUL(const)--> b_scaled_nchw
      ADD(a_nchw, b_scaled_nchw) -> y_nchw
      y_nchw --RESHAPE([N,C,HW])--> r --TRANSPOSE([0,2,1])--> z
      (optional legacy users consume y_nchw)

    Rewrite:
      ADD(a_nhwc, b_scaled_nhwc) -> y_nhwc
      y_nhwc --RESHAPE([N,HW,C])--> z
      (if legacy users exist, keep one adapter TRANSPOSE(0,3,1,2): y_nhwc -> y_nchw)
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_3d_nchw_to_nhwc = [0, 2, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _maybe_swap_reshape_shape_tensor(
        *,
        reshape_op: OperatorIR,
        reshape_idx: int,
        consumers: Dict[str, List[int]],
    ) -> None:
        if len(reshape_op.inputs) < 2:
            return
        shape_name = str(reshape_op.inputs[1])
        shape_tensor = model_ir.tensors.get(shape_name, None)
        if shape_tensor is None or shape_tensor.data is None:
            return
        shape_data = np.asarray(shape_tensor.data)
        if int(shape_data.size) != 3:
            return
        flat = shape_data.reshape(-1).copy()
        flat[[1, 2]] = flat[[2, 1]]
        swapped = flat.reshape(shape_data.shape).astype(shape_data.dtype, copy=False)

        shape_consumers = [int(v) for v in consumers.get(shape_name, [])]
        if any(int(v) != int(reshape_idx) for v in shape_consumers):
            cloned_name = _unique_tensor_name(f"{shape_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(shape_tensor.dtype),
                shape=[int(v) for v in list(swapped.shape)],
                shape_signature=[int(v) for v in list(swapped.shape)],
                data=np.asarray(swapped),
                is_variable=False,
                quantization=_clone_quantization(shape_tensor.quantization),
            )
            reshape_inputs = [str(v) for v in list(reshape_op.inputs)]
            reshape_inputs[1] = str(cloned_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=reshape_inputs,
            )
        else:
            shape_tensor.data = np.asarray(swapped)
            shape_tensor.shape = [int(v) for v in list(swapped.shape)]
            shape_tensor.shape_signature = [int(v) for v in list(swapped.shape)]

    def _maybe_swap_reshape_option_shape(reshape_op: OperatorIR) -> None:
        if not isinstance(reshape_op.options, dict):
            return
        opts = dict(reshape_op.options)
        changed = False
        for key in ["newShape", "onnxRawNewShape"]:
            value = opts.get(key, None)
            if not isinstance(value, list) or len(value) != 3:
                continue
            swapped = [value[0], value[2], value[1]]
            if swapped != value:
                opts[key] = [int(v) for v in swapped]
                changed = True
        if changed:
            reshape_op.options = opts

    def _analyze_input_plan(
        *,
        input_name: str,
        add_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        pre_idx = producers.get(str(input_name), None)
        if pre_idx is not None:
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) == "TRANSPOSE"
                and len(pre_op.inputs) >= 2
                and len(pre_op.outputs) == 1
                and str(pre_op.outputs[0]) == str(input_name)
                and _read_transpose_perm(model_ir, pre_op) == perm_nhwc_to_nchw
                and str(input_name) not in model_outputs
            ):
                users = [int(v) for v in consumers.get(str(input_name), [])]
                if int(add_idx) in users:
                    return {
                        "nhwc_input_name": str(pre_op.inputs[0]),
                        "pre_remove_indices": [int(pre_idx)] if set(users) == {int(add_idx)} else [],
                        "mul_plan": None,
                    }

        mul_idx = producers.get(str(input_name), None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        if str(mul_op.outputs[0]) != str(input_name):
            return None
        if str(input_name) in model_outputs:
            return None
        input_users = [int(v) for v in consumers.get(str(input_name), [])]
        if int(add_idx) not in input_users:
            return None

        data_input_index: Optional[int] = None
        side_input_index: Optional[int] = None
        side_input_name: Optional[str] = None
        data_pre_idx: Optional[int] = None
        data_pre_input_name: Optional[str] = None

        for cand_data_idx, cand_side_idx in [(0, 1), (1, 0)]:
            data_name = str(mul_op.inputs[int(cand_data_idx)])
            side_name = str(mul_op.inputs[int(cand_side_idx)])
            side_tensor = model_ir.tensors.get(side_name, None)
            if side_tensor is None or side_tensor.data is None:
                continue
            side_data = np.asarray(side_tensor.data)
            if int(side_data.size) != 1 and side_data.ndim not in {3, 4}:
                continue

            data_prod_idx = producers.get(data_name, None)
            if data_prod_idx is None:
                continue
            data_prod_op = model_ir.operators[int(data_prod_idx)]
            if (
                str(data_prod_op.op_type) != "TRANSPOSE"
                or len(data_prod_op.inputs) < 2
                or len(data_prod_op.outputs) != 1
                or str(data_prod_op.outputs[0]) != data_name
                or _read_transpose_perm(model_ir, data_prod_op) != perm_nhwc_to_nchw
                or str(data_name) in model_outputs
            ):
                continue
            if set(int(v) for v in consumers.get(data_name, [])) != {int(mul_idx)}:
                continue

            data_tensor = model_ir.tensors.get(data_name, None)
            pre_input_name = str(data_prod_op.inputs[0])
            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            if data_tensor is None or pre_input_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(list(data_tensor.shape))
                or not _is_fully_known_positive_shape(list(pre_input_tensor.shape))
                or len(list(data_tensor.shape)) != 4
                or len(list(pre_input_tensor.shape)) != 4
            ):
                continue

            if int(side_data.size) != 1:
                if _broadcast_static_shapes(list(data_tensor.shape), [int(v) for v in list(side_data.shape)]) is None:
                    continue

            data_input_index = int(cand_data_idx)
            side_input_index = int(cand_side_idx)
            side_input_name = str(side_name)
            data_pre_idx = int(data_prod_idx)
            data_pre_input_name = str(pre_input_name)
            break

        if (
            data_input_index is None
            or side_input_index is None
            or side_input_name is None
            or data_pre_idx is None
            or data_pre_input_name is None
        ):
            return None

        side_tensor = model_ir.tensors.get(side_input_name, None)
        if side_tensor is None or side_tensor.data is None:
            return None
        side_data = np.asarray(side_tensor.data)
        nhwc_side_data: Optional[np.ndarray] = None
        side_needs_update = False
        if int(side_data.size) != 1:
            target_shape = [int(v) for v in list(model_ir.tensors[str(data_pre_input_name)].shape)]
            side_shape = [int(v) for v in list(side_data.shape)]
            is_nchw_channelwise = (
                int(side_data.ndim) == 4
                and len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            )
            if is_nchw_channelwise:
                rotated = np.transpose(side_data, axes=perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is None:
                    return None
                nhwc_side_data = np.asarray(rotated)
                side_needs_update = not np.array_equal(nhwc_side_data, side_data)
            elif _broadcast_static_shapes(target_shape, side_shape) is not None:
                nhwc_side_data = np.asarray(side_data)
            else:
                rotated = np.asarray(side_data)
                found = False
                transpose_perm = perm_nchw_to_nhwc if int(rotated.ndim) == 4 else [1, 2, 0]
                max_rotate = 1 if int(rotated.ndim) == 4 else 3
                for _ in range(int(max_rotate)):
                    rotated = np.transpose(rotated, transpose_perm).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                        nhwc_side_data = np.asarray(rotated)
                        side_needs_update = True
                        found = True
                        break
                if not found:
                    return None

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        side_shared_outside_mul = any(int(v) != int(mul_idx) for v in side_users)

        return {
            "nhwc_input_name": str(input_name),
            "pre_remove_indices": [],
            "mul_plan": {
                "mul_idx": int(mul_idx),
                "mul_data_input_index": int(data_input_index),
                "mul_side_input_index": int(side_input_index),
                "mul_side_input_name": str(side_input_name),
                "mul_out_name": str(input_name),
                "data_pre_idx": int(data_pre_idx),
                "data_pre_input_name": str(data_pre_input_name),
                "side_needs_update": bool(side_needs_update),
                "side_shared_outside_mul": bool(side_shared_outside_mul),
                "side_nhwc_data": (None if nhwc_side_data is None else np.asarray(nhwc_side_data)),
            },
        }

    def _apply_mul_plan(*, plan: Dict[str, Any]) -> List[int]:
        mul_idx = int(plan["mul_idx"])
        mul_data_input_index = int(plan["mul_data_input_index"])
        mul_side_input_index = int(plan["mul_side_input_index"])
        mul_side_input_name = str(plan["mul_side_input_name"])
        mul_out_name = str(plan["mul_out_name"])
        data_pre_idx = int(plan["data_pre_idx"])
        data_pre_input_name = str(plan["data_pre_input_name"])

        side_input_name_for_mul = str(mul_side_input_name)
        if bool(plan.get("side_needs_update", False)):
            nhwc_data = np.asarray(plan.get("side_nhwc_data"))
            side_tensor = model_ir.tensors.get(str(mul_side_input_name), None)
            if side_tensor is None:
                return []
            if bool(plan.get("side_shared_outside_mul", False)):
                side_input_name_for_mul = _unique_tensor_name(f"{mul_side_input_name}_nhwc")
                model_ir.tensors[side_input_name_for_mul] = TensorIR(
                    name=side_input_name_for_mul,
                    dtype=str(side_tensor.dtype),
                    shape=[int(v) for v in list(nhwc_data.shape)],
                    shape_signature=[int(v) for v in list(nhwc_data.shape)],
                    data=np.asarray(nhwc_data),
                    is_variable=False,
                    quantization=_clone_quantization(side_tensor.quantization),
                )
            else:
                side_tensor.data = np.asarray(nhwc_data)
                side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
                side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        mul_op = model_ir.operators[int(mul_idx)]
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        if int(mul_data_input_index) < 0 or int(mul_data_input_index) >= len(mul_inputs):
            return []
        if int(mul_side_input_index) < 0 or int(mul_side_input_index) >= len(mul_inputs):
            return []
        mul_inputs[int(mul_data_input_index)] = str(data_pre_input_name)
        mul_inputs[int(mul_side_input_index)] = str(side_input_name_for_mul)
        _set_operator_inputs(
            model_ir=model_ir,
            op=mul_op,
            new_inputs=mul_inputs,
        )
        _permute_tensor_metadata_if_rank_matches(
            model_ir.tensors.get(mul_out_name, None),
            perm_nchw_to_nhwc,
        )
        return [int(data_pre_idx)]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for add_idx, add_op in enumerate(model_ir.operators):
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            add_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
            if len(add_users) == 0:
                continue

            suffix_reshape_idx: Optional[int] = None
            suffix_transpose_idx: Optional[int] = None
            suffix_output_name: Optional[str] = None
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in add_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "RESHAPE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == add_out_name
                ):
                    reshape_out_name = str(user_op.outputs[0])
                    reshape_users = [
                        int(v)
                        for v in consumers.get(reshape_out_name, [])
                        if int(v) != int(user_idx)
                    ]
                    if len(reshape_users) != 1:
                        legacy_users.append(int(user_idx))
                        continue
                    post_idx = int(reshape_users[0])
                    post_op = model_ir.operators[int(post_idx)]
                    if (
                        str(post_op.op_type) != "TRANSPOSE"
                        or len(post_op.inputs) < 2
                        or len(post_op.outputs) != 1
                        or str(post_op.inputs[0]) != reshape_out_name
                        or _read_transpose_perm(model_ir, post_op) != perm_3d_nchw_to_nhwc
                    ):
                        legacy_users.append(int(user_idx))
                        continue
                    post_out_name = str(post_op.outputs[0])
                    if post_out_name in model_outputs:
                        valid_users = False
                        break
                    if suffix_reshape_idx is not None:
                        valid_users = False
                        break
                    suffix_reshape_idx = int(user_idx)
                    suffix_transpose_idx = int(post_idx)
                    suffix_output_name = str(post_out_name)
                else:
                    legacy_users.append(int(user_idx))
            if not valid_users or suffix_reshape_idx is None or suffix_transpose_idx is None or suffix_output_name is None:
                continue

            input_plans: List[Dict[str, Any]] = []
            rewritable = True
            for input_name in [str(v) for v in list(add_op.inputs)]:
                plan = _analyze_input_plan(
                    input_name=input_name,
                    add_idx=int(add_idx),
                    producers=producers,
                    consumers=consumers,
                    model_outputs=model_outputs,
                )
                if plan is None:
                    rewritable = False
                    break
                input_plans.append(dict(plan))
            if not rewritable:
                continue

            pre_remove_indices: List[int] = []
            for plan in input_plans:
                mul_plan = plan.get("mul_plan", None)
                if mul_plan is not None:
                    pre_remove_indices.extend(_apply_mul_plan(plan=dict(mul_plan)))
                pre_remove_indices.extend([int(v) for v in list(plan.get("pre_remove_indices", []))])

            add_nhwc_name = (
                _unique_tensor_name(f"{add_out_name}_nhwc")
                if len(legacy_users) > 0
                else str(add_out_name)
            )
            new_add_inputs = [str(plan["nhwc_input_name"]) for plan in input_plans]
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=new_add_inputs,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[add_nhwc_name],
            )

            old_add_tensor = model_ir.tensors.get(add_out_name, None)
            if add_nhwc_name != add_out_name:
                if old_add_tensor is not None and add_nhwc_name not in model_ir.tensors:
                    model_ir.tensors[add_nhwc_name] = TensorIR(
                        name=add_nhwc_name,
                        dtype=str(old_add_tensor.dtype),
                        shape=[int(v) for v in list(old_add_tensor.shape)],
                        shape_signature=(
                            [int(v) for v in list(old_add_tensor.shape_signature)]
                            if old_add_tensor.shape_signature is not None
                            else [int(v) for v in list(old_add_tensor.shape)]
                        ),
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(old_add_tensor.quantization),
                    )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add_nhwc_name, None),
                perm_nchw_to_nhwc,
            )

            reshape_op = model_ir.operators[int(suffix_reshape_idx)]
            reshape_inputs = [str(v) for v in list(reshape_op.inputs)]
            reshape_inputs[0] = add_nhwc_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=reshape_inputs,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=reshape_op,
                new_outputs=[str(suffix_output_name)],
            )
            _maybe_swap_reshape_shape_tensor(
                reshape_op=reshape_op,
                reshape_idx=int(suffix_reshape_idx),
                consumers=consumers,
            )
            _maybe_swap_reshape_option_shape(reshape_op)

            remove_indices = [int(suffix_transpose_idx)]
            remove_indices.extend([int(v) for v in pre_remove_indices])
            remove_indices = sorted(list(set(remove_indices)), reverse=True)
            removed_before_add = sum(1 for idx in remove_indices if int(idx) < int(add_idx))
            new_add_idx = int(add_idx) - int(removed_before_add)
            for remove_idx in remove_indices:
                if int(remove_idx) == int(add_idx):
                    continue
                del model_ir.operators[int(remove_idx)]

            if len(legacy_users) > 0:
                adapter_perm_name = _unique_tensor_name(f"{add_out_name}_adapter_perm")
                model_ir.tensors[adapter_perm_name] = TensorIR(
                    name=adapter_perm_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[add_nhwc_name, adapter_perm_name],
                    outputs=[add_out_name],
                    options={},
                )
                model_ir.operators.insert(int(new_add_idx) + 1, adapter_op)

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_pre_add_mulconst_reshape_transpose_suffix_nhwc_chains": int(rewritten)
    }


def _optimize_transpose_pre_unary_reshape_transpose_suffix_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NHWC->NCHW unary wrappers that immediately feed
    [N,C,HW] reshape + [0,2,1] transpose suffixes.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --UNARY(1-input)-> u_nchw
      u_nchw --RESHAPE([N,C,HW])--> r
      r --TRANSPOSE([0,2,1])--> z

    Rewrite:
      x_nhwc --UNARY(1-input)-> u_nhwc
      u_nhwc --RESHAPE([N,HW,C])--> z

    Extended (pseudo-unary Swish):
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      s_nchw = LOGISTIC(x_nchw)
      y_nchw = MUL(x_nchw, s_nchw)
      y_nchw --RESHAPE([N,C,HW])--> r
      r --TRANSPOSE([0,2,1])--> z

    Rewrite:
      s_nhwc = LOGISTIC(x_nhwc)
      y_nhwc = MUL(x_nhwc, s_nhwc)
      y_nhwc --RESHAPE([N,HW,C])--> z
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_3d_nchw_to_nhwc = [0, 2, 1]
    unary_ops = {
        "RELU",
        "RELU6",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
    }

    def _swap_reshape_shape_tensor(reshape_op: OperatorIR) -> bool:
        if len(reshape_op.inputs) < 2:
            return False
        shape_name = str(reshape_op.inputs[1])
        shape_tensor = model_ir.tensors.get(shape_name, None)
        shape_vals = _read_const_ints_from_tensor(shape_tensor)
        if shape_vals is None or len(shape_vals) != 3:
            return False
        swapped = [int(shape_vals[0]), int(shape_vals[2]), int(shape_vals[1])]
        return _write_const_ints_to_tensor(shape_tensor, swapped)

    def _swap_reshape_option_shape(reshape_op: OperatorIR) -> bool:
        if not isinstance(reshape_op.options, dict):
            return False
        opts = dict(reshape_op.options)
        changed = False
        for key in ["newShape", "onnxRawNewShape"]:
            value = opts.get(key, None)
            if not isinstance(value, list) or len(value) != 3:
                continue
            swapped = [int(value[0]), int(value[2]), int(value[1])]
            if swapped != [int(v) for v in value]:
                opts[key] = swapped
                changed = True
        if changed:
            reshape_op.options = opts
        return changed

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape_idx, reshape_op in enumerate(model_ir.operators):
            if str(reshape_op.op_type) != "RESHAPE" or len(reshape_op.inputs) < 1 or len(reshape_op.outputs) != 1:
                continue

            reshape_in_name = str(reshape_op.inputs[0])
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_in_name in model_outputs:
                continue

            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) != 1:
                continue
            post_idx = int(reshape_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != reshape_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_3d_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue

            unary_idx = producers.get(reshape_in_name, None)
            if unary_idx is None:
                continue
            unary_op = model_ir.operators[int(unary_idx)]
            rewrite_kind = ""
            swish_log_idx: Optional[int] = None
            swish_log_op: Optional[OperatorIR] = None
            swish_log_out_name = ""
            swish_mul_data_input_index: Optional[int] = None

            if (
                str(unary_op.op_type) in unary_ops
                and len(unary_op.inputs) == 1
                and len(unary_op.outputs) == 1
                and str(unary_op.outputs[0]) == reshape_in_name
            ):
                rewrite_kind = "unary"
                pre_out_name = str(unary_op.inputs[0])
            elif (
                str(unary_op.op_type) == "MUL"
                and len(unary_op.inputs) == 2
                and len(unary_op.outputs) == 1
                and str(unary_op.outputs[0]) == reshape_in_name
            ):
                mul_inputs = [str(v) for v in unary_op.inputs]
                for log_input_index in [0, 1]:
                    data_input_index = 1 - int(log_input_index)
                    candidate_log_out_name = str(mul_inputs[log_input_index])
                    candidate_pre_out_name = str(mul_inputs[data_input_index])
                    candidate_log_idx = producers.get(candidate_log_out_name, None)
                    if candidate_log_idx is None:
                        continue
                    candidate_log_op = model_ir.operators[int(candidate_log_idx)]
                    if (
                        str(candidate_log_op.op_type) != "LOGISTIC"
                        or len(candidate_log_op.inputs) != 1
                        or len(candidate_log_op.outputs) != 1
                        or str(candidate_log_op.outputs[0]) != candidate_log_out_name
                        or str(candidate_log_op.inputs[0]) != candidate_pre_out_name
                    ):
                        continue
                    candidate_log_users = [int(v) for v in consumers.get(candidate_log_out_name, [])]
                    if set(candidate_log_users) != {int(unary_idx)}:
                        continue
                    rewrite_kind = "swish"
                    pre_out_name = candidate_pre_out_name
                    swish_log_idx = int(candidate_log_idx)
                    swish_log_op = candidate_log_op
                    swish_log_out_name = candidate_log_out_name
                    swish_mul_data_input_index = int(data_input_index)
                    break
            else:
                continue

            if rewrite_kind == "":
                continue

            pre_idx = producers.get(pre_out_name, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != pre_out_name
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_in_name = str(pre_op.inputs[0])

            if (
                pre_in_name in model_outputs
                or pre_out_name in model_outputs
                or reshape_in_name in model_outputs
            ):
                continue

            if rewrite_kind == "unary":
                if set(int(v) for v in consumers.get(pre_out_name, [])) != {int(unary_idx)}:
                    continue
            elif rewrite_kind == "swish":
                expected_pre_users = {int(unary_idx), int(swish_log_idx)} if swish_log_idx is not None else {int(unary_idx)}
                if set(int(v) for v in consumers.get(pre_out_name, [])) != expected_pre_users:
                    continue
                if swish_log_out_name == "":
                    continue
                if set(int(v) for v in consumers.get(swish_log_out_name, [])) != {int(unary_idx)}:
                    continue
            else:
                continue
            if set(int(v) for v in consumers.get(reshape_in_name, [])) != {int(reshape_idx)}:
                continue

            shape_changed = _swap_reshape_shape_tensor(reshape_op)
            opts_changed = _swap_reshape_option_shape(reshape_op)
            if not shape_changed and not opts_changed:
                continue

            if rewrite_kind == "unary":
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_inputs=[pre_in_name],
                )
            else:
                if swish_log_op is None or swish_mul_data_input_index is None:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=swish_log_op,
                    new_inputs=[pre_in_name],
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=unary_op,
                    input_index=int(swish_mul_data_input_index),
                    new_input_name=pre_in_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(swish_log_out_name, None),
                    perm_nchw_to_nhwc,
                )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(reshape_in_name, None),
                perm_nchw_to_nhwc,
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=reshape_op,
                new_outputs=[post_out_name],
            )

            old_reshape_out_tensor = model_ir.tensors.get(reshape_out_name, None)
            canonical_tensor = model_ir.tensors.get(post_out_name, None)
            if old_reshape_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_reshape_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_reshape_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_reshape_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_reshape_out_tensor.shape_signature)]
                    if old_reshape_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_reshape_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_3d_nchw_to_nhwc,
                )

            remove_indices: List[int] = []
            pre_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is pre_op), None)
            if pre_remove_idx is not None:
                remove_indices.append(int(pre_remove_idx))
            post_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is post_op), None)
            if post_remove_idx is not None:
                remove_indices.append(int(post_remove_idx))
            for remove_idx in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_unary_reshape_transpose_suffix_nhwc_chains": int(rewritten)}


def _optimize_transpose_pre_unary_squeeze_transpose_suffix_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NHWC->NCHW unary/Swish wrappers that feed SQUEEZE + [0,2,1] transpose suffix.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --(UNARY | Swish(LOGISTIC+MUL))-> y_nchw
      y_nchw --SQUEEZE(axis in {2,3})--> s
      s --TRANSPOSE([0,2,1])--> z

    Rewrite:
      x_nhwc --(same UNARY/Swish)-> y_nhwc
      y_nhwc --SQUEEZE(mapped axis)-> z
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_3d_nchw_to_nhwc = [0, 2, 1]
    unary_ops = {
        "RELU",
        "RELU6",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
    }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for squeeze_idx, squeeze_op in enumerate(model_ir.operators):
            if str(squeeze_op.op_type) != "SQUEEZE" or len(squeeze_op.inputs) != 1 or len(squeeze_op.outputs) != 1:
                continue

            squeeze_in_name = str(squeeze_op.inputs[0])
            squeeze_out_name = str(squeeze_op.outputs[0])
            if squeeze_in_name in model_outputs or squeeze_out_name in model_outputs:
                continue

            squeeze_users = [int(v) for v in consumers.get(squeeze_out_name, [])]
            if len(squeeze_users) != 1:
                continue
            post_idx = int(squeeze_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != squeeze_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_3d_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])

            unary_idx = producers.get(squeeze_in_name, None)
            if unary_idx is None:
                continue
            unary_op = model_ir.operators[int(unary_idx)]

            rewrite_kind = ""
            swish_log_idx: Optional[int] = None
            swish_log_op: Optional[OperatorIR] = None
            swish_log_out_name = ""
            swish_mul_data_input_index: Optional[int] = None

            if (
                str(unary_op.op_type) in unary_ops
                and len(unary_op.inputs) == 1
                and len(unary_op.outputs) == 1
                and str(unary_op.outputs[0]) == squeeze_in_name
            ):
                rewrite_kind = "unary"
                pre_out_name = str(unary_op.inputs[0])
            elif (
                str(unary_op.op_type) == "MUL"
                and len(unary_op.inputs) == 2
                and len(unary_op.outputs) == 1
                and str(unary_op.outputs[0]) == squeeze_in_name
            ):
                mul_inputs = [str(v) for v in unary_op.inputs]
                for log_input_index in [0, 1]:
                    data_input_index = 1 - int(log_input_index)
                    candidate_log_out_name = str(mul_inputs[log_input_index])
                    candidate_pre_out_name = str(mul_inputs[data_input_index])
                    candidate_log_idx = producers.get(candidate_log_out_name, None)
                    if candidate_log_idx is None:
                        continue
                    candidate_log_op = model_ir.operators[int(candidate_log_idx)]
                    if (
                        str(candidate_log_op.op_type) != "LOGISTIC"
                        or len(candidate_log_op.inputs) != 1
                        or len(candidate_log_op.outputs) != 1
                        or str(candidate_log_op.outputs[0]) != candidate_log_out_name
                        or str(candidate_log_op.inputs[0]) != candidate_pre_out_name
                    ):
                        continue
                    if set(int(v) for v in consumers.get(candidate_log_out_name, [])) != {int(unary_idx)}:
                        continue
                    rewrite_kind = "swish"
                    pre_out_name = candidate_pre_out_name
                    swish_log_idx = int(candidate_log_idx)
                    swish_log_op = candidate_log_op
                    swish_log_out_name = candidate_log_out_name
                    swish_mul_data_input_index = int(data_input_index)
                    break
            else:
                continue

            if rewrite_kind == "":
                continue

            pre_idx = producers.get(pre_out_name, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != pre_out_name
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_in_name = str(pre_op.inputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            if rewrite_kind == "unary":
                if set(int(v) for v in consumers.get(pre_out_name, [])) != {int(unary_idx)}:
                    continue
            elif rewrite_kind == "swish":
                expected_pre_users = {int(unary_idx), int(swish_log_idx)} if swish_log_idx is not None else {int(unary_idx)}
                if set(int(v) for v in consumers.get(pre_out_name, [])) != expected_pre_users:
                    continue
                if swish_log_out_name == "":
                    continue
                if set(int(v) for v in consumers.get(swish_log_out_name, [])) != {int(unary_idx)}:
                    continue
            else:
                continue

            if set(int(v) for v in consumers.get(squeeze_in_name, [])) != {int(squeeze_idx)}:
                continue
            if set(int(v) for v in consumers.get(squeeze_out_name, [])) != {int(post_idx)}:
                continue

            squeeze_in_tensor = model_ir.tensors.get(squeeze_in_name, None)
            squeeze_out_tensor = model_ir.tensors.get(squeeze_out_name, None)
            if (
                squeeze_in_tensor is None
                or squeeze_out_tensor is None
                or not _is_fully_known_positive_shape(squeeze_in_tensor.shape)
                or not _is_fully_known_positive_shape(squeeze_out_tensor.shape)
            ):
                continue
            squeeze_in_shape = [int(v) for v in list(squeeze_in_tensor.shape)]
            squeeze_out_shape = [int(v) for v in list(squeeze_out_tensor.shape)]
            if len(squeeze_in_shape) != 4 or len(squeeze_out_shape) != 3:
                continue

            squeeze_axis_candidates: List[int] = []
            for axis in range(4):
                if int(squeeze_in_shape[axis]) != 1:
                    continue
                candidate_out_shape = [int(v) for i, v in enumerate(squeeze_in_shape) if int(i) != int(axis)]
                if candidate_out_shape == squeeze_out_shape:
                    squeeze_axis_candidates.append(int(axis))
            if len(squeeze_axis_candidates) != 1:
                continue

            squeeze_axis_nchw = int(squeeze_axis_candidates[0])
            if squeeze_axis_nchw == 2:
                squeeze_axis_nhwc = 1
            elif squeeze_axis_nchw == 3:
                squeeze_axis_nhwc = 2
            else:
                continue

            squeeze_in_shape_nhwc = _permute_shape(squeeze_in_shape, perm_nchw_to_nhwc)
            if (
                squeeze_in_shape_nhwc is None
                or len(squeeze_in_shape_nhwc) != 4
                or int(squeeze_in_shape_nhwc[squeeze_axis_nhwc]) != 1
            ):
                continue

            if rewrite_kind == "unary":
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_inputs=[pre_in_name],
                )
            else:
                if swish_log_op is None or swish_mul_data_input_index is None:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=swish_log_op,
                    new_inputs=[pre_in_name],
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=unary_op,
                    input_index=int(swish_mul_data_input_index),
                    new_input_name=pre_in_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(swish_log_out_name, None),
                    perm_nchw_to_nhwc,
                )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(squeeze_in_name, None),
                perm_nchw_to_nhwc,
            )

            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            squeeze_options["squeezeDims"] = [int(squeeze_axis_nhwc)]
            squeeze_op.options = squeeze_options

            _set_operator_outputs(
                model_ir=model_ir,
                op=squeeze_op,
                new_outputs=[post_out_name],
            )

            old_squeeze_out_tensor = model_ir.tensors.get(squeeze_out_name, None)
            canonical_tensor = model_ir.tensors.get(post_out_name, None)
            if old_squeeze_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_squeeze_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_squeeze_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_squeeze_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_squeeze_out_tensor.shape_signature)]
                    if old_squeeze_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_squeeze_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_3d_nchw_to_nhwc,
                )

            remove_indices: List[int] = []
            pre_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is pre_op), None)
            if pre_remove_idx is not None:
                remove_indices.append(int(pre_remove_idx))
            post_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is post_op), None)
            if post_remove_idx is not None:
                remove_indices.append(int(post_remove_idx))
            for remove_idx in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_unary_squeeze_transpose_suffix_nhwc_chains": int(rewritten)}


def _optimize_transpose_reshape_transpose_to_expanddims_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse NHWC->NCHW->RESHAPE->N?HWC wrappers into a single RESHAPE.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --RESHAPE([N,1,C,H,W])--> r
      r      --TRANSPOSE(0,1,3,4,2)--> y

    Rewrite:
      x_nhwc --RESHAPE([N,1,H,W,C])--> y

    This pattern appears in some YOLO heads where layout wrappers are only used
    to insert a singleton anchor axis before 5D decode operations.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_n1chw_to_n1hwc = [0, 1, 3, 4, 2]
    perm_nchw1_to_nhwc1 = [0, 2, 3, 1, 4]
    perm_nhwab_to_nahwb = [0, 3, 1, 2, 4]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) != 1:
                continue
            reshape_idx = int(pre_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 1
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != pre_out_name
            ):
                continue

            reshape_out_name = str(reshape_op.outputs[0])
            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) != 1:
                continue
            post_idx = int(reshape_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != reshape_out_name
            ):
                continue
            post_perm = _read_transpose_perm(model_ir, post_op)
            if post_perm not in [perm_n1chw_to_n1hwc, perm_nchw1_to_nhwc1]:
                continue

            post_out_name = str(post_op.outputs[0])
            if pre_out_name in model_outputs or reshape_out_name in model_outputs:
                continue
            input_tensor = model_ir.tensors.get(pre_in_name, None)
            pre_out_tensor = model_ir.tensors.get(pre_out_name, None)
            reshape_out_tensor = model_ir.tensors.get(reshape_out_name, None)
            post_out_tensor = model_ir.tensors.get(post_out_name, None)
            if (
                input_tensor is None
                or pre_out_tensor is None
                or reshape_out_tensor is None
                or post_out_tensor is None
            ):
                continue

            in_shape = [int(v) for v in list(input_tensor.shape)]
            pre_out_shape = [int(v) for v in list(pre_out_tensor.shape)]
            reshape_out_shape = [int(v) for v in list(reshape_out_tensor.shape)]
            post_out_shape = [int(v) for v in list(post_out_tensor.shape)]
            if len(in_shape) != 4 or len(pre_out_shape) != 4 or len(reshape_out_shape) != 5 or len(post_out_shape) != 5:
                continue

            n, h, w, c = in_shape
            expected_pre_out = [n, c, h, w]
            if not _shapes_match_if_known(pre_out_shape, expected_pre_out):
                continue
            if post_perm == perm_n1chw_to_n1hwc:
                expected_reshape_out = [n, 1, c, h, w]
                expected_post_out = [n, 1, h, w, c]
                reshape_singleton_axis = 1
                post_singleton_axis = 1
            elif post_perm == perm_nchw1_to_nhwc1:
                expected_reshape_out = [n, c, h, w, 1]
                expected_post_out = [n, h, w, c, 1]
                reshape_singleton_axis = 4
                post_singleton_axis = 4
            else:
                continue

            # Case A:
            #   NHWC -> NCHW -> RESHAPE(insert singleton) -> transpose back
            #   => NHWC -> RESHAPE(insert singleton in NHWC order)
            if (
                _shapes_match_if_known(reshape_out_shape, expected_reshape_out)
                and _shapes_match_if_known(post_out_shape, expected_post_out)
                and _dims_compatible(int(reshape_out_shape[int(reshape_singleton_axis)]), 1)
                and _dims_compatible(int(post_out_shape[int(post_singleton_axis)]), 1)
            ):
                target_shape = [int(v) for v in list(post_out_shape)]
                if len(reshape_op.inputs) >= 2:
                    shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
                    if _read_const_ints_from_tensor(shape_tensor) is None:
                        continue
                    if not _write_const_ints_to_tensor(shape_tensor, target_shape):
                        continue
                if isinstance(reshape_op.options, dict):
                    reshape_opts = dict(reshape_op.options)
                    for key in ["newShape", "onnxRawNewShape"]:
                        if isinstance(reshape_opts.get(key, None), list):
                            reshape_opts[key] = [int(v) for v in target_shape]
                    reshape_op.options = reshape_opts

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=reshape_op,
                    new_inputs=[pre_in_name] + [str(v) for v in list(reshape_op.inputs[1:])],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=reshape_op,
                    new_outputs=[post_out_name],
                )

                for remove_idx in sorted([int(post_idx), int(pre_idx)], reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            # Case B (YOLO decode style):
            #   NHWC -> NCHW -> RESHAPE([N,A,B,H,W]) -> TRANSPOSE([0,1,3,4,2]) -> [N,A,H,W,B]
            #   => NHWC -> RESHAPE([N,H,W,A,B]) -> TRANSPOSE([0,3,1,2,4]) -> [N,A,H,W,B]
            if post_perm == perm_n1chw_to_n1hwc:
                a = int(reshape_out_shape[1])
                b = int(reshape_out_shape[2])
                rh = int(reshape_out_shape[3])
                rw = int(reshape_out_shape[4])
                expected_post_out_factorized = [n, a, h, w, b]
                if not _shapes_match_if_known(post_out_shape, expected_post_out_factorized):
                    continue
                if not _dims_compatible(rh, h) or not _dims_compatible(rw, w):
                    continue
                if int(c) > 0 and int(a) > 0 and int(b) > 0 and int(c) != int(a) * int(b):
                    continue

                target_reshape_shape = [int(n), int(h), int(w), int(a), int(b)]
                if len(reshape_op.inputs) < 2:
                    continue
                reshape_shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
                if _read_const_ints_from_tensor(reshape_shape_tensor) is None:
                    continue
                if not _write_const_ints_to_tensor(reshape_shape_tensor, target_reshape_shape):
                    continue

                if isinstance(reshape_op.options, dict):
                    reshape_opts = dict(reshape_op.options)
                    for key in ["newShape", "onnxRawNewShape"]:
                        if isinstance(reshape_opts.get(key, None), list):
                            reshape_opts[key] = [int(v) for v in target_reshape_shape]
                    reshape_op.options = reshape_opts

                if len(post_op.inputs) < 2:
                    continue
                post_perm_tensor = model_ir.tensors.get(str(post_op.inputs[1]), None)
                if _read_const_ints_from_tensor(post_perm_tensor) is None:
                    continue
                if not _write_const_ints_to_tensor(post_perm_tensor, perm_nhwab_to_nahwb):
                    continue

                reshape_shape_norm, reshape_signature_norm = normalize_onnx_shape(target_reshape_shape)
                reshape_out_tensor.shape = [int(v) for v in reshape_shape_norm]
                reshape_out_tensor.shape_signature = [int(v) for v in reshape_signature_norm]

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=reshape_op,
                    new_inputs=[pre_in_name] + [str(v) for v in list(reshape_op.inputs[1:])],
                )
                del model_ir.operators[int(pre_idx)]

                rewritten += 1
                changed = True
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_reshape_transpose_to_expanddims_nhwc_chains": int(rewritten)
    }


def _optimize_transpose_reshape_transpose_to_flatten_hw_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse NHWC->NCHW->RESHAPE->TRANSPOSE flatten wrappers into one RESHAPE.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --RESHAPE([N,C,H*W])--> r
      r      --TRANSPOSE(0,2,1)--> y  (shape [N,H*W,C])

    Rewrite:
      x_nhwc --RESHAPE([N,H*W,C])--> y
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_ncw_to_nwc = [0, 2, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) != 1:
                continue
            reshape_idx = int(pre_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != pre_out_name
            ):
                continue

            reshape_out_name = str(reshape_op.outputs[0])
            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) != 1:
                continue
            post_idx = int(reshape_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != reshape_out_name
            ):
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_ncw_to_nwc:
                continue

            post_out_name = str(post_op.outputs[0])
            if pre_out_name in model_outputs or reshape_out_name in model_outputs:
                continue

            in_tensor = model_ir.tensors.get(pre_in_name, None)
            pre_out_tensor = model_ir.tensors.get(pre_out_name, None)
            reshape_out_tensor = model_ir.tensors.get(reshape_out_name, None)
            post_out_tensor = model_ir.tensors.get(post_out_name, None)
            if (
                in_tensor is None
                or pre_out_tensor is None
                or reshape_out_tensor is None
                or post_out_tensor is None
            ):
                continue

            if (
                not _is_fully_known_positive_shape(in_tensor.shape)
                or not _is_fully_known_positive_shape(pre_out_tensor.shape)
                or not _is_fully_known_positive_shape(reshape_out_tensor.shape)
                or not _is_fully_known_positive_shape(post_out_tensor.shape)
            ):
                continue

            in_shape = [int(v) for v in list(in_tensor.shape)]
            pre_out_shape = [int(v) for v in list(pre_out_tensor.shape)]
            reshape_out_shape = [int(v) for v in list(reshape_out_tensor.shape)]
            post_out_shape = [int(v) for v in list(post_out_tensor.shape)]
            if len(in_shape) != 4 or len(pre_out_shape) != 4 or len(reshape_out_shape) != 3 or len(post_out_shape) != 3:
                continue

            n, h, w, c = in_shape
            if pre_out_shape != [n, c, h, w]:
                continue
            hw = int(h) * int(w)
            if reshape_out_shape != [n, c, hw]:
                continue
            if post_out_shape != [n, hw, c]:
                continue

            shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
            if _read_const_ints_from_tensor(shape_tensor) is None:
                continue
            target_shape = [int(n), int(hw), int(c)]
            if not _write_const_ints_to_tensor(shape_tensor, target_shape):
                continue

            if isinstance(reshape_op.options, dict):
                reshape_opts = dict(reshape_op.options)
                for key in ["newShape", "onnxRawNewShape"]:
                    if isinstance(reshape_opts.get(key, None), list):
                        reshape_opts[key] = [int(v) for v in target_shape]
                reshape_op.options = reshape_opts

            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=[pre_in_name] + [str(v) for v in list(reshape_op.inputs[1:])],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=reshape_op,
                new_outputs=[post_out_name],
            )

            for remove_idx in sorted([int(post_idx), int(pre_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_reshape_transpose_to_flatten_hw_nhwc_chains": int(rewritten)
    }


def _optimize_reshape_transpose_reshape_transpose_to_nhwc_reshape_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse rank-3 -> rank-4 layout shims into a single NHWC reshape.

    Target:
      x --RESHAPE([N,S,C,1])--> r1
        --TRANSPOSE([0,3,2,1])--> t1
        --RESHAPE([N,C,H,W])--> r2
        --TRANSPOSE([0,2,3,1])--> y

    Rewrite:
      x --RESHAPE([N,H,W,C])--> y

    with S == H * W.
    """
    rewritten = 0
    perm_nsc1_to_n1cs = [0, 3, 2, 1]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        serial = 1
        while name in model_ir.tensors:
            name = f"{base}_{serial}"
            serial += 1
        return name

    def _set_reshape_shape_input(
        *,
        reshape_op: OperatorIR,
        reshape_idx: int,
        consumers: Dict[str, List[int]],
        target_shape: List[int],
    ) -> bool:
        if len(reshape_op.inputs) < 2:
            return False
        shape_name = str(reshape_op.inputs[1])
        shape_tensor = model_ir.tensors.get(shape_name, None)
        if shape_tensor is None:
            return False
        shape_users = set(int(v) for v in consumers.get(shape_name, []))
        target = [int(v) for v in list(target_shape)]
        if shape_users == {int(reshape_idx)}:
            return _write_const_ints_to_tensor(shape_tensor, target)

        np_dtype = np.int32
        dtype_name = str(shape_tensor.dtype)
        if shape_tensor.data is not None:
            try:
                np_dtype = np.asarray(shape_tensor.data).dtype
            except Exception:
                np_dtype = np.int32
        new_shape_name = _unique_tensor_name(f"{shape_name}_nhwc")
        model_ir.tensors[new_shape_name] = TensorIR(
            name=new_shape_name,
            dtype=dtype_name,
            shape=[int(len(target))],
            shape_signature=[int(len(target))],
            data=np.asarray(target, dtype=np_dtype),
            is_variable=False,
            quantization=_clone_quantization(shape_tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=reshape_op,
            input_index=1,
            new_input_name=new_shape_name,
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape1_idx, reshape1_op in enumerate(model_ir.operators):
            if str(reshape1_op.op_type) != "RESHAPE" or len(reshape1_op.inputs) < 2 or len(reshape1_op.outputs) != 1:
                continue

            input_name = str(reshape1_op.inputs[0])
            reshape1_out_name = str(reshape1_op.outputs[0])
            if reshape1_out_name in model_outputs:
                continue

            reshape1_users = [int(v) for v in consumers.get(reshape1_out_name, [])]
            if len(reshape1_users) != 1:
                continue
            transpose1_idx = int(reshape1_users[0])
            transpose1_op = model_ir.operators[int(transpose1_idx)]
            if (
                str(transpose1_op.op_type) != "TRANSPOSE"
                or len(transpose1_op.inputs) < 2
                or len(transpose1_op.outputs) != 1
                or str(transpose1_op.inputs[0]) != reshape1_out_name
                or _read_transpose_perm(model_ir, transpose1_op) != perm_nsc1_to_n1cs
            ):
                continue

            transpose1_out_name = str(transpose1_op.outputs[0])
            if transpose1_out_name in model_outputs:
                continue
            transpose1_users = [int(v) for v in consumers.get(transpose1_out_name, [])]
            if len(transpose1_users) != 1:
                continue
            reshape2_idx = int(transpose1_users[0])
            reshape2_op = model_ir.operators[int(reshape2_idx)]
            if (
                str(reshape2_op.op_type) != "RESHAPE"
                or len(reshape2_op.inputs) < 2
                or len(reshape2_op.outputs) != 1
                or str(reshape2_op.inputs[0]) != transpose1_out_name
            ):
                continue

            reshape2_out_name = str(reshape2_op.outputs[0])
            if reshape2_out_name in model_outputs:
                continue
            reshape2_users = [int(v) for v in consumers.get(reshape2_out_name, [])]
            if len(reshape2_users) != 1:
                continue
            transpose2_idx = int(reshape2_users[0])
            transpose2_op = model_ir.operators[int(transpose2_idx)]
            if (
                str(transpose2_op.op_type) != "TRANSPOSE"
                or len(transpose2_op.inputs) < 2
                or len(transpose2_op.outputs) != 1
                or str(transpose2_op.inputs[0]) != reshape2_out_name
                or _read_transpose_perm(model_ir, transpose2_op) != perm_nchw_to_nhwc
            ):
                continue

            output_name = str(transpose2_op.outputs[0])
            in_tensor = model_ir.tensors.get(input_name, None)
            reshape1_out_tensor = model_ir.tensors.get(reshape1_out_name, None)
            transpose1_out_tensor = model_ir.tensors.get(transpose1_out_name, None)
            reshape2_out_tensor = model_ir.tensors.get(reshape2_out_name, None)
            output_tensor = model_ir.tensors.get(output_name, None)
            if (
                in_tensor is None
                or reshape1_out_tensor is None
                or transpose1_out_tensor is None
                or reshape2_out_tensor is None
                or output_tensor is None
            ):
                continue
            if (
                not _is_fully_known_positive_shape(in_tensor.shape)
                or not _is_fully_known_positive_shape(reshape1_out_tensor.shape)
                or not _is_fully_known_positive_shape(transpose1_out_tensor.shape)
                or not _is_fully_known_positive_shape(reshape2_out_tensor.shape)
                or not _is_fully_known_positive_shape(output_tensor.shape)
            ):
                continue

            in_shape = [int(v) for v in list(in_tensor.shape)]
            reshape1_shape = [int(v) for v in list(reshape1_out_tensor.shape)]
            transpose1_shape = [int(v) for v in list(transpose1_out_tensor.shape)]
            reshape2_shape = [int(v) for v in list(reshape2_out_tensor.shape)]
            output_shape = [int(v) for v in list(output_tensor.shape)]
            if (
                len(in_shape) != 3
                or len(reshape1_shape) != 4
                or len(transpose1_shape) != 4
                or len(reshape2_shape) != 4
                or len(output_shape) != 4
            ):
                continue

            n, s, c = [int(v) for v in list(in_shape)]
            if reshape1_shape != [int(n), int(s), int(c), 1]:
                continue
            if transpose1_shape != [int(n), 1, int(c), int(s)]:
                continue

            n2, c2, h, w = [int(v) for v in list(reshape2_shape)]
            if n2 != n or c2 != c:
                continue
            if int(h) <= 0 or int(w) <= 0 or int(s) != int(h) * int(w):
                continue
            if output_shape != [int(n), int(h), int(w), int(c)]:
                continue

            target_shape = [int(n), int(h), int(w), int(c)]
            if not _set_reshape_shape_input(
                reshape_op=reshape1_op,
                reshape_idx=int(reshape1_idx),
                consumers=consumers,
                target_shape=target_shape,
            ):
                continue

            if isinstance(reshape1_op.options, dict):
                reshape_opts = dict(reshape1_op.options)
                for key in ["newShape", "onnxRawNewShape"]:
                    if isinstance(reshape_opts.get(key, None), list):
                        reshape_opts[key] = [int(v) for v in list(target_shape)]
                reshape1_op.options = reshape_opts

            _set_operator_outputs(
                model_ir=model_ir,
                op=reshape1_op,
                new_outputs=[str(output_name)],
            )

            for remove_idx in sorted([int(transpose2_idx), int(reshape2_idx), int(transpose1_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_reshape_transpose_reshape_transpose_to_nhwc_reshape_chains": int(rewritten)}


def _optimize_attention_qkv_reshape_transpose_reshape_to_reshape_transpose_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Reduce QKV-style `RESHAPE -> TRANSPOSE -> RESHAPE` rank adapters.

    Target:
      x[A,1,C] --RESHAPE([A,H,D])--> r
               --TRANSPOSE([1,0,2])--> t
               --RESHAPE([1,H,A,D])--> y
    or
      x[A,1,C] --RESHAPE([A,H,D])--> r
               --TRANSPOSE([1,2,0])--> t
               --RESHAPE([1,H,D,A])--> y

    Rewrite:
      x[A,1,C] --RESHAPE([1,A,H,D])--> r'
               --TRANSPOSE([0,2,1,3] or [0,2,3,1])--> y
    """
    rewritten = 0
    rank3_perm_had = [1, 0, 2]
    rank3_perm_hda = [1, 2, 0]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        serial = 1
        while name in model_ir.tensors:
            name = f"{base}_{serial}"
            serial += 1
        return name

    def _set_const_input(
        *,
        op: OperatorIR,
        op_idx: int,
        input_index: int,
        consumers: Dict[str, List[int]],
        values: List[int],
        base_suffix: str,
    ) -> bool:
        if int(input_index) >= len(op.inputs):
            return False
        input_name = str(op.inputs[int(input_index)])
        input_tensor = model_ir.tensors.get(input_name, None)
        if input_tensor is None:
            return False
        target_vals = [int(v) for v in list(values)]
        input_users = set(int(v) for v in consumers.get(input_name, []))
        if input_users == {int(op_idx)}:
            return _write_const_ints_to_tensor(input_tensor, target_vals)

        np_dtype = np.int32
        if input_tensor.data is not None:
            try:
                np_dtype = np.asarray(input_tensor.data).dtype
            except Exception:
                np_dtype = np.int32
        cloned_name = _unique_tensor_name(f"{input_name}_{base_suffix}")
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(input_tensor.dtype),
            shape=[int(len(target_vals))],
            shape_signature=[int(len(target_vals))],
            data=np.asarray(target_vals, dtype=np_dtype),
            is_variable=False,
            quantization=_clone_quantization(input_tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=cloned_name,
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape1_idx, reshape1_op in enumerate(model_ir.operators):
            if str(reshape1_op.op_type) != "RESHAPE" or len(reshape1_op.inputs) < 2 or len(reshape1_op.outputs) != 1:
                continue

            reshape1_in_name = str(reshape1_op.inputs[0])
            reshape1_out_name = str(reshape1_op.outputs[0])
            if reshape1_out_name in model_outputs:
                continue

            reshape1_users = [int(v) for v in consumers.get(reshape1_out_name, [])]
            if len(reshape1_users) != 1:
                continue
            transpose_idx = int(reshape1_users[0])
            transpose_op = model_ir.operators[int(transpose_idx)]
            if (
                str(transpose_op.op_type) != "TRANSPOSE"
                or len(transpose_op.inputs) < 2
                or len(transpose_op.outputs) != 1
                or str(transpose_op.inputs[0]) != str(reshape1_out_name)
            ):
                continue

            transpose_perm = _read_transpose_perm(model_ir, transpose_op)
            if transpose_perm not in [rank3_perm_had, rank3_perm_hda]:
                continue

            transpose_out_name = str(transpose_op.outputs[0])
            if transpose_out_name in model_outputs:
                continue
            transpose_users = [int(v) for v in consumers.get(transpose_out_name, [])]
            if len(transpose_users) != 1:
                continue
            reshape2_idx = int(transpose_users[0])
            reshape2_op = model_ir.operators[int(reshape2_idx)]
            if (
                str(reshape2_op.op_type) != "RESHAPE"
                or len(reshape2_op.inputs) < 2
                or len(reshape2_op.outputs) != 1
                or str(reshape2_op.inputs[0]) != str(transpose_out_name)
            ):
                continue

            reshape2_out_name = str(reshape2_op.outputs[0])

            reshape1_in_tensor = model_ir.tensors.get(reshape1_in_name, None)
            reshape1_out_tensor = model_ir.tensors.get(reshape1_out_name, None)
            transpose_out_tensor = model_ir.tensors.get(transpose_out_name, None)
            reshape2_out_tensor = model_ir.tensors.get(reshape2_out_name, None)
            if (
                reshape1_in_tensor is None
                or reshape1_out_tensor is None
                or transpose_out_tensor is None
                or reshape2_out_tensor is None
            ):
                continue
            if (
                not _is_fully_known_positive_shape(reshape1_in_tensor.shape)
                or not _is_fully_known_positive_shape(reshape1_out_tensor.shape)
                or not _is_fully_known_positive_shape(transpose_out_tensor.shape)
                or not _is_fully_known_positive_shape(reshape2_out_tensor.shape)
            ):
                continue

            in_shape = [int(v) for v in list(reshape1_in_tensor.shape)]
            reshape1_shape = [int(v) for v in list(reshape1_out_tensor.shape)]
            transpose_shape = [int(v) for v in list(transpose_out_tensor.shape)]
            reshape2_shape = [int(v) for v in list(reshape2_out_tensor.shape)]
            if len(in_shape) != 3 or len(reshape1_shape) != 3 or len(transpose_shape) != 3 or len(reshape2_shape) != 4:
                continue

            a, b, c = [int(v) for v in list(in_shape)]
            if int(b) != 1:
                continue
            if int(reshape1_shape[0]) != int(a):
                continue
            h = int(reshape1_shape[1])
            d = int(reshape1_shape[2])
            if int(h) <= 0 or int(d) <= 0:
                continue
            if int(h) * int(d) != int(c):
                continue

            if transpose_perm == rank3_perm_had:
                if transpose_shape != [int(h), int(a), int(d)]:
                    continue
                if reshape2_shape != [1, int(h), int(a), int(d)]:
                    continue
                new_transpose_perm = [0, 2, 1, 3]
            else:
                if transpose_shape != [int(h), int(d), int(a)]:
                    continue
                if reshape2_shape != [1, int(h), int(d), int(a)]:
                    continue
                new_transpose_perm = [0, 2, 3, 1]

            new_reshape_shape = [1, int(a), int(h), int(d)]
            if not _set_const_input(
                op=reshape1_op,
                op_idx=int(reshape1_idx),
                input_index=1,
                consumers=consumers,
                values=new_reshape_shape,
                base_suffix="qkv_shape",
            ):
                continue
            if not _set_const_input(
                op=transpose_op,
                op_idx=int(transpose_idx),
                input_index=1,
                consumers=consumers,
                values=new_transpose_perm,
                base_suffix="qkv_perm",
            ):
                continue

            if isinstance(reshape1_op.options, dict):
                reshape_opts = dict(reshape1_op.options)
                for key in ["newShape", "onnxRawNewShape"]:
                    if isinstance(reshape_opts.get(key, None), list):
                        reshape_opts[key] = [int(v) for v in list(new_reshape_shape)]
                reshape1_op.options = reshape_opts

            _set_operator_outputs(
                model_ir=model_ir,
                op=transpose_op,
                new_outputs=[str(reshape2_out_name)],
            )

            reshape1_out_tensor.shape = [int(v) for v in list(new_reshape_shape)]
            reshape1_out_tensor.shape_signature = [int(v) for v in list(new_reshape_shape)]
            reshape2_out_tensor.shape = [int(v) for v in list(reshape2_shape)]
            reshape2_out_tensor.shape_signature = [int(v) for v in list(reshape2_shape)]

            del model_ir.operators[int(reshape2_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_qkv_reshape_transpose_reshape_to_reshape_transpose_chains": int(rewritten)}


def _optimize_attention_gather_transpose_reshape_cleanup_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Simplify redundant attention tail chains around Gather/Transpose/Reshape.

    Pattern A:
      x[1,H,T,D] --GATHER(axis=0,index=0)--> [H,T,D]
               --TRANSPOSE([1,0,2])--> [T,H,D]
               --RESHAPE([1,1,T,H*D])--> y
      =>
      x --TRANSPOSE([0,2,1,3])--> [1,T,H,D] --RESHAPE([1,1,T,H*D])--> y

    Pattern B:
      x[1,1,T,C] --GATHER(axis=0,index=0)--> [1,T,C]
                --GATHER(axis=0,index=0)--> [T,C]
                --RESHAPE([1,1,T,C])--> y
      =>
      identity (replace y with x)
    """
    rewritten_pattern_a = 0
    rewritten_pattern_b = 0

    def _normalize_axis(axis: int, rank: int) -> int:
        axis_i = int(axis)
        if axis_i < 0:
            axis_i += int(rank)
        return int(axis_i)

    def _is_zero_index_tensor(tensor: Optional[TensorIR]) -> bool:
        vals = _read_const_ints_from_tensor(tensor)
        if vals is None or len(vals) <= 0:
            return False
        return all(int(v) == 0 for v in list(vals))

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        serial = 1
        while name in model_ir.tensors:
            name = f"{base}_{serial}"
            serial += 1
        return name

    def _set_const_input_with_clone(
        *,
        op: OperatorIR,
        op_idx: int,
        input_index: int,
        consumers: Dict[str, List[int]],
        values: List[int],
        suffix: str,
    ) -> bool:
        if int(input_index) >= len(op.inputs):
            return False
        input_name = str(op.inputs[int(input_index)])
        input_tensor = model_ir.tensors.get(input_name, None)
        if input_tensor is None:
            return False

        target_values = [int(v) for v in list(values)]
        users = set(int(v) for v in consumers.get(input_name, []))
        if users == {int(op_idx)}:
            return _write_const_ints_to_tensor(input_tensor, target_values)

        np_dtype = np.int32
        if input_tensor.data is not None:
            try:
                np_dtype = np.asarray(input_tensor.data).dtype
            except Exception:
                np_dtype = np.int32
        cloned_name = _unique_tensor_name(f"{input_name}_{suffix}")
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(input_tensor.dtype),
            shape=[int(len(target_values))],
            shape_signature=[int(len(target_values))],
            data=np.asarray(target_values, dtype=np_dtype),
            is_variable=False,
            quantization=_clone_quantization(input_tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=cloned_name,
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        # Pattern B: double-gather + reshape back to original shape => identity.
        for gather0_idx, gather0_op in enumerate(model_ir.operators):
            if str(gather0_op.op_type) != "GATHER" or len(gather0_op.inputs) < 2 or len(gather0_op.outputs) != 1:
                continue
            gather0_in_name = str(gather0_op.inputs[0])
            gather0_out_name = str(gather0_op.outputs[0])
            if gather0_out_name in model_outputs:
                continue

            gather0_in_tensor = model_ir.tensors.get(gather0_in_name, None)
            gather0_out_tensor = model_ir.tensors.get(gather0_out_name, None)
            if gather0_in_tensor is None or gather0_out_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(gather0_in_tensor.shape)
                or not _is_fully_known_positive_shape(gather0_out_tensor.shape)
            ):
                continue
            in0_shape = [int(v) for v in list(gather0_in_tensor.shape)]
            out0_shape = [int(v) for v in list(gather0_out_tensor.shape)]
            if len(in0_shape) < 2 or len(out0_shape) != len(in0_shape) - 1:
                continue
            axis0 = _normalize_axis(int(gather0_op.options.get("axis", 0)), len(in0_shape))
            if axis0 != 0 or int(in0_shape[0]) != 1:
                continue
            if not _is_zero_index_tensor(model_ir.tensors.get(str(gather0_op.inputs[1]), None)):
                continue

            gather0_users = [int(v) for v in consumers.get(gather0_out_name, [])]
            if len(gather0_users) != 1:
                continue
            gather1_idx = int(gather0_users[0])
            gather1_op = model_ir.operators[int(gather1_idx)]
            if (
                str(gather1_op.op_type) != "GATHER"
                or len(gather1_op.inputs) < 2
                or len(gather1_op.outputs) != 1
                or str(gather1_op.inputs[0]) != gather0_out_name
            ):
                continue
            gather1_out_name = str(gather1_op.outputs[0])
            if gather1_out_name in model_outputs:
                continue
            gather1_out_tensor = model_ir.tensors.get(gather1_out_name, None)
            if gather1_out_tensor is None or not _is_fully_known_positive_shape(gather1_out_tensor.shape):
                continue
            out1_shape = [int(v) for v in list(gather1_out_tensor.shape)]
            if len(out1_shape) != len(out0_shape) - 1:
                continue
            axis1 = _normalize_axis(int(gather1_op.options.get("axis", 0)), len(out0_shape))
            if axis1 != 0 or int(out0_shape[0]) != 1:
                continue
            if not _is_zero_index_tensor(model_ir.tensors.get(str(gather1_op.inputs[1]), None)):
                continue

            gather1_users = [int(v) for v in consumers.get(gather1_out_name, [])]
            if len(gather1_users) != 1:
                continue
            reshape_idx = int(gather1_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != gather1_out_name
            ):
                continue
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue

            target_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape_op.inputs[1]), None))
            if target_vals is None or len(target_vals) != len(in0_shape):
                continue
            target_shape = [int(v) for v in list(target_vals)]
            if target_shape != [int(v) for v in list(in0_shape)]:
                continue

            _replace_tensor_inputs(model_ir, str(reshape_out_name), str(gather0_in_name))
            for remove_idx in sorted([int(reshape_idx), int(gather1_idx), int(gather0_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]
            rewritten_pattern_b += 1
            changed = True
            break

        if changed:
            continue

        # Pattern A: gather(strip axis-0 singleton) + transpose + reshape -> transpose(rank4) + reshape
        for gather_idx, gather_op in enumerate(model_ir.operators):
            if str(gather_op.op_type) != "GATHER" or len(gather_op.inputs) < 2 or len(gather_op.outputs) != 1:
                continue
            gather_in_name = str(gather_op.inputs[0])
            gather_out_name = str(gather_op.outputs[0])
            if gather_out_name in model_outputs:
                continue

            gather_in_tensor = model_ir.tensors.get(gather_in_name, None)
            gather_out_tensor = model_ir.tensors.get(gather_out_name, None)
            if gather_in_tensor is None or gather_out_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(gather_in_tensor.shape)
                or not _is_fully_known_positive_shape(gather_out_tensor.shape)
            ):
                continue
            in_shape = [int(v) for v in list(gather_in_tensor.shape)]
            out_shape = [int(v) for v in list(gather_out_tensor.shape)]
            if len(in_shape) != 4 or len(out_shape) != 3:
                continue
            axis = _normalize_axis(int(gather_op.options.get("axis", 0)), len(in_shape))
            if axis != 0 or int(in_shape[0]) != 1:
                continue
            if not _is_zero_index_tensor(model_ir.tensors.get(str(gather_op.inputs[1]), None)):
                continue
            if out_shape != [int(in_shape[1]), int(in_shape[2]), int(in_shape[3])]:
                continue

            gather_users = [int(v) for v in consumers.get(gather_out_name, [])]
            if len(gather_users) != 1:
                continue
            transpose_idx = int(gather_users[0])
            transpose_op = model_ir.operators[int(transpose_idx)]
            if (
                str(transpose_op.op_type) != "TRANSPOSE"
                or len(transpose_op.inputs) < 2
                or len(transpose_op.outputs) != 1
                or str(transpose_op.inputs[0]) != gather_out_name
                or _read_transpose_perm(model_ir, transpose_op) != [1, 0, 2]
            ):
                continue

            transpose_out_name = str(transpose_op.outputs[0])
            if transpose_out_name in model_outputs:
                continue
            transpose_out_tensor = model_ir.tensors.get(transpose_out_name, None)
            if transpose_out_tensor is None or not _is_fully_known_positive_shape(transpose_out_tensor.shape):
                continue
            if [int(v) for v in list(transpose_out_tensor.shape)] != [int(in_shape[2]), int(in_shape[1]), int(in_shape[3])]:
                continue

            transpose_users = [int(v) for v in consumers.get(transpose_out_name, [])]
            if len(transpose_users) != 1:
                continue
            reshape_idx = int(transpose_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != transpose_out_name
            ):
                continue

            reshape_shape_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape_op.inputs[1]), None))
            if reshape_shape_vals is None or len(reshape_shape_vals) != 4:
                continue
            expected_last = int(in_shape[1]) * int(in_shape[3])
            expected_shape = [1, 1, int(in_shape[2]), int(expected_last)]
            if [int(v) for v in list(reshape_shape_vals)] != [int(v) for v in list(expected_shape)]:
                continue

            if not _set_const_input_with_clone(
                op=transpose_op,
                op_idx=int(transpose_idx),
                input_index=1,
                consumers=consumers,
                values=[0, 2, 1, 3],
                suffix="qkv_perm4",
            ):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=transpose_op,
                new_inputs=[str(gather_in_name), str(transpose_op.inputs[1])],
            )

            transpose_out_tensor.shape = [1, int(in_shape[2]), int(in_shape[1]), int(in_shape[3])]
            transpose_out_tensor.shape_signature = [1, int(in_shape[2]), int(in_shape[1]), int(in_shape[3])]
            if isinstance(transpose_op.options, dict):
                opts = dict(transpose_op.options)
                for key in ["perm", "onnxPerm"]:
                    if isinstance(opts.get(key, None), list):
                        opts[key] = [0, 2, 1, 3]
                transpose_op.options = opts

            del model_ir.operators[int(gather_idx)]
            rewritten_pattern_a += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_attention_gather_transpose_reshape_cleanup_pattern_a": int(rewritten_pattern_a),
        "optimized_attention_gather_transpose_reshape_cleanup_pattern_b": int(rewritten_pattern_b),
    }


def _optimize_gather_axis0_singleton_to_reshape_input_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove singleton-axis Gather before a lone Reshape.

    Target:
      x[1,...] --GATHER(axis=0, indices={0})--> g
               --RESHAPE(shape=S)--> y

    Rewrite:
      x --RESHAPE(shape=S)--> y

    Safety:
    - Gather output has exactly one consumer (the Reshape).
    - Gather axis is normalized axis 0.
    - Gather indices are compile-time constants and all zeros.
    - Gather input leading dimension is statically 1.
    """
    rewritten = 0

    def _normalize_axis(axis: int, rank: int) -> int:
        axis_i = int(axis)
        if axis_i < 0:
            axis_i += int(rank)
        return int(axis_i)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for gather_idx, gather_op in enumerate(model_ir.operators):
            if str(gather_op.op_type) != "GATHER" or len(gather_op.inputs) < 2 or len(gather_op.outputs) != 1:
                continue
            gather_in_name = str(gather_op.inputs[0])
            gather_out_name = str(gather_op.outputs[0])
            if gather_out_name in model_outputs:
                continue

            gather_in_tensor = model_ir.tensors.get(gather_in_name, None)
            gather_out_tensor = model_ir.tensors.get(gather_out_name, None)
            if gather_in_tensor is None or gather_out_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(gather_in_tensor.shape)
                or not _is_fully_known_positive_shape(gather_out_tensor.shape)
            ):
                continue
            in_shape = [int(v) for v in list(gather_in_tensor.shape)]
            out_shape = [int(v) for v in list(gather_out_tensor.shape)]
            if len(in_shape) <= 0 or len(out_shape) != len(in_shape) - 1:
                continue
            if int(in_shape[0]) != 1:
                continue

            axis = _normalize_axis(int(gather_op.options.get("axis", 0)), len(in_shape))
            if axis != 0:
                continue

            gather_indices = _read_const_ints_from_tensor(model_ir.tensors.get(str(gather_op.inputs[1]), None))
            if gather_indices is None or len(gather_indices) <= 0:
                continue
            if any(int(v) != 0 for v in list(gather_indices)):
                continue

            gather_users = [int(v) for v in consumers.get(gather_out_name, [])]
            if len(gather_users) != 1:
                continue
            reshape_idx = int(gather_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != str(gather_out_name)
            ):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=[str(gather_in_name)] + [str(v) for v in list(reshape_op.inputs[1:])],
            )

            del model_ir.operators[int(gather_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_gather_axis0_singleton_to_reshape_input_chains": int(rewritten)}


def _optimize_attention_preproj_reshape_to_batchmatmul_ranklift_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Lift pre-projection RESHAPE rank before attention Q/K projections.

    Target:
      x[1,1,T,C] --RESHAPE([T,1,C])--> r
      r --BATCH_MATMUL(W_i[C,K_i])--> y_i[T,1,K_i]
      y_i --(ADD/SUB/MUL/DIV + bias)--> z_i[T,1,K_i]
      z_i --RESHAPE([1,T,H,D])--> ...

    Rewrite:
      x --BATCH_MATMUL(W_i)--> y_i'[1,1,T,K_i]
      (same binary op, same tail reshape)
    and remove the leading RESHAPE.
    """
    rewritten = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape_idx, reshape_op in enumerate(model_ir.operators):
            if str(reshape_op.op_type) != "RESHAPE" or len(reshape_op.inputs) < 2 or len(reshape_op.outputs) != 1:
                continue
            reshape_in_name = str(reshape_op.inputs[0])
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue

            reshape_in_tensor = model_ir.tensors.get(reshape_in_name, None)
            reshape_out_tensor = model_ir.tensors.get(reshape_out_name, None)
            if reshape_in_tensor is None or reshape_out_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(reshape_in_tensor.shape)
                or not _is_fully_known_positive_shape(reshape_out_tensor.shape)
            ):
                continue
            in_shape = [int(v) for v in list(reshape_in_tensor.shape)]
            out_shape = [int(v) for v in list(reshape_out_tensor.shape)]
            if len(in_shape) != 4 or len(out_shape) != 3:
                continue
            if int(in_shape[0]) != 1 or int(in_shape[1]) != 1:
                continue

            t_dim = int(in_shape[2])
            c_dim = int(in_shape[3])
            if out_shape != [int(t_dim), 1, int(c_dim)]:
                continue
            target_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape_op.inputs[1]), None))
            if target_vals is None or [int(v) for v in list(target_vals)] != [int(t_dim), 1, int(c_dim)]:
                continue

            bmm_indices = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(bmm_indices) <= 0:
                continue

            branch_specs: List[Tuple[int, int, int]] = []
            valid = True
            for bmm_idx in bmm_indices:
                bmm_op = model_ir.operators[int(bmm_idx)]
                if (
                    str(bmm_op.op_type) != "BATCH_MATMUL"
                    or len(bmm_op.inputs) < 2
                    or len(bmm_op.outputs) != 1
                    or str(bmm_op.inputs[0]) != str(reshape_out_name)
                ):
                    valid = False
                    break
                bmm_out_name = str(bmm_op.outputs[0])
                if bmm_out_name in model_outputs:
                    valid = False
                    break
                bmm_out_tensor = model_ir.tensors.get(bmm_out_name, None)
                rhs_tensor = model_ir.tensors.get(str(bmm_op.inputs[1]), None)
                if (
                    bmm_out_tensor is None
                    or rhs_tensor is None
                    or not _is_fully_known_positive_shape(bmm_out_tensor.shape)
                    or not _is_fully_known_positive_shape(rhs_tensor.shape)
                ):
                    valid = False
                    break
                rhs_shape = [int(v) for v in list(rhs_tensor.shape)]
                bmm_out_shape = [int(v) for v in list(bmm_out_tensor.shape)]
                if len(rhs_shape) != 2 or int(rhs_shape[0]) != int(c_dim):
                    valid = False
                    break
                k_dim = int(rhs_shape[1])
                if bmm_out_shape != [int(t_dim), 1, int(k_dim)]:
                    valid = False
                    break

                bmm_users = [int(v) for v in consumers.get(bmm_out_name, [])]
                if len(bmm_users) != 1:
                    valid = False
                    break
                binary_idx = int(bmm_users[0])
                binary_op = model_ir.operators[int(binary_idx)]
                if (
                    str(binary_op.op_type) not in binary_ops
                    or len(binary_op.inputs) < 2
                    or len(binary_op.outputs) != 1
                ):
                    valid = False
                    break
                if bmm_out_name not in [str(binary_op.inputs[0]), str(binary_op.inputs[1])]:
                    valid = False
                    break
                binary_out_name = str(binary_op.outputs[0])
                if binary_out_name in model_outputs:
                    valid = False
                    break
                binary_out_tensor = model_ir.tensors.get(binary_out_name, None)
                if (
                    binary_out_tensor is None
                    or not _is_fully_known_positive_shape(binary_out_tensor.shape)
                    or [int(v) for v in list(binary_out_tensor.shape)] != [int(t_dim), 1, int(k_dim)]
                ):
                    valid = False
                    break

                binary_users = [int(v) for v in consumers.get(binary_out_name, [])]
                if len(binary_users) != 1:
                    valid = False
                    break
                tail_reshape_idx = int(binary_users[0])
                tail_reshape_op = model_ir.operators[int(tail_reshape_idx)]
                if (
                    str(tail_reshape_op.op_type) != "RESHAPE"
                    or len(tail_reshape_op.inputs) < 2
                    or len(tail_reshape_op.outputs) != 1
                    or str(tail_reshape_op.inputs[0]) != str(binary_out_name)
                ):
                    valid = False
                    break
                tail_target_vals = _read_const_ints_from_tensor(
                    model_ir.tensors.get(str(tail_reshape_op.inputs[1]), None)
                )
                if tail_target_vals is None or len(tail_target_vals) != 4:
                    valid = False
                    break
                tail_target = [int(v) for v in list(tail_target_vals)]
                if tail_target[0] != 1 or tail_target[1] != int(t_dim):
                    valid = False
                    break
                if int(tail_target[2]) * int(tail_target[3]) != int(k_dim):
                    valid = False
                    break

                branch_specs.append((int(bmm_idx), int(binary_idx), int(k_dim)))

            if not valid or len(branch_specs) <= 0:
                continue

            # Rewrite all BATCH_MATMUL inputs and rank metadata.
            for bmm_idx, binary_idx, k_dim in branch_specs:
                bmm_op = model_ir.operators[int(bmm_idx)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=bmm_op,
                    new_inputs=[str(reshape_in_name), str(bmm_op.inputs[1])],
                )
                bmm_out_name = str(bmm_op.outputs[0])
                bmm_out_tensor = model_ir.tensors.get(bmm_out_name, None)
                if bmm_out_tensor is not None:
                    bmm_out_tensor.shape = [1, 1, int(t_dim), int(k_dim)]
                    bmm_out_tensor.shape_signature = [1, 1, int(t_dim), int(k_dim)]

                binary_op = model_ir.operators[int(binary_idx)]
                binary_out_name = str(binary_op.outputs[0])
                binary_out_tensor = model_ir.tensors.get(binary_out_name, None)
                if binary_out_tensor is not None:
                    binary_out_tensor.shape = [1, 1, int(t_dim), int(k_dim)]
                    binary_out_tensor.shape_signature = [1, 1, int(t_dim), int(k_dim)]

            del model_ir.operators[int(reshape_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_preproj_reshape_to_batchmatmul_ranklift_chains": int(rewritten)}


def _optimize_window_partition_reshape_transpose_to_space_to_depth_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Replace Swin-like window-partition wrappers with SPACE_TO_DEPTH.

    Target:
      x[N,H,W,C]
        --RESHAPE([N,OH,BS,OW,BS,C])--> r1
        --TRANSPOSE([0,1,3,2,4,5])--> t1
        --RESHAPE([N*OH*OW, BS*BS, C])--> y

    Rewrite:
      x --SPACE_TO_DEPTH(blockSize=BS)--> s2d [N,OH,OW,BS*BS*C]
      s2d --RESHAPE([N*OH*OW, BS*BS, C])--> y
    """
    rewritten = 0
    target_perm = [0, 1, 3, 2, 4, 5]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape1_idx, reshape1_op in enumerate(model_ir.operators):
            if str(reshape1_op.op_type) != "RESHAPE" or len(reshape1_op.inputs) < 2 or len(reshape1_op.outputs) != 1:
                continue

            input_name = str(reshape1_op.inputs[0])
            reshape1_out_name = str(reshape1_op.outputs[0])
            if reshape1_out_name in model_outputs:
                continue

            input_tensor = model_ir.tensors.get(input_name, None)
            reshape1_out_tensor = model_ir.tensors.get(reshape1_out_name, None)
            if input_tensor is None or reshape1_out_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(input_tensor.shape)
                or not _is_fully_known_positive_shape(reshape1_out_tensor.shape)
            ):
                continue

            in_shape = [int(v) for v in list(input_tensor.shape)]
            r1_shape = [int(v) for v in list(reshape1_out_tensor.shape)]
            if len(in_shape) != 4 or len(r1_shape) != 6:
                continue

            n, h, w, c = [int(v) for v in list(in_shape)]
            n1, oh, bs_h, ow, bs_w, c1 = [int(v) for v in list(r1_shape)]
            if int(n1) != int(n) or int(c1) != int(c):
                continue
            if int(bs_h) <= 1 or int(bs_w) <= 1 or int(bs_h) != int(bs_w):
                continue
            bs = int(bs_h)
            if int(oh) * int(bs) != int(h) or int(ow) * int(bs) != int(w):
                continue

            reshape1_target = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape1_op.inputs[1]), None))
            if reshape1_target is None or [int(v) for v in list(reshape1_target)] != [int(v) for v in list(r1_shape)]:
                continue

            reshape1_users = [int(v) for v in consumers.get(reshape1_out_name, [])]
            if len(reshape1_users) != 1:
                continue
            transpose_idx = int(reshape1_users[0])
            transpose_op = model_ir.operators[int(transpose_idx)]
            if (
                str(transpose_op.op_type) != "TRANSPOSE"
                or len(transpose_op.inputs) < 2
                or len(transpose_op.outputs) != 1
                or str(transpose_op.inputs[0]) != str(reshape1_out_name)
                or _read_transpose_perm(model_ir, transpose_op) != target_perm
            ):
                continue

            transpose_out_name = str(transpose_op.outputs[0])
            if transpose_out_name in model_outputs:
                continue
            transpose_out_tensor = model_ir.tensors.get(transpose_out_name, None)
            if transpose_out_tensor is None or not _is_fully_known_positive_shape(transpose_out_tensor.shape):
                continue
            transpose_shape = [int(v) for v in list(transpose_out_tensor.shape)]
            expected_transpose_shape = [int(n), int(oh), int(ow), int(bs), int(bs), int(c)]
            if transpose_shape != expected_transpose_shape:
                continue

            transpose_users = [int(v) for v in consumers.get(transpose_out_name, [])]
            if len(transpose_users) != 1:
                continue
            reshape2_idx = int(transpose_users[0])
            reshape2_op = model_ir.operators[int(reshape2_idx)]
            if (
                str(reshape2_op.op_type) != "RESHAPE"
                or len(reshape2_op.inputs) < 2
                or len(reshape2_op.outputs) != 1
                or str(reshape2_op.inputs[0]) != str(transpose_out_name)
            ):
                continue

            reshape2_target = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape2_op.inputs[1]), None))
            expected_reshape2 = [int(n) * int(oh) * int(ow), int(bs) * int(bs), int(c)]
            if reshape2_target is None or [int(v) for v in list(reshape2_target)] != expected_reshape2:
                continue

            # Convert TRANSPOSE op in-place to SPACE_TO_DEPTH and bypass first RESHAPE.
            transpose_op.op_type = "SPACE_TO_DEPTH"
            _set_operator_inputs(
                model_ir=model_ir,
                op=transpose_op,
                new_inputs=[str(input_name)],
            )
            transpose_op.options = {"blockSize": int(bs)}

            s2d_shape = [int(n), int(oh), int(ow), int(bs) * int(bs) * int(c)]
            transpose_out_tensor.shape = [int(v) for v in list(s2d_shape)]
            transpose_out_tensor.shape_signature = [int(v) for v in list(s2d_shape)]

            del model_ir.operators[int(reshape1_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_window_partition_reshape_transpose_to_space_to_depth_chains": int(rewritten)}


def _optimize_window_reverse_reshape_transpose_to_depth_to_space_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Replace Swin-like window-reverse wrappers with DEPTH_TO_SPACE.

    Target:
      x[NW, BS*BS, C]
        --RESHAPE([N,OH,OW,BS,BS,C])--> r1
        --TRANSPOSE([0,1,3,2,4,5])--> t1
        --RESHAPE([N,OH*BS,OW*BS,C])--> y

    Rewrite:
      x --RESHAPE([N,OH,OW,BS*BS*C])--> r1'
      r1' --DEPTH_TO_SPACE(blockSize=BS)--> y
    """
    rewritten = 0
    target_perm = [0, 1, 3, 2, 4, 5]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        serial = 1
        while name in model_ir.tensors:
            name = f"{base}_{serial}"
            serial += 1
        return name

    def _set_const_input_with_clone(
        *,
        op: OperatorIR,
        op_idx: int,
        input_index: int,
        consumers: Dict[str, List[int]],
        values: List[int],
        suffix: str,
    ) -> bool:
        if int(input_index) >= len(op.inputs):
            return False
        input_name = str(op.inputs[int(input_index)])
        input_tensor = model_ir.tensors.get(input_name, None)
        if input_tensor is None:
            return False

        target_values = [int(v) for v in list(values)]
        users = set(int(v) for v in consumers.get(input_name, []))
        if users == {int(op_idx)}:
            return _write_const_ints_to_tensor(input_tensor, target_values)

        np_dtype = np.int32
        if input_tensor.data is not None:
            try:
                np_dtype = np.asarray(input_tensor.data).dtype
            except Exception:
                np_dtype = np.int32
        cloned_name = _unique_tensor_name(f"{input_name}_{suffix}")
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(input_tensor.dtype),
            shape=[int(len(target_values))],
            shape_signature=[int(len(target_values))],
            data=np.asarray(target_values, dtype=np_dtype),
            is_variable=False,
            quantization=_clone_quantization(input_tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=cloned_name,
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape1_idx, reshape1_op in enumerate(model_ir.operators):
            if str(reshape1_op.op_type) != "RESHAPE" or len(reshape1_op.inputs) < 2 or len(reshape1_op.outputs) != 1:
                continue

            input_name = str(reshape1_op.inputs[0])
            reshape1_out_name = str(reshape1_op.outputs[0])
            if reshape1_out_name in model_outputs:
                continue

            input_tensor = model_ir.tensors.get(input_name, None)
            reshape1_out_tensor = model_ir.tensors.get(reshape1_out_name, None)
            if input_tensor is None or reshape1_out_tensor is None:
                continue
            if (
                not _is_fully_known_positive_shape(input_tensor.shape)
                or not _is_fully_known_positive_shape(reshape1_out_tensor.shape)
            ):
                continue

            in_shape = [int(v) for v in list(input_tensor.shape)]
            r1_shape = [int(v) for v in list(reshape1_out_tensor.shape)]
            if len(in_shape) != 3 or len(r1_shape) != 6:
                continue

            nw, ws2, c = [int(v) for v in list(in_shape)]
            n, oh, ow, bs_h, bs_w, c1 = [int(v) for v in list(r1_shape)]
            if int(c1) != int(c):
                continue
            if int(bs_h) <= 1 or int(bs_w) <= 1 or int(bs_h) != int(bs_w):
                continue
            bs = int(bs_h)
            if int(ws2) != int(bs) * int(bs):
                continue
            if int(nw) != int(n) * int(oh) * int(ow):
                continue

            reshape1_target = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape1_op.inputs[1]), None))
            if reshape1_target is None or [int(v) for v in list(reshape1_target)] != [int(v) for v in list(r1_shape)]:
                continue

            reshape1_users = [int(v) for v in consumers.get(reshape1_out_name, [])]
            if len(reshape1_users) != 1:
                continue
            transpose_idx = int(reshape1_users[0])
            transpose_op = model_ir.operators[int(transpose_idx)]
            if (
                str(transpose_op.op_type) != "TRANSPOSE"
                or len(transpose_op.inputs) < 2
                or len(transpose_op.outputs) != 1
                or str(transpose_op.inputs[0]) != str(reshape1_out_name)
                or _read_transpose_perm(model_ir, transpose_op) != target_perm
            ):
                continue

            transpose_out_name = str(transpose_op.outputs[0])
            if transpose_out_name in model_outputs:
                continue
            transpose_out_tensor = model_ir.tensors.get(transpose_out_name, None)
            if transpose_out_tensor is None or not _is_fully_known_positive_shape(transpose_out_tensor.shape):
                continue
            transpose_shape = [int(v) for v in list(transpose_out_tensor.shape)]
            expected_transpose_shape = [int(n), int(oh), int(bs), int(ow), int(bs), int(c)]
            if transpose_shape != expected_transpose_shape:
                continue

            transpose_users = [int(v) for v in consumers.get(transpose_out_name, [])]
            if len(transpose_users) != 1:
                continue
            reshape2_idx = int(transpose_users[0])
            reshape2_op = model_ir.operators[int(reshape2_idx)]
            if (
                str(reshape2_op.op_type) != "RESHAPE"
                or len(reshape2_op.inputs) < 2
                or len(reshape2_op.outputs) != 1
                or str(reshape2_op.inputs[0]) != str(transpose_out_name)
            ):
                continue

            reshape2_out_name = str(reshape2_op.outputs[0])
            reshape2_target = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape2_op.inputs[1]), None))
            expected_h = int(oh) * int(bs)
            expected_w = int(ow) * int(bs)
            expected_reshape2 = [int(n), int(expected_h), int(expected_w), int(c)]
            if reshape2_target is None or [int(v) for v in list(reshape2_target)] != expected_reshape2:
                continue

            reshape2_out_tensor = model_ir.tensors.get(reshape2_out_name, None)
            if (
                reshape2_out_tensor is None
                or not _is_fully_known_positive_shape(reshape2_out_tensor.shape)
                or [int(v) for v in list(reshape2_out_tensor.shape)] != expected_reshape2
            ):
                continue

            new_reshape1_shape = [int(n), int(oh), int(ow), int(bs) * int(bs) * int(c)]
            if not _set_const_input_with_clone(
                op=reshape1_op,
                op_idx=int(reshape1_idx),
                input_index=1,
                consumers=consumers,
                values=new_reshape1_shape,
                suffix="d2s_shape",
            ):
                continue

            if isinstance(reshape1_op.options, dict):
                reshape_opts = dict(reshape1_op.options)
                for key in ["newShape", "onnxRawNewShape"]:
                    if isinstance(reshape_opts.get(key, None), list):
                        reshape_opts[key] = [int(v) for v in list(new_reshape1_shape)]
                reshape1_op.options = reshape_opts

            reshape1_out_tensor.shape = [int(v) for v in list(new_reshape1_shape)]
            reshape1_out_tensor.shape_signature = [int(v) for v in list(new_reshape1_shape)]

            transpose_op.op_type = "DEPTH_TO_SPACE"
            _set_operator_inputs(
                model_ir=model_ir,
                op=transpose_op,
                new_inputs=[str(reshape1_out_name)],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=transpose_op,
                new_outputs=[str(reshape2_out_name)],
            )
            transpose_op.options = {"blockSize": int(bs)}

            del model_ir.operators[int(reshape2_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_window_reverse_reshape_transpose_to_depth_to_space_chains": int(rewritten)}


def _optimize_transpose_squeeze_unary_expanddims_transpose_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse Conv1D-shim transpose wrappers around unary activations.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
             --SQUEEZE(axis=k)--> s
             --UNARY--> u
             --EXPAND_DIMS(axis=k)--> e
             --T(0,2,3,1)--> y_nhwc

    Rewrite:
      x_nhwc --UNARY--> y_nhwc

    Safety:
    - Strict linear chain (single consumer on every intermediate).
    - Canonical NHWC<->NCHW transpose pair only.
    - SQUEEZE and EXPAND_DIMS axes must match and round-trip rank-4 shape.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
        "HARD_SWISH",
    }

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            squeeze_idx = int(pre_users[0])
            squeeze_op = model_ir.operators[int(squeeze_idx)]
            if (
                str(squeeze_op.op_type) != "SQUEEZE"
                or len(squeeze_op.inputs) != 1
                or len(squeeze_op.outputs) != 1
                or str(squeeze_op.inputs[0]) != pre_output_name
            ):
                continue
            squeeze_out_name = str(squeeze_op.outputs[0])
            if squeeze_out_name in model_outputs:
                continue

            squeeze_users = [int(v) for v in consumers.get(squeeze_out_name, [])]
            if len(squeeze_users) != 1:
                continue
            unary_idx = int(squeeze_users[0])
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.inputs[0]) != squeeze_out_name
            ):
                continue
            unary_out_name = str(unary_op.outputs[0])
            if unary_out_name in model_outputs:
                continue

            unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
            if len(unary_users) != 1:
                continue
            expand_idx = int(unary_users[0])
            expand_op = model_ir.operators[int(expand_idx)]
            if (
                str(expand_op.op_type) != "EXPAND_DIMS"
                or len(expand_op.inputs) < 2
                or len(expand_op.outputs) != 1
                or str(expand_op.inputs[0]) != unary_out_name
            ):
                continue
            expand_out_name = str(expand_op.outputs[0])
            if expand_out_name in model_outputs:
                continue

            expand_users = [int(v) for v in consumers.get(expand_out_name, [])]
            if len(expand_users) != 1:
                continue
            post_idx = int(expand_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != expand_out_name
            ):
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue
            post_output_name = str(post_op.outputs[0])

            pre_input_shape = _shape_list(pre_input_name)
            pre_output_shape = _shape_list(pre_output_name)
            squeeze_out_shape = _shape_list(squeeze_out_name)
            unary_out_shape = _shape_list(unary_out_name)
            expand_out_shape = _shape_list(expand_out_name)
            post_output_shape = _shape_list(post_output_name)
            if (
                pre_input_shape is None
                or pre_output_shape is None
                or squeeze_out_shape is None
                or unary_out_shape is None
                or expand_out_shape is None
                or post_output_shape is None
            ):
                continue
            if (
                len(pre_input_shape) != 4
                or len(pre_output_shape) != 4
                or len(squeeze_out_shape) != 3
                or len(unary_out_shape) != 3
                or len(expand_out_shape) != 4
                or len(post_output_shape) != 4
            ):
                continue
            if not _shape_compatible(squeeze_out_shape, unary_out_shape):
                continue
            if not _shape_compatible(pre_output_shape, expand_out_shape):
                continue
            if not _shape_compatible(pre_input_shape, post_output_shape):
                continue

            squeeze_axis: Optional[int] = None
            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            if "squeezeDims" in squeeze_options:
                raw_axes = np.asarray(squeeze_options.get("squeezeDims", []), dtype=np.int64).reshape(-1)
                normalized_axes = _normalize_squeeze_axes_for_rank(
                    [int(v) for v in raw_axes.tolist()],
                    4,
                )
                if normalized_axes is None or len(normalized_axes) != 1:
                    continue
                candidate_axis = int(normalized_axes[0])
                if int(pre_output_shape[candidate_axis]) != 1:
                    continue
                candidate_squeezed = [
                    int(pre_output_shape[idx])
                    for idx in range(4)
                    if int(idx) != int(candidate_axis)
                ]
                if not _shape_compatible(candidate_squeezed, squeeze_out_shape):
                    continue
                squeeze_axis = int(candidate_axis)
            else:
                candidate_axes: List[int] = []
                for axis in range(4):
                    if int(pre_output_shape[axis]) != 1:
                        continue
                    candidate_squeezed = [
                        int(pre_output_shape[idx])
                        for idx in range(4)
                        if int(idx) != int(axis)
                    ]
                    if _shape_compatible(candidate_squeezed, squeeze_out_shape):
                        candidate_axes.append(int(axis))
                if len(candidate_axes) != 1:
                    continue
                squeeze_axis = int(candidate_axes[0])

            expand_axis_values = _read_const_ints_from_tensor(
                model_ir.tensors.get(str(expand_op.inputs[1]), None)
            )
            if expand_axis_values is None or len(expand_axis_values) != 1:
                continue
            expand_axis = int(expand_axis_values[0])
            if expand_axis < 0:
                expand_axis += 4
            if expand_axis < 0 or expand_axis >= 4:
                continue
            if int(expand_axis) != int(squeeze_axis):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[pre_input_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=unary_op,
                new_outputs=[post_output_name],
            )

            old_unary_tensor = model_ir.tensors.get(unary_out_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if old_unary_tensor is not None and post_output_tensor is not None:
                post_output_tensor.dtype = str(old_unary_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_unary_tensor.quantization)

            remove_indices = sorted(
                list({int(pre_idx), int(squeeze_idx), int(expand_idx), int(post_idx)}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_squeeze_unary_expanddims_transpose_nhwc_chains": int(rewritten)
    }


def _optimize_transpose_unary_transpose_reshape_expanddims_transpose_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse unary Conv1D-shim wrappers with an intermediate rank-4 transpose.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
             --UNARY--> u_nchw
             --T(0,2,1,3)--> t
             --RESHAPE([H,C,W])--> r3
             --EXPAND_DIMS(axis=2)--> e4
             --T(0,2,3,1)--> y

    Rewrite:
      x_nhwc --UNARY--> u_nhwc
             --RESHAPE([H,W,C])--> r3'
             --EXPAND_DIMS(axis=1)--> y
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nhcw_from_nchw = [0, 2, 1, 3]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_hcw_to_hwc = [0, 2, 1]
    unary_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
        "HARD_SWISH",
    }

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _swap_rank3_shape_in_reshape(reshape_op: OperatorIR) -> bool:
        changed = False
        if len(reshape_op.inputs) >= 2:
            shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
            shape_vals = _read_const_ints_from_tensor(shape_tensor)
            if shape_vals is not None and len(shape_vals) == 3:
                swapped = [int(shape_vals[0]), int(shape_vals[2]), int(shape_vals[1])]
                if swapped != [int(v) for v in list(shape_vals)]:
                    changed = _write_const_ints_to_tensor(shape_tensor, swapped) or changed
        if isinstance(reshape_op.options, dict):
            opts = dict(reshape_op.options)
            opts_changed = False
            for key in ["newShape", "onnxRawNewShape"]:
                value = opts.get(key, None)
                if not isinstance(value, list) or len(value) != 3:
                    continue
                swapped = [int(value[0]), int(value[2]), int(value[1])]
                if swapped != [int(v) for v in list(value)]:
                    opts[key] = swapped
                    opts_changed = True
            if opts_changed:
                reshape_op.options = opts
                changed = True
        return bool(changed)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) != 1:
                continue
            unary_idx = int(pre_users[0])
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.inputs[0]) != pre_out_name
            ):
                continue
            unary_out_name = str(unary_op.outputs[0])
            if unary_out_name in model_outputs:
                continue

            unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
            if len(unary_users) != 1:
                continue
            mid_idx = int(unary_users[0])
            mid_op = model_ir.operators[int(mid_idx)]
            if (
                str(mid_op.op_type) != "TRANSPOSE"
                or len(mid_op.inputs) < 2
                or len(mid_op.outputs) != 1
                or str(mid_op.inputs[0]) != unary_out_name
                or _read_transpose_perm(model_ir, mid_op) != perm_nhcw_from_nchw
            ):
                continue
            mid_out_name = str(mid_op.outputs[0])
            if mid_out_name in model_outputs:
                continue

            mid_users = [int(v) for v in consumers.get(mid_out_name, [])]
            if len(mid_users) != 1:
                continue
            reshape_idx = int(mid_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 1
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != mid_out_name
            ):
                continue
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue

            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) < 1:
                continue
            expand_idx: Optional[int] = None
            post_idx: Optional[int] = None
            expand_op: Optional[OperatorIR] = None
            post_op: Optional[OperatorIR] = None
            expand_out_name = ""
            post_out_name = ""
            for candidate_expand_idx in reshape_users:
                candidate_expand_op = model_ir.operators[int(candidate_expand_idx)]
                if (
                    str(candidate_expand_op.op_type) != "EXPAND_DIMS"
                    or len(candidate_expand_op.inputs) < 2
                    or len(candidate_expand_op.outputs) != 1
                    or str(candidate_expand_op.inputs[0]) != reshape_out_name
                ):
                    continue
                candidate_expand_axis_values = _read_const_ints_from_tensor(
                    model_ir.tensors.get(str(candidate_expand_op.inputs[1]), None)
                )
                if candidate_expand_axis_values is None or len(candidate_expand_axis_values) != 1:
                    continue
                candidate_expand_axis = int(candidate_expand_axis_values[0])
                if candidate_expand_axis < 0:
                    candidate_expand_axis += 4
                if int(candidate_expand_axis) != 2:
                    continue
                candidate_expand_out_name = str(candidate_expand_op.outputs[0])
                if candidate_expand_out_name in model_outputs:
                    continue
                candidate_expand_users = [int(v) for v in consumers.get(candidate_expand_out_name, [])]
                if len(candidate_expand_users) != 1:
                    continue
                candidate_post_idx = int(candidate_expand_users[0])
                candidate_post_op = model_ir.operators[int(candidate_post_idx)]
                if (
                    str(candidate_post_op.op_type) != "TRANSPOSE"
                    or len(candidate_post_op.inputs) < 2
                    or len(candidate_post_op.outputs) != 1
                    or str(candidate_post_op.inputs[0]) != candidate_expand_out_name
                    or _read_transpose_perm(model_ir, candidate_post_op) != perm_nchw_to_nhwc
                ):
                    continue
                expand_idx = int(candidate_expand_idx)
                expand_op = candidate_expand_op
                post_idx = int(candidate_post_idx)
                post_op = candidate_post_op
                expand_out_name = str(candidate_expand_out_name)
                post_out_name = str(candidate_post_op.outputs[0])
                break
            if expand_idx is None or post_idx is None or expand_op is None or post_op is None:
                continue
            non_expand_users = [int(v) for v in reshape_users if int(v) != int(expand_idx)]

            pre_in_shape = _shape_list(pre_in_name)
            pre_out_shape = _shape_list(pre_out_name)
            unary_out_shape = _shape_list(unary_out_name)
            mid_out_shape = _shape_list(mid_out_name)
            reshape_out_shape = _shape_list(reshape_out_name)
            expand_out_shape = _shape_list(expand_out_name)
            post_out_shape = _shape_list(post_out_name)
            if (
                pre_in_shape is None
                or pre_out_shape is None
                or unary_out_shape is None
                or mid_out_shape is None
                or reshape_out_shape is None
                or expand_out_shape is None
                or post_out_shape is None
            ):
                continue
            if (
                len(pre_in_shape) != 4
                or len(pre_out_shape) != 4
                or len(unary_out_shape) != 4
                or len(mid_out_shape) != 4
                or len(reshape_out_shape) != 3
                or len(expand_out_shape) != 4
                or len(post_out_shape) != 4
            ):
                continue

            n, h, w, c = [int(v) for v in list(pre_in_shape)]
            expected_pre_out = [n, c, h, w]
            expected_mid_out = [n, h, c, w]
            expected_reshape_out = [h, c, w]
            expected_expand_out = [h, c, 1, w]
            expected_post_out = [h, 1, w, c]
            if not _shape_compatible(pre_out_shape, expected_pre_out):
                continue
            if not _shape_compatible(unary_out_shape, expected_pre_out):
                continue
            if not _shape_compatible(mid_out_shape, expected_mid_out):
                continue
            if not _shape_compatible(reshape_out_shape, expected_reshape_out):
                continue
            if not _shape_compatible(expand_out_shape, expected_expand_out):
                continue
            if not _shape_compatible(post_out_shape, expected_post_out):
                continue
            if not _dims_compatible(n, 1):
                continue

            reshape_out_legacy_shape = [int(v) for v in list(reshape_out_shape)]
            if not _swap_rank3_shape_in_reshape(reshape_op):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[pre_in_name],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(unary_out_name, None),
                perm_nchw_to_nhwc,
            )

            reshape_inputs = [str(v) for v in list(reshape_op.inputs)]
            reshape_inputs[0] = str(unary_out_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=reshape_inputs,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(reshape_out_name, None),
                perm_hcw_to_hwc,
            )
            if len(non_expand_users) > 0:
                bridge_perm_name = _unique_tensor_name(f"{reshape_out_name}_legacy_hcw_perm")
                bridge_out_name = _unique_tensor_name(f"{reshape_out_name}_legacy_hcw")
                model_ir.tensors[bridge_perm_name] = TensorIR(
                    name=bridge_perm_name,
                    dtype="INT32",
                    shape=[3],
                    shape_signature=[3],
                    data=np.asarray([0, 2, 1], dtype=np.int32),
                    is_variable=False,
                )
                reshape_tensor = model_ir.tensors.get(reshape_out_name, None)
                model_ir.tensors[bridge_out_name] = TensorIR(
                    name=bridge_out_name,
                    dtype=str(reshape_tensor.dtype if reshape_tensor is not None else "FLOAT32"),
                    shape=[int(v) for v in list(reshape_out_legacy_shape)],
                    shape_signature=[int(v) for v in list(reshape_out_legacy_shape)],
                )
                for non_expand_user_idx in non_expand_users:
                    non_expand_user_op = model_ir.operators[int(non_expand_user_idx)]
                    for input_idx, input_name in enumerate(list(non_expand_user_op.inputs)):
                        if str(input_name) == str(reshape_out_name):
                            _replace_operator_input_at(
                                model_ir=model_ir,
                                op=non_expand_user_op,
                                input_index=int(input_idx),
                                new_input_name=str(bridge_out_name),
                            )
                model_ir.operators.append(
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(reshape_out_name), str(bridge_perm_name)],
                        outputs=[str(bridge_out_name)],
                        options={},
                    )
                )

            _write_const_ints_to_tensor(
                model_ir.tensors.get(str(expand_op.inputs[1]), None),
                [1],
            )
            if model_ir.tensors.get(expand_out_name, None) is not None and model_ir.tensors.get(post_out_name, None) is not None:
                post_tensor = model_ir.tensors.get(post_out_name)
                expand_tensor = model_ir.tensors.get(expand_out_name)
                if post_tensor is not None and expand_tensor is not None:
                    expand_tensor.shape = [int(v) for v in list(post_tensor.shape)]
                    expand_tensor.shape_signature = (
                        [int(v) for v in list(post_tensor.shape_signature)]
                        if post_tensor.shape_signature is not None
                        else [int(v) for v in list(post_tensor.shape)]
                    )
                    expand_tensor.dtype = str(post_tensor.dtype)
                    expand_tensor.quantization = _clone_quantization(post_tensor.quantization)

            _replace_tensor_inputs(model_ir, post_out_name, expand_out_name)
            model_ir.outputs = [
                str(expand_out_name) if str(v) == str(post_out_name) else str(v)
                for v in list(model_ir.outputs)
            ]

            for remove_idx in sorted({int(pre_idx), int(mid_idx), int(post_idx)}, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_unary_transpose_reshape_expanddims_transpose_nhwc_chains": int(rewritten)
    }


def _optimize_transpose_squeeze_unary_expanddims_transpose_nhwc_fanout_bypass_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Bypass Conv1D-shim ExpandDims->Transpose branches under unary fanout.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw --SQUEEZE(axis=k)--> s3 --UNARY--> u3
      u3 --EXPAND_DIMS(axis=k)--> e4 --T(0,2,3,1)--> y_nhwc --(...)
      u3 --(other consumers)--> ...

    Rewrite:
      x_nhwc --UNARY--> y_nhwc --T(0,3,1,2)--> x_nchw --SQUEEZE(axis=k)--> u3
      y_nhwc --(...)  # former post-transpose users
      u3 --(other consumers)--> ...

    This keeps the NCHW side branch intact while eliminating only the
    ExpandDims->Transpose wrapper branch.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
        "HARD_SWISH",
    }

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            squeeze_idx = int(pre_users[0])
            squeeze_op = model_ir.operators[int(squeeze_idx)]
            if (
                str(squeeze_op.op_type) != "SQUEEZE"
                or len(squeeze_op.inputs) != 1
                or len(squeeze_op.outputs) != 1
                or str(squeeze_op.inputs[0]) != pre_output_name
            ):
                continue
            squeeze_out_name = str(squeeze_op.outputs[0])

            squeeze_users = [int(v) for v in consumers.get(squeeze_out_name, [])]
            if len(squeeze_users) != 1:
                continue
            unary_idx = int(squeeze_users[0])
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.inputs[0]) != squeeze_out_name
            ):
                continue
            unary_out_name = str(unary_op.outputs[0])

            pre_input_shape = _shape_list(pre_input_name)
            pre_output_shape = _shape_list(pre_output_name)
            squeeze_out_shape = _shape_list(squeeze_out_name)
            unary_out_shape = _shape_list(unary_out_name)
            if (
                pre_input_shape is None
                or pre_output_shape is None
                or squeeze_out_shape is None
                or unary_out_shape is None
            ):
                continue
            if (
                len(pre_input_shape) != 4
                or len(pre_output_shape) != 4
                or len(squeeze_out_shape) != 3
                or len(unary_out_shape) != 3
                or not _shape_compatible(squeeze_out_shape, unary_out_shape)
            ):
                continue

            squeeze_axis: Optional[int] = None
            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            if "squeezeDims" in squeeze_options:
                raw_axes = np.asarray(squeeze_options.get("squeezeDims", []), dtype=np.int64).reshape(-1)
                normalized_axes = _normalize_squeeze_axes_for_rank(
                    [int(v) for v in raw_axes.tolist()],
                    4,
                )
                if normalized_axes is None or len(normalized_axes) != 1:
                    continue
                candidate_axis = int(normalized_axes[0])
                if int(pre_output_shape[candidate_axis]) != 1:
                    continue
                candidate_squeezed = [
                    int(pre_output_shape[idx])
                    for idx in range(4)
                    if int(idx) != int(candidate_axis)
                ]
                if not _shape_compatible(candidate_squeezed, squeeze_out_shape):
                    continue
                squeeze_axis = int(candidate_axis)
            else:
                candidate_axes: List[int] = []
                for axis in range(4):
                    if int(pre_output_shape[axis]) != 1:
                        continue
                    candidate_squeezed = [
                        int(pre_output_shape[idx])
                        for idx in range(4)
                        if int(idx) != int(axis)
                    ]
                    if _shape_compatible(candidate_squeezed, squeeze_out_shape):
                        candidate_axes.append(int(axis))
                if len(candidate_axes) != 1:
                    continue
                squeeze_axis = int(candidate_axes[0])

            unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
            candidate_pairs: List[Tuple[int, int, str]] = []
            for unary_user_idx in unary_users:
                expand_idx = int(unary_user_idx)
                expand_op = model_ir.operators[int(expand_idx)]
                if (
                    str(expand_op.op_type) != "EXPAND_DIMS"
                    or len(expand_op.inputs) < 2
                    or len(expand_op.outputs) != 1
                    or str(expand_op.inputs[0]) != unary_out_name
                ):
                    continue
                expand_axis_values = _read_const_ints_from_tensor(
                    model_ir.tensors.get(str(expand_op.inputs[1]), None)
                )
                if expand_axis_values is None or len(expand_axis_values) != 1:
                    continue
                expand_axis = int(expand_axis_values[0])
                if expand_axis < 0:
                    expand_axis += 4
                if expand_axis < 0 or expand_axis >= 4 or int(expand_axis) != int(squeeze_axis):
                    continue
                expand_out_name = str(expand_op.outputs[0])
                expand_users = [int(v) for v in consumers.get(expand_out_name, [])]
                if len(expand_users) != 1:
                    continue
                post_idx = int(expand_users[0])
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != expand_out_name
                    or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                ):
                    continue
                post_output_name = str(post_op.outputs[0])
                post_output_shape = _shape_list(post_output_name)
                expand_out_shape = _shape_list(expand_out_name)
                if (
                    post_output_shape is None
                    or expand_out_shape is None
                    or len(post_output_shape) != 4
                    or len(expand_out_shape) != 4
                    or not _shape_compatible(pre_input_shape, post_output_shape)
                    or not _shape_compatible(pre_output_shape, expand_out_shape)
                ):
                    continue
                candidate_pairs.append((int(expand_idx), int(post_idx), str(post_output_name)))

            if not candidate_pairs:
                continue

            # Pick first eligible branch and bypass it.
            expand_idx, post_idx, post_output_name = candidate_pairs[0]
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            old_unary_out_tensor = model_ir.tensors.get(unary_out_name, None)
            if post_output_tensor is not None and old_unary_out_tensor is not None:
                post_output_tensor.dtype = str(old_unary_out_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_unary_out_tensor.quantization)

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[pre_input_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=unary_op,
                new_outputs=[post_output_name],
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=pre_op,
                input_index=0,
                new_input_name=str(post_output_name),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=squeeze_op,
                new_outputs=[unary_out_name],
            )

            for remove_idx in sorted({int(expand_idx), int(post_idx)}, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_squeeze_unary_expanddims_transpose_nhwc_fanout_bypass_chains": int(rewritten)
    }


def _optimize_transpose_squeeze_instancenorm_unary_expanddims_transpose_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC<->NCHW wrapper transposes around:
      TRANSPOSE -> SQUEEZE -> RESHAPE -> InstanceNorm(flat) -> RESHAPE
      -> UNARY -> EXPAND_DIMS -> TRANSPOSE

    The middle subgraph operates on [N,1,K] and is permutation-equivariant
    over K in the matched pattern, so we can swap C/W interpretation at the
    two RESHAPE boundaries and drop both transposes.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
        "HARD_SWISH",
    }

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    def _single_consumer_idx(consumers: Dict[str, List[int]], tensor_name: str) -> Optional[int]:
        users = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(users) != 1:
            return None
        return int(users[0])

    def _normalize_axis_list_for_rank(raw_axes: Optional[List[int]], rank: int) -> Optional[List[int]]:
        if raw_axes is None:
            return None
        normalized: List[int] = []
        for axis in raw_axes:
            v = int(axis)
            if v < 0:
                v += int(rank)
            if v < 0 or v >= int(rank):
                return None
            normalized.append(int(v))
        return sorted(set(normalized))

    def _mean_axes_keepdims_ok(mean_op: OperatorIR) -> bool:
        if len(mean_op.inputs) < 2:
            return False
        axes = _read_const_ints_from_tensor(model_ir.tensors.get(str(mean_op.inputs[1]), None))
        axes = _normalize_axis_list_for_rank(axes, 3)
        if axes != [2]:
            return False
        opts = dict(mean_op.options) if isinstance(mean_op.options, dict) else {}
        keep_dims = bool(opts.get("keepDims", opts.get("keep_dims", False)))
        return bool(keep_dims)

    def _find_binary_consumer(
        *,
        consumers: Dict[str, List[int]],
        lhs_tensor: str,
        rhs_tensor: Optional[str],
        op_type: str,
    ) -> Optional[int]:
        for idx in [int(v) for v in consumers.get(str(lhs_tensor), [])]:
            op = model_ir.operators[int(idx)]
            if str(op.op_type) != str(op_type):
                continue
            op_inputs = [str(v) for v in list(op.inputs)]
            if str(lhs_tensor) not in op_inputs:
                continue
            if rhs_tensor is not None and str(rhs_tensor) not in op_inputs:
                continue
            if len(op.outputs) != 1:
                continue
            return int(idx)
        return None

    def _find_unary_consumer(
        *,
        consumers: Dict[str, List[int]],
        input_tensor: str,
        op_type: str,
    ) -> Optional[int]:
        for idx in [int(v) for v in consumers.get(str(input_tensor), [])]:
            op = model_ir.operators[int(idx)]
            if str(op.op_type) != str(op_type):
                continue
            if len(op.inputs) != 1 or len(op.outputs) != 1:
                continue
            if str(op.inputs[0]) != str(input_tensor):
                continue
            return int(idx)
        return None

    def _find_mean_consumer(
        *,
        consumers: Dict[str, List[int]],
        input_tensor: str,
    ) -> Optional[int]:
        for idx in [int(v) for v in consumers.get(str(input_tensor), [])]:
            op = model_ir.operators[int(idx)]
            if str(op.op_type) != "MEAN":
                continue
            if len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            if str(op.inputs[0]) != str(input_tensor):
                continue
            return int(idx)
        return None

    def _update_tensor_shape(name: str, shape: List[int]) -> None:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None:
            return
        tensor.shape = [int(v) for v in list(shape)]
        tensor.shape_signature = [int(v) for v in list(shape)]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            pre_input_shape = _shape_list(pre_input_name)
            pre_output_shape = _shape_list(pre_output_name)
            if (
                pre_input_shape is None
                or pre_output_shape is None
                or len(pre_input_shape) != 4
                or len(pre_output_shape) != 4
                or not _is_fully_known_positive_shape(pre_input_shape)
                or not _is_fully_known_positive_shape(pre_output_shape)
            ):
                continue

            squeeze_idx = _single_consumer_idx(consumers, pre_output_name)
            if squeeze_idx is None:
                continue
            squeeze_op = model_ir.operators[int(squeeze_idx)]
            if (
                str(squeeze_op.op_type) != "SQUEEZE"
                or len(squeeze_op.inputs) != 1
                or len(squeeze_op.outputs) != 1
                or str(squeeze_op.inputs[0]) != pre_output_name
            ):
                continue
            squeeze_out_name = str(squeeze_op.outputs[0])
            squeeze_out_shape = _shape_list(squeeze_out_name)
            if squeeze_out_shape is None or len(squeeze_out_shape) != 3:
                continue

            # Require squeeze axis==2 on pre_output shape [N,C,1,W].
            squeeze_axis: Optional[int] = None
            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            if "squeezeDims" in squeeze_options:
                raw_axes = np.asarray(squeeze_options.get("squeezeDims", []), dtype=np.int64).reshape(-1)
                normalized_axes = _normalize_squeeze_axes_for_rank(
                    [int(v) for v in raw_axes.tolist()],
                    4,
                )
                if normalized_axes is None or len(normalized_axes) != 1:
                    continue
                squeeze_axis = int(normalized_axes[0])
            else:
                candidate_axes: List[int] = []
                for axis in range(4):
                    if int(pre_output_shape[axis]) != 1:
                        continue
                    candidate_squeezed = [
                        int(pre_output_shape[idx])
                        for idx in range(4)
                        if int(idx) != int(axis)
                    ]
                    if _shape_compatible(candidate_squeezed, squeeze_out_shape):
                        candidate_axes.append(int(axis))
                if len(candidate_axes) != 1:
                    continue
                squeeze_axis = int(candidate_axes[0])
            if int(squeeze_axis) != 2:
                continue

            reshape1_idx = _single_consumer_idx(consumers, squeeze_out_name)
            if reshape1_idx is None:
                continue
            reshape1_op = model_ir.operators[int(reshape1_idx)]
            if (
                str(reshape1_op.op_type) != "RESHAPE"
                or len(reshape1_op.inputs) < 2
                or len(reshape1_op.outputs) != 1
                or str(reshape1_op.inputs[0]) != squeeze_out_name
            ):
                continue
            reshape1_out_name = str(reshape1_op.outputs[0])
            reshape1_out_shape = _shape_list(reshape1_out_name)
            if (
                reshape1_out_shape is None
                or len(reshape1_out_shape) != 3
                or not _is_fully_known_positive_shape(reshape1_out_shape)
                or int(reshape1_out_shape[1]) != 1
            ):
                continue

            # Match InstanceNorm(flat) pattern:
            # mean1(x), sub(x,mean1), mul(sub,sub), mean2, add, sqrt, div, mul, mul, add.
            mean1_idx = _find_mean_consumer(
                consumers=consumers,
                input_tensor=reshape1_out_name,
            )
            if mean1_idx is None:
                continue
            mean1_op = model_ir.operators[int(mean1_idx)]
            if not _mean_axes_keepdims_ok(mean1_op):
                continue
            mean1_out_name = str(mean1_op.outputs[0])

            sub_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=reshape1_out_name,
                rhs_tensor=mean1_out_name,
                op_type="SUB",
            )
            if sub_idx is None:
                continue
            sub_op = model_ir.operators[int(sub_idx)]
            sub_out_name = str(sub_op.outputs[0])

            mul_sq_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=sub_out_name,
                rhs_tensor=sub_out_name,
                op_type="MUL",
            )
            if mul_sq_idx is None:
                continue
            mul_sq_out_name = str(model_ir.operators[int(mul_sq_idx)].outputs[0])

            mean2_idx = _find_mean_consumer(
                consumers=consumers,
                input_tensor=mul_sq_out_name,
            )
            if mean2_idx is None:
                continue
            mean2_op = model_ir.operators[int(mean2_idx)]
            if not _mean_axes_keepdims_ok(mean2_op):
                continue
            mean2_out_name = str(mean2_op.outputs[0])

            add_eps_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=mean2_out_name,
                rhs_tensor=None,
                op_type="ADD",
            )
            if add_eps_idx is None:
                continue
            add_eps_out_name = str(model_ir.operators[int(add_eps_idx)].outputs[0])

            sqrt_idx = _find_unary_consumer(
                consumers=consumers,
                input_tensor=add_eps_out_name,
                op_type="SQRT",
            )
            if sqrt_idx is None:
                continue
            sqrt_out_name = str(model_ir.operators[int(sqrt_idx)].outputs[0])

            div_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=sqrt_out_name,
                rhs_tensor=None,
                op_type="DIV",
            )
            if div_idx is None:
                # DIV may have sqrt on input1
                for candidate_idx in [int(v) for v in consumers.get(str(sqrt_out_name), [])]:
                    op = model_ir.operators[int(candidate_idx)]
                    if str(op.op_type) != "DIV" or len(op.outputs) != 1:
                        continue
                    if str(sqrt_out_name) not in [str(v) for v in list(op.inputs)]:
                        continue
                    div_idx = int(candidate_idx)
                    break
            if div_idx is None:
                continue
            div_out_name = str(model_ir.operators[int(div_idx)].outputs[0])

            mul_norm_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=sub_out_name,
                rhs_tensor=div_out_name,
                op_type="MUL",
            )
            if mul_norm_idx is None:
                continue
            mul_norm_out_name = str(model_ir.operators[int(mul_norm_idx)].outputs[0])

            mul_scale_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=mul_norm_out_name,
                rhs_tensor=None,
                op_type="MUL",
            )
            if mul_scale_idx is None:
                continue
            mul_scale_out_name = str(model_ir.operators[int(mul_scale_idx)].outputs[0])

            add_bias_idx = _find_binary_consumer(
                consumers=consumers,
                lhs_tensor=mul_scale_out_name,
                rhs_tensor=None,
                op_type="ADD",
            )
            if add_bias_idx is None:
                continue
            add_bias_out_name = str(model_ir.operators[int(add_bias_idx)].outputs[0])

            reshape2_idx = _single_consumer_idx(consumers, add_bias_out_name)
            if reshape2_idx is None:
                continue
            reshape2_op = model_ir.operators[int(reshape2_idx)]
            if (
                str(reshape2_op.op_type) != "RESHAPE"
                or len(reshape2_op.inputs) < 2
                or len(reshape2_op.outputs) != 1
                or str(reshape2_op.inputs[0]) != add_bias_out_name
            ):
                continue
            reshape2_in_name = str(reshape2_op.inputs[0])
            reshape2_out_name = str(reshape2_op.outputs[0])
            reshape2_in_shape = _shape_list(reshape2_in_name)
            reshape2_out_shape = _shape_list(reshape2_out_name)
            if (
                reshape2_in_shape is None
                or reshape2_out_shape is None
                or len(reshape2_in_shape) != 3
                or len(reshape2_out_shape) != 3
                or not _shape_compatible(reshape2_in_shape, reshape1_out_shape)
                or not _shape_compatible(reshape2_out_shape, squeeze_out_shape)
                or not _is_fully_known_positive_shape(reshape2_out_shape)
            ):
                continue

            unary_idx = _single_consumer_idx(consumers, reshape2_out_name)
            if unary_idx is None:
                continue
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.inputs[0]) != reshape2_out_name
            ):
                continue
            unary_out_name = str(unary_op.outputs[0])
            unary_out_shape = _shape_list(unary_out_name)
            if unary_out_shape is None or len(unary_out_shape) != 3 or not _shape_compatible(unary_out_shape, reshape2_out_shape):
                continue

            expand_idx = _single_consumer_idx(consumers, unary_out_name)
            if expand_idx is None:
                continue
            expand_op = model_ir.operators[int(expand_idx)]
            if (
                str(expand_op.op_type) != "EXPAND_DIMS"
                or len(expand_op.inputs) < 2
                or len(expand_op.outputs) != 1
                or str(expand_op.inputs[0]) != unary_out_name
            ):
                continue
            expand_out_name = str(expand_op.outputs[0])
            expand_axis_values = _read_const_ints_from_tensor(model_ir.tensors.get(str(expand_op.inputs[1]), None))
            if expand_axis_values is None or len(expand_axis_values) != 1:
                continue
            expand_axis = int(expand_axis_values[0])
            if expand_axis < 0:
                expand_axis += 4
            if int(expand_axis) != 2:
                continue

            post_idx = _single_consumer_idx(consumers, expand_out_name)
            if post_idx is None:
                continue
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != expand_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])
            post_out_shape = _shape_list(post_out_name)
            expand_out_shape = _shape_list(expand_out_name)
            if (
                post_out_shape is None
                or expand_out_shape is None
                or len(post_out_shape) != 4
                or len(expand_out_shape) != 4
                or not _shape_compatible(post_out_shape, pre_input_shape)
                or not _shape_compatible(expand_out_shape, pre_output_shape)
            ):
                continue

            n = int(pre_input_shape[0])
            w = int(pre_input_shape[2])
            c = int(pre_input_shape[3])
            new_rank3_shape = [n, w, c]
            new_rank4_shape = [n, 1, w, c]

            # Rewrite axes and reshape target.
            squeeze_opts = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            squeeze_opts["squeezeDims"] = [1]
            squeeze_op.options = squeeze_opts
            _set_operator_inputs(
                model_ir=model_ir,
                op=squeeze_op,
                new_inputs=[pre_input_name],
            )
            _update_tensor_shape(squeeze_out_name, new_rank3_shape)

            reshape2_opts = dict(reshape2_op.options) if isinstance(reshape2_op.options, dict) else {}
            for key in ["newShape", "onnxRawNewShape"]:
                if key in reshape2_opts and isinstance(reshape2_opts[key], list):
                    reshape2_opts[key] = [int(v) for v in list(new_rank3_shape)]
            reshape2_op.options = reshape2_opts
            if len(reshape2_op.inputs) >= 2:
                _write_const_ints_to_tensor(
                    model_ir.tensors.get(str(reshape2_op.inputs[1]), None),
                    [int(v) for v in list(new_rank3_shape)],
                )
            _update_tensor_shape(reshape2_out_name, new_rank3_shape)
            _update_tensor_shape(unary_out_name, new_rank3_shape)

            if len(expand_op.inputs) >= 2:
                _write_const_ints_to_tensor(
                    model_ir.tensors.get(str(expand_op.inputs[1]), None),
                    [1],
                )
            _update_tensor_shape(expand_out_name, new_rank4_shape)

            _replace_tensor_inputs(model_ir, post_out_name, expand_out_name)
            model_ir.outputs = [
                str(expand_out_name) if str(v) == str(post_out_name) else str(v)
                for v in list(model_ir.outputs)
            ]

            for remove_idx in sorted({int(pre_idx), int(post_idx)}, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_squeeze_instancenorm_unary_expanddims_transpose_nhwc_chains": int(rewritten)
    }


def _optimize_tencoder_add_expand_transpose_conv_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove Conv1D-shim post-transposes on dual tencoder branches before residual ADD.

    Target (NCHW-3D bridge):
      lhs4 --T(0,3,1,2)--> lhs4n --SQUEEZE(axis=2)--> lhs3 [N,C,W]
      rhs4 --T(0,3,1,2)--> rhs4n --SQUEEZE(axis=2)--> rhs3 [N,2C,W]
      rhs3 --RESHAPE/InstanceNorm/RESHAPE--> split_src [N,2C,W]
      split_src --SLICE--> a [N,C,W], b [N,C,W] --LOGISTIC/MUL/scale--> rhs_gate [N,C,W]
      lhs3 + rhs_gate -> add_out [N,C,W]
      add_out --EXPAND_DIMS(axis=2)--> add4 --T(0,2,3,1)--> add4_nhwc -> CONV_2D

    Rewrite (NHWC-3D bridge):
      lhs4 --SQUEEZE(axis=1)--> lhs3' [N,W,C]
      rhs4 --SQUEEZE(axis=1)--> rhs3' [N,W,2C]
      reshape/slice/scale constants adjusted to [N,W,*]
      lhs3' + rhs_gate' -> add_out' [N,W,C]
      add_out' --EXPAND_DIMS(axis=1)--> add4' [N,1,W,C] -> CONV_2D
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _update_tensor_shape(name: str, shape: List[int]) -> None:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None:
            return
        tensor.shape = [int(v) for v in list(shape)]
        tensor.shape_signature = [int(v) for v in list(shape)]

    def _resolve_single_squeeze_axis(
        squeeze_op: OperatorIR,
        input_shape: List[int],
        output_shape: List[int],
    ) -> Optional[int]:
        opts = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
        if "squeezeDims" in opts:
            raw_axes = np.asarray(opts.get("squeezeDims", []), dtype=np.int64).reshape(-1)
            normalized_axes = _normalize_squeeze_axes_for_rank(
                [int(v) for v in raw_axes.tolist()],
                int(len(input_shape)),
            )
            if normalized_axes is None or len(normalized_axes) != 1:
                return None
            axis = int(normalized_axes[0])
            if int(input_shape[axis]) != 1:
                return None
            candidate = [int(input_shape[idx]) for idx in range(len(input_shape)) if int(idx) != int(axis)]
            if candidate != [int(v) for v in list(output_shape)]:
                return None
            return int(axis)
        candidates: List[int] = []
        for axis in range(len(input_shape)):
            if int(input_shape[axis]) != 1:
                continue
            candidate = [int(input_shape[idx]) for idx in range(len(input_shape)) if int(idx) != int(axis)]
            if candidate == [int(v) for v in list(output_shape)]:
                candidates.append(int(axis))
        if len(candidates) != 1:
            return None
        return int(candidates[0])

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _clone_or_rewrite_scale_const(
        *,
        const_name: str,
        target_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        const_tensor = model_ir.tensors.get(str(const_name), None)
        if const_tensor is None or const_tensor.data is None:
            return None
        const_data = np.asarray(const_tensor.data)
        if const_data.ndim != 2:
            return None
        if int(const_data.shape[0]) >= 1 and int(const_data.shape[1]) == 1:
            new_data = np.asarray(const_data.T)
        elif int(const_data.shape[0]) == 1 and int(const_data.shape[1]) >= 1:
            new_data = np.asarray(const_data)
        else:
            return None
        users = [int(v) for v in consumers.get(str(const_name), [])]
        shared_outside_target = any(int(v) not in target_users for v in users)
        if shared_outside_target:
            new_name = _unique_tensor_name(f"{const_name}_nhwc")
            new_shape = [int(v) for v in list(new_data.shape)]
            model_ir.tensors[new_name] = TensorIR(
                name=new_name,
                dtype=str(const_tensor.dtype),
                shape=new_shape,
                shape_signature=new_shape,
                data=np.asarray(new_data),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            return str(new_name)
        const_tensor.data = np.asarray(new_data)
        const_tensor.shape = [int(v) for v in list(new_data.shape)]
        const_tensor.shape_signature = [int(v) for v in list(new_data.shape)]
        return str(const_name)

    def _find_upstream_squeeze_transpose(
        *,
        start_tensor: str,
        producers: Dict[str, int],
    ) -> Optional[Tuple[int, int, str, str]]:
        stack: List[str] = [str(start_tensor)]
        visited: set[str] = set()
        while stack:
            current = str(stack.pop())
            if current in visited:
                continue
            visited.add(current)
            producer_idx = producers.get(str(current), None)
            if producer_idx is None:
                continue
            producer_op = model_ir.operators[int(producer_idx)]
            if (
                str(producer_op.op_type) == "SQUEEZE"
                and len(producer_op.inputs) == 1
                and len(producer_op.outputs) == 1
            ):
                squeeze_in = str(producer_op.inputs[0])
                pre_idx = producers.get(str(squeeze_in), None)
                if pre_idx is not None:
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) == "TRANSPOSE"
                        and len(pre_op.inputs) >= 2
                        and len(pre_op.outputs) == 1
                        and _read_transpose_perm(model_ir, pre_op) == perm_nhwc_to_nchw
                        and str(pre_op.outputs[0]) == str(squeeze_in)
                    ):
                        return (
                            int(pre_idx),
                            int(producer_idx),
                            str(pre_op.inputs[0]),
                            str(producer_op.outputs[0]),
                        )
            for in_name in [str(v) for v in list(producer_op.inputs)]:
                if producers.get(str(in_name), None) is not None:
                    stack.append(str(in_name))
        return None

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_in_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])

            expand_idx = producers.get(post_in_name, None)
            if expand_idx is None:
                continue
            expand_op = model_ir.operators[int(expand_idx)]
            if (
                str(expand_op.op_type) != "EXPAND_DIMS"
                or len(expand_op.inputs) < 2
                or len(expand_op.outputs) != 1
                or str(expand_op.outputs[0]) != post_in_name
            ):
                continue
            expand_axis_values = _read_const_ints_from_tensor(
                model_ir.tensors.get(str(expand_op.inputs[1]), None)
            )
            if expand_axis_values is None or len(expand_axis_values) != 1:
                continue
            expand_axis = int(expand_axis_values[0])
            if expand_axis < 0:
                expand_axis += 4
            if int(expand_axis) != 2:
                continue

            add_out_name = str(expand_op.inputs[0])
            if add_out_name in model_outputs:
                continue
            add_idx = producers.get(add_out_name, None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if (
                str(add_op.op_type) != "ADD"
                or len(add_op.inputs) != 2
                or len(add_op.outputs) != 1
                or str(add_op.outputs[0]) != add_out_name
            ):
                continue
            add_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_users) < 1:
                continue

            def _match_simple_branch(out_name: str) -> Optional[Dict[str, Any]]:
                sq_idx = producers.get(str(out_name), None)
                if sq_idx is None:
                    return None
                sq_op = model_ir.operators[int(sq_idx)]
                if (
                    str(sq_op.op_type) != "SQUEEZE"
                    or len(sq_op.inputs) != 1
                    or len(sq_op.outputs) != 1
                    or str(sq_op.outputs[0]) != str(out_name)
                ):
                    return None
                pre_name = str(sq_op.inputs[0])
                pre_idx = producers.get(pre_name, None)
                if pre_idx is None:
                    return None
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != pre_name
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                ):
                    return None
                src4_name = str(pre_op.inputs[0])
                src4_shape = _shape_list(src4_name)
                out_shape = _shape_list(str(out_name))
                if (
                    src4_shape is None
                    or out_shape is None
                    or len(src4_shape) != 4
                    or len(out_shape) != 3
                    or int(src4_shape[1]) != 1
                    or not _is_fully_known_positive_shape(src4_shape)
                    or not _is_fully_known_positive_shape(out_shape)
                ):
                    return None
                axis = _resolve_single_squeeze_axis(sq_op, _shape_list(pre_name) or [], out_shape)
                if axis != 2:
                    return None
                n, _, w, c = [int(v) for v in list(src4_shape)]
                if [n, c, w] != [int(v) for v in list(out_shape)]:
                    return None
                return {
                    "kind": "simple",
                    "pre_idx": int(pre_idx),
                    "sq_idx": int(sq_idx),
                    "src4_name": str(src4_name),
                    "out_name": str(out_name),
                    "n": int(n),
                    "c": int(c),
                    "w": int(w),
                }

            def _match_legacy_rank3_branch(out_name: str) -> Optional[Dict[str, Any]]:
                pre_idx = producers.get(str(out_name), None)
                if pre_idx is None:
                    return None
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(out_name)
                ):
                    return None
                if _read_transpose_perm(model_ir, pre_op) != [0, 2, 1]:
                    return None
                src3_name = str(pre_op.inputs[0])
                src3_shape = _shape_list(src3_name)
                out_shape = _shape_list(str(out_name))
                if (
                    src3_shape is None
                    or out_shape is None
                    or len(src3_shape) != 3
                    or len(out_shape) != 3
                    or not _is_fully_known_positive_shape(src3_shape)
                    or not _is_fully_known_positive_shape(out_shape)
                ):
                    return None
                n, w, c = [int(v) for v in list(src3_shape)]
                if [n, c, w] != [int(v) for v in list(out_shape)]:
                    return None
                return {
                    "kind": "legacy_rank3",
                    "pre_idx": int(pre_idx),
                    "src3_name": str(src3_name),
                    "out_name": str(out_name),
                    "n": int(n),
                    "c": int(c),
                    "w": int(w),
                }

            def _match_lhs_branch(out_name: str) -> Optional[Dict[str, Any]]:
                lhs_branch = _match_simple_branch(out_name)
                if lhs_branch is not None:
                    return lhs_branch
                return _match_legacy_rank3_branch(out_name)

            def _match_gate_branch(out_name: str) -> Optional[Dict[str, Any]]:
                scale_idx = producers.get(str(out_name), None)
                if scale_idx is None:
                    return None
                scale_op = model_ir.operators[int(scale_idx)]
                if (
                    str(scale_op.op_type) != "MUL"
                    or len(scale_op.inputs) != 2
                    or len(scale_op.outputs) != 1
                    or str(scale_op.outputs[0]) != str(out_name)
                ):
                    return None
                scale_inputs = [str(v) for v in list(scale_op.inputs)]
                const_name: Optional[str] = None
                gate_name: Optional[str] = None
                for cand in scale_inputs:
                    tensor = model_ir.tensors.get(str(cand), None)
                    if tensor is not None and tensor.data is not None and np.asarray(tensor.data).ndim == 2:
                        const_name = str(cand)
                    else:
                        gate_name = str(cand)
                if const_name is None or gate_name is None:
                    return None
                const_tensor = model_ir.tensors.get(str(const_name), None)
                if const_tensor is None or const_tensor.data is None:
                    return None
                const_shape = [int(v) for v in list(np.asarray(const_tensor.data).shape)]
                if len(const_shape) != 2 or int(min(const_shape)) != 1:
                    return None
                c = int(max(const_shape))

                gate_idx = producers.get(str(gate_name), None)
                if gate_idx is None:
                    return None
                gate_op = model_ir.operators[int(gate_idx)]
                if (
                    str(gate_op.op_type) != "MUL"
                    or len(gate_op.inputs) != 2
                    or len(gate_op.outputs) != 1
                    or str(gate_op.outputs[0]) != str(gate_name)
                ):
                    return None
                g0, g1 = [str(v) for v in list(gate_op.inputs)]
                g0_prod = producers.get(g0, None)
                g1_prod = producers.get(g1, None)
                logistic_name: Optional[str] = None
                slice0_name: Optional[str] = None
                if g0_prod is not None and str(model_ir.operators[int(g0_prod)].op_type) == "LOGISTIC":
                    logistic_name = str(g0)
                    slice0_name = str(g1)
                elif g1_prod is not None and str(model_ir.operators[int(g1_prod)].op_type) == "LOGISTIC":
                    logistic_name = str(g1)
                    slice0_name = str(g0)
                else:
                    return None
                logistic_idx = producers.get(str(logistic_name), None)
                slice0_idx = producers.get(str(slice0_name), None)
                if logistic_idx is None or slice0_idx is None:
                    return None
                logistic_op = model_ir.operators[int(logistic_idx)]
                slice0_op = model_ir.operators[int(slice0_idx)]
                if (
                    str(logistic_op.op_type) != "LOGISTIC"
                    or len(logistic_op.inputs) != 1
                    or str(slice0_op.op_type) != "SLICE"
                    or len(slice0_op.inputs) < 3
                    or len(slice0_op.outputs) != 1
                    or str(slice0_op.outputs[0]) != str(slice0_name)
                ):
                    return None
                slice1_name = str(logistic_op.inputs[0])
                slice1_idx = producers.get(str(slice1_name), None)
                if slice1_idx is None:
                    return None
                slice1_op = model_ir.operators[int(slice1_idx)]
                if (
                    str(slice1_op.op_type) != "SLICE"
                    or len(slice1_op.inputs) < 3
                    or len(slice1_op.outputs) != 1
                    or str(slice1_op.outputs[0]) != str(slice1_name)
                ):
                    return None
                split_src0 = str(slice0_op.inputs[0])
                split_src1 = str(slice1_op.inputs[0])
                if str(split_src0) != str(split_src1):
                    return None
                split_src = str(split_src0)

                split_src_shape = _shape_list(split_src)
                slice0_shape = _shape_list(slice0_name)
                slice1_shape = _shape_list(slice1_name)
                out_shape = _shape_list(str(out_name))
                if (
                    split_src_shape is None
                    or slice0_shape is None
                    or slice1_shape is None
                    or out_shape is None
                    or len(split_src_shape) != 3
                    or len(slice0_shape) != 3
                    or len(slice1_shape) != 3
                    or len(out_shape) != 3
                    or not _is_fully_known_positive_shape(split_src_shape)
                    or not _is_fully_known_positive_shape(slice0_shape)
                    or not _is_fully_known_positive_shape(out_shape)
                ):
                    return None
                n = int(split_src_shape[0])
                c2 = int(split_src_shape[1])
                w = int(split_src_shape[2])
                if int(c2) != int(c * 2):
                    return None
                if [n, c, w] != [int(v) for v in list(slice0_shape)]:
                    return None
                if [n, c, w] != [int(v) for v in list(slice1_shape)]:
                    return None
                if [n, c, w] != [int(v) for v in list(out_shape)]:
                    return None

                begin0 = _read_const_ints_from_tensor(model_ir.tensors.get(str(slice0_op.inputs[1]), None))
                size0 = _read_const_ints_from_tensor(model_ir.tensors.get(str(slice0_op.inputs[2]), None))
                begin1 = _read_const_ints_from_tensor(model_ir.tensors.get(str(slice1_op.inputs[1]), None))
                size1 = _read_const_ints_from_tensor(model_ir.tensors.get(str(slice1_op.inputs[2]), None))
                if (
                    begin0 is None
                    or size0 is None
                    or begin1 is None
                    or size1 is None
                    or len(begin0) != 3
                    or len(size0) != 3
                    or len(begin1) != 3
                    or len(size1) != 3
                ):
                    return None
                if [int(v) for v in begin0] != [0, 0, 0]:
                    return None
                if [int(v) for v in size0] != [n, c, w]:
                    return None
                if [int(v) for v in begin1] != [0, c, 0]:
                    return None
                if [int(v) for v in size1] != [n, c, w]:
                    return None

                reshape2_idx = producers.get(split_src, None)
                if reshape2_idx is None:
                    return None
                reshape2_op = model_ir.operators[int(reshape2_idx)]
                if (
                    str(reshape2_op.op_type) != "RESHAPE"
                    or len(reshape2_op.inputs) < 2
                    or len(reshape2_op.outputs) != 1
                    or str(reshape2_op.outputs[0]) != split_src
                ):
                    return None

                upstream = _find_upstream_squeeze_transpose(
                    start_tensor=str(reshape2_op.inputs[0]),
                    producers=producers,
                )
                if upstream is None:
                    return None
                pre_idx, sq_idx, src4_name, sq_out_name = upstream
                sq_out_shape = _shape_list(sq_out_name)
                src4_shape = _shape_list(src4_name)
                if (
                    sq_out_shape is None
                    or src4_shape is None
                    or len(sq_out_shape) != 3
                    or len(src4_shape) != 4
                    or int(src4_shape[1]) != 1
                    or [int(src4_shape[0]), int(src4_shape[3]), int(src4_shape[2])] != [n, c2, w]
                    or [int(v) for v in sq_out_shape] != [n, c2, w]
                ):
                    return None
                sq_op = model_ir.operators[int(sq_idx)]
                axis = _resolve_single_squeeze_axis(sq_op, _shape_list(str(sq_op.inputs[0])) or [], sq_out_shape)
                if axis != 2:
                    return None

                return {
                    "scale_idx": int(scale_idx),
                    "scale_const_name": str(const_name),
                    "gate_idx": int(gate_idx),
                    "logistic_idx": int(logistic_idx),
                    "slice0_idx": int(slice0_idx),
                    "slice1_idx": int(slice1_idx),
                    "slice0_name": str(slice0_name),
                    "slice1_name": str(slice1_name),
                    "split_src_name": str(split_src),
                    "reshape2_idx": int(reshape2_idx),
                    "rhs_pre_idx": int(pre_idx),
                    "rhs_sq_idx": int(sq_idx),
                    "rhs_src4_name": str(src4_name),
                    "rhs_sq_out_name": str(sq_out_name),
                    "n": int(n),
                    "c": int(c),
                    "w": int(w),
                    "c2": int(c2),
                    "scale_out_name": str(out_name),
                }

            add_input0 = str(add_op.inputs[0])
            add_input1 = str(add_op.inputs[1])
            lhs = _match_lhs_branch(add_input0)
            rhs = _match_gate_branch(add_input1)
            if lhs is None or rhs is None:
                lhs = _match_lhs_branch(add_input1)
                rhs = _match_gate_branch(add_input0)
            if lhs is None or rhs is None:
                continue

            n = int(lhs["n"])
            c = int(lhs["c"])
            w = int(lhs["w"])
            if int(rhs["n"]) != int(n) or int(rhs["c"]) != int(c) or int(rhs["w"]) != int(w):
                continue

            post_users = [int(v) for v in consumers.get(post_out_name, [])]
            if len(post_users) == 0:
                continue
            can_remove_post_transpose = (
                len(add_users) == 1
                and int(add_users[0]) == int(expand_idx)
                and post_out_name not in model_outputs
            )

            # Rewrite left branch to NHWC-3D.
            if str(lhs.get("kind", "")) == "simple":
                lhs_sq_op = model_ir.operators[int(lhs["sq_idx"])]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=lhs_sq_op,
                    new_inputs=[str(lhs["src4_name"])],
                )
                lhs_sq_opts = dict(lhs_sq_op.options) if isinstance(lhs_sq_op.options, dict) else {}
                lhs_sq_opts["squeezeDims"] = [1]
                lhs_sq_op.options = lhs_sq_opts
                _update_tensor_shape(str(lhs["out_name"]), [n, w, c])
            elif str(lhs.get("kind", "")) == "legacy_rank3":
                lhs_src_name = str(lhs["src3_name"])
                _update_tensor_shape(lhs_src_name, [n, w, c])
                for input_idx, input_name in enumerate(list(add_op.inputs)):
                    if str(input_name) == str(lhs["out_name"]):
                        _replace_operator_input_at(
                            model_ir=model_ir,
                            op=add_op,
                            input_index=int(input_idx),
                            new_input_name=lhs_src_name,
                        )
            else:
                continue

            # Rewrite right branch squeeze + reshape/slice/scale to NHWC-3D.
            rhs_sq_op = model_ir.operators[int(rhs["rhs_sq_idx"])]
            _set_operator_inputs(
                model_ir=model_ir,
                op=rhs_sq_op,
                new_inputs=[str(rhs["rhs_src4_name"])],
            )
            rhs_sq_opts = dict(rhs_sq_op.options) if isinstance(rhs_sq_op.options, dict) else {}
            rhs_sq_opts["squeezeDims"] = [1]
            rhs_sq_op.options = rhs_sq_opts
            _update_tensor_shape(str(rhs["rhs_sq_out_name"]), [n, w, int(rhs["c2"])])

            reshape2_op = model_ir.operators[int(rhs["reshape2_idx"])]
            reshape2_opts = dict(reshape2_op.options) if isinstance(reshape2_op.options, dict) else {}
            for key in ["newShape", "onnxRawNewShape"]:
                if key in reshape2_opts and isinstance(reshape2_opts[key], list):
                    reshape2_opts[key] = [int(v) for v in [n, w, int(rhs["c2"])]]
            reshape2_op.options = reshape2_opts
            if len(reshape2_op.inputs) >= 2:
                _write_const_ints_to_tensor(
                    model_ir.tensors.get(str(reshape2_op.inputs[1]), None),
                    [int(v) for v in [n, w, int(rhs["c2"])]],
                )
            _update_tensor_shape(str(rhs["split_src_name"]), [n, w, int(rhs["c2"])])

            slice0_op = model_ir.operators[int(rhs["slice0_idx"])]
            slice1_op = model_ir.operators[int(rhs["slice1_idx"])]
            _write_const_ints_to_tensor(
                model_ir.tensors.get(str(slice0_op.inputs[1]), None),
                [int(v) for v in [0, 0, 0]],
            )
            _write_const_ints_to_tensor(
                model_ir.tensors.get(str(slice0_op.inputs[2]), None),
                [int(v) for v in [n, w, c]],
            )
            _write_const_ints_to_tensor(
                model_ir.tensors.get(str(slice1_op.inputs[1]), None),
                [int(v) for v in [0, 0, c]],
            )
            _write_const_ints_to_tensor(
                model_ir.tensors.get(str(slice1_op.inputs[2]), None),
                [int(v) for v in [n, w, c]],
            )
            _update_tensor_shape(str(rhs["slice0_name"]), [n, w, c])
            _update_tensor_shape(str(rhs["slice1_name"]), [n, w, c])

            scale_users = {int(rhs["scale_idx"])}
            new_scale_const_name = _clone_or_rewrite_scale_const(
                const_name=str(rhs["scale_const_name"]),
                target_users=scale_users,
                consumers=consumers,
            )
            if new_scale_const_name is None:
                continue
            scale_op = model_ir.operators[int(rhs["scale_idx"])]
            if str(scale_op.inputs[0]) == str(rhs["scale_const_name"]):
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=scale_op,
                    input_index=0,
                    new_input_name=str(new_scale_const_name),
                )
            else:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=scale_op,
                    input_index=1,
                    new_input_name=str(new_scale_const_name),
                )
            _update_tensor_shape(str(rhs["scale_out_name"]), [n, w, c])

            _update_tensor_shape(add_out_name, [n, w, c])

            remove_indices_set = {int(lhs["pre_idx"]), int(rhs["rhs_pre_idx"])}
            if can_remove_post_transpose:
                # Single-consumer form: replace post transpose with direct EXPAND_DIMS output.
                if len(expand_op.inputs) >= 2:
                    _write_const_ints_to_tensor(
                        model_ir.tensors.get(str(expand_op.inputs[1]), None),
                        [1],
                    )
                _update_tensor_shape(post_in_name, [n, 1, w, c])
                _replace_tensor_inputs(model_ir, post_out_name, post_in_name)
                model_ir.outputs = [
                    str(post_in_name) if str(v) == str(post_out_name) else str(v)
                    for v in list(model_ir.outputs)
                ]
                remove_indices_set.add(int(post_idx))
            else:
                # Multi-consumer form:
                # - Keep Conv branch on NHWC-friendly path (drop post-transpose).
                # - Provide legacy [N,C,W] view only for non-Conv users via 3D bridge.
                non_conv_add_users = [int(v) for v in add_users if int(v) != int(expand_idx)]
                if len(non_conv_add_users) == 0:
                    continue
                bridge_out_name = _unique_tensor_name(f"{add_out_name}_legacy_ncw")
                bridge_perm_name = _unique_tensor_name(f"{add_out_name}_legacy_ncw_perm")
                model_ir.tensors[bridge_perm_name] = TensorIR(
                    name=bridge_perm_name,
                    dtype="INT32",
                    shape=[3],
                    shape_signature=[3],
                    data=np.asarray([0, 2, 1], dtype=np.int32),
                    is_variable=False,
                )
                model_ir.tensors[bridge_out_name] = TensorIR(
                    name=bridge_out_name,
                    dtype=str(model_ir.tensors.get(add_out_name).dtype if model_ir.tensors.get(add_out_name) is not None else "FLOAT32"),
                    shape=[n, c, w],
                    shape_signature=[n, c, w],
                )
                for user_idx in non_conv_add_users:
                    user_op = model_ir.operators[int(user_idx)]
                    for input_idx, input_name in enumerate(list(user_op.inputs)):
                        if str(input_name) == str(add_out_name):
                            _replace_operator_input_at(
                                model_ir=model_ir,
                                op=user_op,
                                input_index=int(input_idx),
                                new_input_name=str(bridge_out_name),
                            )
                model_ir.operators.append(
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(add_out_name), str(bridge_perm_name)],
                        outputs=[str(bridge_out_name)],
                        options={},
                    )
                )

                if len(expand_op.inputs) >= 2:
                    _write_const_ints_to_tensor(
                        model_ir.tensors.get(str(expand_op.inputs[1]), None),
                        [1],
                    )
                _update_tensor_shape(post_in_name, [n, 1, w, c])
                _replace_tensor_inputs(model_ir, post_out_name, post_in_name)
                model_ir.outputs = [
                    str(post_in_name) if str(v) == str(post_out_name) else str(v)
                    for v in list(model_ir.outputs)
                ]
                remove_indices_set.add(int(post_idx))

            remove_indices = sorted(list(remove_indices_set), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_tencoder_add_expand_transpose_conv_nhwc_chains": int(rewritten)}


def _optimize_transpose_squeeze_unary_batchmatmul_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC->NCHW transpose before SQUEEZE->UNARY->BATCH_MATMUL tails.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw --SQUEEZE(axis=s)--> s3 --(UNARY)*--> u3
      u3 --BATCH_MATMUL(adjX=old_adj_x, adjY=*)--> y

    Rewrite:
      x_nhwc --SQUEEZE(axis=perm[s])--> s3' --(UNARY)*--> u3'
      u3' --BATCH_MATMUL(adjX=new_adj_x, adjY=*)--> y

    where `new_adj_x` is chosen so matmul semantics are preserved.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
        "HARD_SWISH",
    }

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            squeeze_idx = int(pre_users[0])
            squeeze_op = model_ir.operators[int(squeeze_idx)]
            if (
                str(squeeze_op.op_type) != "SQUEEZE"
                or len(squeeze_op.inputs) != 1
                or len(squeeze_op.outputs) != 1
                or str(squeeze_op.inputs[0]) != pre_output_name
            ):
                continue

            squeeze_out_name = str(squeeze_op.outputs[0])
            if squeeze_out_name in model_outputs:
                continue

            # Collect optional strict unary chain.
            chain_indices: List[int] = []
            chain_output_names: List[str] = []
            current_name = str(squeeze_out_name)
            while True:
                current_users = [int(v) for v in consumers.get(current_name, [])]
                if len(current_users) != 1:
                    break
                unary_idx = int(current_users[0])
                unary_op = model_ir.operators[int(unary_idx)]
                if (
                    str(unary_op.op_type) not in unary_passthrough_ops
                    or len(unary_op.inputs) != 1
                    or len(unary_op.outputs) != 1
                    or str(unary_op.inputs[0]) != current_name
                ):
                    break
                unary_out_name = str(unary_op.outputs[0])
                if unary_out_name in model_outputs:
                    break
                chain_indices.append(int(unary_idx))
                chain_output_names.append(str(unary_out_name))
                current_name = str(unary_out_name)

            tail_users = [int(v) for v in consumers.get(current_name, [])]
            if len(tail_users) != 1:
                continue
            matmul_idx = int(tail_users[0])
            matmul_op = model_ir.operators[int(matmul_idx)]
            if (
                str(matmul_op.op_type) != "BATCH_MATMUL"
                or len(matmul_op.inputs) != 2
                or len(matmul_op.outputs) != 1
            ):
                continue
            matmul_inputs = [str(v) for v in list(matmul_op.inputs)]
            if current_name == matmul_inputs[0]:
                matmul_data_input_index = 0
            elif current_name == matmul_inputs[1]:
                matmul_data_input_index = 1
            else:
                continue
            # Keep this rewrite strict to avoid changing rhs transpose semantics.
            if int(matmul_data_input_index) != 0:
                continue

            pre_input_shape = _shape_list(pre_input_name)
            pre_output_shape = _shape_list(pre_output_name)
            squeeze_shape = _shape_list(squeeze_out_name)
            if (
                pre_input_shape is None
                or pre_output_shape is None
                or squeeze_shape is None
                or len(pre_input_shape) != 4
                or len(pre_output_shape) != 4
                or len(squeeze_shape) != 3
            ):
                continue

            # Determine squeeze axis in transposed domain.
            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            squeeze_axis: Optional[int] = None
            if "squeezeDims" in squeeze_options:
                raw_axes = np.asarray(squeeze_options.get("squeezeDims", []), dtype=np.int64).reshape(-1)
                normalized_axes = _normalize_squeeze_axes_for_rank(
                    [int(v) for v in raw_axes.tolist()],
                    4,
                )
                if normalized_axes is None or len(normalized_axes) != 1:
                    continue
                squeeze_axis = int(normalized_axes[0])
            else:
                continue

            # Old order (after transpose+squeeze): [in[perm[j]] for j != squeeze_axis]
            old_dim_indices = [
                int(perm_nhwc_to_nchw[j])
                for j in range(4)
                if int(j) != int(squeeze_axis)
            ]
            src_axis = int(perm_nhwc_to_nchw[int(squeeze_axis)])
            new_dim_indices = [int(i) for i in range(4) if int(i) != int(src_axis)]

            # Only support identity or last-two swap; BATCH_MATMUL adjX can express this.
            old_adj_x = bool(dict(matmul_op.options).get("adjX", False)) if isinstance(matmul_op.options, dict) else False
            if old_dim_indices == new_dim_indices:
                perm_rank3: Optional[List[int]] = None
                new_adj_x = bool(old_adj_x)
            elif old_dim_indices == [int(new_dim_indices[0]), int(new_dim_indices[2]), int(new_dim_indices[1])]:
                perm_rank3 = [0, 2, 1]
                new_adj_x = bool(not old_adj_x)
            else:
                continue

            expected_old_squeeze_shape = [int(pre_input_shape[idx]) for idx in old_dim_indices]
            if not _shape_compatible(squeeze_shape, expected_old_squeeze_shape):
                continue
            if int(pre_input_shape[src_axis]) >= 0 and int(pre_input_shape[src_axis]) != 1:
                continue

            # Rewire squeeze to consume NHWC tensor and remove pre-transpose.
            _set_operator_inputs(
                model_ir=model_ir,
                op=squeeze_op,
                new_inputs=[pre_input_name],
            )
            squeeze_options["squeezeDims"] = [int(src_axis)]
            squeeze_op.options = squeeze_options

            # Adjust unary metadata to the new rank-3 axis order.
            if perm_rank3 is not None:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(squeeze_out_name, None),
                    perm_rank3,
                )
                for chain_out_name in chain_output_names:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(chain_out_name, None),
                        perm_rank3,
                    )

            # Update BATCH_MATMUL adjX to preserve semantics.
            matmul_options = dict(matmul_op.options) if isinstance(matmul_op.options, dict) else {}
            matmul_options["adjX"] = bool(new_adj_x)
            matmul_op.options = matmul_options

            del model_ir.operators[int(pre_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_squeeze_unary_batchmatmul_nhwc_chains": int(rewritten)}


def _optimize_decoder_batchmatmul_add_expand_transpose_to_nhwc_deconv_input(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove decoder-side EXPAND_DIMS->TRANSPOSE adapters before TRANSPOSE_CONV
    by transposing the preceding BATCH_MATMUL/ADD layout instead.

    Target:
      mm = BATCH_MATMUL(X, W)                  # [N, C, L]
      a  = ADD(mm, b[L])
      e  = EXPAND_DIMS(a, axis=2)              # [N, C, 1, L]
      t  = TRANSPOSE(e, [0,2,3,1])             # [N, 1, L, C]
      y  = TRANSPOSE_CONV(..., t)

    Rewrite (equivalent):
      mm_t = BATCH_MATMUL(W, X, adjX=not old_adjY, adjY=not old_adjX)  # [N, L, C]
      a_t  = ADD(mm_t, b.reshape([1,L,1]))
      e_t  = EXPAND_DIMS(a_t, axis=1)                                   # [N,1,L,C]
      y    = TRANSPOSE_CONV(..., e_t)
    """
    rewritten = 0
    perm_expand_to_deconv = [0, 2, 3, 1]
    rank3_swap = [0, 2, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for transpose_idx, transpose_op in enumerate(model_ir.operators):
            if str(transpose_op.op_type) != "TRANSPOSE" or len(transpose_op.inputs) < 2 or len(transpose_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, transpose_op) != perm_expand_to_deconv:
                continue

            expand_out_name = str(transpose_op.inputs[0])
            transpose_out_name = str(transpose_op.outputs[0])
            if transpose_out_name in model_outputs:
                continue

            transpose_users = [int(v) for v in consumers.get(transpose_out_name, [])]
            if len(transpose_users) != 1:
                continue
            deconv_idx = int(transpose_users[0])
            deconv_op = model_ir.operators[int(deconv_idx)]
            if (
                str(deconv_op.op_type) != "TRANSPOSE_CONV"
                or len(deconv_op.inputs) < 3
                or str(deconv_op.inputs[2]) != transpose_out_name
            ):
                continue

            expand_idx = producers.get(expand_out_name, None)
            if expand_idx is None:
                continue
            expand_op = model_ir.operators[int(expand_idx)]
            if (
                str(expand_op.op_type) != "EXPAND_DIMS"
                or len(expand_op.inputs) < 2
                or len(expand_op.outputs) != 1
                or str(expand_op.outputs[0]) != expand_out_name
            ):
                continue

            expand_in_name = str(expand_op.inputs[0])
            add_idx = producers.get(expand_in_name, None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if (
                str(add_op.op_type) != "ADD"
                or len(add_op.inputs) != 2
                or len(add_op.outputs) != 1
                or str(add_op.outputs[0]) != expand_in_name
            ):
                continue

            add_users = [int(v) for v in consumers.get(expand_in_name, [])]
            if len(add_users) != 1 or int(add_users[0]) != int(expand_idx):
                continue

            add_input0 = str(add_op.inputs[0])
            add_input1 = str(add_op.inputs[1])
            mm_idx = producers.get(add_input0, None)
            bias_name = add_input1
            if mm_idx is None or str(model_ir.operators[int(mm_idx)].op_type) != "BATCH_MATMUL":
                mm_idx = producers.get(add_input1, None)
                bias_name = add_input0
            if mm_idx is None:
                continue
            mm_op = model_ir.operators[int(mm_idx)]
            if (
                str(mm_op.op_type) != "BATCH_MATMUL"
                or len(mm_op.inputs) != 2
                or len(mm_op.outputs) != 1
                or str(mm_op.outputs[0]) not in {add_input0, add_input1}
            ):
                continue

            mm_out_name = str(mm_op.outputs[0])
            mm_users = [int(v) for v in consumers.get(mm_out_name, [])]
            if len(mm_users) != 1 or int(mm_users[0]) != int(add_idx):
                continue

            mm_out_tensor = model_ir.tensors.get(mm_out_name, None)
            add_out_tensor = model_ir.tensors.get(expand_in_name, None)
            expand_out_tensor = model_ir.tensors.get(expand_out_name, None)
            transpose_out_tensor = model_ir.tensors.get(transpose_out_name, None)
            bias_tensor = model_ir.tensors.get(bias_name, None)
            if (
                mm_out_tensor is None
                or add_out_tensor is None
                or expand_out_tensor is None
                or transpose_out_tensor is None
                or bias_tensor is None
                or bias_tensor.data is None
            ):
                continue
            if len(list(mm_out_tensor.shape)) != 3 or len(list(add_out_tensor.shape)) != 3:
                continue
            if len(list(expand_out_tensor.shape)) != 4 or len(list(transpose_out_tensor.shape)) != 4:
                continue

            old_mm_shape = [int(v) for v in list(mm_out_tensor.shape)]
            if len(old_mm_shape) != 3:
                continue
            length_dim = int(old_mm_shape[2])
            if length_dim <= 0:
                continue

            bias_values = np.asarray(bias_tensor.data)
            if int(bias_values.size) != int(length_dim):
                continue

            # Update BATCH_MATMUL to emit transposed rank-3 layout directly.
            mm_options = dict(mm_op.options) if isinstance(mm_op.options, dict) else {}
            old_adj_x = bool(mm_options.get("adjX", False))
            old_adj_y = bool(mm_options.get("adjY", False))
            old_mm_inputs = [str(v) for v in list(mm_op.inputs)]
            _set_operator_inputs(
                model_ir=model_ir,
                op=mm_op,
                new_inputs=[old_mm_inputs[1], old_mm_inputs[0]],
            )
            mm_options["adjX"] = bool(not old_adj_y)
            mm_options["adjY"] = bool(not old_adj_x)
            mm_op.options = mm_options
            _permute_tensor_metadata_if_rank_matches(mm_out_tensor, rank3_swap)
            _permute_tensor_metadata_if_rank_matches(add_out_tensor, rank3_swap)

            # Bias was [L], but after rank-3 transpose it must broadcast on axis=1.
            reshaped_bias = np.asarray(bias_values).reshape(1, int(length_dim), 1)
            bias_tensor.data = reshaped_bias.astype(bias_values.dtype, copy=False)
            bias_tensor.shape = [1, int(length_dim), 1]
            bias_tensor.shape_signature = [1, int(length_dim), 1]

            # EXPAND_DIMS axis: 2 -> 1 for [N,L,C] -> [N,1,L,C].
            expand_axis_tensor = model_ir.tensors.get(str(expand_op.inputs[1]), None)
            if expand_axis_tensor is None:
                continue
            if not _write_const_ints_to_tensor(expand_axis_tensor, [1]):
                continue

            expand_out_tensor.shape = [int(v) for v in list(transpose_out_tensor.shape)]
            expand_out_tensor.shape_signature = (
                [int(v) for v in list(transpose_out_tensor.shape_signature)]
                if transpose_out_tensor.shape_signature is not None
                else [int(v) for v in list(transpose_out_tensor.shape)]
            )

            # Bypass transpose into deconv input.
            _replace_operator_input_at(
                model_ir=model_ir,
                op=deconv_op,
                input_index=2,
                new_input_name=expand_out_name,
            )

            del model_ir.operators[int(transpose_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {
        "optimized_decoder_batchmatmul_add_expand_transpose_to_nhwc_deconv_input": int(rewritten)
    }


def _optimize_transpose_squeeze_mean_squeeze_terminal_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove terminal NHWC->NCHW transpose in:
      x --T(0,3,1,2)--> t --SQUEEZE(axis=2)--> s --MEAN(axis=1,keepDims)--> m --SQUEEZE(axis=1)--> y

    Equivalent NHWC rewrite:
      x --SQUEEZE(axis=1)--> s' --MEAN(axis=2,keepDims)--> m' --SQUEEZE(axis=2)--> y
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    rank3_swap = [0, 2, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for transpose_idx, transpose_op in enumerate(model_ir.operators):
            if str(transpose_op.op_type) != "TRANSPOSE" or len(transpose_op.inputs) < 2 or len(transpose_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, transpose_op) != perm_nhwc_to_nchw:
                continue

            x_name = str(transpose_op.inputs[0])
            t_name = str(transpose_op.outputs[0])
            if t_name in model_outputs:
                continue

            t_users = [int(v) for v in consumers.get(t_name, [])]
            if len(t_users) != 1:
                continue
            squeeze1_idx = int(t_users[0])
            squeeze1_op = model_ir.operators[int(squeeze1_idx)]
            if (
                str(squeeze1_op.op_type) != "SQUEEZE"
                or len(squeeze1_op.inputs) != 1
                or len(squeeze1_op.outputs) != 1
                or str(squeeze1_op.inputs[0]) != t_name
            ):
                continue

            squeeze1_opts = dict(squeeze1_op.options) if isinstance(squeeze1_op.options, dict) else {}
            squeeze1_axes = _normalize_squeeze_axes_for_rank(
                _parse_axes_option(squeeze1_opts.get("squeezeDims", [])),
                4,
            )
            if squeeze1_axes is None or [int(v) for v in squeeze1_axes] != [2]:
                continue
            s_name = str(squeeze1_op.outputs[0])
            if s_name in model_outputs:
                continue

            s_users = [int(v) for v in consumers.get(s_name, [])]
            if len(s_users) != 1:
                continue
            mean_idx = int(s_users[0])
            mean_op = model_ir.operators[int(mean_idx)]
            if (
                str(mean_op.op_type) != "MEAN"
                or len(mean_op.inputs) < 2
                or len(mean_op.outputs) != 1
                or str(mean_op.inputs[0]) != s_name
                or not bool(mean_op.options.get("keepDims", False))
            ):
                continue

            mean_axes_name = str(mean_op.inputs[1])
            mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_vals is None:
                continue
            mean_axes_norm = _normalize_squeeze_axes_for_rank(
                [int(v) for v in mean_axes_vals],
                3,
            )
            if mean_axes_norm is None or [int(v) for v in mean_axes_norm] != [1]:
                continue

            m_name = str(mean_op.outputs[0])
            if m_name in model_outputs:
                continue
            m_users = [int(v) for v in consumers.get(m_name, [])]
            if len(m_users) != 1:
                continue
            squeeze2_idx = int(m_users[0])
            squeeze2_op = model_ir.operators[int(squeeze2_idx)]
            if (
                str(squeeze2_op.op_type) != "SQUEEZE"
                or len(squeeze2_op.inputs) != 1
                or len(squeeze2_op.outputs) != 1
                or str(squeeze2_op.inputs[0]) != m_name
            ):
                continue

            squeeze2_opts = dict(squeeze2_op.options) if isinstance(squeeze2_op.options, dict) else {}
            squeeze2_axes = _normalize_squeeze_axes_for_rank(
                _parse_axes_option(squeeze2_opts.get("squeezeDims", [])),
                3,
            )
            if squeeze2_axes is None or [int(v) for v in squeeze2_axes] != [1]:
                continue

            x_tensor = model_ir.tensors.get(x_name, None)
            t_tensor = model_ir.tensors.get(t_name, None)
            s_tensor = model_ir.tensors.get(s_name, None)
            m_tensor = model_ir.tensors.get(m_name, None)
            if (
                x_tensor is None
                or t_tensor is None
                or s_tensor is None
                or m_tensor is None
                or len(list(x_tensor.shape)) != 4
                or len(list(t_tensor.shape)) != 4
                or len(list(s_tensor.shape)) != 3
                or len(list(m_tensor.shape)) != 3
            ):
                continue

            # Rewire chain onto NHWC source and remap axes.
            _set_operator_inputs(
                model_ir=model_ir,
                op=squeeze1_op,
                new_inputs=[x_name],
            )
            squeeze1_opts["squeezeDims"] = [1]
            squeeze1_op.options = squeeze1_opts

            if mean_axes_tensor is None or not _write_const_ints_to_tensor(mean_axes_tensor, [2]):
                continue

            squeeze2_opts["squeezeDims"] = [2]
            squeeze2_op.options = squeeze2_opts

            # rank-3 tensors between squeeze1 and mean are [N,C,W] -> [N,W,C]
            _permute_tensor_metadata_if_rank_matches(s_tensor, rank3_swap)
            _permute_tensor_metadata_if_rank_matches(m_tensor, rank3_swap)

            del model_ir.operators[int(transpose_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_squeeze_mean_squeeze_terminal_nhwc_chains": int(rewritten)
    }


def _optimize_squeeze_unary_reshape_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant SQUEEZE(axis=0) -> UNARY -> RESHAPE chains.

    Target:
      x[1,H,W,C] --SQUEEZE(axis=0)--> s[H,W,C] --UNARY--> u[H,W,C] --RESHAPE([1,H,W,C])--> y

    Rewrite:
      x[1,H,W,C] --UNARY--> y
    """
    rewritten = 0
    unary_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
        "HARD_SWISH",
    }

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for squeeze_idx, squeeze_op in enumerate(model_ir.operators):
            if str(squeeze_op.op_type) != "SQUEEZE" or len(squeeze_op.inputs) != 1 or len(squeeze_op.outputs) != 1:
                continue

            squeeze_in_name = str(squeeze_op.inputs[0])
            squeeze_out_name = str(squeeze_op.outputs[0])
            squeeze_users = [int(v) for v in consumers.get(squeeze_out_name, [])]
            if len(squeeze_users) != 1:
                continue

            unary_idx = int(squeeze_users[0])
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) not in unary_ops
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.inputs[0]) != squeeze_out_name
            ):
                continue

            unary_out_name = str(unary_op.outputs[0])
            unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
            reshape_user_indices: List[int] = []
            for user_idx in unary_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "RESHAPE"
                    and len(user_op.inputs) >= 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == unary_out_name
                ):
                    reshape_user_indices.append(int(user_idx))
            if len(reshape_user_indices) != 1:
                continue

            reshape_idx = int(reshape_user_indices[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 1
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != unary_out_name
            ):
                continue

            reshape_out_name = str(reshape_op.outputs[0])
            in_shape = _shape_list(squeeze_in_name)
            squeezed_shape = _shape_list(squeeze_out_name)
            reshape_out_shape = _shape_list(reshape_out_name)
            if in_shape is None or squeezed_shape is None or reshape_out_shape is None:
                continue
            if len(in_shape) != len(squeezed_shape) + 1 or len(in_shape) < 1:
                continue
            if int(in_shape[0]) != 1:
                continue
            if not _shape_compatible(squeezed_shape, in_shape[1:]):
                continue
            if not _shape_compatible(reshape_out_shape, in_shape):
                continue

            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            if "squeezeDims" in squeeze_options:
                raw_axes = np.asarray(squeeze_options.get("squeezeDims", []), dtype=np.int64).reshape(-1)
                normalized_axes = _normalize_squeeze_axes_for_rank(
                    [int(v) for v in raw_axes.tolist()],
                    len(in_shape),
                )
                if normalized_axes is None or [int(v) for v in normalized_axes] != [0]:
                    continue

            if len(unary_users) == 1:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_inputs=[squeeze_in_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_outputs=[reshape_out_name],
                )
                for remove_idx in sorted([int(squeeze_idx), int(reshape_idx)], reverse=True):
                    del model_ir.operators[int(remove_idx)]
            else:
                # Fanout case:
                #   squeeze -> unary -> (reshape + other users)
                # Reorder into:
                #   unary(4D) -> squeeze(3D fanout), and remove reshape.
                if unary_out_name in model_outputs:
                    continue
                if reshape_out_name == unary_out_name:
                    continue

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_inputs=[squeeze_in_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=unary_op,
                    new_outputs=[reshape_out_name],
                )
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=squeeze_op,
                    new_inputs=[reshape_out_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=squeeze_op,
                    new_outputs=[unary_out_name],
                )

                # Maintain topological order: unary must run before squeeze.
                squeeze_op_ref = squeeze_op
                del model_ir.operators[int(reshape_idx)]
                adjusted_unary_idx = int(unary_idx)
                if int(reshape_idx) < adjusted_unary_idx:
                    adjusted_unary_idx -= 1
                del model_ir.operators[int(squeeze_idx)]
                if int(squeeze_idx) < adjusted_unary_idx:
                    adjusted_unary_idx -= 1
                model_ir.operators.insert(int(adjusted_unary_idx) + 1, squeeze_op_ref)

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_squeeze_unary_reshape_passthrough_chains": int(rewritten)}


def _optimize_squeeze_reshape_identity_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant SQUEEZE -> RESHAPE chains that round-trip to input shape.

    Target:
      x --SQUEEZE--> s --RESHAPE--> y, where shape(y) == shape(x)

    Rewrite:
      replace all uses of y with x and remove SQUEEZE/RESHAPE.
    """
    rewritten = 0

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for squeeze_idx, squeeze_op in enumerate(model_ir.operators):
            if str(squeeze_op.op_type) != "SQUEEZE" or len(squeeze_op.inputs) != 1 or len(squeeze_op.outputs) != 1:
                continue

            squeeze_in_name = str(squeeze_op.inputs[0])
            squeeze_out_name = str(squeeze_op.outputs[0])
            if squeeze_out_name in model_outputs:
                continue

            squeeze_users = [int(v) for v in consumers.get(squeeze_out_name, [])]
            if len(squeeze_users) != 1:
                continue

            reshape_idx = int(squeeze_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 1
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != squeeze_out_name
            ):
                continue

            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue
            if reshape_out_name == squeeze_in_name:
                continue

            in_shape = _shape_list(squeeze_in_name)
            squeezed_shape = _shape_list(squeeze_out_name)
            reshape_out_shape = _shape_list(reshape_out_name)
            if in_shape is None or squeezed_shape is None or reshape_out_shape is None:
                continue
            if not _shape_compatible(in_shape, reshape_out_shape):
                continue

            squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
            squeeze_axes: List[int]
            if "squeezeDims" in squeeze_options:
                raw_axes = np.asarray(squeeze_options.get("squeezeDims", []), dtype=np.int64).reshape(-1)
                normalized_axes = _normalize_squeeze_axes_for_rank(
                    [int(v) for v in raw_axes.tolist()],
                    len(in_shape),
                )
                if normalized_axes is None:
                    continue
                squeeze_axes = [int(v) for v in normalized_axes]
            else:
                # Keep this conservative for unknown shapes when axes are omitted.
                if any(int(v) < 0 for v in in_shape):
                    continue
                squeeze_axes = [int(idx) for idx, dim in enumerate(in_shape) if int(dim) == 1]

            if len(set(int(v) for v in squeeze_axes)) != len(squeeze_axes):
                continue
            valid_axes = True
            for axis in squeeze_axes:
                if axis < 0 or axis >= len(in_shape):
                    valid_axes = False
                    break
                dim = int(in_shape[int(axis)])
                if dim >= 0 and dim != 1:
                    valid_axes = False
                    break
            if not valid_axes:
                continue

            squeeze_axes_set = set(int(v) for v in squeeze_axes)
            expected_squeezed_shape = [
                int(dim) for idx, dim in enumerate(in_shape) if int(idx) not in squeeze_axes_set
            ]
            if not _shape_compatible(squeezed_shape, expected_squeezed_shape):
                continue

            _replace_tensor_inputs(model_ir, reshape_out_name, squeeze_in_name)
            for remove_idx in sorted([int(reshape_idx), int(squeeze_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_squeeze_reshape_identity_chains": int(rewritten)}


def _optimize_transpose_instancenorm_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NHWC<->NCHW transpose bridges around decomposed InstanceNormalization blocks.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --SQUEEZE(axis=0)--> s0 --RESHAPE--> r0
      r0 --(InstanceNormalization decomposition in NCHW)--> y_nchw
      y_nchw --SQUEEZE(axis=0)--> s1 --UNARY--> u1 --RESHAPE--> r1
      r1 --TRANSPOSE(0,2,3,1)--> z_nhwc

    Rewrite:
      x_nhwc --SQUEEZE(axis=0)--> s0_hwc --RESHAPE(NHWC)--> r0_nhwc
      r0_nhwc --(same decomposition rewritten to NHWC axes/broadcast)--> y_nhwc
      y_nhwc --SQUEEZE(axis=0)--> s1_hwc --UNARY--> u1_hwc --RESHAPE(NHWC)--> z_nhwc
    """
    rewritten = 0
    # Keep this pass conservative for now: broad propagation can degrade
    # intermediate tensor correspondence quality in --report_op_coverage flows.
    max_total_rewrites = 32
    channel_last_hint_names = {
        str(v)
        for v in model_ir.metadata.get("assume_channel_last_layout_tensor_names", [])
        if str(v) != ""
    }
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_3d_chw_to_hwc = [1, 2, 0]
    unary_ops = {
        "RELU",
        "RELU6",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "CAST",
        "FLOOR",
        "CEIL",
        "ROUND",
    }
    tail_nhwc_to_nchw_perm_name = "__instancenorm_tail_nhwc_to_nchw_perm_rank4__"
    tail_hwc_to_chw_perm_name = "__instancenorm_tail_hwc_to_chw_perm_rank3__"

    def _single_consumer_index(consumers: Dict[str, List[int]], tensor_name: str) -> Optional[int]:
        users = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(users) != 1:
            return None
        return int(users[0])

    def _single_producer_index(producers: Dict[str, int], tensor_name: str) -> Optional[int]:
        idx = producers.get(str(tensor_name), None)
        if idx is None:
            return None
        return int(idx)

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _const_ints(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.data is None:
            return None
        arr = np.asarray(tensor.data).reshape(-1)
        if arr.size == 0:
            return []
        return [int(v) for v in arr.tolist()]

    def _is_squeeze_axis0(
        *,
        squeeze_op: OperatorIR,
        input_name: str,
        output_name: str,
    ) -> bool:
        in_tensor = model_ir.tensors.get(str(input_name), None)
        out_tensor = model_ir.tensors.get(str(output_name), None)
        if in_tensor is None or out_tensor is None:
            return False
        in_shape = [int(v) for v in list(in_tensor.shape)]
        out_shape = [int(v) for v in list(out_tensor.shape)]
        if len(in_shape) != 4 or len(out_shape) != 3:
            return False
        if int(in_shape[0]) != 1:
            return False
        if out_shape != [int(in_shape[1]), int(in_shape[2]), int(in_shape[3])]:
            return False
        options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
        if "squeezeDims" in options:
            dims = [int(v) for v in np.asarray(options.get("squeezeDims", [])).reshape(-1).tolist()]
            if dims != [0]:
                return False
        return True

    def _rewrite_reshape_shape_const(
        *,
        reshape_op: OperatorIR,
        reshape_idx: int,
        target_shape: List[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        if len(reshape_op.inputs) < 2:
            return False
        shape_name = str(reshape_op.inputs[1])
        shape_users = set(int(v) for v in consumers.get(shape_name, []))
        if shape_users != {int(reshape_idx)}:
            return False
        shape_tensor = model_ir.tensors.get(shape_name, None)
        if shape_tensor is None or shape_tensor.data is None:
            return False
        dtype = np.asarray(shape_tensor.data).dtype
        if dtype.kind not in {"i", "u"}:
            return False
        new_data = np.asarray([int(v) for v in list(target_shape)], dtype=dtype)
        shape_tensor.data = new_data
        shape_tensor.shape = [int(v) for v in list(new_data.shape)]
        shape_tensor.shape_signature = [int(v) for v in list(new_data.shape)]
        options = dict(reshape_op.options) if isinstance(reshape_op.options, dict) else {}
        options["newShape"] = [int(v) for v in list(target_shape)]
        reshape_op.options = options
        return True

    def _rewrite_reduce_axes_to_nhwc(
        *,
        axes_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        axes_users = set(int(v) for v in consumers.get(str(axes_name), []))
        if axes_users != expected_users:
            return False
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None or axes_tensor.data is None:
            return False
        dtype = np.asarray(axes_tensor.data).dtype
        if dtype.kind not in {"i", "u"}:
            return False
        axes_tensor.data = np.asarray([1, 2], dtype=dtype)
        axes_tensor.shape = [2]
        axes_tensor.shape_signature = [2]
        return True

    def _rewrite_nchw_coeff_to_nhwc(coeff_name: str) -> bool:
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        if coeff_tensor is None or coeff_tensor.data is None:
            return False
        coeff = np.asarray(coeff_tensor.data)
        if coeff.ndim != 4:
            return False
        if int(coeff.shape[0]) != 1 or int(coeff.shape[2]) != 1 or int(coeff.shape[3]) != 1:
            return False
        coeff_tensor.data = np.transpose(coeff, (0, 2, 3, 1))
        coeff_tensor.shape = [int(v) for v in list(coeff_tensor.data.shape)]
        coeff_tensor.shape_signature = [int(v) for v in list(coeff_tensor.data.shape)]
        return True

    while rewritten < max_total_rewrites:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            squeeze1_idx = _single_consumer_index(consumers, pre_out_name)
            if squeeze1_idx is None:
                continue
            squeeze1_op = model_ir.operators[int(squeeze1_idx)]
            if str(squeeze1_op.op_type) != "SQUEEZE" or len(squeeze1_op.inputs) != 1 or len(squeeze1_op.outputs) != 1:
                continue
            squeeze1_out_name = str(squeeze1_op.outputs[0])
            if not _is_squeeze_axis0(
                squeeze_op=squeeze1_op,
                input_name=pre_out_name,
                output_name=squeeze1_out_name,
            ):
                continue

            reshape1_idx = _single_consumer_index(consumers, squeeze1_out_name)
            if reshape1_idx is None:
                continue
            reshape1_op = model_ir.operators[int(reshape1_idx)]
            if str(reshape1_op.op_type) != "RESHAPE" or len(reshape1_op.inputs) < 2 or len(reshape1_op.outputs) != 1:
                continue
            x_name = str(reshape1_op.outputs[0])
            x_users = set(int(v) for v in consumers.get(x_name, []))
            if len(x_users) != 2:
                continue

            mean1_idx = None
            sub_idx = None
            for idx in sorted(list(x_users)):
                op = model_ir.operators[int(idx)]
                if str(op.op_type) == "MEAN" and len(op.inputs) == 2 and str(op.inputs[0]) == x_name:
                    mean1_idx = int(idx)
                elif str(op.op_type) == "SUB" and len(op.inputs) == 2 and x_name in set(str(v) for v in op.inputs):
                    sub_idx = int(idx)
            if mean1_idx is None or sub_idx is None:
                continue
            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0]) if len(mean1_op.outputs) == 1 else ""
            if mean1_out_name == "":
                continue
            if mean1_out_name not in set(str(v) for v in sub_op.inputs):
                continue
            if str(sub_op.outputs[0]) in model_outputs:
                continue
            centered_name = str(sub_op.outputs[0])

            centered_users = set(int(v) for v in consumers.get(centered_name, []))
            if len(centered_users) != 2:
                continue
            mul_square_idx = None
            mul_norm_idx = None
            for idx in sorted(list(centered_users)):
                op = model_ir.operators[int(idx)]
                if str(op.op_type) != "MUL" or len(op.inputs) != 2 or len(op.outputs) != 1:
                    continue
                in0 = str(op.inputs[0])
                in1 = str(op.inputs[1])
                if in0 == centered_name and in1 == centered_name:
                    mul_square_idx = int(idx)
                elif centered_name in {in0, in1}:
                    mul_norm_idx = int(idx)
            if mul_square_idx is None or mul_norm_idx is None:
                continue

            mul_square_op = model_ir.operators[int(mul_square_idx)]
            squared_name = str(mul_square_op.outputs[0])
            mean2_idx = _single_consumer_index(consumers, squared_name)
            if mean2_idx is None:
                continue
            mean2_op = model_ir.operators[int(mean2_idx)]
            if str(mean2_op.op_type) != "MEAN" or len(mean2_op.inputs) != 2 or len(mean2_op.outputs) != 1:
                continue
            if str(mean2_op.inputs[0]) != squared_name:
                continue
            mean2_out_name = str(mean2_op.outputs[0])

            add_eps_idx = _single_consumer_index(consumers, mean2_out_name)
            if add_eps_idx is None:
                continue
            add_eps_op = model_ir.operators[int(add_eps_idx)]
            if str(add_eps_op.op_type) != "ADD" or len(add_eps_op.inputs) != 2 or len(add_eps_op.outputs) != 1:
                continue
            if mean2_out_name not in set(str(v) for v in add_eps_op.inputs):
                continue
            add_eps_out_name = str(add_eps_op.outputs[0])

            sqrt_idx = _single_consumer_index(consumers, add_eps_out_name)
            if sqrt_idx is None:
                continue
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if str(sqrt_op.op_type) != "SQRT" or len(sqrt_op.inputs) != 1 or len(sqrt_op.outputs) != 1:
                continue
            if str(sqrt_op.inputs[0]) != add_eps_out_name:
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])

            div_idx = _single_consumer_index(consumers, sqrt_out_name)
            if div_idx is None:
                continue
            div_op = model_ir.operators[int(div_idx)]
            if str(div_op.op_type) != "DIV" or len(div_op.inputs) != 2 or len(div_op.outputs) != 1:
                continue
            if str(div_op.inputs[1]) != sqrt_out_name:
                continue
            div_out_name = str(div_op.outputs[0])

            mul_norm_op = model_ir.operators[int(mul_norm_idx)]
            if div_out_name not in set(str(v) for v in mul_norm_op.inputs):
                continue
            normalized_name = str(mul_norm_op.outputs[0])

            mul_scale_idx = _single_consumer_index(consumers, normalized_name)
            if mul_scale_idx is None:
                continue
            mul_scale_op = model_ir.operators[int(mul_scale_idx)]
            if str(mul_scale_op.op_type) != "MUL" or len(mul_scale_op.inputs) != 2 or len(mul_scale_op.outputs) != 1:
                continue
            if normalized_name not in set(str(v) for v in mul_scale_op.inputs):
                continue
            scale_const_name = str(mul_scale_op.inputs[0]) if str(mul_scale_op.inputs[1]) == normalized_name else str(mul_scale_op.inputs[1])
            scaled_name = str(mul_scale_op.outputs[0])

            add_bias_idx = _single_consumer_index(consumers, scaled_name)
            if add_bias_idx is None:
                continue
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if str(add_bias_op.op_type) != "ADD" or len(add_bias_op.inputs) != 2 or len(add_bias_op.outputs) != 1:
                continue
            if scaled_name not in set(str(v) for v in add_bias_op.inputs):
                continue
            bias_const_name = str(add_bias_op.inputs[0]) if str(add_bias_op.inputs[1]) == scaled_name else str(add_bias_op.inputs[1])
            inst_out_name = str(add_bias_op.outputs[0])

            tail_mode = "direct_post_transpose"
            side_squeeze_idx: Optional[int] = None
            squeeze2_out_name: Optional[str] = None
            unary_out_name: Optional[str] = None
            tail_add_out_name: Optional[str] = None
            tail_add_residual_input_name: Optional[str] = None
            tail_add_legacy_slots: List[Tuple[int, int]] = []
            tail_add_legacy_meta: Optional[Tuple[str, Optional[Any], List[int], List[int]]] = None
            reshape2_idx: Optional[int] = None
            reshape2_op: Optional[OperatorIR] = None
            reshape2_out_name = str(inst_out_name)
            extra_remove_indices: set[int] = set()

            post_idx = _single_consumer_index(consumers, inst_out_name)
            if post_idx is not None:
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                    or str(post_op.inputs[0]) != inst_out_name
                ):
                    post_idx = None
            if post_idx is None:
                # Variant:
                #   inst_out_nchw -> {post_transpose_nchw_to_nhwc, squeeze(axis=0) side branch}
                # Keep side branch in NCHW via a local adapter transpose and still
                # eliminate the pre/post bridge pair around InstanceNormalization.
                inst_out_users = [int(v) for v in consumers.get(inst_out_name, [])]
                candidate_post_indices: List[int] = []
                for user_idx in inst_out_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == inst_out_name
                        and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    ):
                        candidate_post_indices.append(int(user_idx))
                if len(candidate_post_indices) == 1 and len(inst_out_users) == 2:
                    candidate_post_idx = int(candidate_post_indices[0])
                    candidate_side_idx = (
                        int(inst_out_users[1])
                        if int(inst_out_users[0]) == candidate_post_idx
                        else int(inst_out_users[0])
                    )
                    candidate_side_op = model_ir.operators[int(candidate_side_idx)]
                    candidate_side_out_name = (
                        str(candidate_side_op.outputs[0]) if len(candidate_side_op.outputs) == 1 else ""
                    )
                    if (
                        str(candidate_side_op.op_type) == "SQUEEZE"
                        and len(candidate_side_op.inputs) == 1
                        and len(candidate_side_op.outputs) == 1
                        and _is_squeeze_axis0(
                            squeeze_op=candidate_side_op,
                            input_name=inst_out_name,
                            output_name=candidate_side_out_name,
                        )
                    ):
                        post_idx = int(candidate_post_idx)
                        side_squeeze_idx = int(candidate_side_idx)
                        tail_mode = "direct_post_transpose_side_squeeze"
            if post_idx is None:
                squeeze2_idx = _single_consumer_index(consumers, inst_out_name)
                if squeeze2_idx is not None:
                    candidate_squeeze_op = model_ir.operators[int(squeeze2_idx)]
                    if (
                        str(candidate_squeeze_op.op_type) == "SQUEEZE"
                        and len(candidate_squeeze_op.inputs) == 1
                        and len(candidate_squeeze_op.outputs) == 1
                    ):
                        candidate_squeeze2_out_name = str(candidate_squeeze_op.outputs[0])
                        if _is_squeeze_axis0(
                            squeeze_op=candidate_squeeze_op,
                            input_name=inst_out_name,
                            output_name=candidate_squeeze2_out_name,
                        ):
                            candidate_unary_idx = _single_consumer_index(consumers, candidate_squeeze2_out_name)
                            if candidate_unary_idx is not None:
                                candidate_unary_op = model_ir.operators[int(candidate_unary_idx)]
                                if (
                                    str(candidate_unary_op.op_type) in unary_ops
                                    and len(candidate_unary_op.inputs) == 1
                                    and len(candidate_unary_op.outputs) == 1
                                    and str(candidate_unary_op.inputs[0]) == candidate_squeeze2_out_name
                                ):
                                    candidate_unary_out_name = str(candidate_unary_op.outputs[0])
                                    candidate_reshape2_idx = _single_consumer_index(consumers, candidate_unary_out_name)
                                    if candidate_reshape2_idx is not None:
                                        candidate_reshape2_op = model_ir.operators[int(candidate_reshape2_idx)]
                                        if (
                                            str(candidate_reshape2_op.op_type) == "RESHAPE"
                                            and len(candidate_reshape2_op.inputs) >= 2
                                            and len(candidate_reshape2_op.outputs) == 1
                                        ):
                                            candidate_reshape2_out_name = str(candidate_reshape2_op.outputs[0])
                                            candidate_post_idx = _single_consumer_index(
                                                consumers,
                                                candidate_reshape2_out_name,
                                            )
                                            if candidate_post_idx is not None:
                                                candidate_post_op = model_ir.operators[int(candidate_post_idx)]
                                                if (
                                                    str(candidate_post_op.op_type) == "TRANSPOSE"
                                                    and len(candidate_post_op.inputs) >= 2
                                                    and len(candidate_post_op.outputs) == 1
                                                    and _read_transpose_perm(model_ir, candidate_post_op)
                                                    == perm_nchw_to_nhwc
                                                    and str(candidate_post_op.inputs[0]) == candidate_reshape2_out_name
                                                ):
                                                    tail_mode = "squeeze_unary_reshape"
                                                    squeeze2_out_name = str(candidate_squeeze2_out_name)
                                                    unary_out_name = str(candidate_unary_out_name)
                                                    reshape2_idx = int(candidate_reshape2_idx)
                                                    reshape2_op = candidate_reshape2_op
                                                    reshape2_out_name = str(candidate_reshape2_out_name)
                                                    post_idx = int(candidate_post_idx)

            if post_idx is None:
                # Variant:
                #   inst_out --SQUEEZE(axis=0)--> s
                #   s --ADD(residual_3d)--> a
                #   a --RESHAPE--> r
                #   r --TRANSPOSE(0,2,3,1)--> post_out
                tail_mode = "squeeze_add_reshape"
                squeeze2_idx = _single_consumer_index(consumers, inst_out_name)
                if squeeze2_idx is None:
                    continue
                squeeze2_op = model_ir.operators[int(squeeze2_idx)]
                if str(squeeze2_op.op_type) != "SQUEEZE" or len(squeeze2_op.inputs) != 1 or len(squeeze2_op.outputs) != 1:
                    continue
                squeeze2_out_name = str(squeeze2_op.outputs[0])
                if not _is_squeeze_axis0(
                    squeeze_op=squeeze2_op,
                    input_name=inst_out_name,
                    output_name=squeeze2_out_name,
                ):
                    continue

                add_tail_idx = _single_consumer_index(consumers, squeeze2_out_name)
                if add_tail_idx is None:
                    continue
                add_tail_op = model_ir.operators[int(add_tail_idx)]
                if str(add_tail_op.op_type) != "ADD" or len(add_tail_op.inputs) != 2 or len(add_tail_op.outputs) != 1:
                    continue
                if squeeze2_out_name not in set(str(v) for v in add_tail_op.inputs):
                    continue
                tail_add_residual_input_name = (
                    str(add_tail_op.inputs[0])
                    if str(add_tail_op.inputs[1]) == squeeze2_out_name
                    else str(add_tail_op.inputs[1])
                )
                tail_add_out_name = str(add_tail_op.outputs[0])

                tail_add_users = [int(v) for v in consumers.get(tail_add_out_name, [])]
                reshape2_candidates: List[int] = []
                for user_idx in tail_add_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "RESHAPE"
                        and len(user_op.inputs) >= 1
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == tail_add_out_name
                    ):
                        reshape2_candidates.append(int(user_idx))
                if len(reshape2_candidates) != 1:
                    continue
                reshape2_idx = int(reshape2_candidates[0])
                reshape2_op = model_ir.operators[int(reshape2_idx)]
                if str(reshape2_op.op_type) != "RESHAPE" or len(reshape2_op.inputs) < 2 or len(reshape2_op.outputs) != 1:
                    continue
                if str(reshape2_op.inputs[0]) != tail_add_out_name:
                    continue
                reshape2_out_name = str(reshape2_op.outputs[0])

                tail_add_legacy_slots = []
                for user_idx in tail_add_users:
                    if int(user_idx) == int(reshape2_idx):
                        continue
                    user_op = model_ir.operators[int(user_idx)]
                    input_indices = [
                        int(input_idx)
                        for input_idx, input_name in enumerate(list(user_op.inputs))
                        if str(input_name) == tail_add_out_name
                    ]
                    if len(input_indices) == 0:
                        continue
                    for input_idx in input_indices:
                        tail_add_legacy_slots.append((int(id(user_op)), int(input_idx)))
                if len(tail_add_legacy_slots) > 0:
                    tail_add_tensor = model_ir.tensors.get(tail_add_out_name, None)
                    if tail_add_tensor is None or len(list(tail_add_tensor.shape)) != 3:
                        continue
                    tail_add_legacy_meta = (
                        str(tail_add_tensor.dtype),
                        _clone_quantization(tail_add_tensor.quantization),
                        [int(v) for v in list(tail_add_tensor.shape)],
                        (
                            [int(v) for v in list(tail_add_tensor.shape_signature)]
                            if tail_add_tensor.shape_signature is not None
                            else [int(v) for v in list(tail_add_tensor.shape)]
                        ),
                    )

                post_idx = _single_consumer_index(consumers, reshape2_out_name)
                if post_idx is None:
                    continue
                post_op = model_ir.operators[int(post_idx)]
                if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                    continue
                if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                    continue
                if str(post_op.inputs[0]) != reshape2_out_name:
                    continue

                # If residual side still uses CHW via NHWC->NCHW transpose, lift it to HWC.
                # This removes one extra transpose in residual shortcuts.
                if tail_add_residual_input_name is not None and tail_add_residual_input_name != "":
                    residual_input_name = str(tail_add_residual_input_name)
                    residual_layout_rewritten = False
                    if set(int(v) for v in consumers.get(residual_input_name, [])) != {int(add_tail_idx)}:
                        continue
                    residual_producer_idx = _single_producer_index(producers, residual_input_name)
                    if residual_producer_idx is not None:
                        residual_producer_op = model_ir.operators[int(residual_producer_idx)]
                        if (
                            str(residual_producer_op.op_type) == "TRANSPOSE"
                            and len(residual_producer_op.inputs) >= 2
                            and len(residual_producer_op.outputs) == 1
                            and str(residual_producer_op.outputs[0]) == residual_input_name
                            and _read_transpose_perm(model_ir, residual_producer_op) == [2, 0, 1]
                        ):
                            residual_src_name = str(residual_producer_op.inputs[0])
                            if residual_src_name in model_outputs:
                                continue
                            for input_idx, input_name in enumerate(list(add_tail_op.inputs)):
                                if str(input_name) == residual_input_name:
                                    _replace_operator_input_at(
                                        model_ir=model_ir,
                                        op=add_tail_op,
                                        input_index=int(input_idx),
                                        new_input_name=residual_src_name,
                                    )
                            tail_add_residual_input_name = str(residual_src_name)
                            residual_layout_rewritten = True
                            extra_remove_indices.add(int(residual_producer_idx))
                        else:
                            residual_squeeze_op: Optional[OperatorIR] = None
                            residual_squeeze_idx: Optional[int] = None
                            residual_squeeze_out_name: Optional[str] = None
                            residual_pre_transpose_idx: Optional[int] = None
                            residual_unary_out_name: Optional[str] = None

                            if str(residual_producer_op.op_type) in unary_ops and len(residual_producer_op.inputs) == 1:
                                candidate_squeeze_out = str(residual_producer_op.inputs[0])
                                candidate_squeeze_idx = _single_producer_index(producers, candidate_squeeze_out)
                                if candidate_squeeze_idx is not None:
                                    candidate_squeeze_op = model_ir.operators[int(candidate_squeeze_idx)]
                                    if (
                                        str(candidate_squeeze_op.op_type) == "SQUEEZE"
                                        and len(candidate_squeeze_op.inputs) == 1
                                        and len(candidate_squeeze_op.outputs) == 1
                                        and str(candidate_squeeze_op.outputs[0]) == candidate_squeeze_out
                                    ):
                                        residual_squeeze_op = candidate_squeeze_op
                                        residual_squeeze_idx = int(candidate_squeeze_idx)
                                        residual_squeeze_out_name = str(candidate_squeeze_out)
                                        residual_unary_out_name = str(residual_input_name)
                            elif str(residual_producer_op.op_type) == "SQUEEZE" and len(residual_producer_op.inputs) == 1:
                                residual_squeeze_op = residual_producer_op
                                residual_squeeze_idx = int(residual_producer_idx)
                                residual_squeeze_out_name = str(residual_input_name)

                            if (
                                residual_squeeze_op is not None
                                and residual_squeeze_idx is not None
                                and residual_squeeze_out_name is not None
                            ):
                                residual_squeeze_in_name = str(residual_squeeze_op.inputs[0])
                                if _is_squeeze_axis0(
                                    squeeze_op=residual_squeeze_op,
                                    input_name=residual_squeeze_in_name,
                                    output_name=residual_squeeze_out_name,
                                ):
                                    candidate_pre_idx = _single_producer_index(producers, residual_squeeze_in_name)
                                    if candidate_pre_idx is not None:
                                        candidate_pre_op = model_ir.operators[int(candidate_pre_idx)]
                                        if (
                                            str(candidate_pre_op.op_type) == "TRANSPOSE"
                                            and len(candidate_pre_op.inputs) >= 2
                                            and len(candidate_pre_op.outputs) == 1
                                            and str(candidate_pre_op.outputs[0]) == residual_squeeze_in_name
                                            and _read_transpose_perm(model_ir, candidate_pre_op) == perm_nhwc_to_nchw
                                        ):
                                            pre_out_users = set(int(v) for v in consumers.get(residual_squeeze_in_name, []))
                                            if (
                                                pre_out_users == {int(residual_squeeze_idx)}
                                                and residual_squeeze_in_name not in model_outputs
                                            ):
                                                residual_pre_transpose_idx = int(candidate_pre_idx)
                                                _set_operator_inputs(
                                                    model_ir=model_ir,
                                                    op=residual_squeeze_op,
                                                    new_inputs=[str(candidate_pre_op.inputs[0])],
                                                )
                                                _permute_tensor_metadata_if_rank_matches(
                                                    model_ir.tensors.get(residual_squeeze_out_name, None),
                                                    perm_3d_chw_to_hwc,
                                                )
                                                channel_last_hint_names.add(str(residual_squeeze_out_name))
                                                if residual_unary_out_name is not None:
                                                    _permute_tensor_metadata_if_rank_matches(
                                                        model_ir.tensors.get(str(residual_unary_out_name), None),
                                                        perm_3d_chw_to_hwc,
                                                    )
                                                    channel_last_hint_names.add(str(residual_unary_out_name))
                                                residual_layout_rewritten = True
                            if residual_pre_transpose_idx is not None:
                                extra_remove_indices.add(int(residual_pre_transpose_idx))
                    if not residual_layout_rewritten:
                        continue

            post_op = model_ir.operators[int(post_idx)]
            post_out_name = str(post_op.outputs[0])
            if post_out_name in model_outputs:
                continue

            if not _rewrite_reduce_axes_to_nhwc(
                axes_name=str(mean1_op.inputs[1]),
                expected_users={int(mean1_idx), int(mean2_idx)},
                consumers=consumers,
            ):
                continue
            if str(mean2_op.inputs[1]) != str(mean1_op.inputs[1]):
                if not _rewrite_reduce_axes_to_nhwc(
                    axes_name=str(mean2_op.inputs[1]),
                    expected_users={int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue

            if not _rewrite_nchw_coeff_to_nhwc(scale_const_name):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(bias_const_name):
                continue

            pre_in_tensor = model_ir.tensors.get(pre_in_name, None)
            reshape1_tensor = model_ir.tensors.get(x_name, None)
            if pre_in_tensor is None or reshape1_tensor is None:
                continue
            target_shape1 = [int(v) for v in list(pre_in_tensor.shape)]
            if len(target_shape1) != 4:
                continue
            if not _rewrite_reshape_shape_const(
                reshape_op=reshape1_op,
                reshape_idx=int(reshape1_idx),
                target_shape=target_shape1,
                consumers=consumers,
            ):
                continue

            if tail_mode in {"squeeze_unary_reshape", "squeeze_add_reshape"}:
                if reshape2_op is None or reshape2_idx is None:
                    continue
                reshape2_out_tensor = model_ir.tensors.get(reshape2_out_name, None)
                target_shape2 = None
                if reshape2_out_tensor is not None:
                    target_shape2 = _permute_shape(
                        [int(v) for v in list(reshape2_out_tensor.shape)],
                        perm_nchw_to_nhwc,
                    )
                if target_shape2 is None:
                    post_out_tensor = model_ir.tensors.get(post_out_name, None)
                    if post_out_tensor is not None:
                        target_shape2 = [int(v) for v in list(post_out_tensor.shape)]
                if target_shape2 is None or len(target_shape2) != 4:
                    continue
                if not _rewrite_reshape_shape_const(
                    reshape_op=reshape2_op,
                    reshape_idx=int(reshape2_idx),
                    target_shape=[int(v) for v in list(target_shape2)],
                    consumers=consumers,
                ):
                    continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean1_op,
                input_index=0,
                new_input_name=pre_in_name,
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=squeeze1_op,
                new_inputs=[pre_in_name],
            )
            for input_idx, input_name in enumerate(list(sub_op.inputs)):
                if str(input_name) == x_name:
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=sub_op,
                        input_index=int(input_idx),
                        new_input_name=pre_in_name,
                    )

            rank4_names = [
                x_name,
                mean1_out_name,
                centered_name,
                squared_name,
                mean2_out_name,
                add_eps_out_name,
                sqrt_out_name,
                div_out_name,
                normalized_name,
                scaled_name,
            ]
            if tail_mode != "direct_post_transpose_side_squeeze":
                rank4_names.append(str(inst_out_name))
            if tail_mode in {"squeeze_unary_reshape", "squeeze_add_reshape"} and str(reshape2_out_name) != str(inst_out_name):
                rank4_names.append(str(reshape2_out_name))
            for rank4_name in rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(rank4_name, None),
                    perm_nchw_to_nhwc,
                )
                channel_last_hint_names.add(str(rank4_name))
            rank3_names = [str(squeeze1_out_name)]
            if squeeze2_out_name is not None:
                rank3_names.append(str(squeeze2_out_name))
            if unary_out_name is not None:
                rank3_names.append(str(unary_out_name))
            if tail_mode == "squeeze_add_reshape" and tail_add_out_name is not None:
                rank3_names.append(str(tail_add_out_name))
            for rank3_name in rank3_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(rank3_name, None),
                    perm_3d_chw_to_hwc,
                )
                channel_last_hint_names.add(str(rank3_name))

            channel_last_hint_names.add(str(post_out_name))
            model_ir.metadata["assume_channel_last_layout_tensor_names"] = sorted(
                channel_last_hint_names
            )

            if tail_mode in {"squeeze_unary_reshape", "squeeze_add_reshape"}:
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=reshape2_op,
                    new_outputs=[post_out_name],
                )
            else:
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=add_bias_op,
                    new_outputs=[post_out_name],
                )

            remove_indices_local: set[int] = {int(pre_idx), int(post_idx)} | set(int(v) for v in extra_remove_indices)
            for remove_idx in sorted(list(remove_indices_local), reverse=True):
                del model_ir.operators[int(remove_idx)]

            if (
                tail_mode == "squeeze_add_reshape"
                and tail_add_out_name is not None
                and len(tail_add_legacy_slots) > 0
                and tail_add_legacy_meta is not None
            ):
                if tail_hwc_to_chw_perm_name not in model_ir.tensors:
                    model_ir.tensors[tail_hwc_to_chw_perm_name] = TensorIR(
                        name=tail_hwc_to_chw_perm_name,
                        dtype="INT32",
                        shape=[3],
                        shape_signature=[3],
                        data=np.asarray([2, 0, 1], dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                slot_targets: List[Tuple[int, int]] = []
                for consumer_op_id, input_index in tail_add_legacy_slots:
                    consumer_index = op_index_by_id.get(int(consumer_op_id), None)
                    if consumer_index is None:
                        continue
                    if int(consumer_index) < 0 or int(consumer_index) >= len(model_ir.operators):
                        continue
                    consumer_op = model_ir.operators[int(consumer_index)]
                    if int(input_index) < 0 or int(input_index) >= len(consumer_op.inputs):
                        continue
                    if str(consumer_op.inputs[int(input_index)]) != str(tail_add_out_name):
                        continue
                    slot_targets.append((int(consumer_index), int(input_index)))
                if len(slot_targets) > 0:
                    adapter_name = _unique_tensor_name(f"{tail_add_out_name}_chw_adapter")
                    adapter_dtype, adapter_quant, adapter_shape, adapter_shape_signature = tail_add_legacy_meta
                    model_ir.tensors[adapter_name] = TensorIR(
                        name=adapter_name,
                        dtype=str(adapter_dtype),
                        shape=[int(v) for v in list(adapter_shape)],
                        shape_signature=[int(v) for v in list(adapter_shape_signature)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(adapter_quant),
                    )
                    adapter_op = OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(tail_add_out_name), tail_hwc_to_chw_perm_name],
                        outputs=[adapter_name],
                    )
                    for consumer_index, input_index in slot_targets:
                        consumer_op = model_ir.operators[int(consumer_index)]
                        new_inputs = [str(v) for v in list(consumer_op.inputs)]
                        new_inputs[int(input_index)] = str(adapter_name)
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=consumer_op,
                            new_inputs=new_inputs,
                        )
                    insert_index = min(int(v[0]) for v in slot_targets)
                    model_ir.operators.insert(int(insert_index), adapter_op)

            if tail_mode == "direct_post_transpose_side_squeeze" and side_squeeze_idx is not None:
                if tail_nhwc_to_nchw_perm_name not in model_ir.tensors:
                    model_ir.tensors[tail_nhwc_to_nchw_perm_name] = TensorIR(
                        name=tail_nhwc_to_nchw_perm_name,
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                adapter_insert_idx = int(side_squeeze_idx)
                for removed_idx in sorted(list(remove_indices_local)):
                    if removed_idx < int(side_squeeze_idx):
                        adapter_insert_idx -= 1
                adapter_insert_idx = int(max(0, min(adapter_insert_idx, len(model_ir.operators))))
                model_ir.operators.insert(
                    int(adapter_insert_idx),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(post_out_name), tail_nhwc_to_nchw_perm_name],
                        outputs=[str(inst_out_name)],
                    ),
                )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_instancenorm_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_instancenorm_posttranspose_bias_add_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate strict NHWC<->NCHW transpose bridges around decomposed
    InstanceNormalization where channelwise bias ADD is already on NHWC side.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --(InstanceNormalization decomposition in NCHW, up to MUL(scale))--> s_nchw
      s_nchw --TRANSPOSE(0,2,3,1)--> s_nhwc
      ADD(s_nhwc, bias_nhwc) -> y_nhwc

    Rewrite:
      x_nhwc --(same decomposition rewritten to NHWC axes/broadcast)--> s_nhwc
      ADD(s_nhwc, bias_nhwc) -> y_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    channel_last_hint_names = {
        str(v)
        for v in model_ir.metadata.get("assume_channel_last_layout_tensor_names", [])
        if str(v) != ""
    }

    def _single_consumer_index(consumers: Dict[str, List[int]], tensor_name: str) -> Optional[int]:
        users = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(users) != 1:
            return None
        return int(users[0])

    def _rewrite_axes_to_nhwc(
        *,
        axes_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        axes_users = set(int(v) for v in consumers.get(str(axes_name), []))
        if axes_users != expected_users:
            return False
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None or axes_tensor.data is None:
            return False
        axes_values = _read_const_ints_from_tensor(axes_tensor)
        if axes_values is None:
            return False
        normalized_axes: List[int] = []
        for axis in axes_values:
            axis_value = int(axis)
            if axis_value < 0:
                axis_value += 4
            if axis_value < 0 or axis_value >= 4:
                return False
            normalized_axes.append(int(axis_value))
        if sorted(list(set(normalized_axes))) != [2, 3]:
            return False
        _write_const_ints_to_tensor(axes_tensor, [1, 2])
        return True

    def _rewrite_coeff_to_nhwc(
        *,
        coeff_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        coeff_users = set(int(v) for v in consumers.get(str(coeff_name), []))
        if coeff_users != expected_users:
            return False
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        if coeff_tensor is None or coeff_tensor.data is None:
            return False
        coeff = np.asarray(coeff_tensor.data)
        if int(coeff.size) == 1:
            return True
        if coeff.ndim != 4:
            return False
        shape = [int(v) for v in list(coeff.shape)]
        # Already NHWC broadcast [1,1,1,C].
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
        ):
            return True
        # NCHW broadcast [1,C,1,1] -> NHWC [1,1,1,C].
        if (
            int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            coeff_nhwc = np.transpose(coeff, perm_nchw_to_nhwc).astype(coeff.dtype, copy=False)
            coeff_tensor.data = np.asarray(coeff_nhwc)
            coeff_tensor.shape = [int(v) for v in list(coeff_nhwc.shape)]
            coeff_tensor.shape_signature = [int(v) for v in list(coeff_nhwc.shape)]
            return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue

            pre_users = set(int(v) for v in consumers.get(pre_out_name, []))
            if len(pre_users) != 2:
                continue

            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None
            for user_idx in sorted(list(pre_users)):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_out_name
                ):
                    mean1_idx = int(user_idx)
                elif (
                    str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)
            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name in model_outputs:
                continue
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue
            if set(int(v) for v in consumers.get(mean1_out_name, [])) != {int(sub_idx)}:
                continue

            centered_name = str(sub_op.outputs[0])
            centered_users = set(int(v) for v in consumers.get(centered_name, []))
            if len(centered_users) != 2:
                continue

            mul_square_idx: Optional[int] = None
            mul_norm_idx: Optional[int] = None
            for user_idx in sorted(list(centered_users)):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                if in0 == centered_name and in1 == centered_name:
                    mul_square_idx = int(user_idx)
                elif centered_name in {in0, in1}:
                    mul_norm_idx = int(user_idx)
            if mul_square_idx is None or mul_norm_idx is None:
                continue

            mul_square_op = model_ir.operators[int(mul_square_idx)]
            squared_name = str(mul_square_op.outputs[0])
            mean2_idx = _single_consumer_index(consumers, squared_name)
            if mean2_idx is None:
                continue
            mean2_op = model_ir.operators[int(mean2_idx)]
            if (
                str(mean2_op.op_type) != "MEAN"
                or len(mean2_op.inputs) != 2
                or len(mean2_op.outputs) != 1
                or str(mean2_op.inputs[0]) != squared_name
            ):
                continue
            mean2_out_name = str(mean2_op.outputs[0])

            add_eps_idx = _single_consumer_index(consumers, mean2_out_name)
            if add_eps_idx is None:
                continue
            add_eps_op = model_ir.operators[int(add_eps_idx)]
            if (
                str(add_eps_op.op_type) != "ADD"
                or len(add_eps_op.inputs) != 2
                or len(add_eps_op.outputs) != 1
                or mean2_out_name not in {str(add_eps_op.inputs[0]), str(add_eps_op.inputs[1])}
            ):
                continue
            add_eps_out_name = str(add_eps_op.outputs[0])

            sqrt_idx = _single_consumer_index(consumers, add_eps_out_name)
            if sqrt_idx is None:
                continue
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if (
                str(sqrt_op.op_type) != "SQRT"
                or len(sqrt_op.inputs) != 1
                or len(sqrt_op.outputs) != 1
                or str(sqrt_op.inputs[0]) != add_eps_out_name
            ):
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])

            div_idx = _single_consumer_index(consumers, sqrt_out_name)
            if div_idx is None:
                continue
            div_op = model_ir.operators[int(div_idx)]
            if (
                str(div_op.op_type) != "DIV"
                or len(div_op.inputs) != 2
                or len(div_op.outputs) != 1
                or str(div_op.inputs[1]) != sqrt_out_name
            ):
                continue
            div_out_name = str(div_op.outputs[0])

            mul_norm_op = model_ir.operators[int(mul_norm_idx)]
            if div_out_name not in {str(mul_norm_op.inputs[0]), str(mul_norm_op.inputs[1])}:
                continue
            normalized_name = str(mul_norm_op.outputs[0])

            mul_scale_idx = _single_consumer_index(consumers, normalized_name)
            if mul_scale_idx is None:
                continue
            mul_scale_op = model_ir.operators[int(mul_scale_idx)]
            if (
                str(mul_scale_op.op_type) != "MUL"
                or len(mul_scale_op.inputs) != 2
                or len(mul_scale_op.outputs) != 1
                or normalized_name not in {str(mul_scale_op.inputs[0]), str(mul_scale_op.inputs[1])}
            ):
                continue
            scale_const_name = (
                str(mul_scale_op.inputs[0])
                if str(mul_scale_op.inputs[1]) == normalized_name
                else str(mul_scale_op.inputs[1])
            )
            scaled_name = str(mul_scale_op.outputs[0])

            post_idx = _single_consumer_index(consumers, scaled_name)
            if post_idx is None:
                continue
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != scaled_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])
            if post_out_name in model_outputs:
                continue

            add_bias_idx = _single_consumer_index(consumers, post_out_name)
            if add_bias_idx is None:
                continue
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if (
                str(add_bias_op.op_type) != "ADD"
                or len(add_bias_op.inputs) != 2
                or len(add_bias_op.outputs) != 1
                or post_out_name not in {str(add_bias_op.inputs[0]), str(add_bias_op.inputs[1])}
            ):
                continue
            bias_const_name = (
                str(add_bias_op.inputs[0])
                if str(add_bias_op.inputs[1]) == post_out_name
                else str(add_bias_op.inputs[1])
            )

            mean1_axes_name = str(mean1_op.inputs[1])
            mean2_axes_name = str(mean2_op.inputs[1])
            if mean1_axes_name == mean2_axes_name:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx), int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue
            else:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx)},
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean2_axes_name,
                    expected_users={int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue

            if not _rewrite_coeff_to_nhwc(
                coeff_name=scale_const_name,
                expected_users={int(mul_scale_idx)},
                consumers=consumers,
            ):
                continue
            if not _rewrite_coeff_to_nhwc(
                coeff_name=bias_const_name,
                expected_users={int(add_bias_idx)},
                consumers=consumers,
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean1_op,
                input_index=0,
                new_input_name=pre_in_name,
            )
            sub_inputs = [str(v) for v in list(sub_op.inputs)]
            if sub_inputs[0] == pre_out_name:
                sub_inputs[0] = pre_in_name
            if sub_inputs[1] == pre_out_name:
                sub_inputs[1] = pre_in_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=sub_op,
                new_inputs=sub_inputs,
            )
            add_bias_inputs = [str(v) for v in list(add_bias_op.inputs)]
            if add_bias_inputs[0] == post_out_name:
                add_bias_inputs[0] = scaled_name
            if add_bias_inputs[1] == post_out_name:
                add_bias_inputs[1] = scaled_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_bias_op,
                new_inputs=add_bias_inputs,
            )

            rank4_names = [
                pre_out_name,
                mean1_out_name,
                centered_name,
                squared_name,
                mean2_out_name,
                add_eps_out_name,
                sqrt_out_name,
                div_out_name,
                normalized_name,
                scaled_name,
                post_out_name,
            ]
            for rank4_name in rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(rank4_name), None),
                    perm_nchw_to_nhwc,
                )
                channel_last_hint_names.add(str(rank4_name))
            channel_last_hint_names.add(str(pre_in_name))
            model_ir.metadata["assume_channel_last_layout_tensor_names"] = sorted(
                channel_last_hint_names
            )

            for remove_idx in sorted([int(pre_idx), int(post_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_instancenorm_posttranspose_bias_add_nhwc_chains": int(rewritten)
    }


def _optimize_transpose_instancenorm_pad_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NHWC<->NCHW transpose bridges around decomposed InstanceNormalization
    when the normalized output is padded before returning to NHWC.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --(InstanceNormalization decomposition in NCHW)--> y_nchw
      y_nchw --PAD|MIRROR_PAD(pads_nchw)--> p_nchw
      p_nchw --TRANSPOSE(0,2,3,1)--> z_nhwc

    Rewrite:
      x_nhwc --(same decomposition rewritten to NHWC axes/broadcast)--> y_nhwc
      y_nhwc --PAD|MIRROR_PAD(pads_nhwc)--> z_nhwc

    Notes:
    - If `y_nchw` has non-pad side consumers, preserve them by adding a local
      NHWC->NCHW transpose adapter for those legacy consumers only.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    legacy_adapter_perm_name = "__instancenorm_pad_tail_nhwc_to_nchw_perm_rank4__"

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _normalize_rank4_axes(axes: Optional[List[int]]) -> Optional[List[int]]:
        if axes is None or len(axes) == 0:
            return None
        normalized: List[int] = []
        for axis in axes:
            axis_value = int(axis)
            if axis_value < 0:
                axis_value += 4
            if axis_value < 0 or axis_value >= 4:
                return None
            normalized.append(int(axis_value))
        return sorted(list(set(normalized)))

    def _rewrite_axes_to_nhwc(
        *,
        axes_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        axes_users = set(int(v) for v in consumers.get(str(axes_name), []))
        if axes_users != expected_users:
            return False
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None:
            return False
        axes_values = _read_const_ints_from_tensor(axes_tensor)
        if _normalize_rank4_axes(axes_values) != [2, 3]:
            return False
        _write_const_ints_to_tensor(axes_tensor, [1, 2])
        return True

    def _rewrite_nchw_coeff_to_nhwc(
        *,
        coeff_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        coeff_users = set(int(v) for v in consumers.get(str(coeff_name), []))
        if coeff_users != expected_users:
            return False
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        if coeff_tensor is None or coeff_tensor.data is None:
            return False
        coeff = np.asarray(coeff_tensor.data)
        if int(coeff.size) == 1:
            return True
        if coeff.ndim != 4:
            return False
        shape = [int(v) for v in list(coeff.shape)]
        # Already NHWC broadcast [1,1,1,C].
        if (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
        ):
            return True
        # NCHW broadcast [1,C,1,1] -> NHWC [1,1,1,C].
        if (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            coeff_nhwc = np.transpose(coeff, perm_nchw_to_nhwc).astype(coeff.dtype, copy=False)
            coeff_tensor.data = np.asarray(coeff_nhwc)
            coeff_tensor.shape = [int(v) for v in list(coeff_nhwc.shape)]
            coeff_tensor.shape_signature = [int(v) for v in list(coeff_nhwc.shape)]
            return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            pre_users = set(int(v) for v in consumers.get(pre_out_name, []))
            if len(pre_users) != 2:
                continue

            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None
            for user_idx in sorted(list(pre_users)):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_out_name
                ):
                    mean1_idx = int(user_idx)
                elif (
                    str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)
            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name in model_outputs:
                continue
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue
            if set(int(v) for v in consumers.get(mean1_out_name, [])) != {int(sub_idx)}:
                continue

            centered_name = str(sub_op.outputs[0])
            if centered_name in model_outputs:
                continue
            centered_users = set(int(v) for v in consumers.get(centered_name, []))
            if len(centered_users) != 2:
                continue

            mul_square_idx: Optional[int] = None
            mul_norm_idx: Optional[int] = None
            for user_idx in sorted(list(centered_users)):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                if in0 == centered_name and in1 == centered_name:
                    mul_square_idx = int(user_idx)
                elif centered_name in {in0, in1}:
                    mul_norm_idx = int(user_idx)
            if mul_square_idx is None or mul_norm_idx is None:
                continue

            mul_square_op = model_ir.operators[int(mul_square_idx)]
            squared_name = str(mul_square_op.outputs[0])
            mean2_users = set(int(v) for v in consumers.get(squared_name, []))
            if len(mean2_users) != 1:
                continue
            mean2_idx = int(list(mean2_users)[0])
            mean2_op = model_ir.operators[int(mean2_idx)]
            if (
                str(mean2_op.op_type) != "MEAN"
                or len(mean2_op.inputs) != 2
                or len(mean2_op.outputs) != 1
                or str(mean2_op.inputs[0]) != squared_name
            ):
                continue
            mean2_out_name = str(mean2_op.outputs[0])

            add_eps_users = set(int(v) for v in consumers.get(mean2_out_name, []))
            if len(add_eps_users) != 1:
                continue
            add_eps_idx = int(list(add_eps_users)[0])
            add_eps_op = model_ir.operators[int(add_eps_idx)]
            if (
                str(add_eps_op.op_type) != "ADD"
                or len(add_eps_op.inputs) != 2
                or len(add_eps_op.outputs) != 1
                or mean2_out_name not in {str(add_eps_op.inputs[0]), str(add_eps_op.inputs[1])}
            ):
                continue
            add_eps_out_name = str(add_eps_op.outputs[0])

            sqrt_users = set(int(v) for v in consumers.get(add_eps_out_name, []))
            if len(sqrt_users) != 1:
                continue
            sqrt_idx = int(list(sqrt_users)[0])
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if (
                str(sqrt_op.op_type) != "SQRT"
                or len(sqrt_op.inputs) != 1
                or len(sqrt_op.outputs) != 1
                or str(sqrt_op.inputs[0]) != add_eps_out_name
            ):
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])

            div_users = set(int(v) for v in consumers.get(sqrt_out_name, []))
            if len(div_users) != 1:
                continue
            div_idx = int(list(div_users)[0])
            div_op = model_ir.operators[int(div_idx)]
            if (
                str(div_op.op_type) != "DIV"
                or len(div_op.inputs) != 2
                or len(div_op.outputs) != 1
                or str(div_op.inputs[1]) != sqrt_out_name
            ):
                continue
            div_out_name = str(div_op.outputs[0])

            mul_norm_op = model_ir.operators[int(mul_norm_idx)]
            if div_out_name not in {str(mul_norm_op.inputs[0]), str(mul_norm_op.inputs[1])}:
                continue
            normalized_name = str(mul_norm_op.outputs[0])

            mul_scale_users = set(int(v) for v in consumers.get(normalized_name, []))
            if len(mul_scale_users) != 1:
                continue
            mul_scale_idx = int(list(mul_scale_users)[0])
            mul_scale_op = model_ir.operators[int(mul_scale_idx)]
            if (
                str(mul_scale_op.op_type) != "MUL"
                or len(mul_scale_op.inputs) != 2
                or len(mul_scale_op.outputs) != 1
                or normalized_name not in {str(mul_scale_op.inputs[0]), str(mul_scale_op.inputs[1])}
            ):
                continue
            scale_const_name = (
                str(mul_scale_op.inputs[0])
                if str(mul_scale_op.inputs[1]) == normalized_name
                else str(mul_scale_op.inputs[1])
            )
            scaled_name = str(mul_scale_op.outputs[0])

            add_bias_users = set(int(v) for v in consumers.get(scaled_name, []))
            if len(add_bias_users) != 1:
                continue
            add_bias_idx = int(list(add_bias_users)[0])
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if (
                str(add_bias_op.op_type) != "ADD"
                or len(add_bias_op.inputs) != 2
                or len(add_bias_op.outputs) != 1
                or scaled_name not in {str(add_bias_op.inputs[0]), str(add_bias_op.inputs[1])}
            ):
                continue
            bias_const_name = (
                str(add_bias_op.inputs[0])
                if str(add_bias_op.inputs[1]) == scaled_name
                else str(add_bias_op.inputs[1])
            )
            inst_out_name = str(add_bias_op.outputs[0])
            if inst_out_name in model_outputs:
                continue

            tail_add_idx: Optional[int] = None
            tail_add_new_inputs: Optional[List[str]] = None
            tail_add_output_name: Optional[str] = None
            tail_add_residual_pre_idx: Optional[int] = None
            legacy_adapter_source_name = str(inst_out_name)
            inst_out_users = [int(v) for v in consumers.get(inst_out_name, [])]
            if len(inst_out_users) == 0:
                continue
            pad_candidates: List[int] = []
            for user_idx in inst_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) in {"PAD", "MIRROR_PAD"}
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == inst_out_name
                ):
                    pad_candidates.append(int(user_idx))
            legacy_consumer_slots: List[Tuple[int, int]] = []
            if len(pad_candidates) == 1:
                pad_idx = int(pad_candidates[0])
                pad_op = model_ir.operators[int(pad_idx)]
                for user_idx in inst_out_users:
                    if int(user_idx) == int(pad_idx):
                        continue
                    user_op = model_ir.operators[int(user_idx)]
                    for input_idx, input_name in enumerate(list(user_op.inputs)):
                        if str(input_name) == inst_out_name:
                            legacy_consumer_slots.append((int(id(user_op)), int(input_idx)))
            else:
                # Residual tail variant:
                #   inst_out_nchw --ADD(skip_nchw)--> add_out_nchw --PAD--> ... --T--> nhwc
                if len(inst_out_users) != 1:
                    continue
                candidate_add_idx = int(inst_out_users[0])
                candidate_add_op = model_ir.operators[int(candidate_add_idx)]
                if (
                    str(candidate_add_op.op_type) != "ADD"
                    or len(candidate_add_op.inputs) != 2
                    or len(candidate_add_op.outputs) != 1
                    or inst_out_name not in {str(candidate_add_op.inputs[0]), str(candidate_add_op.inputs[1])}
                ):
                    continue
                candidate_add_out_name = str(candidate_add_op.outputs[0])
                if candidate_add_out_name in model_outputs:
                    continue
                candidate_add_out_users = [int(v) for v in consumers.get(candidate_add_out_name, [])]
                if len(candidate_add_out_users) == 0:
                    continue
                candidate_pad_indices: List[int] = []
                for candidate_user_idx in candidate_add_out_users:
                    candidate_user_op = model_ir.operators[int(candidate_user_idx)]
                    if (
                        str(candidate_user_op.op_type) in {"PAD", "MIRROR_PAD"}
                        and len(candidate_user_op.inputs) >= 2
                        and len(candidate_user_op.outputs) == 1
                        and str(candidate_user_op.inputs[0]) == candidate_add_out_name
                    ):
                        candidate_pad_indices.append(int(candidate_user_idx))
                if len(candidate_pad_indices) != 1:
                    continue
                candidate_pad_idx = int(candidate_pad_indices[0])
                candidate_pad_op = model_ir.operators[int(candidate_pad_idx)]

                residual_input_name = (
                    str(candidate_add_op.inputs[0])
                    if str(candidate_add_op.inputs[1]) == inst_out_name
                    else str(candidate_add_op.inputs[1])
                )
                if residual_input_name in model_outputs:
                    continue
                residual_pre_idx = producers.get(residual_input_name, None)
                if residual_pre_idx is None:
                    continue
                residual_pre_op = model_ir.operators[int(residual_pre_idx)]
                if (
                    str(residual_pre_op.op_type) != "TRANSPOSE"
                    or len(residual_pre_op.inputs) < 2
                    or len(residual_pre_op.outputs) != 1
                    or str(residual_pre_op.outputs[0]) != residual_input_name
                    or _read_transpose_perm(model_ir, residual_pre_op) != perm_nhwc_to_nchw
                ):
                    continue
                if set(int(v) for v in consumers.get(residual_input_name, [])) != {int(candidate_add_idx)}:
                    continue
                residual_src_name = str(residual_pre_op.inputs[0])
                if residual_src_name in model_outputs:
                    continue

                tail_add_idx = int(candidate_add_idx)
                tail_add_output_name = str(candidate_add_out_name)
                tail_add_residual_pre_idx = int(residual_pre_idx)
                tail_add_new_inputs = [str(v) for v in list(candidate_add_op.inputs)]
                for input_idx, input_name in enumerate(list(tail_add_new_inputs)):
                    if str(input_name) == residual_input_name:
                        tail_add_new_inputs[int(input_idx)] = str(residual_src_name)
                for user_idx in candidate_add_out_users:
                    if int(user_idx) == int(candidate_pad_idx):
                        continue
                    user_op = model_ir.operators[int(user_idx)]
                    for input_idx, input_name in enumerate(list(user_op.inputs)):
                        if str(input_name) == candidate_add_out_name:
                            legacy_consumer_slots.append((int(id(user_op)), int(input_idx)))
                legacy_adapter_source_name = str(candidate_add_out_name)

                pad_idx = int(candidate_pad_idx)
                pad_op = candidate_pad_op

            pad_out_name = str(pad_op.outputs[0])
            if pad_out_name in model_outputs:
                continue

            post_users = set(int(v) for v in consumers.get(pad_out_name, []))
            if len(post_users) != 1:
                continue
            post_idx = int(list(post_users)[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != pad_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])

            mean1_axes_name = str(mean1_op.inputs[1])
            mean2_axes_name = str(mean2_op.inputs[1])
            if mean1_axes_name == mean2_axes_name:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx), int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue
            else:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx)},
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean2_axes_name,
                    expected_users={int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue

            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=scale_const_name,
                expected_users={int(mul_scale_idx)},
                consumers=consumers,
            ):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=bias_const_name,
                expected_users={int(add_bias_idx)},
                consumers=consumers,
            ):
                continue

            pads_name = str(pad_op.inputs[1])
            pads_users = set(int(v) for v in consumers.get(pads_name, []))
            if pads_users != {int(pad_idx)}:
                continue
            pads_tensor = model_ir.tensors.get(pads_name, None)
            if pads_tensor is None or pads_tensor.data is None:
                continue
            try:
                pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
            except Exception:
                continue
            if int(pads_pairs.size) != 8:
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean1_op,
                input_index=0,
                new_input_name=pre_in_name,
            )
            sub_inputs = [str(v) for v in list(sub_op.inputs)]
            if sub_inputs[0] == pre_out_name:
                sub_inputs[0] = pre_in_name
            if sub_inputs[1] == pre_out_name:
                sub_inputs[1] = pre_in_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=sub_op,
                new_inputs=sub_inputs,
            )
            if tail_add_idx is not None and tail_add_new_inputs is not None:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=model_ir.operators[int(tail_add_idx)],
                    new_inputs=[str(v) for v in list(tail_add_new_inputs)],
                )

            pads_nhwc = np.asarray(
                [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
                dtype=pads_pairs.dtype,
            )
            pads_tensor.data = np.asarray(pads_nhwc)
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

            _set_operator_outputs(
                model_ir=model_ir,
                op=pad_op,
                new_outputs=[post_out_name],
            )

            adapter_source_tensor = model_ir.tensors.get(str(legacy_adapter_source_name), None)
            adapter_meta: Optional[Tuple[str, Optional[Any], List[int], List[int]]] = None
            if adapter_source_tensor is not None:
                adapter_meta = (
                    str(adapter_source_tensor.dtype),
                    _clone_quantization(adapter_source_tensor.quantization),
                    [int(v) for v in list(adapter_source_tensor.shape)],
                    (
                        [int(v) for v in list(adapter_source_tensor.shape_signature)]
                        if adapter_source_tensor.shape_signature is not None
                        else [int(v) for v in list(adapter_source_tensor.shape)]
                    ),
                )

            rank4_names = [
                pre_out_name,
                mean1_out_name,
                centered_name,
                squared_name,
                mean2_out_name,
                add_eps_out_name,
                sqrt_out_name,
                div_out_name,
                normalized_name,
                scaled_name,
                inst_out_name,
                str(tail_add_output_name) if tail_add_output_name is not None else "",
                pad_out_name,
            ]
            for tensor_name in rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            old_pad_tensor = model_ir.tensors.get(str(pad_out_name), None)
            post_out_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_pad_tensor is not None and post_out_tensor is not None:
                post_out_tensor.dtype = str(old_pad_tensor.dtype)
                post_out_tensor.quantization = _clone_quantization(old_pad_tensor.quantization)
                post_out_tensor.shape = [int(v) for v in list(old_pad_tensor.shape)]
                post_out_tensor.shape_signature = (
                    [int(v) for v in list(old_pad_tensor.shape_signature)]
                    if old_pad_tensor.shape_signature is not None
                    else [int(v) for v in list(old_pad_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_out_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {int(pre_idx), int(post_idx)}
            if tail_add_residual_pre_idx is not None:
                remove_indices.add(int(tail_add_residual_pre_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            if len(legacy_consumer_slots) > 0 and adapter_meta is not None:
                if legacy_adapter_perm_name not in model_ir.tensors:
                    model_ir.tensors[legacy_adapter_perm_name] = TensorIR(
                        name=legacy_adapter_perm_name,
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                op_index_by_id = {
                    int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)
                }
                slot_targets: List[Tuple[int, int]] = []
                for consumer_op_id, input_index in legacy_consumer_slots:
                    consumer_index = op_index_by_id.get(int(consumer_op_id), None)
                    if consumer_index is None:
                        continue
                    consumer_op = model_ir.operators[int(consumer_index)]
                    if int(input_index) < 0 or int(input_index) >= len(consumer_op.inputs):
                        continue
                    if str(consumer_op.inputs[int(input_index)]) != str(legacy_adapter_source_name):
                        continue
                    slot_targets.append((int(consumer_index), int(input_index)))

                if len(slot_targets) > 0:
                    adapter_name = _unique_tensor_name(f"{legacy_adapter_source_name}_nchw_adapter")
                    adapter_dtype, adapter_quant, adapter_shape, adapter_shape_signature = adapter_meta
                    model_ir.tensors[adapter_name] = TensorIR(
                        name=adapter_name,
                        dtype=str(adapter_dtype),
                        shape=[int(v) for v in list(adapter_shape)],
                        shape_signature=[int(v) for v in list(adapter_shape_signature)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(adapter_quant),
                    )
                    for consumer_index, input_index in slot_targets:
                        consumer_op = model_ir.operators[int(consumer_index)]
                        new_inputs = [str(v) for v in list(consumer_op.inputs)]
                        new_inputs[int(input_index)] = str(adapter_name)
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=consumer_op,
                            new_inputs=new_inputs,
                        )
                    insert_index = min(int(v[0]) for v in slot_targets)
                    model_ir.operators.insert(
                        int(insert_index),
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(legacy_adapter_source_name), str(legacy_adapter_perm_name)],
                            outputs=[str(adapter_name)],
                        ),
                    )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_instancenorm_pad_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_flatten_globalnorm_pad_prepost_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC<->NCHW transpose bridges around flattened global-normalization blocks.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw --RESHAPE([N,1,-1])--> r
      r --(global norm decomposition)--> r_out
      r_out --RESHAPE([N,C,H,W])--> y_nchw
      y_nchw --(layout-agnostic tail)*--> t_nchw
      t_nchw --PAD|MIRROR_PAD(pads_nchw)--> p_nchw
      p_nchw --T(0,2,3,1)--> z_nhwc

    Rewrite:
      x_nhwc --RESHAPE([N,1,-1])--> r
      r --(same global norm decomposition)--> r_out
      r_out --RESHAPE([N,H,W,C])--> y_nhwc
      y_nhwc --(same layout-agnostic tail)*--> t_nhwc
      t_nhwc --PAD|MIRROR_PAD(pads_nhwc)--> z_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_tail_ops = {
        "LEAKY_RELU",
        "RELU",
        "RELU6",
        "TANH",
        "LOGISTIC",
        "HARD_SWISH",
        "NEG",
        "ABS",
        "FLOOR",
    }
    binary_tail_ops = {"ADD", "SUB", "MUL", "DIV"}

    def _is_const_broadcast_layout_agnostic(name: str) -> bool:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.data is None:
            return False
        array = np.asarray(tensor.data)
        if int(array.size) == 1:
            return True
        shape = [int(v) for v in list(array.shape)]
        return len(shape) > 0 and all(int(dim) == 1 for dim in shape)

    def _can_rewrite_tail_const_to_nhwc(name: str) -> bool:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.data is None:
            return False
        array = np.asarray(tensor.data)
        if int(array.size) == 1:
            return True
        shape = [int(v) for v in list(array.shape)]
        if len(shape) > 0 and all(int(dim) == 1 for dim in shape):
            return True
        if len(shape) == 3 and int(shape[1]) == 1 and int(shape[2]) == 1:
            return True
        if (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            return True
        return False

    def _rewrite_tail_const_to_nhwc(name: str) -> bool:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.data is None:
            return False
        array = np.asarray(tensor.data)
        if int(array.size) == 1:
            return True
        shape = [int(v) for v in list(array.shape)]
        if len(shape) > 0 and all(int(dim) == 1 for dim in shape):
            return True
        if len(shape) == 3 and int(shape[1]) == 1 and int(shape[2]) == 1:
            nhwc = np.reshape(array, newshape=[1, 1, int(shape[0])]).astype(array.dtype, copy=False)
            tensor.data = np.asarray(nhwc)
            tensor.shape = [int(v) for v in list(nhwc.shape)]
            tensor.shape_signature = [int(v) for v in list(nhwc.shape)]
            return True
        if (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            nhwc = np.transpose(array, perm_nchw_to_nhwc).astype(array.dtype, copy=False)
            tensor.data = np.asarray(nhwc)
            tensor.shape = [int(v) for v in list(nhwc.shape)]
            tensor.shape_signature = [int(v) for v in list(nhwc.shape)]
            return True
        return False

    def _read_reshape_target(op: OperatorIR) -> Optional[List[int]]:
        target: List[int] = []
        if len(op.inputs) >= 2:
            shape_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
            shape_vals = _read_const_ints_from_tensor(shape_tensor)
            if shape_vals is not None:
                target = [int(v) for v in list(shape_vals)]
        if len(target) == 0:
            raw_target = op.options.get("newShape", [])
            try:
                target = [int(v) for v in np.asarray(raw_target).reshape(-1).tolist()]
            except Exception:
                target = []
        return [int(v) for v in list(target)] if len(target) > 0 else None

    def _set_reshape_target(op: OperatorIR, target: List[int]) -> bool:
        updated = False
        shape_values = [int(v) for v in list(target)]
        if len(op.inputs) >= 2:
            shape_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
            if shape_tensor is not None:
                if not _write_const_ints_to_tensor(shape_tensor, shape_values):
                    return False
                updated = True
        if "newShape" in op.options:
            op.options["newShape"] = [int(v) for v in list(shape_values)]
            updated = True
        if "onnxRawNewShape" in op.options and isinstance(op.options["onnxRawNewShape"], list):
            raw_shape = [int(v) for v in list(op.options["onnxRawNewShape"])]
            if (
                len(raw_shape) == len(shape_values)
                and all(int(v) >= 0 for v in raw_shape)
            ):
                op.options["onnxRawNewShape"] = [int(v) for v in list(shape_values)]
                updated = True
        return updated

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(pre_users) != 1:
                continue
            reshape1_idx = int(pre_users[0])
            reshape1_op = model_ir.operators[int(reshape1_idx)]
            if (
                str(reshape1_op.op_type) != "RESHAPE"
                or len(reshape1_op.inputs) < 1
                or len(reshape1_op.outputs) != 1
                or str(reshape1_op.inputs[0]) != pre_out_name
            ):
                continue
            reshape1_out_name = str(reshape1_op.outputs[0])
            reshape1_target = _read_reshape_target(reshape1_op)
            if (
                reshape1_target is None
                or len(reshape1_target) != 3
                or int(reshape1_target[1]) != 1
                or int(reshape1_target[2]) != -1
            ):
                continue

            flat_users = set(int(v) for v in consumers.get(reshape1_out_name, []))
            if len(flat_users) != 2:
                continue
            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None
            for user_idx in sorted(list(flat_users)):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == reshape1_out_name
                ):
                    axes_vals = _read_const_ints_from_tensor(
                        model_ir.tensors.get(str(user_op.inputs[1]), None)
                    )
                    if axes_vals is None or [int(v) for v in list(axes_vals)] != [2]:
                        break
                    mean1_idx = int(user_idx)
                elif (
                    str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and reshape1_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)
            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue
            if set(int(v) for v in consumers.get(mean1_out_name, [])) != {int(sub_idx)}:
                continue

            centered_name = str(sub_op.outputs[0])
            centered_users = set(int(v) for v in consumers.get(centered_name, []))
            if len(centered_users) != 2:
                continue
            mul_square_idx: Optional[int] = None
            mul_norm_idx: Optional[int] = None
            for user_idx in sorted(list(centered_users)):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                if in0 == centered_name and in1 == centered_name:
                    mul_square_idx = int(user_idx)
                elif centered_name in {in0, in1}:
                    mul_norm_idx = int(user_idx)
            if mul_square_idx is None or mul_norm_idx is None:
                continue

            mul_square_op = model_ir.operators[int(mul_square_idx)]
            squared_name = str(mul_square_op.outputs[0])
            mean2_users = set(int(v) for v in consumers.get(squared_name, []))
            if len(mean2_users) != 1:
                continue
            mean2_idx = int(list(mean2_users)[0])
            mean2_op = model_ir.operators[int(mean2_idx)]
            if (
                str(mean2_op.op_type) != "MEAN"
                or len(mean2_op.inputs) != 2
                or len(mean2_op.outputs) != 1
                or str(mean2_op.inputs[0]) != squared_name
            ):
                continue
            axes2_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(mean2_op.inputs[1]), None))
            if axes2_vals is None or [int(v) for v in list(axes2_vals)] != [2]:
                continue

            mean2_out_name = str(mean2_op.outputs[0])
            add_eps_users = set(int(v) for v in consumers.get(mean2_out_name, []))
            if len(add_eps_users) != 1:
                continue
            add_eps_idx = int(list(add_eps_users)[0])
            add_eps_op = model_ir.operators[int(add_eps_idx)]
            if (
                str(add_eps_op.op_type) != "ADD"
                or len(add_eps_op.inputs) != 2
                or len(add_eps_op.outputs) != 1
                or mean2_out_name not in {str(add_eps_op.inputs[0]), str(add_eps_op.inputs[1])}
            ):
                continue
            add_eps_out_name = str(add_eps_op.outputs[0])

            sqrt_users = set(int(v) for v in consumers.get(add_eps_out_name, []))
            if len(sqrt_users) != 1:
                continue
            sqrt_idx = int(list(sqrt_users)[0])
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if (
                str(sqrt_op.op_type) != "SQRT"
                or len(sqrt_op.inputs) != 1
                or len(sqrt_op.outputs) != 1
                or str(sqrt_op.inputs[0]) != add_eps_out_name
            ):
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])

            div_users = set(int(v) for v in consumers.get(sqrt_out_name, []))
            if len(div_users) != 1:
                continue
            div_idx = int(list(div_users)[0])
            div_op = model_ir.operators[int(div_idx)]
            if (
                str(div_op.op_type) != "DIV"
                or len(div_op.inputs) != 2
                or len(div_op.outputs) != 1
                or str(div_op.inputs[1]) != sqrt_out_name
            ):
                continue
            div_out_name = str(div_op.outputs[0])

            mul_norm_op = model_ir.operators[int(mul_norm_idx)]
            if div_out_name not in {str(mul_norm_op.inputs[0]), str(mul_norm_op.inputs[1])}:
                continue
            normalized_name = str(mul_norm_op.outputs[0])

            mul_scale_users = set(int(v) for v in consumers.get(normalized_name, []))
            if len(mul_scale_users) != 1:
                continue
            mul_scale_idx = int(list(mul_scale_users)[0])
            mul_scale_op = model_ir.operators[int(mul_scale_idx)]
            if (
                str(mul_scale_op.op_type) != "MUL"
                or len(mul_scale_op.inputs) != 2
                or len(mul_scale_op.outputs) != 1
                or normalized_name not in {str(mul_scale_op.inputs[0]), str(mul_scale_op.inputs[1])}
            ):
                continue
            scale_const_name = (
                str(mul_scale_op.inputs[0])
                if str(mul_scale_op.inputs[1]) == normalized_name
                else str(mul_scale_op.inputs[1])
            )
            if not _is_const_broadcast_layout_agnostic(scale_const_name):
                continue
            scaled_name = str(mul_scale_op.outputs[0])

            add_bias_users = set(int(v) for v in consumers.get(scaled_name, []))
            if len(add_bias_users) != 1:
                continue
            add_bias_idx = int(list(add_bias_users)[0])
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if (
                str(add_bias_op.op_type) != "ADD"
                or len(add_bias_op.inputs) != 2
                or len(add_bias_op.outputs) != 1
                or scaled_name not in {str(add_bias_op.inputs[0]), str(add_bias_op.inputs[1])}
            ):
                continue
            bias_const_name = (
                str(add_bias_op.inputs[0])
                if str(add_bias_op.inputs[1]) == scaled_name
                else str(add_bias_op.inputs[1])
            )
            if not _is_const_broadcast_layout_agnostic(bias_const_name):
                continue
            inst_flat_name = str(add_bias_op.outputs[0])

            reshape2_users = [int(v) for v in consumers.get(inst_flat_name, [])]
            if len(reshape2_users) != 1:
                continue
            reshape2_idx = int(reshape2_users[0])
            reshape2_op = model_ir.operators[int(reshape2_idx)]
            if (
                str(reshape2_op.op_type) != "RESHAPE"
                or len(reshape2_op.inputs) < 1
                or len(reshape2_op.outputs) != 1
                or str(reshape2_op.inputs[0]) != inst_flat_name
            ):
                continue
            reshape2_target = _read_reshape_target(reshape2_op)
            if (
                reshape2_target is None
                or len(reshape2_target) != 4
                or any(int(v) <= 0 for v in reshape2_target)
            ):
                continue

            tail_rank4_names: List[str] = [str(reshape2_op.outputs[0])]
            tail_const_names: List[str] = []
            cursor_name = str(reshape2_op.outputs[0])
            pad_idx: Optional[int] = None
            pad_op: Optional[OperatorIR] = None
            visited_names: set[str] = {cursor_name}
            while True:
                cursor_users = [int(v) for v in consumers.get(cursor_name, [])]
                if len(cursor_users) != 1:
                    break
                user_idx = int(cursor_users[0])
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if (
                    user_type in {"PAD", "MIRROR_PAD"}
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == cursor_name
                ):
                    pad_idx = int(user_idx)
                    pad_op = user_op
                    tail_rank4_names.append(str(user_op.outputs[0]))
                    break
                if (
                    user_type in unary_tail_ops
                    and len(user_op.inputs) == 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == cursor_name
                ):
                    cursor_name = str(user_op.outputs[0])
                    if cursor_name in visited_names:
                        break
                    visited_names.add(cursor_name)
                    tail_rank4_names.append(cursor_name)
                    continue
                if (
                    user_type in binary_tail_ops
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and cursor_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    other_name = (
                        str(user_op.inputs[0])
                        if str(user_op.inputs[1]) == cursor_name
                        else str(user_op.inputs[1])
                    )
                    if not _can_rewrite_tail_const_to_nhwc(other_name):
                        break
                    tail_const_names.append(str(other_name))
                    cursor_name = str(user_op.outputs[0])
                    if cursor_name in visited_names:
                        break
                    visited_names.add(cursor_name)
                    tail_rank4_names.append(cursor_name)
                    continue
                break
            if pad_idx is None or pad_op is None:
                continue

            pad_out_name = str(pad_op.outputs[0])
            post_users = [int(v) for v in consumers.get(pad_out_name, [])]
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != pad_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])

            pads_name = str(pad_op.inputs[1])
            pads_tensor = model_ir.tensors.get(pads_name, None)
            if pads_tensor is None or pads_tensor.data is None:
                continue
            try:
                pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
            except Exception:
                continue
            if int(pads_pairs.size) != 8:
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=reshape1_op,
                input_index=0,
                new_input_name=pre_in_name,
            )

            nchw_shape = [int(v) for v in list(reshape2_target)]
            nhwc_shape = [
                int(nchw_shape[0]),
                int(nchw_shape[2]),
                int(nchw_shape[3]),
                int(nchw_shape[1]),
            ]
            if not _set_reshape_target(reshape2_op, nhwc_shape):
                continue

            unique_tail_const_names = []
            seen_const_names: set[str] = set()
            for const_name in tail_const_names:
                key = str(const_name)
                if key in seen_const_names:
                    continue
                seen_const_names.add(key)
                unique_tail_const_names.append(key)
            if any(not _rewrite_tail_const_to_nhwc(const_name) for const_name in unique_tail_const_names):
                continue

            pads_nhwc = np.asarray(
                [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
                dtype=pads_pairs.dtype,
            )
            pads_tensor.data = np.asarray(pads_nhwc)
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

            _set_operator_outputs(
                model_ir=model_ir,
                op=pad_op,
                new_outputs=[post_out_name],
            )

            for tensor_name in tail_rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            old_pad_tensor = model_ir.tensors.get(str(pad_out_name), None)
            post_out_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_pad_tensor is not None and post_out_tensor is not None:
                post_out_tensor.dtype = str(old_pad_tensor.dtype)
                post_out_tensor.quantization = _clone_quantization(old_pad_tensor.quantization)
                post_out_tensor.shape = [int(v) for v in list(old_pad_tensor.shape)]
                post_out_tensor.shape_signature = (
                    [int(v) for v in list(old_pad_tensor.shape_signature)]
                    if old_pad_tensor.shape_signature is not None
                    else [int(v) for v in list(old_pad_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_out_tensor,
                    perm_nchw_to_nhwc,
                )

            for remove_idx in sorted([int(pre_idx), int(post_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_flatten_globalnorm_pad_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_instancenorm_residual_add_to_single_post_adapter_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse dual NHWC->NCHW bridges around decomposed InstanceNormalization + residual ADD.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --(InstanceNormalization decomposition in NCHW)--> y_nchw
      r_nhwc --TRANSPOSE(0,3,1,2)--> r_nchw
      ADD(y_nchw, r_nchw) -> a_nchw

    Rewrite:
      x_nhwc --(same decomposition rewritten to NHWC axes/broadcast)--> y_nhwc
      ADD(y_nhwc, r_nhwc) -> a_nhwc
      a_nhwc --TRANSPOSE(0,3,1,2)--> a_nchw

    Safety:
    - Both pre-transpose outputs are consumed only by the rewritten branch.
    - Strict decomposed InstanceNormalization topology.
    - Preserve downstream NCHW users through one post-ADD adapter.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    adapter_perm_name = "__instancenorm_residual_add_nhwc_to_nchw_perm_rank4__"

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _normalize_rank4_axes(axes: Optional[List[int]]) -> Optional[List[int]]:
        if axes is None or len(axes) == 0:
            return None
        normalized: List[int] = []
        for axis in axes:
            axis_value = int(axis)
            if axis_value < 0:
                axis_value += 4
            if axis_value < 0 or axis_value >= 4:
                return None
            normalized.append(int(axis_value))
        return sorted(list(set(normalized)))

    def _rewrite_axes_to_nhwc(
        *,
        axes_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        axes_users = set(int(v) for v in consumers.get(str(axes_name), []))
        if axes_users != expected_users:
            return False
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None:
            return False
        axes_values = _read_const_ints_from_tensor(axes_tensor)
        if _normalize_rank4_axes(axes_values) != [2, 3]:
            return False
        _write_const_ints_to_tensor(axes_tensor, [1, 2])
        return True

    def _rewrite_nchw_coeff_to_nhwc(
        *,
        coeff_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        coeff_users = set(int(v) for v in consumers.get(str(coeff_name), []))
        if coeff_users != expected_users:
            return False
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        if coeff_tensor is None or coeff_tensor.data is None:
            return False
        coeff = np.asarray(coeff_tensor.data)
        if int(coeff.size) == 1:
            return True
        if coeff.ndim != 4:
            return False
        shape = [int(v) for v in list(coeff.shape)]
        if (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
        ):
            return True
        if (
            len(shape) == 4
            and int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            coeff_nhwc = np.transpose(coeff, perm_nchw_to_nhwc).astype(coeff.dtype, copy=False)
            coeff_tensor.data = np.asarray(coeff_nhwc)
            coeff_tensor.shape = [int(v) for v in list(coeff_nhwc.shape)]
            coeff_tensor.shape_signature = [int(v) for v in list(coeff_nhwc.shape)]
            return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            pre_users = set(int(v) for v in consumers.get(pre_out_name, []))
            if len(pre_users) != 2:
                continue

            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None
            for user_idx in sorted(list(pre_users)):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_out_name
                ):
                    mean1_idx = int(user_idx)
                elif (
                    str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)
            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name in model_outputs:
                continue
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue
            if set(int(v) for v in consumers.get(mean1_out_name, [])) != {int(sub_idx)}:
                continue

            centered_name = str(sub_op.outputs[0])
            if centered_name in model_outputs:
                continue
            centered_users = set(int(v) for v in consumers.get(centered_name, []))
            if len(centered_users) != 2:
                continue

            mul_square_idx: Optional[int] = None
            mul_norm_idx: Optional[int] = None
            for user_idx in sorted(list(centered_users)):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                if in0 == centered_name and in1 == centered_name:
                    mul_square_idx = int(user_idx)
                elif centered_name in {in0, in1}:
                    mul_norm_idx = int(user_idx)
            if mul_square_idx is None or mul_norm_idx is None:
                continue

            mul_square_op = model_ir.operators[int(mul_square_idx)]
            squared_name = str(mul_square_op.outputs[0])
            mean2_users = set(int(v) for v in consumers.get(squared_name, []))
            if len(mean2_users) != 1:
                continue
            mean2_idx = int(list(mean2_users)[0])
            mean2_op = model_ir.operators[int(mean2_idx)]
            if (
                str(mean2_op.op_type) != "MEAN"
                or len(mean2_op.inputs) != 2
                or len(mean2_op.outputs) != 1
                or str(mean2_op.inputs[0]) != squared_name
            ):
                continue
            mean2_out_name = str(mean2_op.outputs[0])

            add_eps_users = set(int(v) for v in consumers.get(mean2_out_name, []))
            if len(add_eps_users) != 1:
                continue
            add_eps_idx = int(list(add_eps_users)[0])
            add_eps_op = model_ir.operators[int(add_eps_idx)]
            if (
                str(add_eps_op.op_type) != "ADD"
                or len(add_eps_op.inputs) != 2
                or len(add_eps_op.outputs) != 1
                or mean2_out_name not in {str(add_eps_op.inputs[0]), str(add_eps_op.inputs[1])}
            ):
                continue
            add_eps_out_name = str(add_eps_op.outputs[0])

            sqrt_users = set(int(v) for v in consumers.get(add_eps_out_name, []))
            if len(sqrt_users) != 1:
                continue
            sqrt_idx = int(list(sqrt_users)[0])
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if (
                str(sqrt_op.op_type) != "SQRT"
                or len(sqrt_op.inputs) != 1
                or len(sqrt_op.outputs) != 1
                or str(sqrt_op.inputs[0]) != add_eps_out_name
            ):
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])

            div_users = set(int(v) for v in consumers.get(sqrt_out_name, []))
            if len(div_users) != 1:
                continue
            div_idx = int(list(div_users)[0])
            div_op = model_ir.operators[int(div_idx)]
            if (
                str(div_op.op_type) != "DIV"
                or len(div_op.inputs) != 2
                or len(div_op.outputs) != 1
                or str(div_op.inputs[1]) != sqrt_out_name
            ):
                continue
            div_out_name = str(div_op.outputs[0])

            mul_norm_op = model_ir.operators[int(mul_norm_idx)]
            if div_out_name not in {str(mul_norm_op.inputs[0]), str(mul_norm_op.inputs[1])}:
                continue
            normalized_name = str(mul_norm_op.outputs[0])

            mul_scale_users = set(int(v) for v in consumers.get(normalized_name, []))
            if len(mul_scale_users) != 1:
                continue
            mul_scale_idx = int(list(mul_scale_users)[0])
            mul_scale_op = model_ir.operators[int(mul_scale_idx)]
            if (
                str(mul_scale_op.op_type) != "MUL"
                or len(mul_scale_op.inputs) != 2
                or len(mul_scale_op.outputs) != 1
                or normalized_name not in {str(mul_scale_op.inputs[0]), str(mul_scale_op.inputs[1])}
            ):
                continue
            scale_const_name = (
                str(mul_scale_op.inputs[0])
                if str(mul_scale_op.inputs[1]) == normalized_name
                else str(mul_scale_op.inputs[1])
            )
            scaled_name = str(mul_scale_op.outputs[0])

            add_bias_users = set(int(v) for v in consumers.get(scaled_name, []))
            if len(add_bias_users) != 1:
                continue
            add_bias_idx = int(list(add_bias_users)[0])
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if (
                str(add_bias_op.op_type) != "ADD"
                or len(add_bias_op.inputs) != 2
                or len(add_bias_op.outputs) != 1
                or scaled_name not in {str(add_bias_op.inputs[0]), str(add_bias_op.inputs[1])}
            ):
                continue
            bias_const_name = (
                str(add_bias_op.inputs[0])
                if str(add_bias_op.inputs[1]) == scaled_name
                else str(add_bias_op.inputs[1])
            )
            inst_out_name = str(add_bias_op.outputs[0])
            if inst_out_name in model_outputs:
                continue

            inst_out_users = set(int(v) for v in consumers.get(inst_out_name, []))
            if len(inst_out_users) != 1:
                continue
            tail_add_idx = int(list(inst_out_users)[0])
            tail_add_op = model_ir.operators[int(tail_add_idx)]
            if (
                str(tail_add_op.op_type) != "ADD"
                or len(tail_add_op.inputs) != 2
                or len(tail_add_op.outputs) != 1
                or inst_out_name not in {str(tail_add_op.inputs[0]), str(tail_add_op.inputs[1])}
            ):
                continue
            add_out_name = str(tail_add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            residual_input_name = (
                str(tail_add_op.inputs[0])
                if str(tail_add_op.inputs[1]) == inst_out_name
                else str(tail_add_op.inputs[1])
            )
            residual_pre_idx = producers.get(residual_input_name, None)
            if residual_pre_idx is None:
                continue
            residual_pre_op = model_ir.operators[int(residual_pre_idx)]
            if (
                str(residual_pre_op.op_type) != "TRANSPOSE"
                or len(residual_pre_op.inputs) < 2
                or len(residual_pre_op.outputs) != 1
                or str(residual_pre_op.outputs[0]) != residual_input_name
                or _read_transpose_perm(model_ir, residual_pre_op) != perm_nhwc_to_nchw
            ):
                continue
            if set(int(v) for v in consumers.get(residual_input_name, [])) != {int(tail_add_idx)}:
                continue
            residual_src_name = str(residual_pre_op.inputs[0])
            if residual_src_name in model_outputs:
                continue

            add_out_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_out_users) == 0:
                continue

            mean1_axes_name = str(mean1_op.inputs[1])
            mean2_axes_name = str(mean2_op.inputs[1])
            if mean1_axes_name == mean2_axes_name:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx), int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue
            else:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx)},
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean2_axes_name,
                    expected_users={int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue

            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=scale_const_name,
                expected_users={int(mul_scale_idx)},
                consumers=consumers,
            ):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=bias_const_name,
                expected_users={int(add_bias_idx)},
                consumers=consumers,
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean1_op,
                input_index=0,
                new_input_name=pre_in_name,
            )
            sub_inputs = [str(v) for v in list(sub_op.inputs)]
            if sub_inputs[0] == pre_out_name:
                sub_inputs[0] = pre_in_name
            if sub_inputs[1] == pre_out_name:
                sub_inputs[1] = pre_in_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=sub_op,
                new_inputs=sub_inputs,
            )

            tail_add_inputs = [str(v) for v in list(tail_add_op.inputs)]
            for input_idx, input_name in enumerate(list(tail_add_inputs)):
                if str(input_name) == residual_input_name:
                    tail_add_inputs[int(input_idx)] = str(residual_src_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=tail_add_op,
                new_inputs=tail_add_inputs,
            )

            add_out_nhwc_name = _unique_tensor_name(f"{add_out_name}_nhwc")
            old_add_out_tensor = model_ir.tensors.get(str(add_out_name), None)
            add_out_nhwc_shape: Optional[List[int]] = None
            add_out_nhwc_signature: Optional[List[int]] = None
            add_out_dtype = "FLOAT32"
            add_out_quant = None
            if old_add_out_tensor is not None:
                add_out_dtype = str(old_add_out_tensor.dtype)
                add_out_quant = _clone_quantization(old_add_out_tensor.quantization)
                shape_src = [int(v) for v in list(old_add_out_tensor.shape)]
                signature_src = (
                    [int(v) for v in list(old_add_out_tensor.shape_signature)]
                    if old_add_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_add_out_tensor.shape)]
                )
                add_out_nhwc_shape = _permute_shape(shape_src, perm_nchw_to_nhwc)
                add_out_nhwc_signature = _permute_shape(signature_src, perm_nchw_to_nhwc)

            if add_out_nhwc_shape is None or add_out_nhwc_signature is None:
                continue

            model_ir.tensors[add_out_nhwc_name] = TensorIR(
                name=add_out_nhwc_name,
                dtype=add_out_dtype,
                shape=[int(v) for v in list(add_out_nhwc_shape)],
                shape_signature=[int(v) for v in list(add_out_nhwc_signature)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(add_out_quant),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=tail_add_op,
                new_outputs=[add_out_nhwc_name],
            )

            rank4_names = [
                pre_out_name,
                mean1_out_name,
                centered_name,
                squared_name,
                mean2_out_name,
                add_eps_out_name,
                sqrt_out_name,
                div_out_name,
                normalized_name,
                scaled_name,
                inst_out_name,
            ]
            for tensor_name in rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            if adapter_perm_name not in model_ir.tensors:
                model_ir.tensors[adapter_perm_name] = TensorIR(
                    name=adapter_perm_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
            insert_index = min(int(v) for v in list(add_out_users))
            model_ir.operators.insert(
                int(insert_index),
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(add_out_nhwc_name), str(adapter_perm_name)],
                    outputs=[str(add_out_name)],
                ),
            )

            remove_indices = {int(pre_idx), int(residual_pre_idx)}
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_instancenorm_residual_add_to_single_post_adapter_nhwc_chains": int(rewritten)}


def _optimize_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around decomposed InstanceNorm residual tails.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --(InstanceNorm decomposition in NCHW)--> y_nchw
      y_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc
      ADD(y_nhwc, r_nhwc) -> a_nhwc
      MUL(a_nhwc, c0_nchw[1,C,1,1]) -> m0_nchw
      MUL(a_nhwc, c1_nchw[1,C,1,1]) -> m1_nchw
      CONCAT(axis=1, [m0_nchw, m1_nchw]) -> cat_nchw
      cat_nchw --TRANSPOSE(0,2,3,1)--> cat_nhwc
      CONV_2D(cat_nhwc, ...)

    Rewrite:
      - Rewrite InstanceNorm decomposition to NHWC.
      - Rewrite channelwise constants to NHWC broadcasting.
      - CONCAT axis becomes 3 and directly produces `cat_nhwc`.
      - Remove three transposes (pre/post/tail post).
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _normalize_rank4_axes(axes: Optional[List[int]]) -> Optional[List[int]]:
        if axes is None or len(axes) == 0:
            return None
        normalized: List[int] = []
        for axis in axes:
            axis_value = int(axis)
            if axis_value < 0:
                axis_value += 4
            if axis_value < 0 or axis_value >= 4:
                return None
            normalized.append(int(axis_value))
        return sorted(list(set(normalized)))

    def _rewrite_axes_to_nhwc(
        *,
        axes_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        axes_users = set(int(v) for v in consumers.get(str(axes_name), []))
        if axes_users != expected_users:
            return False
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None:
            return False
        axes_values = _read_const_ints_from_tensor(axes_tensor)
        if _normalize_rank4_axes(axes_values) != [2, 3]:
            return False
        _write_const_ints_to_tensor(axes_tensor, [1, 2])
        return True

    def _rewrite_nchw_coeff_to_nhwc(
        *,
        coeff_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        coeff_users = set(int(v) for v in consumers.get(str(coeff_name), []))
        if coeff_users != expected_users:
            return False
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        if coeff_tensor is None or coeff_tensor.data is None:
            return False
        coeff = np.asarray(coeff_tensor.data)
        if int(coeff.size) == 1:
            return True
        if coeff.ndim != 4:
            return False
        shape = [int(v) for v in list(coeff.shape)]
        # Already NHWC channelwise.
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
        ):
            return True
        # NCHW channelwise -> NHWC channelwise.
        if (
            int(shape[0]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            coeff_nhwc = np.transpose(coeff, perm_nchw_to_nhwc).astype(coeff.dtype, copy=False)
            coeff_tensor.data = np.asarray(coeff_nhwc)
            coeff_tensor.shape = [int(v) for v in list(coeff_nhwc.shape)]
            coeff_tensor.shape_signature = [int(v) for v in list(coeff_nhwc.shape)]
            return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            pre_users = set(int(v) for v in consumers.get(pre_out_name, []))
            if len(pre_users) != 2:
                continue

            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None
            for user_idx in sorted(list(pre_users)):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_out_name
                ):
                    mean1_idx = int(user_idx)
                elif (
                    str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)
            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name in model_outputs:
                continue
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue
            if set(int(v) for v in consumers.get(mean1_out_name, [])) != {int(sub_idx)}:
                continue

            centered_name = str(sub_op.outputs[0])
            if centered_name in model_outputs:
                continue
            centered_users = set(int(v) for v in consumers.get(centered_name, []))
            if len(centered_users) != 2:
                continue

            mul_square_idx: Optional[int] = None
            mul_norm_idx: Optional[int] = None
            for user_idx in sorted(list(centered_users)):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    continue
                in0 = str(user_op.inputs[0])
                in1 = str(user_op.inputs[1])
                if in0 == centered_name and in1 == centered_name:
                    mul_square_idx = int(user_idx)
                elif centered_name in {in0, in1}:
                    mul_norm_idx = int(user_idx)
            if mul_square_idx is None or mul_norm_idx is None:
                continue

            mul_square_op = model_ir.operators[int(mul_square_idx)]
            squared_name = str(mul_square_op.outputs[0])
            mean2_users = set(int(v) for v in consumers.get(squared_name, []))
            if len(mean2_users) != 1:
                continue
            mean2_idx = int(list(mean2_users)[0])
            mean2_op = model_ir.operators[int(mean2_idx)]
            if (
                str(mean2_op.op_type) != "MEAN"
                or len(mean2_op.inputs) != 2
                or len(mean2_op.outputs) != 1
                or str(mean2_op.inputs[0]) != squared_name
            ):
                continue
            mean2_out_name = str(mean2_op.outputs[0])

            add_eps_users = set(int(v) for v in consumers.get(mean2_out_name, []))
            if len(add_eps_users) != 1:
                continue
            add_eps_idx = int(list(add_eps_users)[0])
            add_eps_op = model_ir.operators[int(add_eps_idx)]
            if (
                str(add_eps_op.op_type) != "ADD"
                or len(add_eps_op.inputs) != 2
                or len(add_eps_op.outputs) != 1
                or mean2_out_name not in {str(add_eps_op.inputs[0]), str(add_eps_op.inputs[1])}
            ):
                continue
            add_eps_out_name = str(add_eps_op.outputs[0])

            sqrt_users = set(int(v) for v in consumers.get(add_eps_out_name, []))
            if len(sqrt_users) != 1:
                continue
            sqrt_idx = int(list(sqrt_users)[0])
            sqrt_op = model_ir.operators[int(sqrt_idx)]
            if (
                str(sqrt_op.op_type) != "SQRT"
                or len(sqrt_op.inputs) != 1
                or len(sqrt_op.outputs) != 1
                or str(sqrt_op.inputs[0]) != add_eps_out_name
            ):
                continue
            sqrt_out_name = str(sqrt_op.outputs[0])

            div_users = set(int(v) for v in consumers.get(sqrt_out_name, []))
            if len(div_users) != 1:
                continue
            div_idx = int(list(div_users)[0])
            div_op = model_ir.operators[int(div_idx)]
            if (
                str(div_op.op_type) != "DIV"
                or len(div_op.inputs) != 2
                or len(div_op.outputs) != 1
                or str(div_op.inputs[1]) != sqrt_out_name
            ):
                continue
            div_out_name = str(div_op.outputs[0])

            mul_norm_op = model_ir.operators[int(mul_norm_idx)]
            if div_out_name not in {str(mul_norm_op.inputs[0]), str(mul_norm_op.inputs[1])}:
                continue
            normalized_name = str(mul_norm_op.outputs[0])

            mul_scale_users = set(int(v) for v in consumers.get(normalized_name, []))
            if len(mul_scale_users) != 1:
                continue
            mul_scale_idx = int(list(mul_scale_users)[0])
            mul_scale_op = model_ir.operators[int(mul_scale_idx)]
            if (
                str(mul_scale_op.op_type) != "MUL"
                or len(mul_scale_op.inputs) != 2
                or len(mul_scale_op.outputs) != 1
                or normalized_name not in {str(mul_scale_op.inputs[0]), str(mul_scale_op.inputs[1])}
            ):
                continue
            scale_const_name = (
                str(mul_scale_op.inputs[0])
                if str(mul_scale_op.inputs[1]) == normalized_name
                else str(mul_scale_op.inputs[1])
            )
            scaled_name = str(mul_scale_op.outputs[0])

            add_bias_users = set(int(v) for v in consumers.get(scaled_name, []))
            if len(add_bias_users) != 1:
                continue
            add_bias_idx = int(list(add_bias_users)[0])
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if (
                str(add_bias_op.op_type) != "ADD"
                or len(add_bias_op.inputs) != 2
                or len(add_bias_op.outputs) != 1
                or scaled_name not in {str(add_bias_op.inputs[0]), str(add_bias_op.inputs[1])}
            ):
                continue
            bias_const_name = (
                str(add_bias_op.inputs[0])
                if str(add_bias_op.inputs[1]) == scaled_name
                else str(add_bias_op.inputs[1])
            )
            inst_out_name = str(add_bias_op.outputs[0])
            if inst_out_name in model_outputs:
                continue

            inst_out_users = set(int(v) for v in consumers.get(inst_out_name, []))
            if len(inst_out_users) != 1:
                continue
            post_idx = int(list(inst_out_users)[0])
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                or str(post_op.inputs[0]) != inst_out_name
            ):
                continue
            post_out_name = str(post_op.outputs[0])
            if post_out_name in model_outputs:
                continue

            post_users = set(int(v) for v in consumers.get(post_out_name, []))
            if len(post_users) != 1:
                continue
            tail_add_idx = int(list(post_users)[0])
            tail_add_op = model_ir.operators[int(tail_add_idx)]
            if (
                str(tail_add_op.op_type) != "ADD"
                or len(tail_add_op.inputs) != 2
                or len(tail_add_op.outputs) != 1
                or post_out_name not in {str(tail_add_op.inputs[0]), str(tail_add_op.inputs[1])}
            ):
                continue
            add_out_name = str(tail_add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            add_out_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_out_users) != 2:
                continue

            tail_mul_indices: List[int] = []
            tail_mul_out_names: List[str] = []
            tail_mul_const_names: List[str] = []
            concat_idx: Optional[int] = None
            valid_tail = True
            for user_idx in sorted(add_out_users):
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) != "MUL" or len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                    valid_tail = False
                    break
                if add_out_name not in {str(user_op.inputs[0]), str(user_op.inputs[1])}:
                    valid_tail = False
                    break
                const_name = (
                    str(user_op.inputs[0])
                    if str(user_op.inputs[1]) == add_out_name
                    else str(user_op.inputs[1])
                )
                mul_out_name = str(user_op.outputs[0])
                if mul_out_name in model_outputs:
                    valid_tail = False
                    break
                mul_out_users = set(int(v) for v in consumers.get(mul_out_name, []))
                if len(mul_out_users) != 1:
                    valid_tail = False
                    break
                current_concat_idx = int(list(mul_out_users)[0])
                current_concat_op = model_ir.operators[int(current_concat_idx)]
                if str(current_concat_op.op_type) != "CONCATENATION" or len(current_concat_op.outputs) != 1:
                    valid_tail = False
                    break
                if concat_idx is None:
                    concat_idx = int(current_concat_idx)
                elif int(concat_idx) != int(current_concat_idx):
                    valid_tail = False
                    break
                tail_mul_indices.append(int(user_idx))
                tail_mul_out_names.append(str(mul_out_name))
                tail_mul_const_names.append(str(const_name))
            if not valid_tail or concat_idx is None:
                continue

            concat_op = model_ir.operators[int(concat_idx)]
            concat_axis = int(concat_op.options.get("axis", 1)) if isinstance(concat_op.options, dict) else 1
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue
            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if set(concat_inputs) != set(tail_mul_out_names):
                continue
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue

            concat_users = set(int(v) for v in consumers.get(concat_out_name, []))
            if len(concat_users) != 1:
                continue
            tail_post_idx = int(list(concat_users)[0])
            tail_post_op = model_ir.operators[int(tail_post_idx)]
            if (
                str(tail_post_op.op_type) != "TRANSPOSE"
                or len(tail_post_op.inputs) < 2
                or len(tail_post_op.outputs) != 1
                or _read_transpose_perm(model_ir, tail_post_op) != perm_nchw_to_nhwc
                or str(tail_post_op.inputs[0]) != concat_out_name
            ):
                continue
            tail_post_out_name = str(tail_post_op.outputs[0])

            mean1_axes_name = str(mean1_op.inputs[1])
            mean2_axes_name = str(mean2_op.inputs[1])
            if mean1_axes_name == mean2_axes_name:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx), int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue
            else:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean1_axes_name,
                    expected_users={int(mean1_idx)},
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean2_axes_name,
                    expected_users={int(mean2_idx)},
                    consumers=consumers,
                ):
                    continue

            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=scale_const_name,
                expected_users={int(mul_scale_idx)},
                consumers=consumers,
            ):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=bias_const_name,
                expected_users={int(add_bias_idx)},
                consumers=consumers,
            ):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=tail_mul_const_names[0],
                expected_users={int(tail_mul_indices[0])},
                consumers=consumers,
            ):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=tail_mul_const_names[1],
                expected_users={int(tail_mul_indices[1])},
                consumers=consumers,
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean1_op,
                input_index=0,
                new_input_name=pre_in_name,
            )
            sub_inputs = [str(v) for v in list(sub_op.inputs)]
            if sub_inputs[0] == pre_out_name:
                sub_inputs[0] = pre_in_name
            if sub_inputs[1] == pre_out_name:
                sub_inputs[1] = pre_in_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=sub_op,
                new_inputs=sub_inputs,
            )

            tail_add_inputs = [str(v) for v in list(tail_add_op.inputs)]
            for input_idx, input_name in enumerate(list(tail_add_inputs)):
                if str(input_name) == post_out_name:
                    tail_add_inputs[int(input_idx)] = str(inst_out_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=tail_add_op,
                new_inputs=tail_add_inputs,
            )

            rank4_names = [
                pre_out_name,
                mean1_out_name,
                centered_name,
                squared_name,
                mean2_out_name,
                add_eps_out_name,
                sqrt_out_name,
                div_out_name,
                normalized_name,
                scaled_name,
                inst_out_name,
                tail_mul_out_names[0],
                tail_mul_out_names[1],
                concat_out_name,
            ]
            for tensor_name in rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            concat_opts = dict(concat_op.options) if isinstance(concat_op.options, dict) else {}
            concat_opts["axis"] = 3
            concat_op.options = concat_opts
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[str(tail_post_out_name)],
            )

            for remove_idx in sorted(
                [int(pre_idx), int(post_idx), int(tail_post_idx)],
                reverse=True,
            ):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains": int(rewritten)}


def _optimize_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around dual-stat InstanceNorm-like residual blocks.

    Target (strict topology):
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw --(dual mean/sub/square/mean/sqrt/div branches)--> b0_nchw, b1_nchw
      ADD(b0_nchw, b1_nchw) -> m_nchw
      MUL(m_nchw, gamma_11c11_or_reshape) -> g_nchw
      ADD(g_nchw, beta_11c11_or_reshape) -> y_nchw
      r_nhwc --T(0,3,1,2)--> r_nchw
      ADD(y_nchw, r_nchw) -> a_nchw
      a_nchw --T(0,2,3,1)--> a_nhwc

    Rewrite:
      - Rewrite dual-stat branch tensors to NHWC (axes [2,3] -> [1,2]).
      - Rewrite channelwise coefficients to NHWC broadcast.
      - Bypass channelwise reshapes for gamma/beta when possible.
      - Remove the three transpose bridges.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _normalize_rank4_axes(axes: Optional[List[int]]) -> Optional[List[int]]:
        if axes is None or len(axes) == 0:
            return None
        normalized: List[int] = []
        for axis in axes:
            axis_value = int(axis)
            if axis_value < 0:
                axis_value += 4
            if axis_value < 0 or axis_value >= 4:
                return None
            normalized.append(int(axis_value))
        return sorted(list(set(normalized)))

    def _rewrite_axes_to_nhwc(
        *,
        axes_name: str,
        expected_users: set[int],
        target_axes: List[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        axes_users = set(int(v) for v in consumers.get(str(axes_name), []))
        if axes_users != expected_users:
            return False
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None:
            return False
        _write_const_ints_to_tensor(axes_tensor, [int(v) for v in list(target_axes)])
        return True

    def _axes_match(
        *,
        axes_name: str,
        expected_axes: List[int],
    ) -> bool:
        axes_tensor = model_ir.tensors.get(str(axes_name), None)
        if axes_tensor is None:
            return False
        axes_values = _read_const_ints_from_tensor(axes_tensor)
        return _normalize_rank4_axes(axes_values) == sorted([int(v) for v in list(expected_axes)])

    def _rewrite_nchw_coeff_to_nhwc(
        *,
        coeff_name: str,
        expected_users: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        coeff_users = set(int(v) for v in consumers.get(str(coeff_name), []))
        if coeff_users != expected_users:
            return False
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        if coeff_tensor is None or coeff_tensor.data is None:
            return False
        coeff = np.asarray(coeff_tensor.data)
        if int(coeff.size) == 1:
            return True
        if coeff.ndim != 4:
            return False
        shape = [int(v) for v in list(coeff.shape)]
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
        ):
            return True
        if (
            int(shape[0]) == 1
            and int(shape[1]) > 0
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            coeff_nhwc = np.transpose(coeff, perm_nchw_to_nhwc).astype(coeff.dtype, copy=False)
            coeff_tensor.data = np.asarray(coeff_nhwc)
            coeff_tensor.shape = [int(v) for v in list(coeff_nhwc.shape)]
            coeff_tensor.shape_signature = [int(v) for v in list(coeff_nhwc.shape)]
            return True
        return False

    def _bypass_channelwise_reshape_to_vector(
        *,
        coeff_name: str,
        expected_user: int,
        consumers: Dict[str, List[int]],
        producers: Dict[str, int],
    ) -> Optional[str]:
        coeff_users = set(int(v) for v in consumers.get(str(coeff_name), []))
        if coeff_users != {int(expected_user)}:
            return None
        coeff_prod_idx = producers.get(str(coeff_name), None)
        if coeff_prod_idx is None:
            return None
        coeff_prod_op = model_ir.operators[int(coeff_prod_idx)]
        if (
            str(coeff_prod_op.op_type) != "RESHAPE"
            or len(coeff_prod_op.inputs) != 2
            or len(coeff_prod_op.outputs) != 1
            or str(coeff_prod_op.outputs[0]) != str(coeff_name)
        ):
            return None
        source_name = str(coeff_prod_op.inputs[0])
        coeff_tensor = model_ir.tensors.get(str(coeff_name), None)
        source_tensor = model_ir.tensors.get(str(source_name), None)
        if coeff_tensor is None or source_tensor is None:
            return None
        coeff_shape = [int(v) for v in list(coeff_tensor.shape)]
        source_shape = [int(v) for v in list(source_tensor.shape)]
        if len(coeff_shape) != 4:
            return None
        if not (
            int(coeff_shape[0]) == 1
            and int(coeff_shape[1]) > 0
            and int(coeff_shape[2]) == 1
            and int(coeff_shape[3]) == 1
        ):
            return None
        if len(source_shape) == 2 and int(source_shape[0]) == 1 and int(source_shape[1]) == int(coeff_shape[1]):
            return str(source_name)
        if len(source_shape) == 1 and int(source_shape[0]) == int(coeff_shape[1]):
            return str(source_name)
        return None

    def _match_dual_norm_path(
        *,
        sub_idx: int,
        expected_axes: List[int],
        consumers: Dict[str, List[int]],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        sub_op = model_ir.operators[int(sub_idx)]
        sub_out_name = str(sub_op.outputs[0])
        if sub_out_name in model_outputs:
            return None
        sub_users = set(int(v) for v in consumers.get(sub_out_name, []))
        if len(sub_users) != 2:
            return None

        square_idx: Optional[int] = None
        div_idx: Optional[int] = None
        for user_idx in sorted(list(sub_users)):
            user_op = model_ir.operators[int(user_idx)]
            if (
                str(user_op.op_type) == "MUL"
                and len(user_op.inputs) == 2
                and len(user_op.outputs) == 1
                and str(user_op.inputs[0]) == sub_out_name
                and str(user_op.inputs[1]) == sub_out_name
            ):
                square_idx = int(user_idx)
            elif (
                str(user_op.op_type) == "DIV"
                and len(user_op.inputs) == 2
                and len(user_op.outputs) == 1
                and str(user_op.inputs[0]) == sub_out_name
            ):
                div_idx = int(user_idx)
        if square_idx is None or div_idx is None:
            return None

        square_op = model_ir.operators[int(square_idx)]
        square_out_name = str(square_op.outputs[0])
        square_users = set(int(v) for v in consumers.get(square_out_name, []))
        if len(square_users) != 1:
            return None
        mean2_idx = int(list(square_users)[0])
        mean2_op = model_ir.operators[int(mean2_idx)]
        if (
            str(mean2_op.op_type) != "MEAN"
            or len(mean2_op.inputs) != 2
            or len(mean2_op.outputs) != 1
            or str(mean2_op.inputs[0]) != square_out_name
        ):
            return None
        mean2_axes_name = str(mean2_op.inputs[1])
        if not _axes_match(axes_name=mean2_axes_name, expected_axes=expected_axes):
            return None
        mean2_out_name = str(mean2_op.outputs[0])

        mul_factor_users = set(int(v) for v in consumers.get(mean2_out_name, []))
        if len(mul_factor_users) != 1:
            return None
        mul_factor_idx = int(list(mul_factor_users)[0])
        mul_factor_op = model_ir.operators[int(mul_factor_idx)]
        if (
            str(mul_factor_op.op_type) != "MUL"
            or len(mul_factor_op.inputs) != 2
            or len(mul_factor_op.outputs) != 1
            or mean2_out_name not in {str(mul_factor_op.inputs[0]), str(mul_factor_op.inputs[1])}
        ):
            return None
        mul_factor_out_name = str(mul_factor_op.outputs[0])

        add_eps_users = set(int(v) for v in consumers.get(mul_factor_out_name, []))
        if len(add_eps_users) != 1:
            return None
        add_eps_idx = int(list(add_eps_users)[0])
        add_eps_op = model_ir.operators[int(add_eps_idx)]
        if (
            str(add_eps_op.op_type) != "ADD"
            or len(add_eps_op.inputs) != 2
            or len(add_eps_op.outputs) != 1
            or mul_factor_out_name not in {str(add_eps_op.inputs[0]), str(add_eps_op.inputs[1])}
        ):
            return None
        add_eps_out_name = str(add_eps_op.outputs[0])

        sqrt_users = set(int(v) for v in consumers.get(add_eps_out_name, []))
        if len(sqrt_users) != 1:
            return None
        sqrt_idx = int(list(sqrt_users)[0])
        sqrt_op = model_ir.operators[int(sqrt_idx)]
        if (
            str(sqrt_op.op_type) != "SQRT"
            or len(sqrt_op.inputs) != 1
            or len(sqrt_op.outputs) != 1
            or str(sqrt_op.inputs[0]) != add_eps_out_name
        ):
            return None
        sqrt_out_name = str(sqrt_op.outputs[0])

        div_op = model_ir.operators[int(div_idx)]
        if str(div_op.inputs[1]) != sqrt_out_name:
            return None
        div_out_name = str(div_op.outputs[0])
        div_out_users = set(int(v) for v in consumers.get(div_out_name, []))
        if len(div_out_users) != 1:
            return None
        scale_idx = int(list(div_out_users)[0])
        scale_op = model_ir.operators[int(scale_idx)]
        if (
            str(scale_op.op_type) != "MUL"
            or len(scale_op.inputs) != 2
            or len(scale_op.outputs) != 1
            or div_out_name not in {str(scale_op.inputs[0]), str(scale_op.inputs[1])}
        ):
            return None
        scale_const_name = (
            str(scale_op.inputs[0])
            if str(scale_op.inputs[1]) == div_out_name
            else str(scale_op.inputs[1])
        )
        scaled_out_name = str(scale_op.outputs[0])
        if scaled_out_name in model_outputs:
            return None

        return {
            "sub_idx": int(sub_idx),
            "sub_out_name": str(sub_out_name),
            "square_idx": int(square_idx),
            "square_out_name": str(square_out_name),
            "mean2_idx": int(mean2_idx),
            "mean2_axes_name": str(mean2_axes_name),
            "mean2_out_name": str(mean2_out_name),
            "mul_factor_idx": int(mul_factor_idx),
            "mul_factor_out_name": str(mul_factor_out_name),
            "add_eps_idx": int(add_eps_idx),
            "add_eps_out_name": str(add_eps_out_name),
            "sqrt_idx": int(sqrt_idx),
            "sqrt_out_name": str(sqrt_out_name),
            "div_idx": int(div_idx),
            "div_out_name": str(div_out_name),
            "scale_idx": int(scale_idx),
            "scale_const_name": str(scale_const_name),
            "scaled_out_name": str(scaled_out_name),
        }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            pre_users = sorted(set(int(v) for v in consumers.get(pre_out_name, [])))
            if len(pre_users) != 4:
                continue

            mean_spatial_idx: Optional[int] = None
            mean_global_idx: Optional[int] = None
            mean_spatial_axes_name: Optional[str] = None
            mean_global_axes_name: Optional[str] = None
            sub_candidates: List[int] = []
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_out_name
                ):
                    axes_name = str(user_op.inputs[1])
                    if _axes_match(axes_name=axes_name, expected_axes=[2, 3]):
                        mean_spatial_idx = int(user_idx)
                        mean_spatial_axes_name = str(axes_name)
                    elif _axes_match(axes_name=axes_name, expected_axes=[1, 2, 3]):
                        mean_global_idx = int(user_idx)
                        mean_global_axes_name = str(axes_name)
                elif (
                    str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_candidates.append(int(user_idx))
            if (
                mean_spatial_idx is None
                or mean_global_idx is None
                or len(sub_candidates) != 2
                or mean_spatial_axes_name is None
                or mean_global_axes_name is None
            ):
                continue

            mean_spatial_out_name = str(model_ir.operators[int(mean_spatial_idx)].outputs[0])
            mean_global_out_name = str(model_ir.operators[int(mean_global_idx)].outputs[0])
            if (
                mean_spatial_out_name in model_outputs
                or mean_global_out_name in model_outputs
                or set(int(v) for v in consumers.get(mean_spatial_out_name, [])) == set()
                or set(int(v) for v in consumers.get(mean_global_out_name, [])) == set()
            ):
                continue

            sub_spatial_idx: Optional[int] = None
            sub_global_idx: Optional[int] = None
            for sub_idx in sub_candidates:
                sub_op = model_ir.operators[int(sub_idx)]
                sub_inputs = set(str(v) for v in list(sub_op.inputs))
                if mean_spatial_out_name in sub_inputs:
                    sub_spatial_idx = int(sub_idx)
                elif mean_global_out_name in sub_inputs:
                    sub_global_idx = int(sub_idx)
            if sub_spatial_idx is None or sub_global_idx is None:
                continue
            if set(int(v) for v in consumers.get(mean_spatial_out_name, [])) != {int(sub_spatial_idx)}:
                continue
            if set(int(v) for v in consumers.get(mean_global_out_name, [])) != {int(sub_global_idx)}:
                continue

            path_spatial = _match_dual_norm_path(
                sub_idx=int(sub_spatial_idx),
                expected_axes=[2, 3],
                consumers=consumers,
                model_outputs=model_outputs,
            )
            if path_spatial is None:
                continue
            path_global = _match_dual_norm_path(
                sub_idx=int(sub_global_idx),
                expected_axes=[1, 2, 3],
                consumers=consumers,
                model_outputs=model_outputs,
            )
            if path_global is None:
                continue

            scaled0_name = str(path_spatial["scaled_out_name"])
            scaled1_name = str(path_global["scaled_out_name"])
            scaled0_users = set(int(v) for v in consumers.get(scaled0_name, []))
            scaled1_users = set(int(v) for v in consumers.get(scaled1_name, []))
            if len(scaled0_users) != 1 or len(scaled1_users) != 1 or scaled0_users != scaled1_users:
                continue
            blend_add_idx = int(list(scaled0_users)[0])
            blend_add_op = model_ir.operators[int(blend_add_idx)]
            if (
                str(blend_add_op.op_type) != "ADD"
                or len(blend_add_op.inputs) != 2
                or len(blend_add_op.outputs) != 1
            ):
                continue
            blend_inputs = set(str(v) for v in list(blend_add_op.inputs))
            if blend_inputs != {scaled0_name, scaled1_name}:
                continue
            blend_add_out_name = str(blend_add_op.outputs[0])
            if blend_add_out_name in model_outputs:
                continue

            blend_mul_users = set(int(v) for v in consumers.get(blend_add_out_name, []))
            if len(blend_mul_users) != 1:
                continue
            blend_mul_idx = int(list(blend_mul_users)[0])
            blend_mul_op = model_ir.operators[int(blend_mul_idx)]
            if (
                str(blend_mul_op.op_type) != "MUL"
                or len(blend_mul_op.inputs) != 2
                or len(blend_mul_op.outputs) != 1
                or blend_add_out_name not in {str(blend_mul_op.inputs[0]), str(blend_mul_op.inputs[1])}
            ):
                continue
            gamma_input_index = 0 if str(blend_mul_op.inputs[1]) == blend_add_out_name else 1
            gamma_name = str(blend_mul_op.inputs[int(gamma_input_index)])
            blend_mul_out_name = str(blend_mul_op.outputs[0])
            if blend_mul_out_name in model_outputs:
                continue

            blend_bias_users = set(int(v) for v in consumers.get(blend_mul_out_name, []))
            if len(blend_bias_users) != 1:
                continue
            blend_bias_idx = int(list(blend_bias_users)[0])
            blend_bias_op = model_ir.operators[int(blend_bias_idx)]
            if (
                str(blend_bias_op.op_type) != "ADD"
                or len(blend_bias_op.inputs) != 2
                or len(blend_bias_op.outputs) != 1
                or blend_mul_out_name not in {str(blend_bias_op.inputs[0]), str(blend_bias_op.inputs[1])}
            ):
                continue
            beta_input_index = 0 if str(blend_bias_op.inputs[1]) == blend_mul_out_name else 1
            beta_name = str(blend_bias_op.inputs[int(beta_input_index)])
            inst_out_name = str(blend_bias_op.outputs[0])
            if inst_out_name in model_outputs:
                continue

            inst_out_users = set(int(v) for v in consumers.get(inst_out_name, []))
            if len(inst_out_users) != 1:
                continue
            inst_out_user_idx = int(list(inst_out_users)[0])
            inst_out_user_op = model_ir.operators[int(inst_out_user_idx)]

            use_residual_tail = False
            tail_add_idx: Optional[int] = None
            tail_add_op: Optional[OperatorIR] = None
            post_idx: Optional[int] = None
            post_out_name: Optional[str] = None
            residual_pre_idx: Optional[int] = None
            residual_src_name: Optional[str] = None

            if (
                str(inst_out_user_op.op_type) == "TRANSPOSE"
                and len(inst_out_user_op.inputs) >= 2
                and len(inst_out_user_op.outputs) == 1
                and str(inst_out_user_op.inputs[0]) == inst_out_name
                and _read_transpose_perm(model_ir, inst_out_user_op) == perm_nchw_to_nhwc
            ):
                post_idx = int(inst_out_user_idx)
                post_out_name = str(inst_out_user_op.outputs[0])
            elif (
                str(inst_out_user_op.op_type) == "ADD"
                and len(inst_out_user_op.inputs) == 2
                and len(inst_out_user_op.outputs) == 1
                and inst_out_name in {str(inst_out_user_op.inputs[0]), str(inst_out_user_op.inputs[1])}
            ):
                tail_add_idx = int(inst_out_user_idx)
                tail_add_op = inst_out_user_op
                tail_add_out_name = str(tail_add_op.outputs[0])
                if tail_add_out_name in model_outputs:
                    continue
                tail_add_users2 = set(int(v) for v in consumers.get(tail_add_out_name, []))
                if len(tail_add_users2) != 1:
                    continue
                post_idx = int(list(tail_add_users2)[0])
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != tail_add_out_name
                    or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                ):
                    continue
                post_out_name = str(post_op.outputs[0])
                residual_input_name = (
                    str(tail_add_op.inputs[0])
                    if str(tail_add_op.inputs[1]) == inst_out_name
                    else str(tail_add_op.inputs[1])
                )
                residual_pre_idx = producers.get(str(residual_input_name), None)
                if residual_pre_idx is None:
                    continue
                residual_pre_op = model_ir.operators[int(residual_pre_idx)]
                if (
                    str(residual_pre_op.op_type) != "TRANSPOSE"
                    or len(residual_pre_op.inputs) < 2
                    or len(residual_pre_op.outputs) != 1
                    or str(residual_pre_op.outputs[0]) != residual_input_name
                    or _read_transpose_perm(model_ir, residual_pre_op) != perm_nhwc_to_nchw
                ):
                    continue
                if set(int(v) for v in consumers.get(residual_input_name, [])) != {int(tail_add_idx)}:
                    continue
                residual_src_name = str(residual_pre_op.inputs[0])
                if residual_src_name in model_outputs:
                    continue
                use_residual_tail = True
            else:
                continue

            if post_idx is None or post_out_name is None:
                continue
            if post_out_name in model_outputs:
                continue

            mean2_spatial_axes_name = str(path_spatial["mean2_axes_name"])
            if mean_spatial_axes_name == mean2_spatial_axes_name:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean_spatial_axes_name,
                    expected_users={int(mean_spatial_idx), int(path_spatial["mean2_idx"])},
                    target_axes=[1, 2],
                    consumers=consumers,
                ):
                    continue
            else:
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean_spatial_axes_name,
                    expected_users={int(mean_spatial_idx)},
                    target_axes=[1, 2],
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_axes_to_nhwc(
                    axes_name=mean2_spatial_axes_name,
                    expected_users={int(path_spatial["mean2_idx"])},
                    target_axes=[1, 2],
                    consumers=consumers,
                ):
                    continue

            if not _axes_match(axes_name=mean_global_axes_name, expected_axes=[1, 2, 3]):
                continue
            if not _axes_match(axes_name=str(path_global["mean2_axes_name"]), expected_axes=[1, 2, 3]):
                continue

            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=str(path_spatial["scale_const_name"]),
                expected_users={int(path_spatial["scale_idx"])},
                consumers=consumers,
            ):
                continue
            if not _rewrite_nchw_coeff_to_nhwc(
                coeff_name=str(path_global["scale_const_name"]),
                expected_users={int(path_global["scale_idx"])},
                consumers=consumers,
            ):
                continue

            gamma_bypass_name = _bypass_channelwise_reshape_to_vector(
                coeff_name=str(gamma_name),
                expected_user=int(blend_mul_idx),
                consumers=consumers,
                producers=producers,
            )
            if gamma_bypass_name is not None:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=blend_mul_op,
                    input_index=int(gamma_input_index),
                    new_input_name=str(gamma_bypass_name),
                )
            else:
                if not _rewrite_nchw_coeff_to_nhwc(
                    coeff_name=str(gamma_name),
                    expected_users={int(blend_mul_idx)},
                    consumers=consumers,
                ):
                    continue

            beta_bypass_name = _bypass_channelwise_reshape_to_vector(
                coeff_name=str(beta_name),
                expected_user=int(blend_bias_idx),
                consumers=consumers,
                producers=producers,
            )
            if beta_bypass_name is not None:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=blend_bias_op,
                    input_index=int(beta_input_index),
                    new_input_name=str(beta_bypass_name),
                )
            else:
                if not _rewrite_nchw_coeff_to_nhwc(
                    coeff_name=str(beta_name),
                    expected_users={int(blend_bias_idx)},
                    consumers=consumers,
                ):
                    continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=model_ir.operators[int(mean_spatial_idx)],
                input_index=0,
                new_input_name=str(pre_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=model_ir.operators[int(mean_global_idx)],
                input_index=0,
                new_input_name=str(pre_in_name),
            )
            for sub_idx in [int(sub_spatial_idx), int(sub_global_idx)]:
                sub_op = model_ir.operators[int(sub_idx)]
                sub_inputs = [str(v) for v in list(sub_op.inputs)]
                if sub_inputs[0] == pre_out_name:
                    sub_inputs[0] = pre_in_name
                if sub_inputs[1] == pre_out_name:
                    sub_inputs[1] = pre_in_name
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=sub_op,
                    new_inputs=sub_inputs,
                )

            if use_residual_tail:
                if tail_add_op is None or residual_src_name is None:
                    continue
                tail_add_inputs = [str(v) for v in list(tail_add_op.inputs)]
                for input_idx, input_name in enumerate(list(tail_add_inputs)):
                    if str(input_name) != inst_out_name:
                        tail_add_inputs[int(input_idx)] = str(residual_src_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=tail_add_op,
                    new_inputs=tail_add_inputs,
                )

            rank4_names = [
                pre_out_name,
                mean_spatial_out_name,
                mean_global_out_name,
                str(path_spatial["sub_out_name"]),
                str(path_global["sub_out_name"]),
                str(path_spatial["square_out_name"]),
                str(path_global["square_out_name"]),
                str(path_spatial["mean2_out_name"]),
                str(path_global["mean2_out_name"]),
                str(path_spatial["mul_factor_out_name"]),
                str(path_global["mul_factor_out_name"]),
                str(path_spatial["add_eps_out_name"]),
                str(path_global["add_eps_out_name"]),
                str(path_spatial["sqrt_out_name"]),
                str(path_global["sqrt_out_name"]),
                str(path_spatial["div_out_name"]),
                str(path_global["div_out_name"]),
                str(path_spatial["scaled_out_name"]),
                str(path_global["scaled_out_name"]),
                blend_add_out_name,
                blend_mul_out_name,
                inst_out_name,
            ]
            for tensor_name in rank4_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            if use_residual_tail:
                if tail_add_op is None:
                    continue
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=tail_add_op,
                    new_outputs=[str(post_out_name)],
                )
            else:
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=blend_bias_op,
                    new_outputs=[str(post_out_name)],
                )

            remove_indices = {int(pre_idx), int(post_idx)}
            if use_residual_tail and residual_pre_idx is not None:
                remove_indices.add(int(residual_pre_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains": int(rewritten)}


def _optimize_fold_mul_add_mul_affine_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold strict affine tails:
      x -> MUL(c1) -> ADD(c2) -> MUL(c3)
    into:
      x -> MUL(c1*c3) -> ADD(c2*c3)

    This pass is intentionally conservative:
    - all side inputs must be constants,
    - intermediate tensors must be single-use along the chain.
    """
    optimized = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _extract_data_const(
        *,
        op: OperatorIR,
        data_producer_type: Optional[str],
        producers: Dict[str, int],
    ) -> Optional[Tuple[int, str, int, str]]:
        if len(op.inputs) != 2:
            return None
        in0 = str(op.inputs[0])
        in1 = str(op.inputs[1])
        t0 = model_ir.tensors.get(in0, None)
        t1 = model_ir.tensors.get(in1, None)
        c0 = t0 is not None and t0.data is not None
        c1 = t1 is not None and t1.data is not None
        if c0 == c1:
            return None
        if c0:
            const_idx, const_name = 0, in0
            data_idx, data_name = 1, in1
        else:
            const_idx, const_name = 1, in1
            data_idx, data_name = 0, in0
        if data_producer_type is not None:
            pidx = producers.get(str(data_name), None)
            if pidx is None:
                return None
            pop = model_ir.operators[int(pidx)]
            if str(pop.op_type) != str(data_producer_type):
                return None
            if len(pop.outputs) != 1 or str(pop.outputs[0]) != str(data_name):
                return None
        return int(data_idx), str(data_name), int(const_idx), str(const_name)

    def _assign_const_array(
        *,
        target_name: str,
        new_array: np.ndarray,
        chain_indices: set[int],
        consumers: Dict[str, List[int]],
    ) -> str:
        tensor = model_ir.tensors[str(target_name)]
        users = [int(v) for v in consumers.get(str(target_name), [])]
        shared_outside = any(int(v) not in chain_indices for v in users)
        if not shared_outside:
            tensor.data = np.asarray(new_array)
            tensor.shape = [int(v) for v in list(new_array.shape)]
            tensor.shape_signature = [int(v) for v in list(new_array.shape)]
            return str(target_name)

        cloned_name = _unique_tensor_name(f"{target_name}_folded")
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(tensor.dtype),
            shape=[int(v) for v in list(new_array.shape)],
            shape_signature=[int(v) for v in list(new_array.shape)],
            data=np.asarray(new_array),
            is_variable=False,
            quantization=_clone_quantization(tensor.quantization),
        )
        return str(cloned_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for mul2_idx, mul2_op in enumerate(model_ir.operators):
            if str(mul2_op.op_type) != "MUL" or len(mul2_op.inputs) != 2 or len(mul2_op.outputs) != 1:
                continue
            mul2_out_name = str(mul2_op.outputs[0])

            parsed_mul2 = _extract_data_const(
                op=mul2_op,
                data_producer_type="ADD",
                producers=producers,
            )
            if parsed_mul2 is None:
                continue
            _, add_out_name, mul2_const_idx, mul2_const_name = parsed_mul2
            add_idx = producers.get(str(add_out_name), None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            if set(int(v) for v in consumers.get(str(add_out_name), [])) != {int(mul2_idx)}:
                continue

            parsed_add = _extract_data_const(
                op=add_op,
                data_producer_type="MUL",
                producers=producers,
            )
            if parsed_add is None:
                continue
            _, mul1_out_name, add_const_idx, add_const_name = parsed_add
            mul1_idx = producers.get(str(mul1_out_name), None)
            if mul1_idx is None:
                continue
            mul1_op = model_ir.operators[int(mul1_idx)]
            if str(mul1_op.op_type) != "MUL" or len(mul1_op.inputs) != 2 or len(mul1_op.outputs) != 1:
                continue
            if set(int(v) for v in consumers.get(str(mul1_out_name), [])) != {int(add_idx)}:
                continue

            mul1_inputs = [str(v) for v in list(mul1_op.inputs)]
            t0 = model_ir.tensors.get(str(mul1_inputs[0]), None)
            t1 = model_ir.tensors.get(str(mul1_inputs[1]), None)
            c0 = t0 is not None and t0.data is not None
            c1 = t1 is not None and t1.data is not None
            if c0 == c1:
                continue
            mul1_const_idx = 0 if c0 else 1
            mul1_const_name = str(mul1_inputs[int(mul1_const_idx)])

            mul1_const_tensor = model_ir.tensors.get(str(mul1_const_name), None)
            add_const_tensor = model_ir.tensors.get(str(add_const_name), None)
            mul2_const_tensor = model_ir.tensors.get(str(mul2_const_name), None)
            if (
                mul1_const_tensor is None
                or add_const_tensor is None
                or mul2_const_tensor is None
                or mul1_const_tensor.data is None
                or add_const_tensor.data is None
                or mul2_const_tensor.data is None
            ):
                continue

            try:
                c1_arr = np.asarray(mul1_const_tensor.data)
                c2_arr = np.asarray(add_const_tensor.data)
                c3_arr = np.asarray(mul2_const_tensor.data)
                c1_new = np.asarray(c1_arr * c3_arr, dtype=c1_arr.dtype)
                c2_new = np.asarray(c2_arr * c3_arr, dtype=c2_arr.dtype)
            except Exception:
                continue

            chain_indices = {int(mul1_idx), int(add_idx), int(mul2_idx)}

            c1_target_name = _assign_const_array(
                target_name=str(mul1_const_name),
                new_array=np.asarray(c1_new),
                chain_indices=chain_indices,
                consumers=consumers,
            )
            c2_target_name = _assign_const_array(
                target_name=str(add_const_name),
                new_array=np.asarray(c2_new),
                chain_indices=chain_indices,
                consumers=consumers,
            )

            if c1_target_name != str(mul1_const_name):
                mul1_new_inputs = [str(v) for v in list(mul1_op.inputs)]
                mul1_new_inputs[int(mul1_const_idx)] = str(c1_target_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul1_op,
                    new_inputs=mul1_new_inputs,
                )
            if c2_target_name != str(add_const_name):
                add_new_inputs = [str(v) for v in list(add_op.inputs)]
                add_new_inputs[int(add_const_idx)] = str(c2_target_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_inputs=add_new_inputs,
                )

            add_out_tensor = model_ir.tensors.get(str(add_out_name), None)
            mul2_out_tensor = model_ir.tensors.get(str(mul2_out_name), None)
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[str(mul2_out_name)],
            )
            if add_out_tensor is not None and mul2_out_tensor is not None:
                mul2_out_tensor.dtype = str(add_out_tensor.dtype)
                mul2_out_tensor.quantization = _clone_quantization(add_out_tensor.quantization)
                mul2_out_tensor.shape = [int(v) for v in list(add_out_tensor.shape)]
                mul2_out_tensor.shape_signature = (
                    [int(v) for v in list(add_out_tensor.shape_signature)]
                    if add_out_tensor.shape_signature is not None
                    else [int(v) for v in list(add_out_tensor.shape)]
                )

            del model_ir.operators[int(mul2_idx)]
            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_fold_mul_add_mul_affine_chains": int(optimized)}


def _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around linear TRANSPOSE->MUL(const)->ADD(const)->TRANSPOSE chains.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --MUL(const)--> m_nchw --ADD(const)--> y_nchw
      y_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      x_nhwc --MUL(const_nhwc)--> m_nhwc --ADD(const_nhwc)--> y_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_binary_const_to_nhwc(
        *,
        op_idx: int,
        side_input_index: int,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False

        nhwc_data: Optional[np.ndarray] = None
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is not None:
            side_shape = [int(v) for v in list(side_data.shape)]
            is_nchw_channelwise = (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            )
            if is_nchw_channelwise:
                rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is None:
                    return False
                nhwc_data = np.asarray(rotated)
                if np.array_equal(nhwc_data, side_data):
                    return True
            elif _broadcast_static_shapes(target_shape, side_shape) is not None:
                # Already NHWC-compatible. Keep idempotent across repeated sweeps.
                return True
            else:
                rotated = np.asarray(side_data)
                for _ in range(3):
                    rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                    rotated_shape = [int(v) for v in list(rotated.shape)]
                    if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                        nhwc_data = np.asarray(rotated)
                        break
                if nhwc_data is None:
                    return False
        else:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 1
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                nhwc_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            else:
                return True

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)

        target_name = str(side_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{side_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
        else:
            side_tensor.data = np.asarray(nhwc_data)
            side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        _replace_operator_input_at(
            model_ir=model_ir,
            op=model_ir.operators[int(op_idx)],
            input_index=int(side_input_index),
            new_input_name=str(target_name),
        )
        return True

    def _can_rewrite_binary_const_to_nhwc(
        *,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is None:
            return True
        side_shape = [int(v) for v in list(side_data.shape)]
        is_nchw_channelwise = (
            len(side_shape) == 4
            and int(side_shape[0]) == 1
            and int(side_shape[1]) > 0
            and int(side_shape[2]) == 1
            and int(side_shape[3]) == 1
        )
        if is_nchw_channelwise:
            rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            return _broadcast_static_shapes(target_shape, rotated_shape) is not None
        if _broadcast_static_shapes(target_shape, side_shape) is not None:
            return True
        rotated = np.asarray(side_data)
        for _ in range(3):
            rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_outputs or pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue
            for mul_idx in pre_users:
                mul_op = model_ir.operators[int(mul_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    continue
                mul_in0 = str(mul_op.inputs[0])
                mul_in1 = str(mul_op.inputs[1])
                if mul_in0 == pre_output_name:
                    mul_side_input_index = 1
                    mul_side_name = mul_in1
                elif mul_in1 == pre_output_name:
                    mul_side_input_index = 0
                    mul_side_name = mul_in0
                else:
                    continue

                add_input_name = str(mul_op.outputs[0])
                if add_input_name in model_outputs:
                    continue
                add_users = [int(v) for v in consumers.get(add_input_name, [])]
                if len(add_users) != 1:
                    continue
                add_idx = int(add_users[0])
                add_op = model_ir.operators[add_idx]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    continue
                add_in0 = str(add_op.inputs[0])
                add_in1 = str(add_op.inputs[1])
                if add_in0 == add_input_name:
                    add_side_input_index = 1
                    add_side_name = add_in1
                elif add_in1 == add_input_name:
                    add_side_input_index = 0
                    add_side_name = add_in0
                else:
                    continue

                add_output_name = str(add_op.outputs[0])
                if add_output_name in model_outputs:
                    continue

                bridge_output_name = str(add_output_name)
                bridge_output_producer_idx = int(add_idx)
                bridge_output_producer_op = add_op
                prelu_idx: Optional[int] = None
                prelu_side_input_index: Optional[int] = None
                prelu_side_name: Optional[str] = None

                add_output_users = [int(v) for v in consumers.get(add_output_name, []) if int(v) != int(add_idx)]
                if len(add_output_users) == 0:
                    continue
                if False and len(add_output_users) == 1:
                    candidate_idx = int(add_output_users[0])
                    candidate_op = model_ir.operators[int(candidate_idx)]
                    if (
                        str(candidate_op.op_type) == "PRELU"
                        and len(candidate_op.inputs) == 2
                        and len(candidate_op.outputs) == 1
                        and str(candidate_op.inputs[0]) == add_output_name
                    ):
                        candidate_out_name = str(candidate_op.outputs[0])
                        if candidate_out_name in model_outputs:
                            continue
                        prelu_idx = int(candidate_idx)
                        prelu_side_input_index = 1
                        prelu_side_name = str(candidate_op.inputs[1])
                        bridge_output_name = str(candidate_out_name)
                        bridge_output_producer_idx = int(candidate_idx)
                        bridge_output_producer_op = candidate_op

                bridge_output_users = [
                    int(v)
                    for v in consumers.get(bridge_output_name, [])
                    if int(v) != int(bridge_output_producer_idx)
                ]
                if len(bridge_output_users) == 0:
                    continue

                post_indices: List[int] = []
                post_output_names: List[str] = []
                legacy_users: List[int] = []
                valid_posts = True
                for user_idx in bridge_output_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == bridge_output_name
                        and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        post_indices.append(int(user_idx))
                        post_output_names.append(str(user_op.outputs[0]))
                    else:
                        legacy_users.append(int(user_idx))
                if not valid_posts or len(post_indices) == 0:
                    continue
                if len(legacy_users) > 0:
                    continue

                chain_index_set = {int(mul_idx), int(add_idx)}
                if prelu_idx is not None:
                    chain_index_set.add(int(prelu_idx))
                target_nhwc_shape = (
                    list(model_ir.tensors[pre_input_name].shape)
                    if pre_input_name in model_ir.tensors
                    else None
                )
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(mul_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    continue
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(add_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    continue
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(mul_idx),
                    side_input_index=int(mul_side_input_index),
                    side_input_name=str(mul_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(add_idx),
                    side_input_index=int(add_side_input_index),
                    side_input_name=str(add_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    continue
                if prelu_idx is not None:
                    if prelu_side_input_index is None or prelu_side_name is None:
                        continue
                    if not _can_rewrite_binary_const_to_nhwc(
                        side_input_name=str(prelu_side_name),
                        target_broadcast_shape=target_nhwc_shape,
                    ):
                        continue
                    if not _rewrite_binary_const_to_nhwc(
                        op_idx=int(prelu_idx),
                        side_input_index=int(prelu_side_input_index),
                        side_input_name=str(prelu_side_name),
                        target_broadcast_shape=target_nhwc_shape,
                        chain_index_set=chain_index_set,
                        consumers=consumers,
                    ):
                        continue

                # Bypass only this branch input from pre_output to pre_input.
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                if mul_inputs[0] == pre_output_name:
                    mul_inputs[0] = pre_input_name
                elif mul_inputs[1] == pre_output_name:
                    mul_inputs[1] = pre_input_name
                else:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )

                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(add_input_name, None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(add_output_name, None),
                    perm_nchw_to_nhwc,
                )
                if prelu_idx is not None:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(bridge_output_name, None),
                        perm_nchw_to_nhwc,
                    )

                canonical_post_output_name = str(post_output_names[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=bridge_output_producer_op,
                    new_outputs=[canonical_post_output_name],
                )
                for alias_name in post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

                old_add_out_tensor = model_ir.tensors.get(bridge_output_name, None)
                canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
                if old_add_out_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_add_out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_add_out_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_add_out_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_add_out_tensor.shape_signature)]
                        if old_add_out_tensor.shape_signature is not None
                        else [int(v) for v in list(old_add_out_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        canonical_tensor,
                        perm_nchw_to_nhwc,
                    )

                post_remove_indices = [int(v) for v in post_indices]

                remove_indices = set(int(v) for v in post_remove_indices)
                # Drop pre transpose only when this rewrite consumed its last user.
                pre_remaining_users = [int(v) for v in pre_users if int(v) != int(mul_idx)]
                if len(pre_remaining_users) == 0:
                    remove_indices.add(int(pre_idx))
                for remove_idx in sorted(remove_indices, reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_mul_add_const_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_pad_mul_posttranspose_add_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate strict NHWC<->NCHW round-trips around:
      T(NHWC->NCHW) -> PAD|MIRROR_PAD -> MUL(const) -> T(NCHW->NHWC) -> ADD(const)

    Rewrite:
      PAD|MIRROR_PAD(NHWC pads) -> MUL(const_nhwc) -> ADD(const_nhwc)
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_rank4_mul_const_to_nhwc(
        *,
        mul_op: OperatorIR,
        mul_idx: int,
        const_input_index: int,
        const_input_name: str,
        target_nhwc_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        const_tensor = model_ir.tensors.get(str(const_input_name), None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        target_shape = (
            [int(v) for v in list(target_nhwc_shape)]
            if _is_fully_known_positive_shape(target_nhwc_shape)
            else None
        )

        rotated_data: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            as_is_shape = [int(v) for v in list(const_data.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, as_is_shape) is not None:
                return True
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is not None:
                if _broadcast_static_shapes(target_shape, rotated_shape) is None:
                    return False
            rotated_data = np.asarray(rotated_candidate)
        else:
            if target_shape is None:
                return False
            side_shape = [int(v) for v in list(const_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is None:
                return False
            return True

        if rotated_data is None:
            return False

        const_users = [int(v) for v in consumers.get(str(const_input_name), [])]
        shared_outside_chain = any(int(v) not in chain_index_set for v in const_users)

        target_name = str(const_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{const_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated_data.shape)],
                shape_signature=[int(v) for v in list(rotated_data.shape)],
                data=np.asarray(rotated_data),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
        else:
            const_tensor.data = np.asarray(rotated_data)
            const_tensor.shape = [int(v) for v in list(rotated_data.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated_data.shape)]

        if target_name != str(const_input_name):
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            mul_inputs[int(const_input_index)] = str(target_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        return True

    def _rewrite_pad_spec_to_nhwc(
        *,
        pad_op: OperatorIR,
        pad_idx: int,
        pads_input_name: str,
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        pads_tensor = model_ir.tensors.get(str(pads_input_name), None)
        if pads_tensor is None or pads_tensor.data is None:
            return False
        try:
            pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
        except Exception:
            return False
        if int(pads_pairs.size) != 8:
            return False
        pads_nhwc = np.asarray(
            [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
            dtype=pads_pairs.dtype,
        )

        pads_users = [int(v) for v in consumers.get(str(pads_input_name), [])]
        shared_outside_chain = any(int(v) not in chain_index_set for v in pads_users)
        target_name = str(pads_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{pads_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(pads_tensor.dtype),
                shape=[4, 2],
                shape_signature=[4, 2],
                data=np.asarray(pads_nhwc),
                is_variable=False,
                quantization=_clone_quantization(pads_tensor.quantization),
            )
        else:
            pads_tensor.data = np.asarray(pads_nhwc)
            pads_tensor.shape = [4, 2]
            pads_tensor.shape_signature = [4, 2]

        if target_name != str(pads_input_name):
            pad_inputs = [str(v) for v in list(pad_op.inputs)]
            pad_inputs[1] = str(target_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=pad_op,
                new_inputs=pad_inputs,
            )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            mul_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if mul_output_name in model_outputs or post_output_name in model_outputs:
                continue

            mul_idx = producers.get(str(mul_output_name), None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if (
                str(mul_op.op_type) != "MUL"
                or len(mul_op.inputs) != 2
                or len(mul_op.outputs) != 1
                or str(mul_op.outputs[0]) != str(mul_output_name)
            ):
                continue
            if set(int(v) for v in consumers.get(str(mul_output_name), [])) != {int(post_idx)}:
                continue

            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            data_input_index: Optional[int] = None
            const_input_index: Optional[int] = None
            data_input_name: Optional[str] = None
            const_input_name: Optional[str] = None
            for idx, input_name in enumerate(mul_inputs):
                input_producer_idx = producers.get(str(input_name), None)
                if input_producer_idx is not None:
                    input_producer_op = model_ir.operators[int(input_producer_idx)]
                    if (
                        str(input_producer_op.op_type) in {"PAD", "MIRROR_PAD"}
                        and len(input_producer_op.outputs) == 1
                        and str(input_producer_op.outputs[0]) == str(input_name)
                    ):
                        data_input_index = int(idx)
                        data_input_name = str(input_name)
                        continue
                candidate_const = model_ir.tensors.get(str(input_name), None)
                if candidate_const is not None and candidate_const.data is not None:
                    const_input_index = int(idx)
                    const_input_name = str(input_name)

            if (
                data_input_index is None
                or data_input_name is None
                or const_input_index is None
                or const_input_name is None
            ):
                continue

            pad_idx = producers.get(str(data_input_name), None)
            if pad_idx is None:
                continue
            pad_op = model_ir.operators[int(pad_idx)]
            if (
                str(pad_op.op_type) not in {"PAD", "MIRROR_PAD"}
                or len(pad_op.inputs) < 2
                or len(pad_op.outputs) != 1
                or str(pad_op.outputs[0]) != str(data_input_name)
            ):
                continue
            if set(int(v) for v in consumers.get(str(data_input_name), [])) != {int(mul_idx)}:
                continue

            pre_nchw_name = str(pad_op.inputs[0])
            pre_idx = producers.get(str(pre_nchw_name), None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != str(pre_nchw_name)
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                or pre_nchw_name in model_outputs
            ):
                continue
            if set(int(v) for v in consumers.get(str(pre_nchw_name), [])) != {int(pad_idx)}:
                continue

            post_users = [int(v) for v in consumers.get(str(post_output_name), [])]
            if len(post_users) != 1:
                continue
            add_idx = int(post_users[0])
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_inputs = [str(v) for v in list(add_op.inputs)]
            add_side_name: Optional[str] = None
            if add_inputs[0] == str(post_output_name):
                add_side_name = str(add_inputs[1])
            elif add_inputs[1] == str(post_output_name):
                add_side_name = str(add_inputs[0])
            if add_side_name is None:
                continue
            add_side_tensor = model_ir.tensors.get(str(add_side_name), None)
            if add_side_tensor is None or add_side_tensor.data is None:
                continue

            target_nhwc_shape = (
                [int(v) for v in list(model_ir.tensors[str(post_output_name)].shape)]
                if str(post_output_name) in model_ir.tensors
                else None
            )
            if _is_fully_known_positive_shape(target_nhwc_shape):
                if _broadcast_static_shapes(
                    [int(v) for v in list(target_nhwc_shape)],
                    [int(v) for v in list(add_side_tensor.shape)],
                ) is None:
                    continue

            chain_index_set = {
                int(pre_idx),
                int(pad_idx),
                int(mul_idx),
                int(post_idx),
                int(add_idx),
            }

            if not _rewrite_rank4_mul_const_to_nhwc(
                mul_op=mul_op,
                mul_idx=int(mul_idx),
                const_input_index=int(const_input_index),
                const_input_name=str(const_input_name),
                target_nhwc_shape=target_nhwc_shape,
                chain_index_set=chain_index_set,
                consumers=consumers,
            ):
                continue

            if not _rewrite_pad_spec_to_nhwc(
                pad_op=pad_op,
                pad_idx=int(pad_idx),
                pads_input_name=str(pad_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            ):
                continue

            pre_input_name = str(pre_op.inputs[0])
            pad_inputs = [str(v) for v in list(pad_op.inputs)]
            pad_inputs[0] = str(pre_input_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=pad_op,
                new_inputs=pad_inputs,
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(data_input_name), None),
                perm_nchw_to_nhwc,
            )

            old_mul_tensor = model_ir.tensors.get(str(mul_output_name), None)
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[str(post_output_name)],
            )
            post_output_tensor = model_ir.tensors.get(str(post_output_name), None)
            if old_mul_tensor is not None and post_output_tensor is not None:
                post_output_tensor.dtype = str(old_mul_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                post_output_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                post_output_tensor.shape_signature = (
                    [int(v) for v in list(old_mul_tensor.shape_signature)]
                    if old_mul_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {int(post_idx)}
            pre_remaining_users = [int(v) for v in consumers.get(str(pre_nchw_name), []) if int(v) != int(pad_idx)]
            if len(pre_remaining_users) == 0:
                remove_indices.add(int(pre_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pad_mul_posttranspose_add_nhwc_chains": int(rewritten)}


def _optimize_transpose_mul_posttranspose_add_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate strict NHWC<->NCHW round-trips around:
      TRANSPOSE(NHWC->NCHW) -> MUL(const) -> TRANSPOSE(NCHW->NHWC) -> ADD(const)

    Rewrite:
      MUL(const_nhwc) -> ADD(const)
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_op: OperatorIR,
        const_input_index: int,
        const_input_name: str,
        target_nhwc_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        const_tensor = model_ir.tensors.get(str(const_input_name), None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        target_shape = (
            [int(v) for v in list(target_nhwc_shape)]
            if _is_fully_known_positive_shape(target_nhwc_shape)
            else None
        )

        rotated_data: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            as_is_shape = [int(v) for v in list(const_data.shape)]
            is_nchw_channelwise = (
                len(as_is_shape) == 4
                and int(as_is_shape[0]) == 1
                and int(as_is_shape[1]) > 0
                and int(as_is_shape[2]) == 1
                and int(as_is_shape[3]) == 1
            )
            if target_shape is not None and _broadcast_static_shapes(target_shape, as_is_shape) is not None:
                return True
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is not None:
                if _broadcast_static_shapes(target_shape, rotated_shape) is None:
                    return False
            elif not is_nchw_channelwise:
                return False
            rotated_data = np.asarray(rotated_candidate)
        else:
            if target_shape is None:
                return False
            side_shape = [int(v) for v in list(const_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is not None:
                return True
            return False

        if rotated_data is None:
            return False

        const_users = [int(v) for v in consumers.get(str(const_input_name), [])]
        shared_outside_chain = any(int(v) not in chain_index_set for v in const_users)
        target_name = str(const_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{const_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated_data.shape)],
                shape_signature=[int(v) for v in list(rotated_data.shape)],
                data=np.asarray(rotated_data),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
        else:
            const_tensor.data = np.asarray(rotated_data)
            const_tensor.shape = [int(v) for v in list(rotated_data.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated_data.shape)]

        if target_name != str(const_input_name):
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            mul_inputs[int(const_input_index)] = str(target_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            mul_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if mul_output_name in model_outputs or post_output_name in model_outputs:
                continue

            mul_idx = producers.get(str(mul_output_name), None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if (
                str(mul_op.op_type) != "MUL"
                or len(mul_op.inputs) != 2
                or len(mul_op.outputs) != 1
                or str(mul_op.outputs[0]) != str(mul_output_name)
            ):
                continue
            if set(int(v) for v in consumers.get(str(mul_output_name), [])) != {int(post_idx)}:
                continue

            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            pre_input_index: Optional[int] = None
            pre_idx: Optional[int] = None
            pre_op: Optional[OperatorIR] = None
            mul_const_input_index: Optional[int] = None
            mul_const_name: Optional[str] = None
            for i, input_name in enumerate(mul_inputs):
                producer_idx = producers.get(str(input_name), None)
                if producer_idx is not None:
                    producer_op = model_ir.operators[int(producer_idx)]
                    if (
                        str(producer_op.op_type) == "TRANSPOSE"
                        and len(producer_op.inputs) >= 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == str(input_name)
                        and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                    ):
                        if pre_idx is not None:
                            pre_idx = None
                            break
                        pre_input_index = int(i)
                        pre_idx = int(producer_idx)
                        pre_op = producer_op
                        continue
                const_tensor = model_ir.tensors.get(str(input_name), None)
                if const_tensor is not None and const_tensor.data is not None:
                    mul_const_input_index = int(i)
                    mul_const_name = str(input_name)

            if (
                pre_idx is None
                or pre_op is None
                or pre_input_index is None
                or mul_const_input_index is None
                or mul_const_name is None
            ):
                continue

            pre_output_name = str(pre_op.outputs[0])
            pre_input_name = str(pre_op.inputs[0])
            if pre_output_name in model_outputs:
                continue

            add_users = [int(v) for v in consumers.get(str(post_output_name), [])]
            if len(add_users) == 0:
                continue

            valid_add_tail = True
            for add_idx in add_users:
                add_op = model_ir.operators[int(add_idx)]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    valid_add_tail = False
                    break
                add_inputs = [str(v) for v in list(add_op.inputs)]
                if str(post_output_name) == add_inputs[0]:
                    side_name = str(add_inputs[1])
                elif str(post_output_name) == add_inputs[1]:
                    side_name = str(add_inputs[0])
                else:
                    valid_add_tail = False
                    break
                side_tensor = model_ir.tensors.get(side_name, None)
                if side_tensor is None or side_tensor.data is None:
                    valid_add_tail = False
                    break
                side_data = np.asarray(side_tensor.data)
                if int(side_data.size) != 1:
                    side_shape = [int(v) for v in list(side_data.shape)]
                    if len(side_shape) != 4:
                        valid_add_tail = False
                        break
                    if not (
                        int(side_shape[0]) == 1
                        and int(side_shape[1]) == 1
                        and int(side_shape[2]) == 1
                        and int(side_shape[3]) > 0
                    ):
                        valid_add_tail = False
                        break
            if not valid_add_tail:
                continue

            target_nhwc_shape: Optional[List[int]] = None
            pre_input_tensor = model_ir.tensors.get(str(pre_input_name), None)
            if pre_input_tensor is not None and len(list(pre_input_tensor.shape)) == 4:
                target_nhwc_shape = [int(v) for v in list(pre_input_tensor.shape)]

            chain_index_set = {int(mul_idx), int(post_idx), int(pre_idx)}
            chain_index_set.update(int(v) for v in add_users)
            if not _rewrite_mul_const_to_nhwc_if_needed(
                mul_op=mul_op,
                const_input_index=int(mul_const_input_index),
                const_input_name=str(mul_const_name),
                target_nhwc_shape=target_nhwc_shape,
                chain_index_set=chain_index_set,
                consumers=consumers,
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=int(pre_input_index),
                new_input_name=str(pre_input_name),
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul_output_name), None),
                perm_nchw_to_nhwc,
            )

            for add_idx in add_users:
                add_op = model_ir.operators[int(add_idx)]
                add_inputs = [str(v) for v in list(add_op.inputs)]
                if str(post_output_name) == add_inputs[0]:
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=add_op,
                        input_index=0,
                        new_input_name=str(mul_output_name),
                    )
                elif str(post_output_name) == add_inputs[1]:
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=add_op,
                        input_index=1,
                        new_input_name=str(mul_output_name),
                    )

            remove_indices = {int(post_idx)}
            pre_remaining_users = [int(v) for v in consumers.get(str(pre_output_name), []) if int(v) != int(mul_idx)]
            if len(pre_remaining_users) == 0:
                remove_indices.add(int(pre_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_mul_posttranspose_add_nhwc_chains": int(rewritten)}


def _optimize_sinet_shuffle_residual_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Optimize a strict ShuffleNet-like residual pattern observed in SiNet blocks:

      T13: a_nhwc -> a_nchw
      T26/T36: x_nhwc,y_nhwc -> x_nchw,y_nchw
      ADD -> MUL(const) -> ADD(const) -> PRELU -> T42 (to NHWC conv branch)
      CONCAT(a_nchw, prelu_nchw) -> MUL(const) -> ADD(const) -> PRELU -> T49

    Rewrite to NHWC core and drop redundant adapters:
      - remove T26/T36 and run first residual chain in NHWC
      - remove T42 by exposing first PRELU output as NHWC directly
      - rewrite CONCAT to NHWC (axis=3), consume a_nhwc directly (drop T13)
      - run second residual chain in NHWC and remove T49
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _ensure_channel_last_const(
        *,
        tensor_name: str,
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        data = np.asarray(tensor.data)
        if int(data.size) == 1:
            return str(tensor_name)
        if data.ndim != 4:
            return None
        shape = [int(v) for v in list(data.shape)]
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) > 0
        ):
            nhwc_data = np.asarray(data)
        elif (
            int(shape[0]) == 1
            and int(shape[1]) > 0
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            nhwc_data = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
        else:
            return None

        side_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        if shared_outside_chain:
            cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            return str(cloned_name)

        tensor.data = np.asarray(nhwc_data)
        tensor.shape = [int(v) for v in list(nhwc_data.shape)]
        tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
        return str(tensor_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post2_idx, post2_op in enumerate(model_ir.operators):
            if (
                str(post2_op.op_type) != "TRANSPOSE"
                or len(post2_op.inputs) < 2
                or len(post2_op.outputs) != 1
                or _read_transpose_perm(model_ir, post2_op) != perm_nchw_to_nhwc
            ):
                continue
            post2_input_name = str(post2_op.inputs[0])
            post2_output_name = str(post2_op.outputs[0])
            if post2_output_name in model_outputs:
                continue

            prelu2_idx = producers.get(post2_input_name, None)
            if prelu2_idx is None:
                continue
            prelu2_op = model_ir.operators[int(prelu2_idx)]
            if (
                str(prelu2_op.op_type) != "PRELU"
                or len(prelu2_op.inputs) != 2
                or len(prelu2_op.outputs) != 1
                or str(prelu2_op.outputs[0]) != post2_input_name
            ):
                continue

            add2_idx = producers.get(str(prelu2_op.inputs[0]), None)
            if add2_idx is None:
                continue
            add2_op = model_ir.operators[int(add2_idx)]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_out_name = str(add2_op.outputs[0])
            if str(prelu2_op.inputs[0]) != add2_out_name:
                continue
            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])

            mul2_idx = producers.get(add2_in0, None)
            if mul2_idx is None:
                mul2_idx = producers.get(add2_in1, None)
            if mul2_idx is None:
                continue
            mul2_op = model_ir.operators[int(mul2_idx)]
            if str(mul2_op.op_type) != "MUL" or len(mul2_op.inputs) != 2 or len(mul2_op.outputs) != 1:
                continue
            mul2_out_name = str(mul2_op.outputs[0])
            if mul2_out_name == add2_in0:
                add2_side_input_index = 1
                add2_side_name = add2_in1
            elif mul2_out_name == add2_in1:
                add2_side_input_index = 0
                add2_side_name = add2_in0
            else:
                continue

            mul2_in0 = str(mul2_op.inputs[0])
            mul2_in1 = str(mul2_op.inputs[1])
            if str(mul2_out_name) != str(prelu2_op.inputs[0]) and str(mul2_out_name) != str(add2_in0) and str(mul2_out_name) != str(add2_in1):
                continue
            if str(mul2_out_name) in model_outputs:
                continue

            if mul2_in0 == str(mul2_in1):
                continue
            if mul2_in0 in model_outputs or mul2_in1 in model_outputs:
                continue

            if mul2_in0 in producers and str(model_ir.operators[int(producers[mul2_in0])].op_type) == "CONCATENATION":
                concat2_data_name = str(mul2_in0)
                mul2_side_input_index = 1
                mul2_side_name = str(mul2_in1)
            elif mul2_in1 in producers and str(model_ir.operators[int(producers[mul2_in1])].op_type) == "CONCATENATION":
                concat2_data_name = str(mul2_in1)
                mul2_side_input_index = 0
                mul2_side_name = str(mul2_in0)
            else:
                continue

            concat2_idx = producers.get(concat2_data_name, None)
            if concat2_idx is None:
                continue
            concat2_op = model_ir.operators[int(concat2_idx)]
            if (
                str(concat2_op.op_type) != "CONCATENATION"
                or len(concat2_op.inputs) != 2
                or len(concat2_op.outputs) != 1
                or int(concat2_op.options.get("axis", 1)) != 1
                or str(concat2_op.outputs[0]) != concat2_data_name
            ):
                continue

            concat2_inputs = [str(v) for v in list(concat2_op.inputs)]
            pre13_idx: Optional[int] = None
            pre13_input_name: Optional[str] = None
            prelu1_out_name: Optional[str] = None
            for input_name in concat2_inputs:
                input_prod_idx = producers.get(str(input_name), None)
                if input_prod_idx is None:
                    continue
                input_prod_op = model_ir.operators[int(input_prod_idx)]
                if (
                    str(input_prod_op.op_type) == "TRANSPOSE"
                    and len(input_prod_op.inputs) >= 2
                    and len(input_prod_op.outputs) == 1
                    and str(input_prod_op.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, input_prod_op) == perm_nhwc_to_nchw
                ):
                    pre13_idx = int(input_prod_idx)
                    pre13_input_name = str(input_prod_op.inputs[0])
                elif str(input_prod_op.op_type) == "PRELU":
                    prelu1_out_name = str(input_name)
            if pre13_idx is None or pre13_input_name is None or prelu1_out_name is None:
                continue

            prelu1_idx = producers.get(str(prelu1_out_name), None)
            if prelu1_idx is None:
                continue
            prelu1_op = model_ir.operators[int(prelu1_idx)]
            if (
                str(prelu1_op.op_type) != "PRELU"
                or len(prelu1_op.inputs) != 2
                or len(prelu1_op.outputs) != 1
                or str(prelu1_op.outputs[0]) != str(prelu1_out_name)
            ):
                continue
            if str(prelu1_out_name) in model_outputs:
                continue

            prelu1_users = [int(v) for v in consumers.get(str(prelu1_out_name), []) if int(v) != int(prelu1_idx)]
            if len(prelu1_users) != 2 or int(concat2_idx) not in prelu1_users:
                continue
            post1_idx = None
            for u in prelu1_users:
                if int(u) == int(concat2_idx):
                    continue
                u_op = model_ir.operators[int(u)]
                if (
                    str(u_op.op_type) == "TRANSPOSE"
                    and len(u_op.inputs) >= 2
                    and len(u_op.outputs) == 1
                    and str(u_op.inputs[0]) == str(prelu1_out_name)
                    and _read_transpose_perm(model_ir, u_op) == perm_nchw_to_nhwc
                    and str(u_op.outputs[0]) not in model_outputs
                ):
                    post1_idx = int(u)
            if post1_idx is None:
                continue
            post1_op = model_ir.operators[int(post1_idx)]
            post1_output_name = str(post1_op.outputs[0])

            add1_idx = producers.get(str(prelu1_op.inputs[0]), None)
            if add1_idx is None:
                continue
            add1_op = model_ir.operators[int(add1_idx)]
            if str(add1_op.op_type) != "ADD" or len(add1_op.inputs) != 2 or len(add1_op.outputs) != 1:
                continue
            add1_in0 = str(add1_op.inputs[0])
            add1_in1 = str(add1_op.inputs[1])

            mul1_idx = producers.get(add1_in0, None)
            if mul1_idx is None:
                mul1_idx = producers.get(add1_in1, None)
            if mul1_idx is None:
                continue
            mul1_op = model_ir.operators[int(mul1_idx)]
            if str(mul1_op.op_type) != "MUL" or len(mul1_op.inputs) != 2 or len(mul1_op.outputs) != 1:
                continue
            mul1_out_name = str(mul1_op.outputs[0])
            if mul1_out_name == add1_in0:
                add1_side_input_index = 1
                add1_side_name = add1_in1
            elif mul1_out_name == add1_in1:
                add1_side_input_index = 0
                add1_side_name = add1_in0
            else:
                continue

            mul1_in0 = str(mul1_op.inputs[0])
            mul1_in1 = str(mul1_op.inputs[1])
            add0_out_name: Optional[str] = None
            mul1_side_input_index: Optional[int] = None
            mul1_side_name: Optional[str] = None
            for data_idx, side_idx in [(0, 1), (1, 0)]:
                cand_data = str(mul1_op.inputs[int(data_idx)])
                cand_side = str(mul1_op.inputs[int(side_idx)])
                cand_prod_idx = producers.get(cand_data, None)
                if cand_prod_idx is not None:
                    cand_prod_op = model_ir.operators[int(cand_prod_idx)]
                    if str(cand_prod_op.op_type) == "ADD":
                        add0_out_name = str(cand_data)
                        mul1_side_input_index = int(side_idx)
                        mul1_side_name = str(cand_side)
                        break
            if add0_out_name is None or mul1_side_input_index is None or mul1_side_name is None:
                continue

            add0_idx = producers.get(str(add0_out_name), None)
            if add0_idx is None:
                continue
            add0_op = model_ir.operators[int(add0_idx)]
            if str(add0_op.op_type) != "ADD" or len(add0_op.inputs) != 2 or len(add0_op.outputs) != 1:
                continue
            if str(add0_op.outputs[0]) != str(add0_out_name):
                continue

            pre_add_indices: List[int] = []
            pre_add_inputs_nhwc: List[str] = []
            pre_add_out_names = [str(v) for v in list(add0_op.inputs)]
            valid_pre_add = True
            for pre_add_out in pre_add_out_names:
                pre_add_idx = producers.get(str(pre_add_out), None)
                if pre_add_idx is None:
                    valid_pre_add = False
                    break
                pre_add_op = model_ir.operators[int(pre_add_idx)]
                if (
                    str(pre_add_op.op_type) != "TRANSPOSE"
                    or len(pre_add_op.inputs) < 2
                    or len(pre_add_op.outputs) != 1
                    or str(pre_add_op.outputs[0]) != str(pre_add_out)
                    or _read_transpose_perm(model_ir, pre_add_op) != perm_nhwc_to_nchw
                ):
                    valid_pre_add = False
                    break
                if set(int(v) for v in consumers.get(str(pre_add_out), [])) != {int(add0_idx)}:
                    valid_pre_add = False
                    break
                pre_add_indices.append(int(pre_add_idx))
                pre_add_inputs_nhwc.append(str(pre_add_op.inputs[0]))
            if not valid_pre_add or len(pre_add_indices) != 2:
                continue

            if set(int(v) for v in consumers.get(str(concat2_inputs[0]), [])) != {int(concat2_idx)} and set(
                int(v) for v in consumers.get(str(concat2_inputs[1]), [])
            ) != {int(concat2_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(pre13_input_name if pre13_input_name else ""), [])):
                pass
            if set(int(v) for v in consumers.get(str(model_ir.operators[int(pre13_idx)].outputs[0]), [])) != {int(concat2_idx)}:
                continue

            chain_index_set = {
                int(add0_idx),
                int(mul1_idx),
                int(add1_idx),
                int(prelu1_idx),
                int(concat2_idx),
                int(mul2_idx),
                int(add2_idx),
                int(prelu2_idx),
                int(pre13_idx),
                int(post1_idx),
                int(post2_idx),
                int(pre_add_indices[0]),
                int(pre_add_indices[1]),
            }
            mul1_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mul1_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            add1_side_for_op = _ensure_channel_last_const(
                tensor_name=str(add1_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            prelu1_side_for_op = _ensure_channel_last_const(
                tensor_name=str(prelu1_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            mul2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mul2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            add2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(add2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            prelu2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(prelu2_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            if any(
                v is None
                for v in [
                    mul1_side_for_op,
                    add1_side_for_op,
                    prelu1_side_for_op,
                    mul2_side_for_op,
                    add2_side_for_op,
                    prelu2_side_for_op,
                ]
            ):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=add0_op,
                new_inputs=[str(pre_add_inputs_nhwc[0]), str(pre_add_inputs_nhwc[1])],
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul1_op,
                input_index=int(mul1_side_input_index),
                new_input_name=str(mul1_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add1_op,
                input_index=int(add1_side_input_index),
                new_input_name=str(add1_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu1_op,
                input_index=1,
                new_input_name=str(prelu1_side_for_op),
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu1_op,
                new_outputs=[str(post1_output_name)],
            )

            new_concat_inputs = []
            for inp in concat2_inputs:
                if str(inp) == str(model_ir.operators[int(pre13_idx)].outputs[0]):
                    new_concat_inputs.append(str(pre13_input_name))
                elif str(inp) == str(prelu1_out_name):
                    new_concat_inputs.append(str(post1_output_name))
                else:
                    new_concat_inputs.append(str(inp))
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat2_op,
                new_inputs=[str(v) for v in new_concat_inputs],
            )
            concat2_op.options["axis"] = 3

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul2_op,
                input_index=int(mul2_side_input_index),
                new_input_name=str(mul2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add2_op,
                input_index=int(add2_side_input_index),
                new_input_name=str(add2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu2_op,
                input_index=1,
                new_input_name=str(prelu2_side_for_op),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu2_op,
                new_outputs=[str(post2_output_name)],
            )

            for t_name in [add0_out_name, mul1_out_name, str(add1_op.outputs[0]), str(post1_output_name)]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(t_name), None),
                    perm_nchw_to_nhwc,
                )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat2_data_name), None),
                perm_nchw_to_nhwc,
            )
            for t_name in [mul2_out_name, add2_out_name, str(post2_output_name)]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(t_name), None),
                    perm_nchw_to_nhwc,
                )

            remove_indices = set([int(post1_idx), int(post2_idx), int(pre13_idx)])
            remove_indices.update(set(int(v) for v in pre_add_indices))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_sinet_shuffle_residual_transpose_chains": int(optimized)}


def _optimize_sinet_concat_resize_affine_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Optimize strict SiNet branch merge blocks with three NHWC->NCHW adapters:

      pre_nhwc --T(0,3,1,2)--> pre_nchw
      branch0_nhwc --T(0,3,1,2)--> b0_nchw
      resize_nhwc --T(0,3,1,2)--> rz_nchw --MUL(const)--ADD(const)--> b1_nchw
      CONCAT(axis=1, [b0_nchw, b1_nchw]) -> c_nchw
      ADD(pre_nchw, c_nchw) -> a_nchw --MUL(const)--ADD(const)--PRELU--> p_nchw
      p_nchw --T(0,2,3,1)--> p_nhwc

    Rewrite to NHWC core and remove the four transpose adapters.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _ensure_channel_last_const(
        *,
        tensor_name: str,
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        data = np.asarray(tensor.data)
        if int(data.size) == 1:
            return str(tensor_name)
        if data.ndim != 4:
            return None
        shape = [int(v) for v in list(data.shape)]
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) > 0
        ):
            nhwc_data = np.asarray(data)
        elif (
            int(shape[0]) == 1
            and int(shape[1]) > 0
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            nhwc_data = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
        else:
            return None

        side_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        if shared_outside_chain:
            cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            return str(cloned_name)

        tensor.data = np.asarray(nhwc_data)
        tensor.shape = [int(v) for v in list(nhwc_data.shape)]
        tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
        return str(tensor_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue
            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue

            add0_idx = int(pre_users[0])
            add0_op = model_ir.operators[int(add0_idx)]
            if str(add0_op.op_type) != "ADD" or len(add0_op.inputs) != 2 or len(add0_op.outputs) != 1:
                continue
            add0_inputs = [str(v) for v in list(add0_op.inputs)]
            if add0_inputs[0] == pre_output_name:
                concat_name = str(add0_inputs[1])
            elif add0_inputs[1] == pre_output_name:
                concat_name = str(add0_inputs[0])
            else:
                continue
            add0_out_name = str(add0_op.outputs[0])
            if add0_out_name in model_outputs:
                continue

            concat_idx = producers.get(concat_name, None)
            if concat_idx is None:
                continue
            concat_op = model_ir.operators[int(concat_idx)]
            if (
                str(concat_op.op_type) != "CONCATENATION"
                or len(concat_op.inputs) != 2
                or len(concat_op.outputs) != 1
                or str(concat_op.outputs[0]) != concat_name
                or int(concat_op.options.get("axis", 1)) != 1
            ):
                continue
            if set(int(v) for v in consumers.get(concat_name, [])) != {int(add0_idx)}:
                continue

            trans_a_idx: Optional[int] = None
            trans_a_in_name: Optional[str] = None
            trans_a_out_name: Optional[str] = None
            addb_idx: Optional[int] = None
            addb_out_name: Optional[str] = None
            for c_in in [str(v) for v in list(concat_op.inputs)]:
                c_prod_idx = producers.get(str(c_in), None)
                if c_prod_idx is None:
                    continue
                c_prod_op = model_ir.operators[int(c_prod_idx)]
                if (
                    str(c_prod_op.op_type) == "TRANSPOSE"
                    and len(c_prod_op.inputs) >= 2
                    and len(c_prod_op.outputs) == 1
                    and str(c_prod_op.outputs[0]) == str(c_in)
                    and _read_transpose_perm(model_ir, c_prod_op) == perm_nhwc_to_nchw
                    and set(int(v) for v in consumers.get(str(c_in), [])) == {int(concat_idx)}
                    and str(c_in) not in model_outputs
                ):
                    trans_a_idx = int(c_prod_idx)
                    trans_a_in_name = str(c_prod_op.inputs[0])
                    trans_a_out_name = str(c_in)
                elif (
                    str(c_prod_op.op_type) == "ADD"
                    and len(c_prod_op.inputs) == 2
                    and len(c_prod_op.outputs) == 1
                    and str(c_prod_op.outputs[0]) == str(c_in)
                    and set(int(v) for v in consumers.get(str(c_in), [])) == {int(concat_idx)}
                    and str(c_in) not in model_outputs
                ):
                    addb_idx = int(c_prod_idx)
                    addb_out_name = str(c_in)
            if (
                trans_a_idx is None
                or trans_a_in_name is None
                or trans_a_out_name is None
                or addb_idx is None
                or addb_out_name is None
            ):
                continue

            addb_op = model_ir.operators[int(addb_idx)]
            addb_in0 = str(addb_op.inputs[0])
            addb_in1 = str(addb_op.inputs[1])
            mulb_idx = producers.get(addb_in0, None)
            if mulb_idx is None:
                mulb_idx = producers.get(addb_in1, None)
            if mulb_idx is None:
                continue
            mulb_op = model_ir.operators[int(mulb_idx)]
            if str(mulb_op.op_type) != "MUL" or len(mulb_op.inputs) != 2 or len(mulb_op.outputs) != 1:
                continue
            mulb_out_name = str(mulb_op.outputs[0])
            if mulb_out_name == addb_in0:
                addb_side_index = 1
                addb_side_name = str(addb_in1)
            elif mulb_out_name == addb_in1:
                addb_side_index = 0
                addb_side_name = str(addb_in0)
            else:
                continue

            mulb_inputs = [str(v) for v in list(mulb_op.inputs)]
            if mulb_inputs[0] == mulb_inputs[1]:
                continue
            trans_b_idx: Optional[int] = None
            trans_b_in_name: Optional[str] = None
            mulb_data_index: Optional[int] = None
            mulb_side_name: Optional[str] = None
            for data_idx, side_idx in [(0, 1), (1, 0)]:
                trans_idx = producers.get(str(mulb_inputs[int(data_idx)]), None)
                if trans_idx is None:
                    continue
                trans_op = model_ir.operators[int(trans_idx)]
                if (
                    str(trans_op.op_type) != "TRANSPOSE"
                    or len(trans_op.inputs) < 2
                    or len(trans_op.outputs) != 1
                    or str(trans_op.outputs[0]) != str(mulb_inputs[int(data_idx)])
                    or _read_transpose_perm(model_ir, trans_op) != perm_nhwc_to_nchw
                    or set(int(v) for v in consumers.get(str(mulb_inputs[int(data_idx)]), [])) != {int(mulb_idx)}
                    or str(mulb_inputs[int(data_idx)]) in model_outputs
                ):
                    continue
                trans_b_idx = int(trans_idx)
                trans_b_in_name = str(trans_op.inputs[0])
                mulb_data_index = int(data_idx)
                mulb_side_name = str(mulb_inputs[int(side_idx)])
                break
            if trans_b_idx is None or trans_b_in_name is None or mulb_data_index is None or mulb_side_name is None:
                continue

            add0_users = [int(v) for v in consumers.get(str(add0_out_name), [])]
            if len(add0_users) != 1:
                continue
            mul2_idx = int(add0_users[0])
            mul2_op = model_ir.operators[int(mul2_idx)]
            if str(mul2_op.op_type) != "MUL" or len(mul2_op.inputs) != 2 or len(mul2_op.outputs) != 1:
                continue
            mul2_in0 = str(mul2_op.inputs[0])
            mul2_in1 = str(mul2_op.inputs[1])
            if mul2_in0 == add0_out_name:
                mul2_side_index = 1
                mul2_side_name = str(mul2_in1)
            elif mul2_in1 == add0_out_name:
                mul2_side_index = 0
                mul2_side_name = str(mul2_in0)
            else:
                continue
            mul2_out_name = str(mul2_op.outputs[0])
            if mul2_out_name in model_outputs:
                continue

            mul2_users = [int(v) for v in consumers.get(str(mul2_out_name), [])]
            if len(mul2_users) != 1:
                continue
            add2_idx = int(mul2_users[0])
            add2_op = model_ir.operators[int(add2_idx)]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])
            if add2_in0 == mul2_out_name:
                add2_side_index = 1
                add2_side_name = str(add2_in1)
            elif add2_in1 == mul2_out_name:
                add2_side_index = 0
                add2_side_name = str(add2_in0)
            else:
                continue
            add2_out_name = str(add2_op.outputs[0])
            if add2_out_name in model_outputs:
                continue

            add2_users = [int(v) for v in consumers.get(str(add2_out_name), [])]
            if len(add2_users) != 1:
                continue
            prelu2_idx = int(add2_users[0])
            prelu2_op = model_ir.operators[int(prelu2_idx)]
            if (
                str(prelu2_op.op_type) != "PRELU"
                or len(prelu2_op.inputs) != 2
                or len(prelu2_op.outputs) != 1
                or str(prelu2_op.inputs[0]) != str(add2_out_name)
            ):
                continue
            prelu2_out_name = str(prelu2_op.outputs[0])
            if prelu2_out_name in model_outputs:
                continue

            prelu2_users = [int(v) for v in consumers.get(str(prelu2_out_name), [])]
            if len(prelu2_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            for u_idx in prelu2_users:
                u_op = model_ir.operators[int(u_idx)]
                if (
                    str(u_op.op_type) == "TRANSPOSE"
                    and len(u_op.inputs) >= 2
                    and len(u_op.outputs) == 1
                    and str(u_op.inputs[0]) == str(prelu2_out_name)
                    and _read_transpose_perm(model_ir, u_op) == perm_nchw_to_nhwc
                    and str(u_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(u_idx))
                    post_output_names.append(str(u_op.outputs[0]))
                else:
                    legacy_users.append(int(u_idx))
            if len(post_indices) == 0:
                continue
            post_out_name = str(post_output_names[0])

            chain_index_set = {
                int(pre_idx),
                int(trans_a_idx),
                int(trans_b_idx),
                int(addb_idx),
                int(mulb_idx),
                int(concat_idx),
                int(add0_idx),
                int(mul2_idx),
                int(add2_idx),
                int(prelu2_idx),
                *[int(v) for v in list(post_indices)],
            }
            mulb_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mulb_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            addb_side_for_op = _ensure_channel_last_const(
                tensor_name=str(addb_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            mul2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mul2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            add2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(add2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            prelu2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(prelu2_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            if any(
                v is None
                for v in [
                    mulb_side_for_op,
                    addb_side_for_op,
                    mul2_side_for_op,
                    add2_side_for_op,
                    prelu2_side_for_op,
                ]
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mulb_op,
                input_index=int(mulb_data_index),
                new_input_name=str(trans_b_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mulb_op,
                input_index=int(1 - int(mulb_data_index)),
                new_input_name=str(mulb_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=addb_op,
                input_index=int(addb_side_index),
                new_input_name=str(addb_side_for_op),
            )

            new_concat_inputs = [
                str(trans_a_in_name) if str(v) == str(trans_a_out_name) else str(v)
                for v in list(concat_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in new_concat_inputs],
            )
            concat_op.options["axis"] = 3

            new_add0_inputs = [str(v) for v in add0_inputs]
            if new_add0_inputs[0] == str(pre_output_name):
                new_add0_inputs[0] = str(pre_input_name)
            elif new_add0_inputs[1] == str(pre_output_name):
                new_add0_inputs[1] = str(pre_input_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=add0_op,
                new_inputs=[str(v) for v in new_add0_inputs],
            )

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul2_op,
                input_index=int(mul2_side_index),
                new_input_name=str(mul2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add2_op,
                input_index=int(add2_side_index),
                new_input_name=str(add2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu2_op,
                input_index=1,
                new_input_name=str(prelu2_side_for_op),
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu2_op,
                new_outputs=[str(post_out_name)],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), str(post_out_name))

            for t_name in [
                str(mulb_out_name),
                str(addb_out_name),
                str(concat_name),
                str(add0_out_name),
                str(mul2_out_name),
                str(add2_out_name),
                str(post_out_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(t_name), None),
                    perm_nchw_to_nhwc,
                )

            old_prelu2_tensor = model_ir.tensors.get(str(prelu2_out_name), None)
            post_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_prelu2_tensor is not None and post_tensor is not None:
                post_tensor.dtype = str(old_prelu2_tensor.dtype)
                post_tensor.quantization = _clone_quantization(old_prelu2_tensor.quantization)
                post_tensor.shape = [int(v) for v in list(old_prelu2_tensor.shape)]
                post_tensor.shape_signature = (
                    [int(v) for v in list(old_prelu2_tensor.shape_signature)]
                    if old_prelu2_tensor.shape_signature is not None
                    else [int(v) for v in list(old_prelu2_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                # Preserve NCHW legacy fan-out from the rewritten NHWC output.
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(post_out_name), str(pre_op.inputs[1])],
                    outputs=[str(prelu2_out_name)],
                    options={},
                )
                op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                live_legacy_indices = [
                    int(op_index_by_id[int(id(model_ir.operators[int(v)]))])
                    for v in legacy_users
                    if int(v) < len(model_ir.operators)
                ]
                insert_index = int(min(live_legacy_indices)) if len(live_legacy_indices) > 0 else int(len(model_ir.operators))
                model_ir.operators.insert(int(insert_index), adapter_op)

            remove_indices = {
                int(pre_idx),
                int(trans_a_idx),
                int(trans_b_idx),
                *[int(v) for v in list(post_indices)],
            }
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_sinet_concat_resize_affine_transpose_chains": int(optimized)}


def _optimize_sinet_dual_resize_affine_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Optimize strict SiNet branch-merge blocks where BOTH concat branches are:
      RESIZE -> TRANSPOSE(NHWC->NCHW) -> MUL(const) -> ADD(const)

    Target:
      pre_nhwc --T(0,3,1,2)--> pre_nchw
      b0_nhwc --Resize--> r0_nhwc --T(0,3,1,2)--> r0_nchw --MUL(c0)--ADD(c1)--> b0_nchw
      b1_nhwc --Resize--> r1_nhwc --T(0,3,1,2)--> r1_nchw --MUL(c2)--ADD(c3)--> b1_nchw
      CONCAT(axis=1, [b0_nchw, b1_nchw]) -> c_nchw
      ADD(pre_nchw, c_nchw) -> a_nchw --MUL(const)--ADD(const)--PRELU--> p_nchw
      p_nchw --T(0,2,3,1)--> p_nhwc

    Rewrite to NHWC core and remove the four transpose adapters.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _ensure_channel_last_const(
        *,
        tensor_name: str,
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        data = np.asarray(tensor.data)
        if int(data.size) == 1:
            return str(tensor_name)
        if data.ndim != 4:
            return None
        shape = [int(v) for v in list(data.shape)]
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) > 0
        ):
            nhwc_data = np.asarray(data)
        elif (
            int(shape[0]) == 1
            and int(shape[1]) > 0
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            nhwc_data = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
        else:
            return None

        side_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        if shared_outside_chain:
            cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            return str(cloned_name)

        tensor.data = np.asarray(nhwc_data)
        tensor.shape = [int(v) for v in list(nhwc_data.shape)]
        tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
        return str(tensor_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue
            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue

            add0_idx = int(pre_users[0])
            add0_op = model_ir.operators[int(add0_idx)]
            if str(add0_op.op_type) != "ADD" or len(add0_op.inputs) != 2 or len(add0_op.outputs) != 1:
                continue
            add0_inputs = [str(v) for v in list(add0_op.inputs)]
            if add0_inputs[0] == pre_output_name:
                concat_name = str(add0_inputs[1])
            elif add0_inputs[1] == pre_output_name:
                concat_name = str(add0_inputs[0])
            else:
                continue
            add0_out_name = str(add0_op.outputs[0])
            if add0_out_name in model_outputs:
                continue

            concat_idx = producers.get(concat_name, None)
            if concat_idx is None:
                continue
            concat_op = model_ir.operators[int(concat_idx)]
            if (
                str(concat_op.op_type) != "CONCATENATION"
                or len(concat_op.inputs) != 2
                or len(concat_op.outputs) != 1
                or str(concat_op.outputs[0]) != concat_name
                or int(concat_op.options.get("axis", 1)) != 1
            ):
                continue
            if set(int(v) for v in consumers.get(concat_name, [])) != {int(add0_idx)}:
                continue

            branch_plans: List[Dict[str, Any]] = []

            def _match_resize_affine_branch(add_out_name: str) -> Optional[Dict[str, Any]]:
                add_idx = producers.get(str(add_out_name), None)
                if add_idx is None:
                    return None
                add_op = model_ir.operators[int(add_idx)]
                if (
                    str(add_op.op_type) != "ADD"
                    or len(add_op.inputs) != 2
                    or len(add_op.outputs) != 1
                    or str(add_op.outputs[0]) != str(add_out_name)
                ):
                    return None
                if set(int(v) for v in consumers.get(str(add_out_name), [])) != {int(concat_idx)}:
                    return None

                add_in0 = str(add_op.inputs[0])
                add_in1 = str(add_op.inputs[1])
                mul_idx = producers.get(add_in0, None)
                if mul_idx is None:
                    mul_idx = producers.get(add_in1, None)
                if mul_idx is None:
                    return None
                mul_op = model_ir.operators[int(mul_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    return None

                mul_out_name = str(mul_op.outputs[0])
                if mul_out_name == add_in0:
                    add_side_index = 1
                    add_side_name = str(add_in1)
                elif mul_out_name == add_in1:
                    add_side_index = 0
                    add_side_name = str(add_in0)
                else:
                    return None

                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                if mul_inputs[0] == mul_inputs[1]:
                    return None

                trans_idx: Optional[int] = None
                trans_in_name: Optional[str] = None
                mul_data_index: Optional[int] = None
                mul_side_name: Optional[str] = None

                for data_idx, side_idx in [(0, 1), (1, 0)]:
                    candidate_trans_idx = producers.get(str(mul_inputs[int(data_idx)]), None)
                    if candidate_trans_idx is None:
                        continue
                    candidate_trans_op = model_ir.operators[int(candidate_trans_idx)]
                    if (
                        str(candidate_trans_op.op_type) != "TRANSPOSE"
                        or len(candidate_trans_op.inputs) < 2
                        or len(candidate_trans_op.outputs) != 1
                        or str(candidate_trans_op.outputs[0]) != str(mul_inputs[int(data_idx)])
                        or _read_transpose_perm(model_ir, candidate_trans_op) != perm_nhwc_to_nchw
                        or set(int(v) for v in consumers.get(str(mul_inputs[int(data_idx)]), []))
                        != {int(mul_idx)}
                        or str(mul_inputs[int(data_idx)]) in model_outputs
                    ):
                        continue
                    candidate_resize_out_name = str(candidate_trans_op.inputs[0])
                    resize_idx = producers.get(candidate_resize_out_name, None)
                    if resize_idx is None:
                        continue
                    resize_op = model_ir.operators[int(resize_idx)]
                    if str(resize_op.op_type) not in {"RESIZE_BILINEAR", "RESIZE_NEAREST_NEIGHBOR"}:
                        continue
                    if len(resize_op.outputs) != 1 or str(resize_op.outputs[0]) != candidate_resize_out_name:
                        continue
                    trans_idx = int(candidate_trans_idx)
                    trans_in_name = candidate_resize_out_name
                    mul_data_index = int(data_idx)
                    mul_side_name = str(mul_inputs[int(side_idx)])
                    break

                if trans_idx is None or trans_in_name is None or mul_data_index is None or mul_side_name is None:
                    return None

                return {
                    "add_idx": int(add_idx),
                    "add_op": add_op,
                    "add_out_name": str(add_out_name),
                    "add_side_index": int(add_side_index),
                    "add_side_name": str(add_side_name),
                    "mul_idx": int(mul_idx),
                    "mul_op": mul_op,
                    "mul_out_name": str(mul_out_name),
                    "mul_data_index": int(mul_data_index),
                    "mul_side_name": str(mul_side_name),
                    "trans_idx": int(trans_idx),
                    "trans_in_name": str(trans_in_name),
                }

            valid_branches = True
            for concat_input_name in [str(v) for v in list(concat_op.inputs)]:
                plan = _match_resize_affine_branch(str(concat_input_name))
                if plan is None:
                    valid_branches = False
                    break
                branch_plans.append(plan)
            if not valid_branches or len(branch_plans) != 2:
                continue

            add0_users = [int(v) for v in consumers.get(str(add0_out_name), [])]
            if len(add0_users) != 1:
                continue
            mul2_idx = int(add0_users[0])
            mul2_op = model_ir.operators[int(mul2_idx)]
            if str(mul2_op.op_type) != "MUL" or len(mul2_op.inputs) != 2 or len(mul2_op.outputs) != 1:
                continue
            mul2_in0 = str(mul2_op.inputs[0])
            mul2_in1 = str(mul2_op.inputs[1])
            if mul2_in0 == add0_out_name:
                mul2_side_index = 1
                mul2_side_name = str(mul2_in1)
            elif mul2_in1 == add0_out_name:
                mul2_side_index = 0
                mul2_side_name = str(mul2_in0)
            else:
                continue
            mul2_out_name = str(mul2_op.outputs[0])
            if mul2_out_name in model_outputs:
                continue

            mul2_users = [int(v) for v in consumers.get(str(mul2_out_name), [])]
            if len(mul2_users) != 1:
                continue
            add2_idx = int(mul2_users[0])
            add2_op = model_ir.operators[int(add2_idx)]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])
            if add2_in0 == mul2_out_name:
                add2_side_index = 1
                add2_side_name = str(add2_in1)
            elif add2_in1 == mul2_out_name:
                add2_side_index = 0
                add2_side_name = str(add2_in0)
            else:
                continue
            add2_out_name = str(add2_op.outputs[0])
            if add2_out_name in model_outputs:
                continue

            add2_users = [int(v) for v in consumers.get(str(add2_out_name), [])]
            if len(add2_users) != 1:
                continue
            prelu2_idx = int(add2_users[0])
            prelu2_op = model_ir.operators[int(prelu2_idx)]
            if (
                str(prelu2_op.op_type) != "PRELU"
                or len(prelu2_op.inputs) != 2
                or len(prelu2_op.outputs) != 1
                or str(prelu2_op.inputs[0]) != str(add2_out_name)
            ):
                continue
            prelu2_out_name = str(prelu2_op.outputs[0])
            if prelu2_out_name in model_outputs:
                continue

            prelu2_users = [int(v) for v in consumers.get(str(prelu2_out_name), [])]
            if len(prelu2_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            for u_idx in prelu2_users:
                u_op = model_ir.operators[int(u_idx)]
                if (
                    str(u_op.op_type) == "TRANSPOSE"
                    and len(u_op.inputs) >= 2
                    and len(u_op.outputs) == 1
                    and str(u_op.inputs[0]) == str(prelu2_out_name)
                    and _read_transpose_perm(model_ir, u_op) == perm_nchw_to_nhwc
                    and str(u_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(u_idx))
                    post_output_names.append(str(u_op.outputs[0]))
                else:
                    legacy_users.append(int(u_idx))
            if len(post_indices) == 0:
                continue
            post_out_name = str(post_output_names[0])

            chain_index_set = {
                int(pre_idx),
                int(concat_idx),
                int(add0_idx),
                int(mul2_idx),
                int(add2_idx),
                int(prelu2_idx),
                *[int(v) for v in list(post_indices)],
            }
            for plan in branch_plans:
                chain_index_set.add(int(plan["trans_idx"]))
                chain_index_set.add(int(plan["mul_idx"]))
                chain_index_set.add(int(plan["add_idx"]))

            branch_mul_side_names: List[str] = []
            branch_add_side_names: List[str] = []
            for plan in branch_plans:
                branch_mul_side = _ensure_channel_last_const(
                    tensor_name=str(plan["mul_side_name"]),
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                )
                branch_add_side = _ensure_channel_last_const(
                    tensor_name=str(plan["add_side_name"]),
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                )
                if branch_mul_side is None or branch_add_side is None:
                    branch_mul_side_names = []
                    branch_add_side_names = []
                    break
                branch_mul_side_names.append(str(branch_mul_side))
                branch_add_side_names.append(str(branch_add_side))
            if len(branch_mul_side_names) != 2 or len(branch_add_side_names) != 2:
                continue

            mul2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mul2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            add2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(add2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            prelu2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(prelu2_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            if any(v is None for v in [mul2_side_for_op, add2_side_for_op, prelu2_side_for_op]):
                continue

            for plan, branch_mul_side, branch_add_side in zip(
                branch_plans,
                branch_mul_side_names,
                branch_add_side_names,
            ):
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=plan["mul_op"],
                    input_index=int(plan["mul_data_index"]),
                    new_input_name=str(plan["trans_in_name"]),
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=plan["mul_op"],
                    input_index=int(1 - int(plan["mul_data_index"])),
                    new_input_name=str(branch_mul_side),
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=plan["add_op"],
                    input_index=int(plan["add_side_index"]),
                    new_input_name=str(branch_add_side),
                )

            concat_op.options["axis"] = 3

            new_add0_inputs = [str(v) for v in add0_inputs]
            if new_add0_inputs[0] == str(pre_output_name):
                new_add0_inputs[0] = str(pre_input_name)
            elif new_add0_inputs[1] == str(pre_output_name):
                new_add0_inputs[1] = str(pre_input_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=add0_op,
                new_inputs=[str(v) for v in new_add0_inputs],
            )

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul2_op,
                input_index=int(mul2_side_index),
                new_input_name=str(mul2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add2_op,
                input_index=int(add2_side_index),
                new_input_name=str(add2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu2_op,
                input_index=1,
                new_input_name=str(prelu2_side_for_op),
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu2_op,
                new_outputs=[str(post_out_name)],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), str(post_out_name))

            for plan in branch_plans:
                for t_name in [str(plan["mul_out_name"]), str(plan["add_out_name"])]:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(t_name), None),
                        perm_nchw_to_nhwc,
                    )
            for t_name in [
                str(concat_name),
                str(add0_out_name),
                str(mul2_out_name),
                str(add2_out_name),
                str(post_out_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(t_name), None),
                    perm_nchw_to_nhwc,
                )

            old_prelu2_tensor = model_ir.tensors.get(str(prelu2_out_name), None)
            post_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_prelu2_tensor is not None and post_tensor is not None:
                post_tensor.dtype = str(old_prelu2_tensor.dtype)
                post_tensor.quantization = _clone_quantization(old_prelu2_tensor.quantization)
                post_tensor.shape = [int(v) for v in list(old_prelu2_tensor.shape)]
                post_tensor.shape_signature = (
                    [int(v) for v in list(old_prelu2_tensor.shape_signature)]
                    if old_prelu2_tensor.shape_signature is not None
                    else [int(v) for v in list(old_prelu2_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[str(post_out_name), str(pre_op.inputs[1])],
                    outputs=[str(prelu2_out_name)],
                    options={},
                )
                op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                live_legacy_indices = [
                    int(op_index_by_id[int(id(model_ir.operators[int(v)]))])
                    for v in legacy_users
                    if int(v) < len(model_ir.operators)
                ]
                insert_index = int(min(live_legacy_indices)) if len(live_legacy_indices) > 0 else int(len(model_ir.operators))
                model_ir.operators.insert(int(insert_index), adapter_op)

            remove_indices = {
                int(pre_idx),
                *[int(v) for v in list(post_indices)],
                *[int(plan["trans_idx"]) for plan in branch_plans],
            }
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_sinet_dual_resize_affine_transpose_chains": int(optimized)}


def _optimize_sinet_concat_resize_affine_tail_concat_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Optimize a late SiNet residual variant:

      skip_nhwc --T(0,3,1,2)--> skip_nchw
      pre_nhwc --T(0,3,1,2)--> pre_nchw
      b0_nhwc --T(0,3,1,2)--> b0_nchw
      rz_nhwc --T(0,3,1,2)--> rz_nchw --MUL(const)--ADD(const)--> b1_nchw
      CONCAT(axis=1,[b0_nchw,b1_nchw]) -> c1_nchw
      ADD(pre_nchw,c1_nchw)->MUL(const)->ADD(const)->PRELU -> p1_nchw
      CONCAT(axis=1,[skip_nchw,p1_nchw]) -> c2_nchw
      c2_nchw -> MUL(const) -> ADD(const) -> PRELU -> T(0,2,3,1)

    Rewrite the whole block to NHWC and remove redundant transpose bridges.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _ensure_channel_last_const(
        *,
        tensor_name: str,
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        data = np.asarray(tensor.data)
        if int(data.size) == 1:
            return str(tensor_name)
        if data.ndim != 4:
            return None
        shape = [int(v) for v in list(data.shape)]
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) > 0
        ):
            nhwc_data = np.asarray(data)
        elif (
            int(shape[0]) == 1
            and int(shape[1]) > 0
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            nhwc_data = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
        else:
            return None

        side_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        if shared_outside_chain:
            cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            return str(cloned_name)

        tensor.data = np.asarray(nhwc_data)
        tensor.shape = [int(v) for v in list(nhwc_data.shape)]
        tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
        return str(tensor_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post2_idx, post2_op in enumerate(model_ir.operators):
            if (
                str(post2_op.op_type) != "TRANSPOSE"
                or len(post2_op.inputs) < 2
                or len(post2_op.outputs) != 1
                or _read_transpose_perm(model_ir, post2_op) != perm_nchw_to_nhwc
            ):
                continue
            post2_input_name = str(post2_op.inputs[0])
            post2_output_name = str(post2_op.outputs[0])
            if post2_output_name in model_outputs:
                continue

            prelu2_idx = producers.get(post2_input_name, None)
            if prelu2_idx is None:
                continue
            prelu2_op = model_ir.operators[int(prelu2_idx)]
            if (
                str(prelu2_op.op_type) != "PRELU"
                or len(prelu2_op.inputs) != 2
                or len(prelu2_op.outputs) != 1
                or str(prelu2_op.outputs[0]) != str(post2_input_name)
            ):
                continue

            add2_idx = producers.get(str(prelu2_op.inputs[0]), None)
            if add2_idx is None:
                continue
            add2_op = model_ir.operators[int(add2_idx)]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_out_name = str(add2_op.outputs[0])
            if str(prelu2_op.inputs[0]) != str(add2_out_name):
                continue

            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])
            if add2_in0 == add2_in1:
                continue
            if add2_out_name in model_outputs:
                continue

            mul2_idx = producers.get(add2_in0, None)
            if mul2_idx is None:
                mul2_idx = producers.get(add2_in1, None)
            if mul2_idx is None:
                continue
            mul2_op = model_ir.operators[int(mul2_idx)]
            if str(mul2_op.op_type) != "MUL" or len(mul2_op.inputs) != 2 or len(mul2_op.outputs) != 1:
                continue
            mul2_out_name = str(mul2_op.outputs[0])
            if mul2_out_name == add2_in0:
                add2_side_index = 1
                add2_side_name = str(add2_in1)
            elif mul2_out_name == add2_in1:
                add2_side_index = 0
                add2_side_name = str(add2_in0)
            else:
                continue

            mul2_in0 = str(mul2_op.inputs[0])
            mul2_in1 = str(mul2_op.inputs[1])
            if mul2_in0 == mul2_in1:
                continue
            if mul2_out_name in model_outputs:
                continue

            if mul2_in0 in producers and str(model_ir.operators[int(producers[mul2_in0])].op_type) == "CONCATENATION":
                concat2_name = str(mul2_in0)
                mul2_side_index = 1
                mul2_side_name = str(mul2_in1)
            elif mul2_in1 in producers and str(model_ir.operators[int(producers[mul2_in1])].op_type) == "CONCATENATION":
                concat2_name = str(mul2_in1)
                mul2_side_index = 0
                mul2_side_name = str(mul2_in0)
            else:
                continue

            concat2_idx = producers.get(concat2_name, None)
            if concat2_idx is None:
                continue
            concat2_op = model_ir.operators[int(concat2_idx)]
            if (
                str(concat2_op.op_type) != "CONCATENATION"
                or len(concat2_op.inputs) != 2
                or len(concat2_op.outputs) != 1
                or str(concat2_op.outputs[0]) != str(concat2_name)
                or int(concat2_op.options.get("axis", 1)) != 1
            ):
                continue
            if set(int(v) for v in consumers.get(str(concat2_name), [])) != {int(mul2_idx)}:
                continue

            pre13_idx: Optional[int] = None
            pre13_in_name: Optional[str] = None
            pre13_out_name: Optional[str] = None
            prelu1_idx: Optional[int] = None
            prelu1_out_name: Optional[str] = None
            for c_in in [str(v) for v in list(concat2_op.inputs)]:
                prod_idx = producers.get(str(c_in), None)
                if prod_idx is None:
                    continue
                prod_op = model_ir.operators[int(prod_idx)]
                if (
                    str(prod_op.op_type) == "TRANSPOSE"
                    and len(prod_op.inputs) >= 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(c_in)
                    and _read_transpose_perm(model_ir, prod_op) == perm_nhwc_to_nchw
                    and set(int(v) for v in consumers.get(str(c_in), [])) == {int(concat2_idx)}
                    and str(c_in) not in model_outputs
                ):
                    pre13_idx = int(prod_idx)
                    pre13_in_name = str(prod_op.inputs[0])
                    pre13_out_name = str(c_in)
                elif str(prod_op.op_type) == "PRELU":
                    prelu1_idx = int(prod_idx)
                    prelu1_out_name = str(c_in)
            if (
                pre13_idx is None
                or pre13_in_name is None
                or pre13_out_name is None
                or prelu1_idx is None
                or prelu1_out_name is None
            ):
                continue

            prelu1_op = model_ir.operators[int(prelu1_idx)]
            if (
                str(prelu1_op.op_type) != "PRELU"
                or len(prelu1_op.inputs) != 2
                or len(prelu1_op.outputs) != 1
                or str(prelu1_op.outputs[0]) != str(prelu1_out_name)
                or str(prelu1_out_name) in model_outputs
            ):
                continue
            # strict variant: prelu1 only feeds concat2
            if set(int(v) for v in consumers.get(str(prelu1_out_name), [])) != {int(concat2_idx)}:
                continue

            add1_idx = producers.get(str(prelu1_op.inputs[0]), None)
            if add1_idx is None:
                continue
            add1_op = model_ir.operators[int(add1_idx)]
            if str(add1_op.op_type) != "ADD" or len(add1_op.inputs) != 2 or len(add1_op.outputs) != 1:
                continue
            add1_out_name = str(add1_op.outputs[0])
            if str(prelu1_op.inputs[0]) != str(add1_out_name):
                continue
            add1_in0 = str(add1_op.inputs[0])
            add1_in1 = str(add1_op.inputs[1])

            mul1_idx = producers.get(add1_in0, None)
            if mul1_idx is None:
                mul1_idx = producers.get(add1_in1, None)
            if mul1_idx is None:
                continue
            mul1_op = model_ir.operators[int(mul1_idx)]
            if str(mul1_op.op_type) != "MUL" or len(mul1_op.inputs) != 2 or len(mul1_op.outputs) != 1:
                continue
            mul1_out_name = str(mul1_op.outputs[0])
            if mul1_out_name == add1_in0:
                add1_side_index = 1
                add1_side_name = str(add1_in1)
            elif mul1_out_name == add1_in1:
                add1_side_index = 0
                add1_side_name = str(add1_in0)
            else:
                continue

            add0_name: Optional[str] = None
            mul1_side_index: Optional[int] = None
            mul1_side_name: Optional[str] = None
            for data_idx, side_idx in [(0, 1), (1, 0)]:
                cand_data = str(mul1_op.inputs[int(data_idx)])
                cand_side = str(mul1_op.inputs[int(side_idx)])
                cand_prod_idx = producers.get(cand_data, None)
                if cand_prod_idx is None:
                    continue
                cand_prod_op = model_ir.operators[int(cand_prod_idx)]
                if str(cand_prod_op.op_type) == "ADD":
                    add0_name = str(cand_data)
                    mul1_side_index = int(side_idx)
                    mul1_side_name = str(cand_side)
                    break
            if add0_name is None or mul1_side_index is None or mul1_side_name is None:
                continue

            add0_idx = producers.get(str(add0_name), None)
            if add0_idx is None:
                continue
            add0_op = model_ir.operators[int(add0_idx)]
            if (
                str(add0_op.op_type) != "ADD"
                or len(add0_op.inputs) != 2
                or len(add0_op.outputs) != 1
                or str(add0_op.outputs[0]) != str(add0_name)
                or str(add0_name) in model_outputs
            ):
                continue

            # add0 inputs: pre transpose + concat(axis=1)
            pre_idx: Optional[int] = None
            pre_in_name: Optional[str] = None
            pre_out_name: Optional[str] = None
            concat1_idx: Optional[int] = None
            concat1_name: Optional[str] = None
            for a_in in [str(v) for v in list(add0_op.inputs)]:
                a_prod_idx = producers.get(str(a_in), None)
                if a_prod_idx is None:
                    continue
                a_prod_op = model_ir.operators[int(a_prod_idx)]
                if (
                    str(a_prod_op.op_type) == "TRANSPOSE"
                    and len(a_prod_op.inputs) >= 2
                    and len(a_prod_op.outputs) == 1
                    and str(a_prod_op.outputs[0]) == str(a_in)
                    and _read_transpose_perm(model_ir, a_prod_op) == perm_nhwc_to_nchw
                    and set(int(v) for v in consumers.get(str(a_in), [])) == {int(add0_idx)}
                    and str(a_in) not in model_outputs
                ):
                    pre_idx = int(a_prod_idx)
                    pre_in_name = str(a_prod_op.inputs[0])
                    pre_out_name = str(a_in)
                elif (
                    str(a_prod_op.op_type) == "CONCATENATION"
                    and len(a_prod_op.inputs) == 2
                    and len(a_prod_op.outputs) == 1
                    and str(a_prod_op.outputs[0]) == str(a_in)
                    and int(a_prod_op.options.get("axis", 1)) == 1
                    and set(int(v) for v in consumers.get(str(a_in), [])) == {int(add0_idx)}
                    and str(a_in) not in model_outputs
                ):
                    concat1_idx = int(a_prod_idx)
                    concat1_name = str(a_in)
            if (
                pre_idx is None
                or pre_in_name is None
                or pre_out_name is None
                or concat1_idx is None
                or concat1_name is None
            ):
                continue

            concat1_op = model_ir.operators[int(concat1_idx)]
            trans_a_idx: Optional[int] = None
            trans_a_in_name: Optional[str] = None
            trans_a_out_name: Optional[str] = None
            addb_idx: Optional[int] = None
            addb_out_name: Optional[str] = None
            for c1_in in [str(v) for v in list(concat1_op.inputs)]:
                c1_prod_idx = producers.get(str(c1_in), None)
                if c1_prod_idx is None:
                    continue
                c1_prod_op = model_ir.operators[int(c1_prod_idx)]
                if (
                    str(c1_prod_op.op_type) == "TRANSPOSE"
                    and len(c1_prod_op.inputs) >= 2
                    and len(c1_prod_op.outputs) == 1
                    and str(c1_prod_op.outputs[0]) == str(c1_in)
                    and _read_transpose_perm(model_ir, c1_prod_op) == perm_nhwc_to_nchw
                    and set(int(v) for v in consumers.get(str(c1_in), [])) == {int(concat1_idx)}
                    and str(c1_in) not in model_outputs
                ):
                    trans_a_idx = int(c1_prod_idx)
                    trans_a_in_name = str(c1_prod_op.inputs[0])
                    trans_a_out_name = str(c1_in)
                elif (
                    str(c1_prod_op.op_type) == "ADD"
                    and len(c1_prod_op.inputs) == 2
                    and len(c1_prod_op.outputs) == 1
                    and str(c1_prod_op.outputs[0]) == str(c1_in)
                    and set(int(v) for v in consumers.get(str(c1_in), [])) == {int(concat1_idx)}
                    and str(c1_in) not in model_outputs
                ):
                    addb_idx = int(c1_prod_idx)
                    addb_out_name = str(c1_in)
            if (
                trans_a_idx is None
                or trans_a_in_name is None
                or trans_a_out_name is None
                or addb_idx is None
                or addb_out_name is None
            ):
                continue

            addb_op = model_ir.operators[int(addb_idx)]
            addb_in0 = str(addb_op.inputs[0])
            addb_in1 = str(addb_op.inputs[1])
            mulb_idx = producers.get(addb_in0, None)
            if mulb_idx is None:
                mulb_idx = producers.get(addb_in1, None)
            if mulb_idx is None:
                continue
            mulb_op = model_ir.operators[int(mulb_idx)]
            if str(mulb_op.op_type) != "MUL" or len(mulb_op.inputs) != 2 or len(mulb_op.outputs) != 1:
                continue
            mulb_out_name = str(mulb_op.outputs[0])
            if mulb_out_name == addb_in0:
                addb_side_index = 1
                addb_side_name = str(addb_in1)
            elif mulb_out_name == addb_in1:
                addb_side_index = 0
                addb_side_name = str(addb_in0)
            else:
                continue

            trans_b_idx: Optional[int] = None
            trans_b_in_name: Optional[str] = None
            trans_b_out_name: Optional[str] = None
            mulb_side_index: Optional[int] = None
            mulb_side_name: Optional[str] = None
            for data_idx, side_idx in [(0, 1), (1, 0)]:
                cand_data = str(mulb_op.inputs[int(data_idx)])
                cand_side = str(mulb_op.inputs[int(side_idx)])
                t_idx = producers.get(str(cand_data), None)
                if t_idx is None:
                    continue
                t_op = model_ir.operators[int(t_idx)]
                if (
                    str(t_op.op_type) != "TRANSPOSE"
                    or len(t_op.inputs) < 2
                    or len(t_op.outputs) != 1
                    or str(t_op.outputs[0]) != str(cand_data)
                    or _read_transpose_perm(model_ir, t_op) != perm_nhwc_to_nchw
                    or set(int(v) for v in consumers.get(str(cand_data), [])) != {int(mulb_idx)}
                    or str(cand_data) in model_outputs
                ):
                    continue
                rz_name = str(t_op.inputs[0])
                rz_idx = producers.get(rz_name, None)
                if rz_idx is None:
                    continue
                rz_op = model_ir.operators[int(rz_idx)]
                if str(rz_op.op_type) not in {"RESIZE_BILINEAR", "RESIZE_NEAREST_NEIGHBOR"}:
                    continue
                if len(rz_op.outputs) != 1 or str(rz_op.outputs[0]) != rz_name:
                    continue
                trans_b_idx = int(t_idx)
                trans_b_in_name = str(rz_name)
                trans_b_out_name = str(cand_data)
                mulb_side_index = int(side_idx)
                mulb_side_name = str(cand_side)
                break
            if (
                trans_b_idx is None
                or trans_b_in_name is None
                or trans_b_out_name is None
                or mulb_side_index is None
                or mulb_side_name is None
            ):
                continue

            chain_index_set = {
                int(pre13_idx),
                int(pre_idx),
                int(trans_a_idx),
                int(trans_b_idx),
                int(concat1_idx),
                int(addb_idx),
                int(mulb_idx),
                int(add0_idx),
                int(mul1_idx),
                int(add1_idx),
                int(prelu1_idx),
                int(concat2_idx),
                int(mul2_idx),
                int(add2_idx),
                int(prelu2_idx),
                int(post2_idx),
            }
            mulb_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mulb_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            addb_side_for_op = _ensure_channel_last_const(
                tensor_name=str(addb_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            mul1_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mul1_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            add1_side_for_op = _ensure_channel_last_const(
                tensor_name=str(add1_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            prelu1_side_for_op = _ensure_channel_last_const(
                tensor_name=str(prelu1_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            mul2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(mul2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            add2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(add2_side_name),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            prelu2_side_for_op = _ensure_channel_last_const(
                tensor_name=str(prelu2_op.inputs[1]),
                chain_index_set=chain_index_set,
                consumers=consumers,
            )
            if any(
                v is None
                for v in [
                    mulb_side_for_op,
                    addb_side_for_op,
                    mul1_side_for_op,
                    add1_side_for_op,
                    prelu1_side_for_op,
                    mul2_side_for_op,
                    add2_side_for_op,
                    prelu2_side_for_op,
                ]
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mulb_op,
                input_index=int(1 - int(mulb_side_index)),
                new_input_name=str(trans_b_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mulb_op,
                input_index=int(mulb_side_index),
                new_input_name=str(mulb_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=addb_op,
                input_index=int(addb_side_index),
                new_input_name=str(addb_side_for_op),
            )

            new_concat1_inputs = [
                str(trans_a_in_name) if str(v) == str(trans_a_out_name) else str(v)
                for v in list(concat1_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat1_op,
                new_inputs=[str(v) for v in new_concat1_inputs],
            )
            concat1_op.options["axis"] = 3

            _replace_operator_input_at(
                model_ir=model_ir,
                op=add0_op,
                input_index=0 if str(add0_op.inputs[0]) == str(pre_out_name) else 1,
                new_input_name=str(pre_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul1_op,
                input_index=int(mul1_side_index),
                new_input_name=str(mul1_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add1_op,
                input_index=int(add1_side_index),
                new_input_name=str(add1_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu1_op,
                input_index=1,
                new_input_name=str(prelu1_side_for_op),
            )

            new_concat2_inputs = [
                str(pre13_in_name) if str(v) == str(pre13_out_name) else str(v)
                for v in list(concat2_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat2_op,
                new_inputs=[str(v) for v in new_concat2_inputs],
            )
            concat2_op.options["axis"] = 3

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul2_op,
                input_index=int(mul2_side_index),
                new_input_name=str(mul2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add2_op,
                input_index=int(add2_side_index),
                new_input_name=str(add2_side_for_op),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=prelu2_op,
                input_index=1,
                new_input_name=str(prelu2_side_for_op),
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu2_op,
                new_outputs=[str(post2_output_name)],
            )

            for t_name in [
                str(mulb_out_name),
                str(addb_out_name),
                str(concat1_name),
                str(add0_name),
                str(mul1_out_name),
                str(add1_out_name),
                str(prelu1_out_name),
                str(concat2_name),
                str(mul2_out_name),
                str(add2_out_name),
                str(post2_output_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(t_name), None),
                    perm_nchw_to_nhwc,
                )

            old_prelu2_tensor = model_ir.tensors.get(str(post2_input_name), None)
            post2_tensor = model_ir.tensors.get(str(post2_output_name), None)
            if old_prelu2_tensor is not None and post2_tensor is not None:
                post2_tensor.dtype = str(old_prelu2_tensor.dtype)
                post2_tensor.quantization = _clone_quantization(old_prelu2_tensor.quantization)
                post2_tensor.shape = [int(v) for v in list(old_prelu2_tensor.shape)]
                post2_tensor.shape_signature = (
                    [int(v) for v in list(old_prelu2_tensor.shape_signature)]
                    if old_prelu2_tensor.shape_signature is not None
                    else [int(v) for v in list(old_prelu2_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post2_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {
                int(pre13_idx),
                int(pre_idx),
                int(trans_a_idx),
                int(trans_b_idx),
                int(post2_idx),
            }
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_sinet_concat_resize_affine_tail_concat_transpose_chains": int(optimized)}


def _optimize_sinet_softmax_mask_residual_nhwc_tail_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Optimize late SiNet tail blocks where NCHW adapters remain around
    Softmax+mask residual composition.

    Target sketch:
      side_nhwc --T(0,3,1,2)--> side_nchw --PRELU--> side_prelu_nchw
      main_nhwc --T(0,3,1,2)--> main_nchw --MUL(const)--ADD(const)--> bn_nchw
      bn_nchw --T(0,3,2,1)--> s_in --SOFTMAX--> s_out --T(0,3,2,1)--> soft_nchw
      REDUCE_MAX(axis=1)->SUB->RESHAPE([N,1,H,W])->MUL(ones)
      MUL(side_prelu_nchw, mask_nchw) + bn_nchw -> sum_nchw --T(0,2,3,1)--> sum_nhwc

    Rewrite:
      - keep side/main branches in NHWC
      - remove the three NCHW bridge transposes plus terminal pre-resize transpose
      - run SOFTMAX directly on NHWC tensor and remap reduce axis/reshape/consts
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nchw_to_nwhc = [0, 3, 2, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _set_const_int_data(
        *,
        tensor_name: str,
        new_values: List[int],
        target_consumer_idx: int,
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        users = [int(v) for v in consumers.get(str(tensor_name), [])]
        new_data = np.asarray([int(v) for v in list(new_values)], dtype=np.int32)
        if len(users) == 1 and int(users[0]) == int(target_consumer_idx):
            tensor.data = new_data
            tensor.dtype = "INT32"
            tensor.shape = [int(len(new_values))]
            tensor.shape_signature = [int(len(new_values))]
            return str(tensor_name)

        cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype="INT32",
            shape=[int(len(new_values))],
            shape_signature=[int(len(new_values))],
            data=new_data,
            is_variable=False,
            quantization=None,
        )
        return str(cloned_name)

    def _ensure_const_matches_target_nhwc(
        *,
        tensor_name: str,
        target_shape: Optional[List[int]],
        target_consumer_idx: int,
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        data = np.asarray(tensor.data)
        if int(data.size) == 1:
            return str(tensor_name)
        if data.ndim != 4:
            return None

        shape = [int(v) for v in list(data.shape)]
        nhwc_data: Optional[np.ndarray] = None
        if _is_fully_known_positive_shape(target_shape):
            target = [int(v) for v in list(target_shape)]
            if _broadcast_static_shapes(target, shape) is not None:
                nhwc_data = np.asarray(data)
            else:
                rotated = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target, rotated_shape) is None:
                    return None
                nhwc_data = np.asarray(rotated)
        else:
            # Unknown target: preserve idempotence and only rotate strict NCHW-per-channel constants.
            if (
                int(shape[0]) == 1
                and int(shape[1]) > 0
                and int(shape[2]) == 1
                and int(shape[3]) == 1
            ):
                nhwc_data = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
            elif (
                int(shape[0]) == 1
                and int(shape[1]) == 1
                and int(shape[2]) == 1
                and int(shape[3]) > 0
            ):
                nhwc_data = np.asarray(data)
            else:
                return None

        users = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(users) == 1 and int(users[0]) == int(target_consumer_idx):
            tensor.data = np.asarray(nhwc_data)
            tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
            return str(tensor_name)

        cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(tensor.dtype),
            shape=[int(v) for v in list(nhwc_data.shape)],
            shape_signature=[int(v) for v in list(nhwc_data.shape)],
            data=np.asarray(nhwc_data),
            is_variable=False,
            quantization=_clone_quantization(tensor.quantization),
        )
        return str(cloned_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue

            post_in = str(post_op.inputs[0])
            post_out = str(post_op.outputs[0])
            if post_out in model_outputs:
                continue

            add_res_idx = producers.get(str(post_in), None)
            if add_res_idx is None:
                continue
            add_res_op = model_ir.operators[int(add_res_idx)]
            if (
                str(add_res_op.op_type) != "ADD"
                or len(add_res_op.inputs) != 2
                or len(add_res_op.outputs) != 1
                or str(add_res_op.outputs[0]) != str(post_in)
            ):
                continue

            # Parse add_res inputs:
            #   - main_nchw (ADD from BN affine)
            #   - side_mask_nchw (MUL from side-prelu and mask)
            add_res_inputs = [str(v) for v in list(add_res_op.inputs)]
            main_name: Optional[str] = None
            side_mask_name: Optional[str] = None
            main_add_idx: Optional[int] = None
            side_mask_idx: Optional[int] = None
            for inp in add_res_inputs:
                inp_prod_idx = producers.get(str(inp), None)
                if inp_prod_idx is None:
                    continue
                inp_prod_op = model_ir.operators[int(inp_prod_idx)]
                if str(inp_prod_op.op_type) == "ADD":
                    # Candidate: BN affine result
                    a_in0 = str(inp_prod_op.inputs[0])
                    a_in1 = str(inp_prod_op.inputs[1])
                    m_idx = producers.get(a_in0, None)
                    if m_idx is None:
                        m_idx = producers.get(a_in1, None)
                    if m_idx is None:
                        continue
                    m_op = model_ir.operators[int(m_idx)]
                    if str(m_op.op_type) != "MUL" or len(m_op.inputs) != 2:
                        continue
                    if str(m_op.outputs[0]) not in {a_in0, a_in1}:
                        continue
                    main_name = str(inp)
                    main_add_idx = int(inp_prod_idx)
                elif str(inp_prod_op.op_type) == "MUL":
                    side_mask_name = str(inp)
                    side_mask_idx = int(inp_prod_idx)
            if (
                main_name is None
                or side_mask_name is None
                or main_add_idx is None
                or side_mask_idx is None
            ):
                continue

            main_add_op = model_ir.operators[int(main_add_idx)]
            main_add_in0 = str(main_add_op.inputs[0])
            main_add_in1 = str(main_add_op.inputs[1])
            main_mul_idx = producers.get(main_add_in0, None)
            if main_mul_idx is None:
                main_mul_idx = producers.get(main_add_in1, None)
            if main_mul_idx is None:
                continue
            main_mul_op = model_ir.operators[int(main_mul_idx)]
            if str(main_mul_op.op_type) != "MUL" or len(main_mul_op.inputs) != 2:
                continue
            if str(main_mul_op.outputs[0]) == str(main_add_in0):
                main_add_side_idx = 1
                main_add_side_name = str(main_add_in1)
            elif str(main_mul_op.outputs[0]) == str(main_add_in1):
                main_add_side_idx = 0
                main_add_side_name = str(main_add_in0)
            else:
                continue

            main_pre_idx: Optional[int] = None
            main_pre_in_name: Optional[str] = None
            main_pre_out_name: Optional[str] = None
            main_mul_side_idx: Optional[int] = None
            main_mul_side_name: Optional[str] = None
            for data_idx, side_idx in [(0, 1), (1, 0)]:
                cand_data = str(main_mul_op.inputs[int(data_idx)])
                cand_side = str(main_mul_op.inputs[int(side_idx)])
                cand_pre_idx = producers.get(str(cand_data), None)
                if cand_pre_idx is None:
                    continue
                cand_pre_op = model_ir.operators[int(cand_pre_idx)]
                if (
                    str(cand_pre_op.op_type) != "TRANSPOSE"
                    or len(cand_pre_op.inputs) < 2
                    or len(cand_pre_op.outputs) != 1
                    or str(cand_pre_op.outputs[0]) != str(cand_data)
                    or _read_transpose_perm(model_ir, cand_pre_op) != perm_nhwc_to_nchw
                    or set(int(v) for v in consumers.get(str(cand_data), [])) != {int(main_mul_idx)}
                    or str(cand_data) in model_outputs
                ):
                    continue
                main_pre_idx = int(cand_pre_idx)
                main_pre_in_name = str(cand_pre_op.inputs[0])
                main_pre_out_name = str(cand_data)
                main_mul_side_idx = int(side_idx)
                main_mul_side_name = str(cand_side)
                break
            if (
                main_pre_idx is None
                or main_pre_in_name is None
                or main_pre_out_name is None
                or main_mul_side_idx is None
                or main_mul_side_name is None
            ):
                continue

            # Softmax chain on main_name.
            main_users = set(int(v) for v in consumers.get(str(main_name), []))
            if int(add_res_idx) not in main_users or len(main_users) != 2:
                continue
            soft_pre_idx = next((int(v) for v in main_users if int(v) != int(add_res_idx)), None)
            if soft_pre_idx is None:
                continue
            soft_pre_op = model_ir.operators[int(soft_pre_idx)]
            if (
                str(soft_pre_op.op_type) != "TRANSPOSE"
                or len(soft_pre_op.inputs) < 2
                or len(soft_pre_op.outputs) != 1
                or str(soft_pre_op.inputs[0]) != str(main_name)
                or _read_transpose_perm(model_ir, soft_pre_op) != perm_nchw_to_nwhc
            ):
                continue
            soft_pre_out = str(soft_pre_op.outputs[0])
            softmax_idx = consumers.get(str(soft_pre_out), [])
            if len(softmax_idx) != 1:
                continue
            softmax_idx = int(softmax_idx[0])
            softmax_op = model_ir.operators[int(softmax_idx)]
            if (
                str(softmax_op.op_type) != "SOFTMAX"
                or len(softmax_op.inputs) != 1
                or len(softmax_op.outputs) != 1
                or str(softmax_op.inputs[0]) != str(soft_pre_out)
            ):
                continue
            softmax_out = str(softmax_op.outputs[0])
            soft_post_users = [int(v) for v in consumers.get(str(softmax_out), [])]
            if len(soft_post_users) != 1:
                continue
            soft_post_idx = int(soft_post_users[0])
            soft_post_op = model_ir.operators[int(soft_post_idx)]
            if (
                str(soft_post_op.op_type) != "TRANSPOSE"
                or len(soft_post_op.inputs) < 2
                or len(soft_post_op.outputs) != 1
                or str(soft_post_op.inputs[0]) != str(softmax_out)
                or _read_transpose_perm(model_ir, soft_post_op) != perm_nchw_to_nwhc
            ):
                continue
            soft_back_name = str(soft_post_op.outputs[0])

            reduce_users = [int(v) for v in consumers.get(str(soft_back_name), [])]
            if len(reduce_users) != 1:
                continue
            reduce_idx = int(reduce_users[0])
            reduce_op = model_ir.operators[int(reduce_idx)]
            if (
                str(reduce_op.op_type) != "REDUCE_MAX"
                or len(reduce_op.inputs) < 2
                or len(reduce_op.outputs) != 1
                or str(reduce_op.inputs[0]) != str(soft_back_name)
                or bool(reduce_op.options.get("keepDims", False))
            ):
                continue
            reduce_axis_name = str(reduce_op.inputs[1])
            reduce_axis_tensor = model_ir.tensors.get(str(reduce_axis_name), None)
            reduce_axis_vals = _read_const_ints_from_tensor(reduce_axis_tensor)
            if reduce_axis_vals is None or list(reduce_axis_vals) != [1]:
                continue
            reduce_out = str(reduce_op.outputs[0])

            sub_users = [int(v) for v in consumers.get(str(reduce_out), [])]
            if len(sub_users) != 1:
                continue
            sub_idx = int(sub_users[0])
            sub_op = model_ir.operators[int(sub_idx)]
            if str(sub_op.op_type) != "SUB" or len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue
            sub_in0 = str(sub_op.inputs[0])
            sub_in1 = str(sub_op.inputs[1])
            if sub_in1 == str(reduce_out):
                if not _is_singleton_constant_tensor(model_ir, sub_in0):
                    continue
            elif sub_in0 == str(reduce_out):
                if not _is_singleton_constant_tensor(model_ir, sub_in1):
                    continue
            else:
                continue
            sub_out = str(sub_op.outputs[0])

            reshape_users = [int(v) for v in consumers.get(str(sub_out), [])]
            if len(reshape_users) != 1:
                continue
            reshape_idx = int(reshape_users[0])
            reshape_op = model_ir.operators[int(reshape_idx)]
            if (
                str(reshape_op.op_type) != "RESHAPE"
                or len(reshape_op.inputs) < 2
                or len(reshape_op.outputs) != 1
                or str(reshape_op.inputs[0]) != str(sub_out)
            ):
                continue
            reshape_shape_name = str(reshape_op.inputs[1])
            reshape_shape_tensor = model_ir.tensors.get(str(reshape_shape_name), None)
            reshape_vals = _read_const_ints_from_tensor(reshape_shape_tensor)
            if reshape_vals is None or len(reshape_vals) != 4:
                continue
            if int(reshape_vals[0]) != 1 or int(reshape_vals[1]) != 1:
                continue
            reshape_out = str(reshape_op.outputs[0])

            expand_users = [int(v) for v in consumers.get(str(reshape_out), [])]
            if len(expand_users) != 1:
                continue
            expand_idx = int(expand_users[0])
            expand_op = model_ir.operators[int(expand_idx)]
            if str(expand_op.op_type) != "MUL" or len(expand_op.inputs) != 2 or len(expand_op.outputs) != 1:
                continue
            expand_in0 = str(expand_op.inputs[0])
            expand_in1 = str(expand_op.inputs[1])
            if expand_in0 == str(reshape_out):
                expand_const_idx = 1
                expand_const_name = str(expand_in1)
            elif expand_in1 == str(reshape_out):
                expand_const_idx = 0
                expand_const_name = str(expand_in0)
            else:
                continue
            expand_const_tensor = model_ir.tensors.get(str(expand_const_name), None)
            if expand_const_tensor is None or expand_const_tensor.data is None:
                continue
            expand_out = str(expand_op.outputs[0])

            side_mask_op = model_ir.operators[int(side_mask_idx)]
            side_mask_in0 = str(side_mask_op.inputs[0])
            side_mask_in1 = str(side_mask_op.inputs[1])
            if side_mask_in0 == str(expand_out):
                side_tensor_name = str(side_mask_in1)
            elif side_mask_in1 == str(expand_out):
                side_tensor_name = str(side_mask_in0)
            else:
                continue

            side_prelu_idx = producers.get(str(side_tensor_name), None)
            if side_prelu_idx is None:
                continue
            side_prelu_op = model_ir.operators[int(side_prelu_idx)]
            if (
                str(side_prelu_op.op_type) != "PRELU"
                or len(side_prelu_op.inputs) != 2
                or len(side_prelu_op.outputs) != 1
                or str(side_prelu_op.outputs[0]) != str(side_tensor_name)
            ):
                continue
            side_pre_in = str(side_prelu_op.inputs[0])
            side_pre_idx = producers.get(str(side_pre_in), None)
            if side_pre_idx is None:
                continue
            side_pre_op = model_ir.operators[int(side_pre_idx)]
            if (
                str(side_pre_op.op_type) != "TRANSPOSE"
                or len(side_pre_op.inputs) < 2
                or len(side_pre_op.outputs) != 1
                or str(side_pre_op.outputs[0]) != str(side_pre_in)
                or _read_transpose_perm(model_ir, side_pre_op) != perm_nhwc_to_nchw
                or set(int(v) for v in consumers.get(str(side_pre_in), [])) != {int(side_prelu_idx)}
                or str(side_pre_in) in model_outputs
            ):
                continue
            side_pre_in_name = str(side_pre_op.inputs[0])

            # Locality and output guards.
            if any(
                name in model_outputs
                for name in [
                    str(main_name),
                    str(side_tensor_name),
                    str(soft_pre_out),
                    str(softmax_out),
                    str(soft_back_name),
                    str(post_in),
                ]
            ):
                continue

            main_nhwc_shape = (
                list(model_ir.tensors[str(main_pre_in_name)].shape)
                if str(main_pre_in_name) in model_ir.tensors
                else None
            )
            side_nhwc_shape = (
                list(model_ir.tensors[str(side_pre_in_name)].shape)
                if str(side_pre_in_name) in model_ir.tensors
                else None
            )

            main_mul_const = _ensure_const_matches_target_nhwc(
                tensor_name=str(main_mul_side_name),
                target_shape=main_nhwc_shape,
                target_consumer_idx=int(main_mul_idx),
                consumers=consumers,
            )
            main_add_const = _ensure_const_matches_target_nhwc(
                tensor_name=str(main_add_side_name),
                target_shape=main_nhwc_shape,
                target_consumer_idx=int(main_add_idx),
                consumers=consumers,
            )
            side_prelu_const = _ensure_const_matches_target_nhwc(
                tensor_name=str(side_prelu_op.inputs[1]),
                target_shape=side_nhwc_shape,
                target_consumer_idx=int(side_prelu_idx),
                consumers=consumers,
            )
            expand_const_nhwc = _ensure_const_matches_target_nhwc(
                tensor_name=str(expand_const_name),
                target_shape=side_nhwc_shape,
                target_consumer_idx=int(expand_idx),
                consumers=consumers,
            )
            if any(v is None for v in [main_mul_const, main_add_const, side_prelu_const, expand_const_nhwc]):
                continue

            new_reduce_axis_name = _set_const_int_data(
                tensor_name=str(reduce_axis_name),
                new_values=[3],
                target_consumer_idx=int(reduce_idx),
                consumers=consumers,
            )
            if new_reduce_axis_name is None:
                continue
            new_reshape_shape_name = _set_const_int_data(
                tensor_name=str(reshape_shape_name),
                new_values=[int(reshape_vals[0]), int(reshape_vals[2]), int(reshape_vals[3]), int(reshape_vals[1])],
                target_consumer_idx=int(reshape_idx),
                consumers=consumers,
            )
            if new_reshape_shape_name is None:
                continue

            # Rewrite side branch to NHWC.
            _set_operator_inputs(
                model_ir=model_ir,
                op=side_prelu_op,
                new_inputs=[str(side_pre_in_name), str(side_prelu_const)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(side_tensor_name), None),
                perm_nchw_to_nhwc,
            )

            # Rewrite main branch to NHWC.
            _replace_operator_input_at(
                model_ir=model_ir,
                op=main_mul_op,
                input_index=int(1 - int(main_mul_side_idx)),
                new_input_name=str(main_pre_in_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=main_mul_op,
                input_index=int(main_mul_side_idx),
                new_input_name=str(main_mul_const),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=main_add_op,
                input_index=int(main_add_side_idx),
                new_input_name=str(main_add_const),
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(main_mul_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(main_name), None),
                perm_nchw_to_nhwc,
            )

            # Remove softmax wrapper transposes.
            _set_operator_inputs(
                model_ir=model_ir,
                op=softmax_op,
                new_inputs=[str(main_name)],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=softmax_op,
                new_outputs=[str(soft_back_name)],
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=reduce_op,
                input_index=1,
                new_input_name=str(new_reduce_axis_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=reshape_op,
                input_index=1,
                new_input_name=str(new_reshape_shape_name),
            )
            reshape_opts = dict(reshape_op.options) if isinstance(reshape_op.options, dict) else {}
            for shape_key in ["newShape", "onnxRawNewShape"]:
                shape_val = reshape_opts.get(shape_key, None)
                if isinstance(shape_val, list) and len(shape_val) == 4:
                    reshape_opts[shape_key] = [
                        int(shape_val[0]),
                        int(shape_val[2]),
                        int(shape_val[3]),
                        int(shape_val[1]),
                    ]
            reshape_op.options = reshape_opts
            _replace_operator_input_at(
                model_ir=model_ir,
                op=expand_op,
                input_index=int(expand_const_idx),
                new_input_name=str(expand_const_nhwc),
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(soft_back_name), None),
                perm_nchw_to_nhwc,
            )

            # Keep mask/residual in NHWC and bypass terminal transpose.
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(expand_out), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(side_mask_name), None),
                perm_nchw_to_nhwc,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_res_op,
                new_outputs=[str(post_out)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(post_out), None),
                perm_nchw_to_nhwc,
            )

            remove_indices = {
                int(side_pre_idx),
                int(main_pre_idx),
                int(soft_pre_idx),
                int(soft_post_idx),
                int(post_idx),
            }
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_sinet_softmax_mask_residual_nhwc_tail_chains": int(optimized)}


def _optimize_transpose_mul_add_const_prelu_prepost_nhwc_terminal_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Terminal rewrite for:
      TRANSPOSE(NHWC->NCHW) -> MUL(const) -> ADD(const) -> PRELU(const) -> TRANSPOSE(NCHW->NHWC)

    This pass is intentionally strict and runs only in the terminal stage to avoid
    interacting with earlier graph-wide transpose optimizers.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _ensure_const_nhwc_channel_last(
        *,
        tensor_name: str,
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        data = np.asarray(tensor.data)
        if int(data.size) == 1:
            return str(tensor_name)
        if data.ndim != 4:
            return None
        shape = [int(v) for v in list(data.shape)]
        nhwc_data: Optional[np.ndarray] = None
        if (
            int(shape[0]) == 1
            and int(shape[1]) == 1
            and int(shape[2]) == 1
            and int(shape[3]) > 0
        ):
            nhwc_data = np.asarray(data)
        elif (
            int(shape[0]) == 1
            and int(shape[1]) > 0
            and int(shape[2]) == 1
            and int(shape[3]) == 1
        ):
            nhwc_data = np.transpose(data, perm_nchw_to_nhwc).astype(data.dtype, copy=False)
        else:
            return None

        side_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)
        if shared_outside_chain:
            cloned_name = _unique_tensor_name(f"{tensor_name}_nhwc")
            model_ir.tensors[cloned_name] = TensorIR(
                name=cloned_name,
                dtype=str(tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            return str(cloned_name)

        tensor.data = np.asarray(nhwc_data)
        tensor.shape = [int(v) for v in list(nhwc_data.shape)]
        tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]
        return str(tensor_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_outputs or pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue
            for mul_idx in pre_users:
                mul_op = model_ir.operators[int(mul_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    continue
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                if mul_inputs[0] == pre_output_name:
                    mul_side_input_index = 1
                    mul_side_name = str(mul_inputs[1])
                elif mul_inputs[1] == pre_output_name:
                    mul_side_input_index = 0
                    mul_side_name = str(mul_inputs[0])
                else:
                    continue

                mul_out_name = str(mul_op.outputs[0])
                if mul_out_name in model_outputs:
                    continue
                mul_out_users = [int(v) for v in consumers.get(mul_out_name, []) if int(v) != int(mul_idx)]
                if len(mul_out_users) != 1:
                    continue
                add_idx = int(mul_out_users[0])
                add_op = model_ir.operators[int(add_idx)]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    continue
                add_inputs = [str(v) for v in list(add_op.inputs)]
                if add_inputs[0] == mul_out_name:
                    add_side_input_index = 1
                    add_side_name = str(add_inputs[1])
                elif add_inputs[1] == mul_out_name:
                    add_side_input_index = 0
                    add_side_name = str(add_inputs[0])
                else:
                    continue

                add_out_name = str(add_op.outputs[0])
                if add_out_name in model_outputs:
                    continue
                add_out_users = [int(v) for v in consumers.get(add_out_name, []) if int(v) != int(add_idx)]
                if len(add_out_users) != 1:
                    continue
                prelu_idx = int(add_out_users[0])
                prelu_op = model_ir.operators[int(prelu_idx)]
                if (
                    str(prelu_op.op_type) != "PRELU"
                    or len(prelu_op.inputs) != 2
                    or len(prelu_op.outputs) != 1
                    or str(prelu_op.inputs[0]) != add_out_name
                ):
                    continue
                prelu_side_name = str(prelu_op.inputs[1])
                prelu_out_name = str(prelu_op.outputs[0])
                if prelu_out_name in model_outputs:
                    continue

                prelu_users = [int(v) for v in consumers.get(prelu_out_name, []) if int(v) != int(prelu_idx)]
                if len(prelu_users) == 0:
                    continue
                post_indices: List[int] = []
                post_output_names: List[str] = []
                legacy_users: List[int] = []
                for user_idx in prelu_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == prelu_out_name
                        and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        post_indices.append(int(user_idx))
                        post_output_names.append(str(user_op.outputs[0]))
                    else:
                        legacy_users.append(int(user_idx))
                if len(post_indices) == 0:
                    continue

                chain_index_set = {int(mul_idx), int(add_idx), int(prelu_idx)}
                mul_side_for_op = _ensure_const_nhwc_channel_last(
                    tensor_name=str(mul_side_name),
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                )
                if mul_side_for_op is None:
                    continue
                add_side_for_op = _ensure_const_nhwc_channel_last(
                    tensor_name=str(add_side_name),
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                )
                if add_side_for_op is None:
                    continue
                prelu_side_for_op = _ensure_const_nhwc_channel_last(
                    tensor_name=str(prelu_side_name),
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                )
                if prelu_side_for_op is None:
                    continue

                mul_inputs[int(mul_side_input_index)] = str(mul_side_for_op)
                if mul_inputs[0] == pre_output_name:
                    mul_inputs[0] = str(pre_input_name)
                elif mul_inputs[1] == pre_output_name:
                    mul_inputs[1] = str(pre_input_name)
                else:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=add_op,
                    input_index=int(add_side_input_index),
                    new_input_name=str(add_side_for_op),
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=prelu_op,
                    input_index=1,
                    new_input_name=str(prelu_side_for_op),
                )

                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(mul_out_name, None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(add_out_name, None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(prelu_out_name, None),
                    perm_nchw_to_nhwc,
                )

                canonical_post_output_name = str(post_output_names[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=prelu_op,
                    new_outputs=[canonical_post_output_name],
                )
                for alias_name in post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

                old_prelu_tensor = model_ir.tensors.get(prelu_out_name, None)
                canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
                if old_prelu_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_prelu_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_prelu_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_prelu_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_prelu_tensor.shape_signature)]
                        if old_prelu_tensor.shape_signature is not None
                        else [int(v) for v in list(old_prelu_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        canonical_tensor,
                        perm_nchw_to_nhwc,
                    )

                if len(legacy_users) > 0:
                    keep_post_idx = int(post_indices[0])
                    keep_post_op = model_ir.operators[int(keep_post_idx)]
                    keep_perm_name = str(keep_post_op.inputs[1])
                    keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                    if keep_perm_tensor is not None:
                        keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=keep_post_op,
                        new_inputs=[canonical_post_output_name, keep_perm_name],
                    )
                    _set_operator_outputs(
                        model_ir=model_ir,
                        op=keep_post_op,
                        new_outputs=[prelu_out_name],
                    )
                    post_remove_indices = [int(v) for v in post_indices[1:]]
                else:
                    post_remove_indices = [int(v) for v in post_indices]

                remove_indices = set(int(v) for v in post_remove_indices)
                pre_remaining_users = [int(v) for v in pre_users if int(v) != int(mul_idx)]
                if len(pre_remaining_users) == 0:
                    remove_indices.add(int(pre_idx))
                for remove_idx in sorted(list(remove_indices), reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_mul_add_const_prelu_prepost_nhwc_terminal_chains": int(rewritten)}


def _optimize_transpose_mean_mul_add_const_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around TRANSPOSE->MEAN->MUL(const)->ADD(const)->TRANSPOSE chains.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --MEAN(axes_nchw, keepDims=True)--> m_nchw
      m_nchw --MUL(const)--> u_nchw --ADD(const)--> y_nchw
      y_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      x_nhwc --MEAN(axes_nhwc, keepDims=True)--> m_nhwc
      m_nhwc --MUL(const_nhwc)--> u_nhwc --ADD(const_nhwc)--> y_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _rewrite_binary_const_to_nhwc(
        *,
        op_idx: int,
        side_input_index: int,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
        chain_index_set: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False

        nhwc_data: Optional[np.ndarray] = None
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is not None:
            side_shape = [int(v) for v in list(side_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is not None:
                # Already NHWC-compatible. Keep idempotent across repeated sweeps.
                return True
            rotated = np.asarray(side_data)
            for _ in range(3):
                rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                    nhwc_data = np.asarray(rotated)
                    break
            if nhwc_data is None:
                return False
        else:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 1
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                nhwc_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            else:
                return True

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)

        target_name = str(side_input_name)
        if shared_outside_chain:
            target_name = _unique_tensor_name(f"{side_input_name}_nhwc")
            model_ir.tensors[target_name] = TensorIR(
                name=target_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
        else:
            side_tensor.data = np.asarray(nhwc_data)
            side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        _replace_operator_input_at(
            model_ir=model_ir,
            op=model_ir.operators[int(op_idx)],
            input_index=int(side_input_index),
            new_input_name=str(target_name),
        )
        return True

    def _can_rewrite_binary_const_to_nhwc(
        *,
        side_input_name: str,
        target_broadcast_shape: Optional[List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False
        target_shape = (
            [int(v) for v in list(target_broadcast_shape)]
            if _is_fully_known_positive_shape(target_broadcast_shape)
            else None
        )
        if target_shape is None:
            return True
        side_shape = [int(v) for v in list(side_data.shape)]
        if _broadcast_static_shapes(target_shape, side_shape) is not None:
            return True
        rotated = np.asarray(side_data)
        for _ in range(3):
            rotated = np.transpose(rotated, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated.shape)]
            if _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                return True
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_outputs or pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue

            for mean_idx in pre_users:
                mean_op = model_ir.operators[int(mean_idx)]
                if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                    continue
                if str(mean_op.inputs[0]) != pre_output_name:
                    continue
                if not bool(mean_op.options.get("keepDims", False)):
                    continue

                mean_axes_name = str(mean_op.inputs[1])
                mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                mean_axes = _read_const_ints_from_tensor(mean_axes_tensor)
                if mean_axes is None or len(mean_axes) == 0:
                    continue

                mean_output_name = str(mean_op.outputs[0])
                if mean_output_name in model_outputs:
                    continue
                mul_users = [int(v) for v in consumers.get(mean_output_name, [])]
                if len(mul_users) != 1:
                    continue
                mul_idx = int(mul_users[0])
                mul_op = model_ir.operators[mul_idx]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    continue

                mul_in0 = str(mul_op.inputs[0])
                mul_in1 = str(mul_op.inputs[1])
                if mul_in0 == mean_output_name:
                    mul_side_input_index = 1
                    mul_side_name = mul_in1
                elif mul_in1 == mean_output_name:
                    mul_side_input_index = 0
                    mul_side_name = mul_in0
                else:
                    continue

                add_input_name = str(mul_op.outputs[0])
                if add_input_name in model_outputs:
                    continue
                add_users = [int(v) for v in consumers.get(add_input_name, [])]
                if len(add_users) != 1:
                    continue
                add_idx = int(add_users[0])
                add_op = model_ir.operators[add_idx]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    continue

                add_in0 = str(add_op.inputs[0])
                add_in1 = str(add_op.inputs[1])
                if add_in0 == add_input_name:
                    add_side_input_index = 1
                    add_side_name = add_in1
                elif add_in1 == add_input_name:
                    add_side_input_index = 0
                    add_side_name = add_in0
                else:
                    continue

                add_output_name = str(add_op.outputs[0])
                if add_output_name in model_outputs:
                    continue
                add_output_users = [int(v) for v in consumers.get(add_output_name, [])]
                if len(add_output_users) == 0:
                    continue

                post_indices: List[int] = []
                post_output_names: List[str] = []
                valid_posts = True
                for user_idx in add_output_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == add_output_name
                        and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        post_indices.append(int(user_idx))
                        post_output_names.append(str(user_op.outputs[0]))
                    else:
                        valid_posts = False
                        break
                if not valid_posts or len(post_indices) == 0:
                    continue

                # Normalize NCHW-side axes and map them to NHWC-side axes.
                mapped_axes: List[int] = []
                valid_axes = True
                for axis in mean_axes:
                    a = int(axis)
                    if a < 0:
                        a += 4
                    if a < 0 or a >= 4:
                        valid_axes = False
                        break
                    mapped_axes.append(int(perm_nchw_to_nhwc[int(a)]))
                if not valid_axes or len(mapped_axes) == 0:
                    continue

                chain_index_set = {int(mean_idx), int(mul_idx), int(add_idx)}
                target_nhwc_shape = (
                    list(model_ir.tensors[pre_input_name].shape)
                    if pre_input_name in model_ir.tensors
                    else None
                )
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(mul_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    continue
                if not _can_rewrite_binary_const_to_nhwc(
                    side_input_name=str(add_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                ):
                    continue
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(mul_idx),
                    side_input_index=int(mul_side_input_index),
                    side_input_name=str(mul_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    continue
                if not _rewrite_binary_const_to_nhwc(
                    op_idx=int(add_idx),
                    side_input_index=int(add_side_input_index),
                    side_input_name=str(add_side_name),
                    target_broadcast_shape=target_nhwc_shape,
                    chain_index_set=chain_index_set,
                    consumers=consumers,
                ):
                    continue

                if not _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes]):
                    # Keep going even when unchanged; rewrite still valid.
                    pass

                mean_inputs = [str(v) for v in list(mean_op.inputs)]
                mean_inputs[0] = pre_input_name
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mean_op,
                    new_inputs=mean_inputs,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(mean_output_name, None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(add_input_name, None),
                    perm_nchw_to_nhwc,
                )

                canonical_post_output_name = str(post_output_names[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_outputs=[canonical_post_output_name],
                )
                for alias_name in post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

                old_add_out_tensor = model_ir.tensors.get(add_output_name, None)
                canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
                if old_add_out_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_add_out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_add_out_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_add_out_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_add_out_tensor.shape_signature)]
                        if old_add_out_tensor.shape_signature is not None
                        else [int(v) for v in list(old_add_out_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        canonical_tensor,
                        perm_nchw_to_nhwc,
                    )

                remove_indices = set(int(v) for v in post_indices)
                # Drop pre transpose only when this rewrite consumed its last user.
                pre_remaining_users = [int(v) for v in pre_users if int(v) != int(mean_idx)]
                if len(pre_remaining_users) == 0:
                    remove_indices.add(int(pre_idx))
                for remove_idx in sorted(remove_indices, reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_mean_mul_add_const_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around TRANSPOSE->MEAN->TRANSPOSE chains.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --MEAN(axes_nchw, keepDims=True)--> m_nchw
      m_nchw --TRANSPOSE(0,2,3,1)--> m_nhwc

    Rewrite:
      x_nhwc --MEAN(axes_nhwc, keepDims=True)--> m_nhwc

    Safety:
    - MEAN output is consumed only by inverse post-transpose adapters.
    - Rank-4 and constant MEAN axes.
    - When pre-transpose has extra consumers, keep that transpose and bypass only the MEAN branch.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue

            for mean_idx in pre_users:
                mean_op = model_ir.operators[int(mean_idx)]
                if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                    continue
                if str(mean_op.inputs[0]) != pre_output_name:
                    continue
                if not bool(mean_op.options.get("keepDims", False)):
                    continue

                mean_axes_name = str(mean_op.inputs[1])
                mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
                if mean_axes_vals is None or len(mean_axes_vals) == 0:
                    continue

                rank = 4
                pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
                if pre_input_tensor is not None and pre_input_tensor.shape is not None and len(pre_input_tensor.shape) > 0:
                    rank = int(len(list(pre_input_tensor.shape)))
                if rank != 4:
                    continue

                normalized_axes: List[int] = []
                valid_axes = True
                for axis in mean_axes_vals:
                    a = int(axis)
                    if a < 0:
                        a += int(rank)
                    if a < 0 or a >= int(rank):
                        valid_axes = False
                        break
                    normalized_axes.append(int(a))
                if not valid_axes:
                    continue
                mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]

                mean_out_name = str(mean_op.outputs[0])
                if mean_out_name in model_outputs:
                    continue
                mean_users = [int(v) for v in consumers.get(mean_out_name, [])]
                if len(mean_users) == 0:
                    continue

                post_indices: List[int] = []
                post_output_names: List[str] = []
                valid_posts = True
                for user_idx in mean_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) != "TRANSPOSE"
                        or len(user_op.inputs) < 2
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != mean_out_name
                        or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    ):
                        valid_posts = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                if not valid_posts or len(post_indices) == 0:
                    continue

                # Commit axis remap only when this rewrite is actually applicable.
                _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes])

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mean_op,
                    new_inputs=[pre_input_name, mean_axes_name],
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(mean_out_name, None),
                    perm_nchw_to_nhwc,
                )

                representative_output_name = str(post_output_names[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=mean_op,
                    new_outputs=[representative_output_name],
                )
                for alias_name in post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, str(alias_name), representative_output_name)

                old_mean_tensor = model_ir.tensors.get(mean_out_name, None)
                representative_tensor = model_ir.tensors.get(representative_output_name, None)
                if old_mean_tensor is not None and representative_tensor is not None:
                    representative_tensor.dtype = str(old_mean_tensor.dtype)
                    representative_tensor.quantization = _clone_quantization(old_mean_tensor.quantization)
                    representative_tensor.shape = [int(v) for v in list(old_mean_tensor.shape)]
                    representative_tensor.shape_signature = (
                        [int(v) for v in list(old_mean_tensor.shape_signature)]
                        if old_mean_tensor.shape_signature is not None
                        else [int(v) for v in list(old_mean_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        representative_tensor,
                        perm_nchw_to_nhwc,
                    )

                remove_indices = set(int(v) for v in post_indices)
                pre_remaining_users = [int(v) for v in pre_users if int(v) != int(mean_idx)]
                if len(pre_remaining_users) == 0:
                    remove_indices.add(int(pre_idx))
                for remove_idx in sorted(list(remove_indices), reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_mean_prepost_nhwc_passthrough_chains": int(rewritten)}


def _optimize_transpose_layernorm_stats_nhwc_propagation_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Re-propagate NHWC through LayerNorm-statistics branches after token-mixing residual paths.

    Target (rank-4):
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      m1 = MEAN(x_nchw, axes=[1], keepDims=True)
      d  = SUB(x_nchw, m1)
      sq = MUL(d, d)
      v  = MEAN(sq, axes=[1], keepDims=True)

    Rewrite:
      m1 = MEAN(x_nhwc, axes=[3], keepDims=True)
      d  = SUB(x_nhwc, m1)
      v  = MEAN(sq, axes=[3], keepDims=True)

    Notes:
    - Pre-transpose is removed only when it became dead.
    - This pass focuses on the statistics path and keeps other residual users intact.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _map_mean_axes_to_nhwc(
        *,
        axes_tensor: Optional[TensorIR],
        rank: int,
    ) -> Optional[List[int]]:
        if axes_tensor is None:
            return None
        axes_vals = _read_const_ints_from_tensor(axes_tensor)
        if axes_vals is None or len(axes_vals) == 0:
            return None
        normalized_axes: List[int] = []
        for axis in axes_vals:
            a = int(axis)
            if a < 0:
                a += int(rank)
            if a < 0 or a >= int(rank):
                return None
            normalized_axes.append(int(a))
        return [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            rank = int(len(list(pre_input_tensor.shape))) if pre_input_tensor is not None and pre_input_tensor.shape is not None else 4
            if int(rank) != 4:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue

            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None

            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    mean1_idx is None
                    and str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_output_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean1_idx = int(user_idx)
                    continue
                if (
                    sub_idx is None
                    and str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_output_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)

            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name in model_outputs:
                continue
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue

            mean1_axes_name = str(mean1_op.inputs[1])
            mean1_axes_tensor = model_ir.tensors.get(mean1_axes_name, None)
            mapped_axes1 = _map_mean_axes_to_nhwc(
                axes_tensor=mean1_axes_tensor,
                rank=rank,
            )
            if mapped_axes1 is None:
                continue

            sub_out_name = str(sub_op.outputs[0])
            if sub_out_name in model_outputs:
                continue
            sub_users = [int(v) for v in consumers.get(sub_out_name, [])]
            if len(sub_users) == 0:
                continue

            sq_idx: Optional[int] = None
            for user_idx in sub_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MUL"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == sub_out_name
                    and str(user_op.inputs[1]) == sub_out_name
                ):
                    sq_idx = int(user_idx)
                    break
            if sq_idx is None:
                continue
            # Only rewrite variance-statistics branches.
            # If SUB output feeds non-square branches (e.g. InstanceNorm affine path),
            # keep original layout to avoid breaking downstream broadcast semantics.
            sub_other_users = [
                int(v) for v in sub_users if int(v) != int(sq_idx)
            ]
            if len(sub_other_users) > 0:
                continue

            sq_op = model_ir.operators[int(sq_idx)]
            sq_out_name = str(sq_op.outputs[0])
            if sq_out_name in model_outputs:
                continue

            sq_users = [int(v) for v in consumers.get(sq_out_name, [])]
            mean2_idx: Optional[int] = None
            for user_idx in sq_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == sq_out_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean2_idx = int(user_idx)
                    break
            if mean2_idx is None:
                continue

            mean2_op = model_ir.operators[int(mean2_idx)]
            mean2_out_name = str(mean2_op.outputs[0])
            if mean2_out_name in model_outputs:
                continue

            mean2_axes_name = str(mean2_op.inputs[1])
            mean2_axes_tensor = model_ir.tensors.get(mean2_axes_name, None)
            mapped_axes2 = _map_mean_axes_to_nhwc(
                axes_tensor=mean2_axes_tensor,
                rank=rank,
            )
            if mapped_axes2 is None:
                continue

            _write_const_ints_to_tensor(mean1_axes_tensor, [int(v) for v in mapped_axes1])
            _set_operator_inputs(
                model_ir=model_ir,
                op=mean1_op,
                new_inputs=[pre_input_name, mean1_axes_name],
            )

            new_sub_inputs = [
                pre_input_name if str(v) == pre_output_name else str(v)
                for v in list(sub_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=sub_op,
                new_inputs=new_sub_inputs,
            )

            _write_const_ints_to_tensor(mean2_axes_tensor, [int(v) for v in mapped_axes2])
            _set_operator_inputs(
                model_ir=model_ir,
                op=mean2_op,
                new_inputs=[sq_out_name, mean2_axes_name],
            )

            for tensor_name in [
                mean1_out_name,
                sub_out_name,
                sq_out_name,
                mean2_out_name,
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            pre_remaining_users = [
                int(v) for v in pre_users if int(v) not in {int(mean1_idx), int(sub_idx)}
            ]
            if len(pre_remaining_users) == 0:
                del model_ir.operators[int(pre_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_layernorm_stats_nhwc_propagation_chains": int(rewritten)}


def _optimize_layernorm_stats_via_existing_post_transpose_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Re-propagate NHWC through LayerNorm statistics when an NCHW tensor already has
    an existing post-transpose (NCHW->NHWC) projection.

    Target:
      x_nchw --TRANSPOSE(0,2,3,1)--> x_nhwc
      m1 = MEAN(x_nchw, axes=[1], keepDims=True)
      d  = SUB(x_nchw, m1)
      sq = MUL(d, d)
      v  = MEAN(sq, axes=[1], keepDims=True)

    Rewrite:
      m1 = MEAN(x_nhwc, axes=[3], keepDims=True)
      d  = SUB(x_nhwc, m1)
      v  = MEAN(sq, axes=[3], keepDims=True)
    """
    rewritten = 0
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    axis_map_nchw_to_nhwc = _invert_perm(perm_nchw_to_nhwc)
    if axis_map_nchw_to_nhwc is None:
        return {"optimized_layernorm_stats_via_existing_post_transpose_nhwc_chains": 0}

    def _map_axes(
        *,
        axes_tensor: Optional[TensorIR],
        rank: int,
    ) -> Optional[List[int]]:
        if axes_tensor is None:
            return None
        axes_vals = _read_const_ints_from_tensor(axes_tensor)
        if axes_vals is None or len(axes_vals) == 0:
            return None
        normalized_axes: List[int] = []
        for axis in axes_vals:
            a = int(axis)
            if a < 0:
                a += int(rank)
            if a < 0 or a >= int(rank):
                return None
            normalized_axes.append(int(a))
        return [int(axis_map_nchw_to_nhwc[int(axis)]) for axis in normalized_axes]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            x_nchw_name = str(post_op.inputs[0])
            x_nhwc_name = str(post_op.outputs[0])
            if x_nchw_name in model_outputs:
                continue

            x_tensor = model_ir.tensors.get(x_nchw_name, None)
            rank = int(len(list(x_tensor.shape))) if x_tensor is not None and x_tensor.shape is not None else 4
            if int(rank) != 4:
                continue

            x_users = [int(v) for v in consumers.get(x_nchw_name, [])]
            if len(x_users) == 0:
                continue

            mean1_idx: Optional[int] = None
            sub_idx: Optional[int] = None
            for user_idx in x_users:
                if int(user_idx) == int(post_idx):
                    continue
                user_op = model_ir.operators[int(user_idx)]
                if (
                    mean1_idx is None
                    and str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == x_nchw_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean1_idx = int(user_idx)
                    continue
                if (
                    sub_idx is None
                    and str(user_op.op_type) == "SUB"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and x_nchw_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                ):
                    sub_idx = int(user_idx)

            if mean1_idx is None or sub_idx is None:
                continue

            mean1_op = model_ir.operators[int(mean1_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            mean1_out_name = str(mean1_op.outputs[0])
            if mean1_out_name in model_outputs:
                continue
            if mean1_out_name not in {str(sub_op.inputs[0]), str(sub_op.inputs[1])}:
                continue

            mean1_axes_name = str(mean1_op.inputs[1])
            mean1_axes_tensor = model_ir.tensors.get(mean1_axes_name, None)
            mapped_axes1 = _map_axes(
                axes_tensor=mean1_axes_tensor,
                rank=rank,
            )
            if mapped_axes1 is None:
                continue

            sub_out_name = str(sub_op.outputs[0])
            if sub_out_name in model_outputs:
                continue
            sq_users = [int(v) for v in consumers.get(sub_out_name, [])]
            sq_idx: Optional[int] = None
            for user_idx in sq_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MUL"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == sub_out_name
                    and str(user_op.inputs[1]) == sub_out_name
                ):
                    sq_idx = int(user_idx)
                    break
            if sq_idx is None:
                continue
            # Same safety guard as the transpose-driven variant above.
            sub_other_users = [
                int(v) for v in sq_users if int(v) != int(sq_idx)
            ]
            if len(sub_other_users) > 0:
                continue

            sq_op = model_ir.operators[int(sq_idx)]
            sq_out_name = str(sq_op.outputs[0])
            if sq_out_name in model_outputs:
                continue

            mean2_idx: Optional[int] = None
            for user_idx in consumers.get(sq_out_name, []):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == sq_out_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean2_idx = int(user_idx)
                    break
            if mean2_idx is None:
                continue

            mean2_op = model_ir.operators[int(mean2_idx)]
            mean2_out_name = str(mean2_op.outputs[0])
            if mean2_out_name in model_outputs:
                continue

            mean2_axes_name = str(mean2_op.inputs[1])
            mean2_axes_tensor = model_ir.tensors.get(mean2_axes_name, None)
            mapped_axes2 = _map_axes(
                axes_tensor=mean2_axes_tensor,
                rank=rank,
            )
            if mapped_axes2 is None:
                continue

            _write_const_ints_to_tensor(mean1_axes_tensor, [int(v) for v in mapped_axes1])
            _set_operator_inputs(
                model_ir=model_ir,
                op=mean1_op,
                new_inputs=[x_nhwc_name, mean1_axes_name],
            )
            _write_const_ints_to_tensor(mean2_axes_tensor, [int(v) for v in mapped_axes2])
            _set_operator_inputs(
                model_ir=model_ir,
                op=mean2_op,
                new_inputs=[sq_out_name, mean2_axes_name],
            )

            new_sub_inputs = [
                x_nhwc_name if str(v) == x_nchw_name else str(v)
                for v in list(sub_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=sub_op,
                new_inputs=new_sub_inputs,
            )

            for tensor_name in [
                mean1_out_name,
                sub_out_name,
                sq_out_name,
                mean2_out_name,
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_layernorm_stats_via_existing_post_transpose_nhwc_chains": int(rewritten)}


def _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate terminal NHWC->NCHW adapters before unary/MEAN suffixes.

    Target (strict chain):
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --(UNARY)*--> u_nchw
      u_nchw --MEAN(axes_nchw, keepDims=True)--> m_nchw
      m_nchw --RESHAPE--> y

    Rewrite:
      x_nhwc --(same UNARY)*--> u_nhwc
      u_nhwc --MEAN(axes_nhwc, keepDims=True)--> m_nhwc
      m_nhwc --(same RESHAPE)--> y

    Safety:
    - Pre-transpose direction is canonical NHWC->NCHW.
    - Unary chain is linear and layout-agnostic.
    - MEAN rank is 4 with constant axes.
    - Skip chains already handled by pre/post MEAN transpose elimination.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "FLOOR",
        "CEIL",
    }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) == 0:
                continue

            for first_user_idx in pre_users:
                first_user_op = model_ir.operators[int(first_user_idx)]
                first_user_type = str(first_user_op.op_type)
                if (
                    first_user_type != "MEAN"
                    and first_user_type not in unary_passthrough_ops
                ):
                    continue

                unary_chain_indices: List[int] = []
                unary_chain_ops: List[OperatorIR] = []
                current_tensor_name = pre_output_name
                current_op_idx = int(first_user_idx)
                mean_idx: Optional[int] = None
                valid_path = True

                while True:
                    current_op = model_ir.operators[int(current_op_idx)]
                    current_type = str(current_op.op_type)
                    if current_type in unary_passthrough_ops:
                        if len(current_op.inputs) != 1 or len(current_op.outputs) != 1:
                            valid_path = False
                            break
                        if str(current_op.inputs[0]) != current_tensor_name:
                            valid_path = False
                            break
                        unary_out_name = str(current_op.outputs[0])
                        if unary_out_name in model_outputs:
                            valid_path = False
                            break
                        unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
                        if len(unary_users) != 1:
                            valid_path = False
                            break
                        unary_chain_indices.append(int(current_op_idx))
                        unary_chain_ops.append(current_op)
                        current_tensor_name = unary_out_name
                        current_op_idx = int(unary_users[0])
                        continue
                    if current_type == "MEAN":
                        mean_idx = int(current_op_idx)
                        break
                    valid_path = False
                    break
                if not valid_path or mean_idx is None:
                    continue

                mean_op = model_ir.operators[int(mean_idx)]
                if len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                    continue
                if str(mean_op.inputs[0]) != current_tensor_name:
                    continue
                keep_dims = bool(
                    mean_op.options.get(
                        "keepDims",
                        mean_op.options.get(
                            "keep_dims",
                            mean_op.options.get("keepdims", False),
                        ),
                    )
                )

                rank = 4
                pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
                if (
                    pre_input_tensor is not None
                    and pre_input_tensor.shape is not None
                    and len(pre_input_tensor.shape) > 0
                ):
                    rank = int(len(list(pre_input_tensor.shape)))
                if rank != 4:
                    continue

                mean_axes_name = str(mean_op.inputs[1])
                mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                mean_axes = _read_const_ints_from_tensor(mean_axes_tensor)
                if mean_axes is None or len(mean_axes) == 0:
                    continue

                normalized_axes: List[int] = []
                valid_axes = True
                for axis in mean_axes:
                    a = int(axis)
                    if a < 0:
                        a += int(rank)
                    if a < 0 or a >= int(rank):
                        valid_axes = False
                        break
                    normalized_axes.append(int(a))
                if not valid_axes:
                    continue
                mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]

                mean_out_name = str(mean_op.outputs[0])
                if not keep_dims:
                    # Some lowered MEAN ops omit keepDims in options while still
                    # producing rank-preserving tensors. Accept those by shape.
                    mean_in_tensor = model_ir.tensors.get(current_tensor_name, None)
                    mean_out_tensor = model_ir.tensors.get(mean_out_name, None)
                    mean_in_shape = (
                        [int(v) for v in list(mean_in_tensor.shape)]
                        if mean_in_tensor is not None and mean_in_tensor.shape is not None
                        else None
                    )
                    mean_out_shape = (
                        [int(v) for v in list(mean_out_tensor.shape)]
                        if mean_out_tensor is not None and mean_out_tensor.shape is not None
                        else None
                    )
                    if (
                        mean_in_shape is None
                        or mean_out_shape is None
                        or len(mean_in_shape) != len(mean_out_shape)
                    ):
                        continue
                    keep_dims_by_shape = True
                    for axis in normalized_axes:
                        if int(axis) < 0 or int(axis) >= len(mean_out_shape):
                            keep_dims_by_shape = False
                            break
                        out_dim = int(mean_out_shape[int(axis)])
                        if out_dim > 0 and out_dim != 1:
                            keep_dims_by_shape = False
                            break
                    if not keep_dims_by_shape:
                        continue

                mean_users = [int(v) for v in consumers.get(mean_out_name, [])]
                if len(mean_users) != 1:
                    continue

                tail_idx = int(mean_users[0])
                tail_op = model_ir.operators[int(tail_idx)]
                tail_type = str(tail_op.op_type)
                if (
                    tail_type == "TRANSPOSE"
                    and len(tail_op.inputs) >= 2
                    and len(tail_op.outputs) == 1
                    and str(tail_op.inputs[0]) == mean_out_name
                    and _read_transpose_perm(model_ir, tail_op) == perm_nchw_to_nhwc
                ):
                    # Leave paired pre/post transpose chains to dedicated rewrite.
                    continue
                if tail_type != "RESHAPE":
                    continue
                if len(tail_op.inputs) < 2 or len(tail_op.outputs) != 1:
                    continue
                if str(tail_op.inputs[0]) != mean_out_name:
                    continue
                if _read_const_ints_from_tensor(model_ir.tensors.get(str(tail_op.inputs[1]), None)) is None:
                    continue

                _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes])

                if len(unary_chain_ops) > 0:
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=unary_chain_ops[0],
                        new_inputs=[pre_input_name],
                    )
                    for unary_op in unary_chain_ops:
                        _permute_tensor_metadata_if_rank_matches(
                            model_ir.tensors.get(str(unary_op.outputs[0]), None),
                            perm_nchw_to_nhwc,
                        )
                else:
                    mean_inputs = [str(v) for v in list(mean_op.inputs)]
                    mean_inputs[0] = pre_input_name
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=mean_op,
                        new_inputs=mean_inputs,
                    )

                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(mean_out_name, None),
                    perm_nchw_to_nhwc,
                )

                pre_remaining_users = [int(v) for v in pre_users if int(v) != int(first_user_idx)]
                remove_indices: List[int] = []
                if len(pre_remaining_users) == 0:
                    remove_indices.append(int(pre_idx))
                for remove_idx in sorted(remove_indices, reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_unary_mean_terminal_nhwc_chains": int(rewritten)}


def _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips in EfficientNet-style SE conv gating blocks.

    Target (constrained):
      x_nhwc --T(0,3,1,2)--> x_nchw --LOGISTIC--> s1
      MUL(x_nchw, s1) -> sw_nchw --MEAN(axes=[2,3],keepDims=1)--> m_nchw --T(0,2,3,1)--> m_nhwc
      m_nhwc -> ... -> g_nhwc --T(0,3,1,2)--> g_nchw --GATE--> s2
      MUL(sw_nchw, s2) -> y_nchw --T(0,2,3,1)--> y_nhwc

    GATE can be either:
      - LOGISTIC(g_nchw), or
      - ADD(g_nchw, const) -> MUL|DIV(const)

    Rewrite:
      - Bypass the two NHWC->NCHW pre-transposes and NHWC inverse post-transposes.
      - Keep LOGISTIC/ADD/MUL/DIV in NHWC (layout-agnostic).
      - Remap MEAN axes from NCHW to NHWC.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 2:
                continue

            log1_idx: Optional[int] = None
            mul1_idx: Optional[int] = None
            for user_idx in pre_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if (
                    user_type == "LOGISTIC"
                    and len(user_op.inputs) == 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == pre_output_name
                ):
                    log1_idx = int(user_idx)
                elif (
                    user_type == "MUL"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and pre_output_name in {str(v) for v in list(user_op.inputs)}
                ):
                    mul1_idx = int(user_idx)
            if log1_idx is None or mul1_idx is None:
                continue

            log1_op = model_ir.operators[int(log1_idx)]
            log1_out_name = str(log1_op.outputs[0])
            if set(int(v) for v in consumers.get(log1_out_name, [])) != {int(mul1_idx)}:
                continue

            mul1_op = model_ir.operators[int(mul1_idx)]
            mul1_inputs = [str(v) for v in list(mul1_op.inputs)]
            if pre_output_name not in mul1_inputs or log1_out_name not in mul1_inputs:
                continue
            mul1_data_input_index = int(0 if mul1_inputs[0] == pre_output_name else 1)
            mul1_out_name = str(mul1_op.outputs[0])
            if mul1_out_name in model_outputs:
                continue

            mul1_users = [int(v) for v in consumers.get(mul1_out_name, []) if int(v) != int(mul1_idx)]
            if len(mul1_users) != 2:
                continue

            mean_idx: Optional[int] = None
            mul2_idx: Optional[int] = None
            for user_idx in mul1_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if (
                    user_type == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == mul1_out_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean_idx = int(user_idx)
                elif (
                    user_type == "MUL"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                    and mul1_out_name in {str(v) for v in list(user_op.inputs)}
                ):
                    mul2_idx = int(user_idx)
            if mean_idx is None or mul2_idx is None:
                continue

            mean_op = model_ir.operators[int(mean_idx)]
            mean_out_name = str(mean_op.outputs[0])
            if mean_out_name in model_outputs:
                continue
            mean_axes_name = str(mean_op.inputs[1])
            mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_vals is None or len(mean_axes_vals) == 0:
                continue

            x_tensor = model_ir.tensors.get(pre_input_name, None)
            x_shape = list(x_tensor.shape) if x_tensor is not None else []
            rank = int(len(x_shape)) if len(x_shape) > 0 else 4
            if rank != 4:
                continue
            normalized_axes: List[int] = []
            valid_axes = True
            for axis in mean_axes_vals:
                a = int(axis)
                if a < 0:
                    a += int(rank)
                if a < 0 or a >= int(rank):
                    valid_axes = False
                    break
                normalized_axes.append(int(a))
            if not valid_axes or sorted(normalized_axes) != [2, 3]:
                continue
            mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]
            if sorted(mapped_axes) != [1, 2]:
                continue

            mean_users = [int(v) for v in consumers.get(mean_out_name, []) if int(v) != int(mean_idx)]
            if len(mean_users) == 0:
                continue
            post_mean_indices: List[int] = []
            post_mean_outputs: List[str] = []
            mean_passthrough_via_squeeze = False
            mean_squeeze_idx: Optional[int] = None
            mean_squeeze_op: Optional[OperatorIR] = None
            # Legacy path: MEAN output is wrapped by NCHW->NHWC transpose adapter(s).
            valid_post_mean = True
            for user_idx in mean_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != mean_out_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    valid_post_mean = False
                    break
                post_mean_indices.append(int(user_idx))
                post_mean_outputs.append(str(user_op.outputs[0]))
            # EfficientDet-like variant: MEAN output is consumed by SQUEEZE directly
            # and converted back to NHWC with RESHAPE (no explicit post-transpose).
            if not valid_post_mean or len(post_mean_indices) == 0:
                post_mean_indices = []
                post_mean_outputs = []
                if len(mean_users) != 1:
                    continue
                candidate_idx = int(mean_users[0])
                candidate_op = model_ir.operators[int(candidate_idx)]
                if (
                    str(candidate_op.op_type) != "SQUEEZE"
                    or len(candidate_op.inputs) != 1
                    or len(candidate_op.outputs) != 1
                    or str(candidate_op.inputs[0]) != mean_out_name
                    or str(candidate_op.outputs[0]) in model_outputs
                ):
                    continue
                mean_passthrough_via_squeeze = True
                mean_squeeze_idx = int(candidate_idx)
                mean_squeeze_op = candidate_op

            mul2_op = model_ir.operators[int(mul2_idx)]
            mul2_inputs = [str(v) for v in list(mul2_op.inputs)]
            if mul2_inputs[0] == mul1_out_name:
                gate_out_name = str(mul2_inputs[1])
            elif mul2_inputs[1] == mul1_out_name:
                gate_out_name = str(mul2_inputs[0])
            else:
                continue

            gate_pre_idx: Optional[int] = None
            gate_pre_op: Optional[OperatorIR] = None
            gate_rewrite_op: Optional[OperatorIR] = None
            gate_rewrite_input_index: Optional[int] = None
            gate_nchw_tensor_names: List[str] = []
            gate_pre_needs_removal = False
            gate_pre_is_reshape = False

            if set(int(v) for v in consumers.get(gate_out_name, [])) != {int(mul2_idx)}:
                continue

            gate_prod_idx = producers.get(gate_out_name, None)
            if gate_prod_idx is None:
                continue
            gate_prod_op = model_ir.operators[int(gate_prod_idx)]
            gate_prod_type = str(gate_prod_op.op_type)

            if (
                gate_prod_type == "LOGISTIC"
                and len(gate_prod_op.inputs) == 1
                and len(gate_prod_op.outputs) == 1
                and str(gate_prod_op.outputs[0]) == gate_out_name
            ):
                gate_pre_out_name = str(gate_prod_op.inputs[0])
                gate_pre_idx = producers.get(gate_pre_out_name, None)
                if gate_pre_idx is None:
                    continue
                gate_pre_op = model_ir.operators[int(gate_pre_idx)]
                # Legacy path: gate LOGISTIC input is wrapped by NHWC->NCHW transpose.
                if (
                    str(gate_pre_op.op_type) != "TRANSPOSE"
                    or len(gate_pre_op.inputs) < 2
                    or len(gate_pre_op.outputs) != 1
                    or str(gate_pre_op.outputs[0]) != gate_pre_out_name
                    or _read_transpose_perm(model_ir, gate_pre_op) != perm_nhwc_to_nchw
                    or gate_pre_out_name in model_outputs
                ):
                    # EfficientDet-like path: gate LOGISTIC input is RESHAPE to NCHW
                    # from NHWC tensor. Keep RESHAPE op and rewrite shape to NHWC.
                    if (
                        str(gate_pre_op.op_type) != "RESHAPE"
                        or len(gate_pre_op.inputs) < 2
                        or len(gate_pre_op.outputs) != 1
                        or str(gate_pre_op.outputs[0]) != gate_pre_out_name
                        or gate_pre_out_name in model_outputs
                    ):
                        continue
                    gate_shape_tensor = model_ir.tensors.get(str(gate_pre_op.inputs[1]), None)
                    gate_shape_vals = _read_const_ints_from_tensor(gate_shape_tensor)
                    if gate_shape_vals is None or len(gate_shape_vals) != 4:
                        continue
                    if set(int(v) for v in consumers.get(gate_pre_out_name, [])) != {int(gate_prod_idx)}:
                        continue
                    swapped_shape = [
                        int(gate_shape_vals[0]),
                        int(gate_shape_vals[2]),
                        int(gate_shape_vals[3]),
                        int(gate_shape_vals[1]),
                    ]
                    _write_const_ints_to_tensor(gate_shape_tensor, swapped_shape)
                    gate_pre_is_reshape = True
                if set(int(v) for v in consumers.get(gate_pre_out_name, [])) != {int(gate_prod_idx)}:
                    continue

                gate_rewrite_op = gate_prod_op
                if not gate_pre_is_reshape:
                    gate_rewrite_input_index = 0
                    gate_pre_needs_removal = True
                gate_nchw_tensor_names = [str(gate_pre_out_name), str(gate_out_name)]

            elif (
                gate_prod_type in {"MUL", "DIV"}
                and len(gate_prod_op.inputs) == 2
                and len(gate_prod_op.outputs) == 1
                and str(gate_prod_op.outputs[0]) == gate_out_name
            ):
                gate_scale_inputs = [str(v) for v in list(gate_prod_op.inputs)]
                gate_add_idx: Optional[int] = None
                gate_add_out_name: Optional[str] = None
                gate_scale_side_name: Optional[str] = None
                for gate_scale_input_name in gate_scale_inputs:
                    gate_scale_input_prod = producers.get(str(gate_scale_input_name), None)
                    if gate_scale_input_prod is not None:
                        gate_scale_input_op = model_ir.operators[int(gate_scale_input_prod)]
                        if (
                            str(gate_scale_input_op.op_type) == "ADD"
                            and len(gate_scale_input_op.inputs) == 2
                            and len(gate_scale_input_op.outputs) == 1
                            and str(gate_scale_input_op.outputs[0]) == str(gate_scale_input_name)
                        ):
                            gate_add_idx = int(gate_scale_input_prod)
                            gate_add_out_name = str(gate_scale_input_name)
                            continue
                    gate_scale_side_name = str(gate_scale_input_name)
                if (
                    gate_add_idx is None
                    or gate_add_out_name is None
                    or gate_scale_side_name is None
                    or not _is_singleton_constant_tensor(model_ir, gate_scale_side_name)
                ):
                    continue
                if set(int(v) for v in consumers.get(gate_add_out_name, [])) != {int(gate_prod_idx)}:
                    continue

                gate_add_op = model_ir.operators[int(gate_add_idx)]
                gate_add_inputs = [str(v) for v in list(gate_add_op.inputs)]
                gate_add_data_input_index: Optional[int] = None
                gate_pre_out_name: Optional[str] = None
                for gate_add_input_index, gate_add_input_name in enumerate(gate_add_inputs):
                    gate_add_input_prod = producers.get(str(gate_add_input_name), None)
                    if gate_add_input_prod is not None:
                        gate_add_input_op = model_ir.operators[int(gate_add_input_prod)]
                        if (
                            str(gate_add_input_op.op_type) == "TRANSPOSE"
                            and len(gate_add_input_op.inputs) >= 2
                            and len(gate_add_input_op.outputs) == 1
                            and str(gate_add_input_op.outputs[0]) == str(gate_add_input_name)
                            and _read_transpose_perm(model_ir, gate_add_input_op) == perm_nhwc_to_nchw
                            and str(gate_add_input_name) not in model_outputs
                        ):
                            gate_pre_idx = int(gate_add_input_prod)
                            gate_pre_op = gate_add_input_op
                            gate_pre_out_name = str(gate_add_input_name)
                            gate_add_data_input_index = int(gate_add_input_index)
                            continue
                    if not _is_singleton_constant_tensor(model_ir, str(gate_add_input_name)):
                        gate_add_data_input_index = None
                        break
                if (
                    gate_pre_idx is None
                    or gate_pre_op is None
                    or gate_pre_out_name is None
                    or gate_add_data_input_index is None
                ):
                    continue
                if set(int(v) for v in consumers.get(gate_pre_out_name, [])) != {int(gate_add_idx)}:
                    continue

                gate_rewrite_op = gate_add_op
                gate_rewrite_input_index = int(gate_add_data_input_index)
                gate_nchw_tensor_names = [str(gate_add_out_name), str(gate_out_name)]
                gate_pre_needs_removal = True

            else:
                continue

            if (
                gate_pre_idx is None
                or gate_pre_op is None
                or gate_rewrite_op is None
                or (gate_rewrite_input_index is None and not gate_pre_is_reshape)
            ):
                continue

            mul2_out_name = str(mul2_op.outputs[0])
            if mul2_out_name in model_outputs:
                continue
            mul2_users = [int(v) for v in consumers.get(mul2_out_name, []) if int(v) != int(mul2_idx)]
            if len(mul2_users) == 0:
                continue
            post_out_indices: List[int] = []
            post_out_outputs: List[str] = []
            valid_post_out = True
            for user_idx in mul2_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != mul2_out_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    valid_post_out = False
                    break
                post_out_indices.append(int(user_idx))
                post_out_outputs.append(str(user_op.outputs[0]))
            if not valid_post_out or len(post_out_indices) == 0:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=log1_op,
                new_inputs=[pre_input_name],
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul1_op,
                input_index=int(mul1_data_input_index),
                new_input_name=pre_input_name,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(log1_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mul1_out_name, None),
                perm_nchw_to_nhwc,
            )

            _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes])
            if isinstance(mean_op.options, dict):
                mean_options = dict(mean_op.options)
                for key in ["axis", "axes", "onnxRawAxes"]:
                    value = mean_options.get(key, None)
                    if isinstance(value, list) and len(value) == len(mapped_axes):
                        mean_options[key] = [int(v) for v in mapped_axes]
                mean_op.options = mean_options
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mean_out_name, None),
                perm_nchw_to_nhwc,
            )
            for post_mean_out_name in post_mean_outputs:
                _replace_tensor_inputs(model_ir, post_mean_out_name, mean_out_name)
            if mean_passthrough_via_squeeze and mean_squeeze_op is not None:
                squeeze_dims = list(mean_squeeze_op.options.get("squeezeDims", []))
                remapped_dims: List[int] = []
                rank4 = 4
                for axis in squeeze_dims:
                    a = int(axis)
                    if a < 0:
                        a += int(rank4)
                    if a < 0 or a >= int(rank4):
                        remapped_dims = []
                        break
                    remapped_dims.append(int(perm_nhwc_to_nchw[int(a)]))
                if len(remapped_dims) > 0:
                    mean_squeeze_op.options = dict(mean_squeeze_op.options)
                    mean_squeeze_op.options["squeezeDims"] = sorted([int(v) for v in remapped_dims])
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(mean_squeeze_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )

            if gate_rewrite_input_index is not None:
                gate_pre_input_name = str(gate_pre_op.inputs[0])
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=gate_rewrite_op,
                    input_index=int(gate_rewrite_input_index),
                    new_input_name=gate_pre_input_name,
                )
            for gate_nchw_tensor_name in gate_nchw_tensor_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(gate_nchw_tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            representative_output_name = str(post_out_outputs[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul2_op,
                new_outputs=[representative_output_name],
            )
            for alias_name in post_out_outputs[1:]:
                _replace_tensor_inputs(model_ir, alias_name, representative_output_name)

            old_mul2_tensor = model_ir.tensors.get(mul2_out_name, None)
            representative_tensor = model_ir.tensors.get(representative_output_name, None)
            if old_mul2_tensor is not None and representative_tensor is not None:
                representative_tensor.dtype = str(old_mul2_tensor.dtype)
                representative_tensor.quantization = _clone_quantization(old_mul2_tensor.quantization)
                representative_tensor.shape = [int(v) for v in list(old_mul2_tensor.shape)]
                representative_tensor.shape_signature = (
                    [int(v) for v in list(old_mul2_tensor.shape_signature)]
                    if old_mul2_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul2_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    representative_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = set([int(pre_idx)])
            if gate_pre_needs_removal and gate_pre_idx is not None:
                remove_indices.add(int(gate_pre_idx))
            remove_indices.update(int(v) for v in post_mean_indices)
            remove_indices.update(int(v) for v in post_out_indices)
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_se_conv_mul_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips in SE-like gating blocks.

    Target:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nhwc --AVG_POOL_2D--> p_nhwc --T(0,3,1,2)--> p_nchw
      p_nchw --RESHAPE--> r --FULLY_CONNECTED--> f --PRELU--> g --RESHAPE--> s_nchw
      MUL(x_nchw, s_nchw) -> m_nchw --T(0,2,3,1)--> m_nhwc

    Rewrite:
      - Bypass both pre-transposes and the post-transpose.
      - Keep AVG_POOL/FC/PRELU chain as-is, but rewrite final gate reshape to NHWC.
      - Feed MUL directly with x_nhwc and gate_nhwc.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    gate_head_passthrough_unary = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LOGISTIC",
        "TANH",
        "HARD_SWISH",
    }
    output_bridge_passthrough_unary = {
        "GELU",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LOGISTIC",
        "TANH",
        "HARD_SWISH",
    }

    def _clone_const_tensor_for_operator_input(
        *,
        op: OperatorIR,
        input_index: int,
        tag: str,
    ) -> str:
        input_name = str(op.inputs[int(input_index)])
        tensor = model_ir.tensors.get(input_name, None)
        if tensor is None or tensor.data is None:
            return input_name
        cloned_name = f"{input_name}{tag}"
        suffix = 1
        while cloned_name in model_ir.tensors:
            cloned_name = f"{input_name}{tag}_{suffix}"
            suffix += 1
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(tensor.dtype),
            shape=[int(v) for v in list(tensor.shape)] if tensor.shape is not None else [],
            shape_signature=(
                [int(v) for v in list(tensor.shape_signature)]
                if tensor.shape_signature is not None
                else (
                    [int(v) for v in list(tensor.shape)]
                    if tensor.shape is not None
                    else []
                )
            ),
            data=np.array(tensor.data, copy=True),
            is_variable=bool(tensor.is_variable),
            quantization=_clone_quantization(tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=cloned_name,
        )
        return str(cloned_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_mul_idx, pre_mul_op in enumerate(model_ir.operators):
            if str(pre_mul_op.op_type) != "TRANSPOSE" or len(pre_mul_op.inputs) < 2 or len(pre_mul_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_mul_op) != perm_nhwc_to_nchw:
                continue

            x_nhwc_name = str(pre_mul_op.inputs[0])
            x_nchw_name = str(pre_mul_op.outputs[0])
            if x_nchw_name in model_outputs:
                continue

            x_nchw_users = [int(v) for v in consumers.get(x_nchw_name, [])]
            if len(x_nchw_users) == 0:
                continue

            for mul_idx in x_nchw_users:
                mul_op = model_ir.operators[int(mul_idx)]
                if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                    continue

                mul_in0 = str(mul_op.inputs[0])
                mul_in1 = str(mul_op.inputs[1])
                if mul_in0 == x_nchw_name:
                    gate_name = str(mul_in1)
                    x_side_index = 0
                elif mul_in1 == x_nchw_name:
                    gate_name = str(mul_in0)
                    x_side_index = 1
                else:
                    continue

                gate_post_name = str(gate_name)
                gate_post_idx = producers.get(gate_post_name, None)
                if gate_post_idx is None:
                    continue
                gate_post_op = model_ir.operators[int(gate_post_idx)]
                gate_reshape_idx = int(gate_post_idx)
                gate_reshape_op = gate_post_op
                gate_requires_post_unary = False
                if str(gate_reshape_op.op_type) != "RESHAPE":
                    if (
                        str(gate_post_op.op_type) in gate_head_passthrough_unary
                        and len(gate_post_op.inputs) == 1
                        and len(gate_post_op.outputs) == 1
                        and str(gate_post_op.outputs[0]) == gate_post_name
                    ):
                        gate_reshape_input_name = str(gate_post_op.inputs[0])
                        gate_reshape_idx_candidate = producers.get(gate_reshape_input_name, None)
                        if gate_reshape_idx_candidate is None:
                            continue
                        gate_reshape_op = model_ir.operators[int(gate_reshape_idx_candidate)]
                        if (
                            str(gate_reshape_op.op_type) != "RESHAPE"
                            or len(gate_reshape_op.inputs) < 1
                            or len(gate_reshape_op.outputs) != 1
                            or str(gate_reshape_op.outputs[0]) != gate_reshape_input_name
                        ):
                            continue
                        gate_reshape_idx = int(gate_reshape_idx_candidate)
                        gate_requires_post_unary = True
                    else:
                        continue
                if gate_post_name in model_outputs:
                    continue

                # Allow stacked gate heads:
                #   RESHAPE -> ((FULLY_CONNECTED | 1x1 CONV_2D | 1x1 DEPTHWISE_CONV_2D) -> PRELU/UNARY)* -> RESHAPE(gate)
                gate_mlp_input_name = str(gate_reshape_op.inputs[0])
                flat_idx: Optional[int] = None
                flat_op: Optional[OperatorIR] = None
                saw_fc = False
                max_hops = 16
                hop = 0
                while hop < max_hops:
                    producer_idx = producers.get(gate_mlp_input_name, None)
                    if producer_idx is None:
                        break
                    producer_op = model_ir.operators[int(producer_idx)]
                    producer_type = str(producer_op.op_type)
                    if (
                        producer_type == "RESHAPE"
                        and len(producer_op.inputs) >= 1
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == gate_mlp_input_name
                    ):
                        flat_idx = int(producer_idx)
                        flat_op = producer_op
                        break
                    if (
                        producer_type == "PRELU"
                        and len(producer_op.inputs) == 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == gate_mlp_input_name
                    ):
                        gate_mlp_input_name = str(producer_op.inputs[0])
                        hop += 1
                        continue
                    if (
                        producer_type in gate_head_passthrough_unary
                        and len(producer_op.inputs) == 1
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == gate_mlp_input_name
                    ):
                        gate_mlp_input_name = str(producer_op.inputs[0])
                        hop += 1
                        continue
                    if (
                        producer_type == "FULLY_CONNECTED"
                        and len(producer_op.inputs) >= 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == gate_mlp_input_name
                    ):
                        saw_fc = True
                        gate_mlp_input_name = str(producer_op.inputs[0])
                        hop += 1
                        continue
                    if (
                        producer_type in {"CONV_2D", "DEPTHWISE_CONV_2D"}
                        and len(producer_op.inputs) >= 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == gate_mlp_input_name
                    ):
                        conv_kernel_tensor = model_ir.tensors.get(str(producer_op.inputs[1]), None)
                        conv_kernel_shape = list(conv_kernel_tensor.shape) if conv_kernel_tensor is not None else []
                        is_1x1_conv = (
                            producer_type == "CONV_2D"
                            and len(conv_kernel_shape) == 4
                            and int(conv_kernel_shape[1]) == 1
                            and int(conv_kernel_shape[2]) == 1
                        )
                        is_1x1_dwconv = (
                            producer_type == "DEPTHWISE_CONV_2D"
                            and len(conv_kernel_shape) == 4
                            and int(conv_kernel_shape[1]) == 1
                            and int(conv_kernel_shape[2]) == 1
                        )
                        if not (is_1x1_conv or is_1x1_dwconv):
                            break
                        saw_fc = True
                        gate_mlp_input_name = str(producer_op.inputs[0])
                        hop += 1
                        continue
                    break
                if flat_idx is None or flat_op is None or not saw_fc:
                    continue

                pre_gate_name = str(flat_op.inputs[0])
                pre_gate_idx = producers.get(pre_gate_name, None)
                if pre_gate_idx is None:
                    continue
                pre_gate_op = model_ir.operators[int(pre_gate_idx)]

                pool_out_name: Optional[str] = None
                pool_idx: Optional[int] = None
                removable_pre_gate_idx: Optional[int] = None
                rewrite_pool_mean_to_nhwc = False
                mapped_mean_axes: Optional[List[int]] = None

                # Legacy path: pooled output is wrapped by NHWC->NCHW transpose.
                if (
                    str(pre_gate_op.op_type) == "TRANSPOSE"
                    and len(pre_gate_op.inputs) >= 2
                    and len(pre_gate_op.outputs) == 1
                    and str(pre_gate_op.outputs[0]) == pre_gate_name
                    and _read_transpose_perm(model_ir, pre_gate_op) == perm_nhwc_to_nchw
                    and pre_gate_name not in model_outputs
                ):
                    pool_out_name = str(pre_gate_op.inputs[0])
                    pool_idx = producers.get(pool_out_name, None)
                    removable_pre_gate_idx = int(pre_gate_idx)
                # Already optimized path: flatten consumes pooled NHWC output directly.
                elif (
                    str(pre_gate_op.op_type) in {"AVERAGE_POOL_2D", "MEAN"}
                    and len(pre_gate_op.inputs) >= 1
                    and len(pre_gate_op.outputs) == 1
                    and str(pre_gate_op.outputs[0]) == pre_gate_name
                ):
                    pool_out_name = str(pre_gate_name)
                    pool_idx = int(pre_gate_idx)
                    removable_pre_gate_idx = None
                else:
                    continue

                if pool_out_name is None or pool_idx is None:
                    continue
                pool_op = model_ir.operators[int(pool_idx)]
                pool_op_type = str(pool_op.op_type)
                if pool_op_type == "AVERAGE_POOL_2D":
                    if len(pool_op.inputs) != 1 or len(pool_op.outputs) != 1:
                        continue
                    if str(pool_op.outputs[0]) != str(pool_out_name):
                        continue
                    if str(pool_op.inputs[0]) != x_nhwc_name:
                        continue
                elif pool_op_type == "MEAN":
                    if len(pool_op.inputs) < 2 or len(pool_op.outputs) != 1:
                        continue
                    if str(pool_op.outputs[0]) != str(pool_out_name):
                        continue
                    pool_keep_dims = bool(
                        pool_op.options.get(
                            "keepDims",
                            pool_op.options.get(
                                "keep_dims",
                                pool_op.options.get("keepdims", False),
                            ),
                        )
                    )
                    if not pool_keep_dims:
                        continue
                    pool_mean_input_name = str(pool_op.inputs[0])
                    mean_axes_name = str(pool_op.inputs[1])
                    mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                    mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
                    if mean_axes_vals is None or len(mean_axes_vals) == 0:
                        continue
                    normalized_axes: List[int] = []
                    axes_valid = True
                    for axis in mean_axes_vals:
                        a = int(axis)
                        if a < 0:
                            a += 4
                        if a < 0 or a >= 4:
                            axes_valid = False
                            break
                        normalized_axes.append(int(a))
                    if not axes_valid:
                        continue
                    if pool_mean_input_name == x_nhwc_name:
                        if sorted(normalized_axes) != [1, 2]:
                            continue
                    elif pool_mean_input_name == x_nchw_name:
                        if sorted(normalized_axes) != [2, 3]:
                            continue
                        mapped_mean_axes = [int(perm_nhwc_to_nchw[int(v)]) for v in normalized_axes]
                        if sorted(mapped_mean_axes) != [1, 2]:
                            continue
                        rewrite_pool_mean_to_nhwc = True
                    else:
                        continue
                else:
                    continue

                mul_out_name = str(mul_op.outputs[0])
                if mul_out_name in model_outputs:
                    continue
                mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
                if len(mul_users) == 0:
                    continue

                post_indices: List[int] = []
                post_output_names: List[str] = []
                post_bridge_op: OperatorIR = mul_op
                post_bridge_output_name = str(mul_out_name)
                post_bridge_via_unary = False
                post_bridge_unary_idx: Optional[int] = None
                valid_posts = True
                for user_idx in mul_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) != "TRANSPOSE"
                    ):
                        if (
                            str(user_op.op_type) in output_bridge_passthrough_unary
                            and len(user_op.inputs) == 1
                            and len(user_op.outputs) == 1
                            and str(user_op.inputs[0]) == mul_out_name
                            and post_bridge_unary_idx is None
                        ):
                            unary_out_name = str(user_op.outputs[0])
                            if unary_out_name in model_outputs:
                                valid_posts = False
                                break
                            unary_users = [int(v) for v in consumers.get(unary_out_name, [])]
                            if len(unary_users) == 0:
                                valid_posts = False
                                break
                            valid_unary_tail = True
                            for unary_user_idx in unary_users:
                                unary_user_op = model_ir.operators[int(unary_user_idx)]
                                if (
                                    str(unary_user_op.op_type) != "TRANSPOSE"
                                    or len(unary_user_op.inputs) < 2
                                    or len(unary_user_op.outputs) != 1
                                    or str(unary_user_op.inputs[0]) != unary_out_name
                                    or _read_transpose_perm(model_ir, unary_user_op) != perm_nchw_to_nhwc
                                    or str(unary_user_op.outputs[0]) in model_outputs
                                ):
                                    valid_unary_tail = False
                                    break
                            if not valid_unary_tail:
                                valid_posts = False
                                break
                            post_bridge_via_unary = True
                            post_bridge_unary_idx = int(user_idx)
                            post_bridge_op = user_op
                            post_bridge_output_name = str(unary_out_name)
                            post_indices.extend([int(v) for v in unary_users])
                            post_output_names.extend(
                                [str(model_ir.operators[int(v)].outputs[0]) for v in unary_users]
                            )
                            continue
                        valid_posts = False
                        break
                    if (
                        len(user_op.inputs) < 2
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != mul_out_name
                        or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                        or str(user_op.outputs[0]) in model_outputs
                    ):
                        valid_posts = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                if not valid_posts or len(post_indices) == 0:
                    continue

                # Rewrite gate reshape shape [N,C,H,W] -> [N,H,W,C].
                shape_rewritten = False
                if len(gate_reshape_op.inputs) >= 2:
                    gate_shape_name = str(gate_reshape_op.inputs[1])
                    gate_shape_tensor = model_ir.tensors.get(gate_shape_name, None)
                    gate_shape_vals = _read_const_ints_from_tensor(gate_shape_tensor)
                    if gate_shape_vals is not None and len(gate_shape_vals) == 4:
                        swapped_shape = [
                            int(gate_shape_vals[0]),
                            int(gate_shape_vals[2]),
                            int(gate_shape_vals[3]),
                            int(gate_shape_vals[1]),
                        ]
                        _write_const_ints_to_tensor(gate_shape_tensor, swapped_shape)
                        shape_rewritten = True

                if isinstance(gate_reshape_op.options, dict):
                    opts = dict(gate_reshape_op.options)
                    options_changed = False
                    for key in ["newShape", "onnxRawNewShape"]:
                        value = opts.get(key, None)
                        if not isinstance(value, list) or len(value) != 4:
                            continue
                        swapped_value = [
                            int(value[0]),
                            int(value[2]),
                            int(value[3]),
                            int(value[1]),
                        ]
                        if [int(v) for v in list(value)] != swapped_value:
                            opts[key] = swapped_value
                            options_changed = True
                            shape_rewritten = True
                    if options_changed:
                        gate_reshape_op.options = opts

                if not shape_rewritten:
                    continue

                # Bypass pre-gate transpose into flatten reshape when present.
                if removable_pre_gate_idx is not None:
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=flat_op,
                        input_index=0,
                        new_input_name=pool_out_name,
                    )
                if rewrite_pool_mean_to_nhwc:
                    pool_inputs = [str(v) for v in list(pool_op.inputs)]
                    pool_inputs[0] = str(x_nhwc_name)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=pool_op,
                        new_inputs=pool_inputs,
                    )
                    if mapped_mean_axes is not None:
                        mean_axes_name = _clone_const_tensor_for_operator_input(
                            op=pool_op,
                            input_index=1,
                            tag="__se_fc_nhwc_axes",
                        )
                        mean_axes_tensor = model_ir.tensors.get(str(mean_axes_name), None)
                        if not _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_mean_axes]):
                            continue
                        if isinstance(pool_op.options, dict):
                            mean_opts = dict(pool_op.options)
                            for key in ["axis", "axes", "onnxRawAxes"]:
                                value = mean_opts.get(key, None)
                                if isinstance(value, list) and len(value) == len(mapped_mean_axes):
                                    mean_opts[key] = [int(v) for v in mapped_mean_axes]
                            pool_op.options = mean_opts

                # Bypass pre-mul transpose into MUL.
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                mul_inputs[int(x_side_index)] = str(x_nhwc_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )

                # Gate/MUL metadata: convert from NCHW to NHWC.
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(gate_post_name, None),
                    perm_nchw_to_nhwc,
                )
                if gate_requires_post_unary:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(gate_reshape_op.outputs[0]), None),
                        perm_nchw_to_nhwc,
                    )
                if rewrite_pool_mean_to_nhwc:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(pool_out_name), None),
                        perm_nchw_to_nhwc,
                    )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(mul_out_name, None),
                    perm_nchw_to_nhwc,
                )
                if post_bridge_via_unary and post_bridge_unary_idx is not None:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(post_bridge_output_name), None),
                        perm_nchw_to_nhwc,
                    )

                # Remove post transpose by letting bridge producer emit NHWC directly.
                representative_output_name = str(post_output_names[0])
                _set_operator_outputs(model_ir=model_ir, op=post_bridge_op, new_outputs=[representative_output_name])
                for alias_name in post_output_names[1:]:
                    _replace_tensor_inputs(model_ir, alias_name, representative_output_name)

                old_mul_tensor = model_ir.tensors.get(
                    str(post_bridge_output_name),
                    None,
                )
                representative_tensor = model_ir.tensors.get(representative_output_name, None)
                if old_mul_tensor is not None and representative_tensor is not None:
                    representative_tensor.dtype = str(old_mul_tensor.dtype)
                    representative_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                    representative_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                    representative_tensor.shape_signature = (
                        [int(v) for v in list(old_mul_tensor.shape_signature)]
                        if old_mul_tensor.shape_signature is not None
                        else [int(v) for v in list(old_mul_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        representative_tensor,
                        perm_nchw_to_nhwc,
                    )

                remove_indices = set(int(v) for v in post_indices)
                pre_mul_exempt_users = {int(mul_idx)}
                if rewrite_pool_mean_to_nhwc and str(pool_op.op_type) == "MEAN":
                    pre_mul_exempt_users.add(int(pool_idx))
                pre_mul_remaining_users = [
                    int(v)
                    for v in x_nchw_users
                    if int(v) not in pre_mul_exempt_users
                ]
                if len(pre_mul_remaining_users) == 0:
                    remove_indices.add(int(pre_mul_idx))

                if removable_pre_gate_idx is not None:
                    pre_gate_users = [int(v) for v in consumers.get(pre_gate_name, [])]
                    pre_gate_remaining_users = [int(v) for v in pre_gate_users if int(v) != int(flat_idx)]
                    if len(pre_gate_remaining_users) == 0:
                        remove_indices.add(int(removable_pre_gate_idx))

                for remove_idx in sorted(list(remove_indices), reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

            # Alternate float SE-like path:
            #   x_nhwc --T--> x_nchw --ADD(c)--> a --(MUL|DIV)(c)--> b --MUL(x_nchw,b)--> xh_nchw
            #   xh_nchw --MEAN--> m_nchw --T--> m_nhwc --CONV--CONV--> g_nhwc --T--> g_nchw
            #   g_nchw --MUL(alpha)--ADD(beta)--RELU_0_TO_1--> h_nchw
            #   MUL(xh_nchw, h_nchw) -> y_nchw --T--> y_nhwc
            if not changed:
                add0_idx: Optional[int] = None
                main_mul_idx: Optional[int] = None
                add0_out_name: Optional[str] = None
                main_mul_gate_name: Optional[str] = None

                for user_idx in x_nchw_users:
                    user_op = model_ir.operators[int(user_idx)]
                    user_type = str(user_op.op_type)
                    if user_type == "ADD":
                        if len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                            continue
                        add_inputs = [str(v) for v in list(user_op.inputs)]
                        if x_nchw_name == add_inputs[0]:
                            side_name = str(add_inputs[1])
                        elif x_nchw_name == add_inputs[1]:
                            side_name = str(add_inputs[0])
                        else:
                            continue
                        if not _is_singleton_constant_tensor(model_ir, side_name):
                            continue
                        if add0_idx is not None:
                            add0_idx = None
                            break
                        add0_idx = int(user_idx)
                        add0_out_name = str(user_op.outputs[0])
                    elif user_type == "MUL":
                        if len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                            continue
                        mul_inputs = [str(v) for v in list(user_op.inputs)]
                        if x_nchw_name == mul_inputs[0]:
                            other_name = str(mul_inputs[1])
                        elif x_nchw_name == mul_inputs[1]:
                            other_name = str(mul_inputs[0])
                        else:
                            continue
                        if main_mul_idx is not None:
                            main_mul_idx = None
                            break
                        main_mul_idx = int(user_idx)
                        main_mul_gate_name = str(other_name)

                if (
                    add0_idx is not None
                    and main_mul_idx is not None
                    and add0_out_name is not None
                    and main_mul_gate_name is not None
                ):
                    add0_op = model_ir.operators[int(add0_idx)]
                    main_mul_op = model_ir.operators[int(main_mul_idx)]
                    main_mul_out_name = str(main_mul_op.outputs[0])
                    if main_mul_out_name not in model_outputs:
                        scale_idx = producers.get(str(main_mul_gate_name), None)
                        if scale_idx is not None:
                            scale_op = model_ir.operators[int(scale_idx)]
                            scale_type = str(scale_op.op_type)
                            if (
                                scale_type in {"MUL", "DIV"}
                                and len(scale_op.inputs) == 2
                                and len(scale_op.outputs) == 1
                                and str(scale_op.outputs[0]) == str(main_mul_gate_name)
                            ):
                                scale_inputs = [str(v) for v in list(scale_op.inputs)]
                                if str(add0_out_name) == scale_inputs[0]:
                                    scale_side_name = str(scale_inputs[1])
                                elif str(add0_out_name) == scale_inputs[1]:
                                    scale_side_name = str(scale_inputs[0])
                                else:
                                    scale_side_name = ""
                                if (
                                    len(scale_side_name) > 0
                                    and _is_singleton_constant_tensor(model_ir, scale_side_name)
                                ):
                                    main_out_users = [int(v) for v in consumers.get(main_mul_out_name, [])]
                                    mean_idx: Optional[int] = None
                                    out_mul_idx: Optional[int] = None
                                    for user_idx in main_out_users:
                                        user_op = model_ir.operators[int(user_idx)]
                                        user_type = str(user_op.op_type)
                                        if (
                                            user_type == "MEAN"
                                            and len(user_op.inputs) >= 2
                                            and len(user_op.outputs) == 1
                                            and str(user_op.inputs[0]) == str(main_mul_out_name)
                                            and bool(user_op.options.get("keepDims", False))
                                        ):
                                            if mean_idx is not None:
                                                mean_idx = None
                                                break
                                            mean_idx = int(user_idx)
                                        elif (
                                            user_type == "MUL"
                                            and len(user_op.inputs) == 2
                                            and len(user_op.outputs) == 1
                                            and str(main_mul_out_name) in {str(v) for v in list(user_op.inputs)}
                                        ):
                                            if out_mul_idx is not None:
                                                out_mul_idx = None
                                                break
                                            out_mul_idx = int(user_idx)
                                    if mean_idx is not None and out_mul_idx is not None:
                                        mean_op = model_ir.operators[int(mean_idx)]
                                        mean_out_name = str(mean_op.outputs[0])
                                        if mean_out_name not in model_outputs:
                                            mean_axes_name = str(mean_op.inputs[1])
                                            mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                                            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
                                            if mean_axes_vals is not None and len(mean_axes_vals) > 0:
                                                mean_users = [int(v) for v in consumers.get(mean_out_name, [])]
                                                if len(mean_users) == 1:
                                                    post_mean_idx = int(mean_users[0])
                                                    post_mean_op = model_ir.operators[int(post_mean_idx)]
                                                    if (
                                                        str(post_mean_op.op_type) == "TRANSPOSE"
                                                        and len(post_mean_op.inputs) >= 2
                                                        and len(post_mean_op.outputs) == 1
                                                        and str(post_mean_op.inputs[0]) == str(mean_out_name)
                                                        and _read_transpose_perm(model_ir, post_mean_op) == perm_nchw_to_nhwc
                                                    ):
                                                        post_mean_out_name = str(post_mean_op.outputs[0])
                                                        post_mean_out_users = [int(v) for v in consumers.get(post_mean_out_name, [])]
                                                        if len(post_mean_out_users) == 1:
                                                            conv1_idx = int(post_mean_out_users[0])
                                                            conv1_op = model_ir.operators[int(conv1_idx)]
                                                            if (
                                                                str(conv1_op.op_type) in {"CONV_2D", "DEPTHWISE_CONV_2D"}
                                                                and len(conv1_op.inputs) >= 1
                                                                and len(conv1_op.outputs) == 1
                                                                and str(conv1_op.inputs[0]) == str(post_mean_out_name)
                                                            ):
                                                                conv1_out_name = str(conv1_op.outputs[0])
                                                                conv1_out_users = [int(v) for v in consumers.get(conv1_out_name, [])]
                                                                if len(conv1_out_users) == 1:
                                                                    conv2_idx = int(conv1_out_users[0])
                                                                    conv2_op = model_ir.operators[int(conv2_idx)]
                                                                    if (
                                                                        str(conv2_op.op_type) in {"CONV_2D", "DEPTHWISE_CONV_2D"}
                                                                        and len(conv2_op.inputs) >= 1
                                                                        and len(conv2_op.outputs) == 1
                                                                        and str(conv2_op.inputs[0]) == str(conv1_out_name)
                                                                    ):
                                                                        conv2_out_name = str(conv2_op.outputs[0])
                                                                        conv2_out_users = [int(v) for v in consumers.get(conv2_out_name, [])]
                                                                        if len(conv2_out_users) == 1:
                                                                            pre_gate_idx = int(conv2_out_users[0])
                                                                            pre_gate_op = model_ir.operators[int(pre_gate_idx)]
                                                                            if (
                                                                                str(pre_gate_op.op_type) == "TRANSPOSE"
                                                                                and len(pre_gate_op.inputs) >= 2
                                                                                and len(pre_gate_op.outputs) == 1
                                                                                and str(pre_gate_op.inputs[0]) == str(conv2_out_name)
                                                                                and _read_transpose_perm(model_ir, pre_gate_op) == perm_nhwc_to_nchw
                                                                            ):
                                                                                pre_gate_out_name = str(pre_gate_op.outputs[0])
                                                                                pre_gate_users = [int(v) for v in consumers.get(pre_gate_out_name, [])]
                                                                                if len(pre_gate_users) == 1:
                                                                                    hs_mul_idx = int(pre_gate_users[0])
                                                                                    hs_mul_op = model_ir.operators[int(hs_mul_idx)]
                                                                                    if (
                                                                                        str(hs_mul_op.op_type) == "MUL"
                                                                                        and len(hs_mul_op.inputs) == 2
                                                                                        and len(hs_mul_op.outputs) == 1
                                                                                    ):
                                                                                        hs_mul_inputs = [str(v) for v in list(hs_mul_op.inputs)]
                                                                                        if pre_gate_out_name == hs_mul_inputs[0]:
                                                                                            hs_mul_side_name = str(hs_mul_inputs[1])
                                                                                        elif pre_gate_out_name == hs_mul_inputs[1]:
                                                                                            hs_mul_side_name = str(hs_mul_inputs[0])
                                                                                        else:
                                                                                            hs_mul_side_name = ""
                                                                                        hs_mul_out_name = str(hs_mul_op.outputs[0])
                                                                                        if (
                                                                                            len(hs_mul_side_name) > 0
                                                                                            and _is_singleton_constant_tensor(model_ir, hs_mul_side_name)
                                                                                        ):
                                                                                            hs_mul_users = [int(v) for v in consumers.get(hs_mul_out_name, [])]
                                                                                            if len(hs_mul_users) == 1:
                                                                                                hs_add_idx = int(hs_mul_users[0])
                                                                                                hs_add_op = model_ir.operators[int(hs_add_idx)]
                                                                                                if (
                                                                                                    str(hs_add_op.op_type) == "ADD"
                                                                                                    and len(hs_add_op.inputs) == 2
                                                                                                    and len(hs_add_op.outputs) == 1
                                                                                                ):
                                                                                                    hs_add_inputs = [str(v) for v in list(hs_add_op.inputs)]
                                                                                                    if hs_mul_out_name == hs_add_inputs[0]:
                                                                                                        hs_add_side_name = str(hs_add_inputs[1])
                                                                                                    elif hs_mul_out_name == hs_add_inputs[1]:
                                                                                                        hs_add_side_name = str(hs_add_inputs[0])
                                                                                                    else:
                                                                                                        hs_add_side_name = ""
                                                                                                    hs_add_out_name = str(hs_add_op.outputs[0])
                                                                                                    if (
                                                                                                        len(hs_add_side_name) > 0
                                                                                                        and _is_singleton_constant_tensor(model_ir, hs_add_side_name)
                                                                                                    ):
                                                                                                        hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
                                                                                                        if len(hs_add_users) == 1:
                                                                                                            hs_out_name: Optional[str] = None
                                                                                                            hs_gate_op = model_ir.operators[int(hs_add_users[0])]
                                                                                                            if (
                                                                                                                str(hs_gate_op.op_type) == "RELU_0_TO_1"
                                                                                                                and len(hs_gate_op.inputs) == 1
                                                                                                                and len(hs_gate_op.outputs) == 1
                                                                                                                and str(hs_gate_op.inputs[0]) == str(hs_add_out_name)
                                                                                                            ):
                                                                                                                hs_out_name = str(hs_gate_op.outputs[0])
                                                                                                            elif (
                                                                                                                str(hs_gate_op.op_type) == "MAXIMUM"
                                                                                                                and len(hs_gate_op.inputs) == 2
                                                                                                                and len(hs_gate_op.outputs) == 1
                                                                                                            ):
                                                                                                                hs_max_inputs = [str(v) for v in list(hs_gate_op.inputs)]
                                                                                                                if hs_add_out_name == hs_max_inputs[0]:
                                                                                                                    hs_max_side_name = str(hs_max_inputs[1])
                                                                                                                elif hs_add_out_name == hs_max_inputs[1]:
                                                                                                                    hs_max_side_name = str(hs_max_inputs[0])
                                                                                                                else:
                                                                                                                    hs_max_side_name = ""
                                                                                                                hs_max_out_name = str(hs_gate_op.outputs[0])
                                                                                                                if (
                                                                                                                    len(hs_max_side_name) > 0
                                                                                                                    and _is_singleton_constant_tensor(model_ir, hs_max_side_name)
                                                                                                                ):
                                                                                                                    hs_max_users = [int(v) for v in consumers.get(hs_max_out_name, [])]
                                                                                                                    if len(hs_max_users) == 1:
                                                                                                                        hs_min_op = model_ir.operators[int(hs_max_users[0])]
                                                                                                                        if (
                                                                                                                            str(hs_min_op.op_type) == "MINIMUM"
                                                                                                                            and len(hs_min_op.inputs) == 2
                                                                                                                            and len(hs_min_op.outputs) == 1
                                                                                                                        ):
                                                                                                                            hs_min_inputs = [str(v) for v in list(hs_min_op.inputs)]
                                                                                                                            if hs_max_out_name == hs_min_inputs[0]:
                                                                                                                                hs_min_side_name = str(hs_min_inputs[1])
                                                                                                                            elif hs_max_out_name == hs_min_inputs[1]:
                                                                                                                                hs_min_side_name = str(hs_min_inputs[0])
                                                                                                                            else:
                                                                                                                                hs_min_side_name = ""
                                                                                                                            if (
                                                                                                                                len(hs_min_side_name) > 0
                                                                                                                                and _is_singleton_constant_tensor(model_ir, hs_min_side_name)
                                                                                                                            ):
                                                                                                                                hs_out_name = str(hs_min_op.outputs[0])
                                                                                                            if hs_out_name is not None:
                                                                                                                out_mul_op = model_ir.operators[int(out_mul_idx)]
                                                                                                                out_mul_inputs = [str(v) for v in list(out_mul_op.inputs)]
                                                                                                                if (
                                                                                                                    str(main_mul_out_name) in out_mul_inputs
                                                                                                                    and str(hs_out_name) in out_mul_inputs
                                                                                                                ):
                                                                                                                    out_mul_out_name = str(out_mul_op.outputs[0])
                                                                                                                    if out_mul_out_name not in model_outputs:
                                                                                                                        out_mul_users = [int(v) for v in consumers.get(out_mul_out_name, [])]
                                                                                                                        post_indices: List[int] = []
                                                                                                                        post_output_names: List[str] = []
                                                                                                                        valid_posts = True
                                                                                                                        for user_idx in out_mul_users:
                                                                                                                            user_op = model_ir.operators[int(user_idx)]
                                                                                                                            if (
                                                                                                                                str(user_op.op_type) != "TRANSPOSE"
                                                                                                                                or len(user_op.inputs) < 2
                                                                                                                                or len(user_op.outputs) != 1
                                                                                                                                or str(user_op.inputs[0]) != str(out_mul_out_name)
                                                                                                                                or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                                                                                                                                or str(user_op.outputs[0]) in model_outputs
                                                                                                                            ):
                                                                                                                                valid_posts = False
                                                                                                                                break
                                                                                                                            post_indices.append(int(user_idx))
                                                                                                                            post_output_names.append(str(user_op.outputs[0]))
                                                                                                                        if valid_posts and len(post_indices) > 0:
                                                                                                                            rank = 4
                                                                                                                            x_nhwc_tensor = model_ir.tensors.get(x_nhwc_name, None)
                                                                                                                            if (
                                                                                                                                x_nhwc_tensor is not None
                                                                                                                                and x_nhwc_tensor.shape is not None
                                                                                                                                and len(x_nhwc_tensor.shape) > 0
                                                                                                                            ):
                                                                                                                                rank = int(len(list(x_nhwc_tensor.shape)))
                                                                                                                            if rank == 4:
                                                                                                                                normalized_axes: List[int] = []
                                                                                                                                valid_axes = True
                                                                                                                                for axis in mean_axes_vals:
                                                                                                                                    a = int(axis)
                                                                                                                                    if a < 0:
                                                                                                                                        a += int(rank)
                                                                                                                                    if a < 0 or a >= int(rank):
                                                                                                                                        valid_axes = False
                                                                                                                                        break
                                                                                                                                    normalized_axes.append(int(a))
                                                                                                                                if valid_axes and len(normalized_axes) > 0:
                                                                                                                                    mapped_axes = [
                                                                                                                                        int(perm_nhwc_to_nchw[int(axis)])
                                                                                                                                        for axis in normalized_axes
                                                                                                                                    ]
                                                                                                                                    if (
                                                                                                                                        _write_const_ints_to_tensor(
                                                                                                                                            mean_axes_tensor,
                                                                                                                                            [int(v) for v in mapped_axes],
                                                                                                                                        )
                                                                                                                                        or _read_const_ints_from_tensor(mean_axes_tensor)
                                                                                                                                        == [int(v) for v in mapped_axes]
                                                                                                                                    ):
                                                                                                                                        add0_inputs = [str(v) for v in list(add0_op.inputs)]
                                                                                                                                        add0_inputs = [
                                                                                                                                            str(x_nhwc_name) if str(v) == str(x_nchw_name) else str(v)
                                                                                                                                            for v in add0_inputs
                                                                                                                                        ]
                                                                                                                                        _set_operator_inputs(
                                                                                                                                            model_ir=model_ir,
                                                                                                                                            op=add0_op,
                                                                                                                                            new_inputs=add0_inputs,
                                                                                                                                        )

                                                                                                                                        main_mul_inputs = [str(v) for v in list(main_mul_op.inputs)]
                                                                                                                                        main_mul_inputs = [
                                                                                                                                            str(x_nhwc_name) if str(v) == str(x_nchw_name) else str(v)
                                                                                                                                            for v in main_mul_inputs
                                                                                                                                        ]
                                                                                                                                        _set_operator_inputs(
                                                                                                                                            model_ir=model_ir,
                                                                                                                                            op=main_mul_op,
                                                                                                                                            new_inputs=main_mul_inputs,
                                                                                                                                        )

                                                                                                                                        conv1_inputs = [str(v) for v in list(conv1_op.inputs)]
                                                                                                                                        conv1_inputs = [
                                                                                                                                            str(mean_out_name) if str(v) == str(post_mean_out_name) else str(v)
                                                                                                                                            for v in conv1_inputs
                                                                                                                                        ]
                                                                                                                                        _set_operator_inputs(
                                                                                                                                            model_ir=model_ir,
                                                                                                                                            op=conv1_op,
                                                                                                                                            new_inputs=conv1_inputs,
                                                                                                                                        )

                                                                                                                                        hs_mul_inputs_updated = [str(v) for v in list(hs_mul_op.inputs)]
                                                                                                                                        hs_mul_inputs_updated = [
                                                                                                                                            str(conv2_out_name) if str(v) == str(pre_gate_out_name) else str(v)
                                                                                                                                            for v in hs_mul_inputs_updated
                                                                                                                                        ]
                                                                                                                                        _set_operator_inputs(
                                                                                                                                            model_ir=model_ir,
                                                                                                                                            op=hs_mul_op,
                                                                                                                                            new_inputs=hs_mul_inputs_updated,
                                                                                                                                        )

                                                                                                                                        for tensor_name in [
                                                                                                                                            str(add0_out_name),
                                                                                                                                            str(main_mul_gate_name),
                                                                                                                                            str(main_mul_out_name),
                                                                                                                                            str(mean_out_name),
                                                                                                                                            str(hs_mul_out_name),
                                                                                                                                            str(hs_add_out_name),
                                                                                                                                            str(hs_out_name),
                                                                                                                                            str(out_mul_out_name),
                                                                                                                                        ]:
                                                                                                                                            _permute_tensor_metadata_if_rank_matches(
                                                                                                                                                model_ir.tensors.get(tensor_name, None),
                                                                                                                                                perm_nchw_to_nhwc,
                                                                                                                                            )

                                                                                                                                        canonical_post_output_name = str(post_output_names[0])
                                                                                                                                        _set_operator_outputs(
                                                                                                                                            model_ir=model_ir,
                                                                                                                                            op=out_mul_op,
                                                                                                                                            new_outputs=[canonical_post_output_name],
                                                                                                                                        )
                                                                                                                                        for alias_name in post_output_names[1:]:
                                                                                                                                            _replace_tensor_inputs(
                                                                                                                                                model_ir,
                                                                                                                                                str(alias_name),
                                                                                                                                                canonical_post_output_name,
                                                                                                                                            )

                                                                                                                                        old_out_mul_tensor = model_ir.tensors.get(out_mul_out_name, None)
                                                                                                                                        canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
                                                                                                                                        if old_out_mul_tensor is not None and canonical_tensor is not None:
                                                                                                                                            canonical_tensor.dtype = str(old_out_mul_tensor.dtype)
                                                                                                                                            canonical_tensor.quantization = _clone_quantization(old_out_mul_tensor.quantization)
                                                                                                                                            canonical_tensor.shape = [int(v) for v in list(old_out_mul_tensor.shape)]
                                                                                                                                            canonical_tensor.shape_signature = (
                                                                                                                                                [int(v) for v in list(old_out_mul_tensor.shape_signature)]
                                                                                                                                                if old_out_mul_tensor.shape_signature is not None
                                                                                                                                                else [int(v) for v in list(old_out_mul_tensor.shape)]
                                                                                                                                            )
                                                                                                                                            _permute_tensor_metadata_if_rank_matches(
                                                                                                                                                canonical_tensor,
                                                                                                                                                perm_nchw_to_nhwc,
                                                                                                                                            )

                                                                                                                                        remove_indices = set(int(v) for v in post_indices)
                                                                                                                                        remove_indices.add(int(post_mean_idx))
                                                                                                                                        remove_indices.add(int(pre_gate_idx))
                                                                                                                                        pre_mul_remaining_users = [
                                                                                                                                            int(v)
                                                                                                                                            for v in x_nchw_users
                                                                                                                                            if int(v) not in {int(add0_idx), int(main_mul_idx)}
                                                                                                                                        ]
                                                                                                                                        if len(pre_mul_remaining_users) == 0:
                                                                                                                                            remove_indices.add(int(pre_mul_idx))
                                                                                                                                        for remove_idx in sorted(list(remove_indices), reverse=True):
                                                                                                                                            del model_ir.operators[int(remove_idx)]

                                                                                                                                        rewritten += 1
                                                                                                                                        changed = True

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_se_fc_mul_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_sum_logistic_muladd_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around SUM/LOGISTIC/SUB + MUL + ADD merge blocks.

    Target:
      a_nhwc --T(0,3,1,2)--> a_nchw
      b_nhwc --T(0,3,1,2)--> b_nchw
      c_nhwc --T(0,3,1,2)--> c_nchw --SUM(axis=1, keepDims=True)--> s_nchw --LOGISTIC--> sig_nchw
      one-sigmoid branch: SUB(one, sig_nchw) -> sub_nchw (optional)
      MUL(sub_or_sig_nchw, a_nchw) -> m0_nchw
      MUL(sig_nchw, b_nchw)        -> m1_nchw
      ADD(m0_nchw, m1_nchw) -> y_nchw
      y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - Bypass all three pre-transposes (a/b/c) to NHWC.
      - Remap SUM reduction axes from NCHW to NHWC.
      - Keep SUM/LOGISTIC/SUB/MUL/ADD in NHWC.
      - Remove post-transpose by making ADD produce NHWC output directly.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        def _extract_mul_pre_gate(
            mul_idx: int,
        ) -> Optional[Tuple[int, str, str, str]]:
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                return None
            pre_idx: Optional[int] = None
            pre_input_name: Optional[str] = None
            pre_output_name: Optional[str] = None
            gate_name: Optional[str] = None
            for input_name in list(mul_op.inputs):
                prod_idx = producers.get(str(input_name), None)
                if prod_idx is not None:
                    prod_op = model_ir.operators[int(prod_idx)]
                    if (
                        str(prod_op.op_type) == "TRANSPOSE"
                        and len(prod_op.inputs) >= 2
                        and len(prod_op.outputs) == 1
                        and str(prod_op.outputs[0]) == str(input_name)
                        and _read_transpose_perm(model_ir, prod_op) == perm_nhwc_to_nchw
                        and str(prod_op.inputs[0]) not in model_outputs
                        and str(prod_op.outputs[0]) not in model_outputs
                    ):
                        if pre_idx is not None:
                            return None
                        pre_idx = int(prod_idx)
                        pre_input_name = str(prod_op.inputs[0])
                        pre_output_name = str(prod_op.outputs[0])
                        continue
                if gate_name is not None:
                    return None
                gate_name = str(input_name)
            if pre_idx is None or pre_input_name is None or pre_output_name is None or gate_name is None:
                return None
            return int(pre_idx), str(pre_input_name), str(pre_output_name), str(gate_name)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue
            add_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if add_output_name in model_outputs or post_output_name in model_outputs:
                continue

            add_idx = producers.get(add_output_name, None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            if str(add_op.outputs[0]) != add_output_name:
                continue

            add_inputs = [str(v) for v in list(add_op.inputs)]
            mul0_idx = producers.get(add_inputs[0], None)
            mul1_idx = producers.get(add_inputs[1], None)
            if mul0_idx is None or mul1_idx is None:
                continue
            if int(mul0_idx) == int(mul1_idx):
                continue

            mul0_info = _extract_mul_pre_gate(int(mul0_idx))
            mul1_info = _extract_mul_pre_gate(int(mul1_idx))
            if mul0_info is None or mul1_info is None:
                continue

            pre0_idx, pre0_in_name, pre0_out_name, gate0_name = mul0_info
            pre1_idx, pre1_in_name, pre1_out_name, gate1_name = mul1_info
            if int(pre0_idx) == int(pre1_idx):
                continue

            # Identify LOGISTIC gate and optional SUB(one, sigmoid) gate.
            logistic_idx: Optional[int] = None
            logistic_out_name: Optional[str] = None
            sub_idx: Optional[int] = None
            sub_out_name: Optional[str] = None

            def _try_resolve_gate_pair(
                g_logistic: str,
                g_other: str,
            ) -> Optional[Tuple[int, str, Optional[int], Optional[str]]]:
                g_logistic_prod = producers.get(str(g_logistic), None)
                if g_logistic_prod is None:
                    return None
                g_logistic_op = model_ir.operators[int(g_logistic_prod)]
                if str(g_logistic_op.op_type) != "LOGISTIC" or len(g_logistic_op.outputs) != 1:
                    return None
                if str(g_logistic_op.outputs[0]) != str(g_logistic):
                    return None
                if str(g_other) == str(g_logistic):
                    return int(g_logistic_prod), str(g_logistic), None, None
                g_other_prod = producers.get(str(g_other), None)
                if g_other_prod is None:
                    return None
                g_other_op = model_ir.operators[int(g_other_prod)]
                if str(g_other_op.op_type) != "SUB" or len(g_other_op.inputs) != 2 or len(g_other_op.outputs) != 1:
                    return None
                if str(g_other_op.outputs[0]) != str(g_other):
                    return None
                if str(g_logistic) not in {str(v) for v in list(g_other_op.inputs)}:
                    return None
                return int(g_logistic_prod), str(g_logistic), int(g_other_prod), str(g_other)

            resolved = _try_resolve_gate_pair(gate0_name, gate1_name)
            if resolved is None:
                resolved = _try_resolve_gate_pair(gate1_name, gate0_name)
            if resolved is None:
                continue
            logistic_idx, logistic_out_name, sub_idx, sub_out_name = resolved

            logistic_op = model_ir.operators[int(logistic_idx)]
            if len(logistic_op.inputs) != 1:
                continue
            sum_output_name = str(logistic_op.inputs[0])
            if sum_output_name in model_outputs:
                continue

            sum_idx = producers.get(sum_output_name, None)
            if sum_idx is None:
                continue
            sum_op = model_ir.operators[int(sum_idx)]
            if str(sum_op.op_type) != "SUM" or len(sum_op.inputs) < 2 or len(sum_op.outputs) != 1:
                continue
            if str(sum_op.outputs[0]) != sum_output_name:
                continue
            if str(sum_op.inputs[0]) in model_outputs:
                continue

            pre2_out_name = str(sum_op.inputs[0])
            pre2_idx = producers.get(pre2_out_name, None)
            if pre2_idx is None:
                continue
            pre2_op = model_ir.operators[int(pre2_idx)]
            if (
                str(pre2_op.op_type) != "TRANSPOSE"
                or len(pre2_op.inputs) < 2
                or len(pre2_op.outputs) != 1
                or str(pre2_op.outputs[0]) != pre2_out_name
                or _read_transpose_perm(model_ir, pre2_op) != perm_nhwc_to_nchw
            ):
                continue
            pre2_in_name = str(pre2_op.inputs[0])
            if pre2_in_name in model_outputs or pre2_out_name in model_outputs:
                continue

            # Pre-transpose outputs must be local to this chain.
            if set(int(v) for v in consumers.get(pre0_out_name, [])) != {int(mul0_idx)} and set(
                int(v) for v in consumers.get(pre0_out_name, [])
            ) != {int(mul1_idx)}:
                continue
            if set(int(v) for v in consumers.get(pre1_out_name, [])) != {int(mul0_idx)} and set(
                int(v) for v in consumers.get(pre1_out_name, [])
            ) != {int(mul1_idx)}:
                continue
            if set(int(v) for v in consumers.get(pre2_out_name, [])) != {int(sum_idx)}:
                continue

            # ADD output must fan out only to inverse post-transposes.
            add_out_users = [int(v) for v in consumers.get(add_output_name, [])]
            if len(add_out_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            valid_posts = True
            for user_idx in add_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == add_output_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    valid_posts = False
                    break
            if not valid_posts or len(post_indices) == 0:
                continue

            # Remap SUM axes from NCHW to NHWC.
            sum_axes_name = str(sum_op.inputs[1])
            sum_axes_tensor = model_ir.tensors.get(sum_axes_name, None)
            sum_axes_vals = _read_const_ints_from_tensor(sum_axes_tensor)
            if sum_axes_vals is None or len(sum_axes_vals) == 0:
                continue
            rank = 4
            pre2_in_tensor = model_ir.tensors.get(pre2_in_name, None)
            if pre2_in_tensor is not None and pre2_in_tensor.shape is not None and len(pre2_in_tensor.shape) > 0:
                rank = int(len(pre2_in_tensor.shape))
            if rank != 4:
                continue
            old_to_new_axis = _invert_perm(perm_nchw_to_nhwc)
            if old_to_new_axis is None:
                continue
            mapped_axes: List[int] = []
            valid_axes = True
            for axis in sum_axes_vals:
                a = int(axis)
                if a < 0:
                    a += int(rank)
                if a < 0 or a >= int(rank):
                    valid_axes = False
                    break
                mapped_axes.append(int(old_to_new_axis[int(a)]))
            if not valid_axes or len(mapped_axes) == 0:
                continue
            _write_const_ints_to_tensor(sum_axes_tensor, [int(v) for v in mapped_axes])

            # Rewire SUM to NHWC source.
            sum_inputs = [str(v) for v in list(sum_op.inputs)]
            sum_inputs[0] = str(pre2_in_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=sum_op,
                new_inputs=sum_inputs,
            )

            # Rewire both MULs from transpose outputs to transpose inputs.
            for mul_idx, pre_out_name, pre_in_name in [
                (int(mul0_idx), str(pre0_out_name), str(pre0_in_name)),
                (int(mul1_idx), str(pre1_out_name), str(pre1_in_name)),
            ]:
                mul_op = model_ir.operators[int(mul_idx)]
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                replaced = False
                for in_idx, in_name in enumerate(list(mul_inputs)):
                    if str(in_name) == str(pre_out_name):
                        mul_inputs[in_idx] = str(pre_in_name)
                        replaced = True
                if not replaced:
                    valid_posts = False
                    break
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )
            if not valid_posts:
                continue

            # Update metadata to NHWC across the rewritten chain.
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(sum_output_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(logistic_out_name), None),
                perm_nchw_to_nhwc,
            )
            if sub_out_name is not None:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(sub_out_name), None),
                    perm_nchw_to_nhwc,
                )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(model_ir.operators[int(mul0_idx)].outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(model_ir.operators[int(mul1_idx)].outputs[0]), None),
                perm_nchw_to_nhwc,
            )

            # Canonicalize ADD output to post-transpose output tensor.
            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            old_add_out_tensor = model_ir.tensors.get(add_output_name, None)
            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_add_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_add_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_add_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_add_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_add_out_tensor.shape_signature)]
                    if old_add_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_add_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {
                int(pre0_idx),
                int(pre1_idx),
                int(pre2_idx),
                *[int(v) for v in post_indices],
            }
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_sum_logistic_muladd_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_weighted_add_swish_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around weighted ADD + Swish chains.

    Target:
      a_nhwc --T(0,3,1,2)--> a_nchw --MUL(w0_scalar_like)--> wa_nchw
      b_nhwc --T(0,3,1,2)--> b_nchw --MUL(w1_scalar_like)--> wb_nchw
      ADD(wa_nchw, wb_nchw) -> s_nchw --LOGISTIC--> sig_nchw
      MUL(s_nchw, sig_nchw) -> y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - Bypass both pre-transposes to NHWC.
      - Keep ADD/LOGISTIC/MUL in NHWC.
      - Remove inverse post-transpose by making MUL emit NHWC directly.
      - Keep one adapter transpose only when legacy NCHW users exist.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        def _extract_weighted_mul(
            *,
            mul_idx: int,
            expected_output_name: str,
            add_idx: int,
        ) -> Optional[Tuple[int, str, str, str, str]]:
            mul_op = model_ir.operators[int(mul_idx)]
            if (
                str(mul_op.op_type) != "MUL"
                or len(mul_op.inputs) != 2
                or len(mul_op.outputs) != 1
                or str(mul_op.outputs[0]) != str(expected_output_name)
            ):
                return None

            inputs = [str(v) for v in list(mul_op.inputs)]
            pre_idx: Optional[int] = None
            pre_in_name: Optional[str] = None
            pre_out_name: Optional[str] = None
            side_name: Optional[str] = None
            for input_name in inputs:
                prod_idx = producers.get(str(input_name), None)
                if prod_idx is None:
                    if side_name is not None:
                        return None
                    side_name = str(input_name)
                    continue
                prod_op = model_ir.operators[int(prod_idx)]
                if (
                    str(prod_op.op_type) == "TRANSPOSE"
                    and len(prod_op.inputs) >= 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, prod_op) == perm_nhwc_to_nchw
                    and str(prod_op.inputs[0]) not in model_outputs
                    and str(prod_op.outputs[0]) not in model_outputs
                ):
                    if pre_idx is not None:
                        return None
                    pre_idx = int(prod_idx)
                    pre_in_name = str(prod_op.inputs[0])
                    pre_out_name = str(prod_op.outputs[0])
                    continue
                if side_name is not None:
                    return None
                side_name = str(input_name)

            if (
                pre_idx is None
                or pre_in_name is None
                or pre_out_name is None
                or side_name is None
                or not _is_scalar_like_tensor(model_ir, str(side_name))
            ):
                return None

            if set(int(v) for v in consumers.get(str(expected_output_name), [])) != {int(add_idx)}:
                return None
            if int(mul_idx) not in set(int(v) for v in consumers.get(str(pre_out_name), [])):
                return None

            return int(pre_idx), str(pre_in_name), str(pre_out_name), str(side_name), str(expected_output_name)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            swish_out_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if swish_out_name in model_outputs or post_output_name in model_outputs:
                continue

            swish_mul_idx = producers.get(swish_out_name, None)
            if swish_mul_idx is None:
                continue
            swish_mul_op = model_ir.operators[int(swish_mul_idx)]
            if (
                str(swish_mul_op.op_type) != "MUL"
                or len(swish_mul_op.inputs) != 2
                or len(swish_mul_op.outputs) != 1
                or str(swish_mul_op.outputs[0]) != swish_out_name
            ):
                continue

            swish_inputs = [str(v) for v in list(swish_mul_op.inputs)]
            logistic_idx: Optional[int] = None
            logistic_out_name: Optional[str] = None
            add_idx: Optional[int] = None
            add_out_name: Optional[str] = None

            for input_name in swish_inputs:
                prod_idx = producers.get(str(input_name), None)
                if prod_idx is None:
                    continue
                prod_op = model_ir.operators[int(prod_idx)]
                if (
                    str(prod_op.op_type) == "LOGISTIC"
                    and len(prod_op.inputs) == 1
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(input_name)
                ):
                    if logistic_idx is not None:
                        logistic_idx = None
                        break
                    logistic_idx = int(prod_idx)
                    logistic_out_name = str(input_name)
                    add_out_name = str(prod_op.inputs[0])
                else:
                    add_out_name = str(input_name)

            if logistic_idx is None or logistic_out_name is None or add_out_name is None:
                continue
            if add_out_name in model_outputs:
                continue

            add_idx = producers.get(add_out_name, None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if (
                str(add_op.op_type) != "ADD"
                or len(add_op.inputs) != 2
                or len(add_op.outputs) != 1
                or str(add_op.outputs[0]) != str(add_out_name)
            ):
                continue

            # Strict swish topology only.
            if set(int(v) for v in consumers.get(str(logistic_out_name), [])) != {int(swish_mul_idx)}:
                continue
            add_out_users = set(int(v) for v in consumers.get(str(add_out_name), []))
            if add_out_users != {int(logistic_idx), int(swish_mul_idx)}:
                continue

            add_inputs = [str(v) for v in list(add_op.inputs)]
            mul0_idx = producers.get(add_inputs[0], None)
            mul1_idx = producers.get(add_inputs[1], None)
            if mul0_idx is None or mul1_idx is None or int(mul0_idx) == int(mul1_idx):
                continue

            mul0_info = _extract_weighted_mul(
                mul_idx=int(mul0_idx),
                expected_output_name=str(add_inputs[0]),
                add_idx=int(add_idx),
            )
            mul1_info = _extract_weighted_mul(
                mul_idx=int(mul1_idx),
                expected_output_name=str(add_inputs[1]),
                add_idx=int(add_idx),
            )
            if mul0_info is None or mul1_info is None:
                continue

            pre0_idx, pre0_in_name, pre0_out_name, _side0_name, mul0_out_name = mul0_info
            pre1_idx, pre1_in_name, pre1_out_name, _side1_name, mul1_out_name = mul1_info

            if int(pre0_idx) == int(pre1_idx):
                continue

            swish_out_users = [int(v) for v in consumers.get(swish_out_name, []) if int(v) != int(swish_mul_idx)]
            if len(swish_out_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_posts = True
            for user_idx in swish_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(swish_out_name)
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))
            if not valid_posts or len(post_indices) == 0:
                continue

            # Rewire weighted MULs to NHWC sources.
            mul0_op = model_ir.operators[int(mul0_idx)]
            mul0_inputs = [str(v) for v in list(mul0_op.inputs)]
            for idx, input_name in enumerate(list(mul0_inputs)):
                if str(input_name) == str(pre0_out_name):
                    mul0_inputs[idx] = str(pre0_in_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul0_op,
                new_inputs=mul0_inputs,
            )

            mul1_op = model_ir.operators[int(mul1_idx)]
            mul1_inputs = [str(v) for v in list(mul1_op.inputs)]
            for idx, input_name in enumerate(list(mul1_inputs)):
                if str(input_name) == str(pre1_out_name):
                    mul1_inputs[idx] = str(pre1_in_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul1_op,
                new_inputs=mul1_inputs,
            )

            # LOGISTIC is layout-agnostic; keep its input/output names and update metadata.
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(logistic_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul0_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul1_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(add_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(swish_out_name), None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=swish_mul_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            old_swish_out_tensor = model_ir.tensors.get(str(swish_out_name), None)
            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_swish_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_swish_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_swish_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_swish_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_swish_out_tensor.shape_signature)]
                    if old_swish_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_swish_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[swish_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = set(int(v) for v in post_remove_indices)
            pre0_remaining_users = [int(v) for v in consumers.get(str(pre0_out_name), []) if int(v) != int(mul0_idx)]
            if len(pre0_remaining_users) == 0:
                remove_indices.add(int(pre0_idx))
            pre1_remaining_users = [int(v) for v in consumers.get(str(pre1_out_name), []) if int(v) != int(mul1_idx)]
            if len(pre1_remaining_users) == 0:
                remove_indices.add(int(pre1_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_weighted_add_swish_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_nested_weighted_add_swish_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around nested weighted-ADD trees followed by Swish.

    Target (constrained):
      leaves: MUL(w_scalar_like, T(0,3,1,2)(x_nhwc)) ...
      nested ADD tree over leaves -> sum_nchw
      LOGISTIC(sum_nchw) + MUL(sum_nchw, sigmoid) -> y_nchw
      y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - Rewire each leaf MUL to consume NHWC tensor directly.
      - Keep nested ADD + LOGISTIC + MUL in NHWC.
      - Remove post inverse transpose (or keep one adapter for legacy users).
      - Remove now-unused leaf pre-transposes.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            swish_out_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if swish_out_name in model_outputs or post_output_name in model_outputs:
                continue

            swish_mul_idx = producers.get(swish_out_name, None)
            if swish_mul_idx is None:
                continue
            swish_mul_op = model_ir.operators[int(swish_mul_idx)]
            if (
                str(swish_mul_op.op_type) != "MUL"
                or len(swish_mul_op.inputs) != 2
                or len(swish_mul_op.outputs) != 1
                or str(swish_mul_op.outputs[0]) != str(swish_out_name)
            ):
                continue

            swish_inputs = [str(v) for v in list(swish_mul_op.inputs)]
            logistic_idx: Optional[int] = None
            logistic_out_name: Optional[str] = None
            root_add_out_name: Optional[str] = None
            for input_name in swish_inputs:
                prod_idx = producers.get(str(input_name), None)
                if prod_idx is None:
                    continue
                prod_op = model_ir.operators[int(prod_idx)]
                if (
                    str(prod_op.op_type) == "LOGISTIC"
                    and len(prod_op.inputs) == 1
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(input_name)
                ):
                    if logistic_idx is not None:
                        logistic_idx = None
                        break
                    logistic_idx = int(prod_idx)
                    logistic_out_name = str(input_name)
                    root_add_out_name = str(prod_op.inputs[0])
                else:
                    root_add_out_name = str(input_name)
            if logistic_idx is None or logistic_out_name is None or root_add_out_name is None:
                continue
            if root_add_out_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(str(logistic_out_name), [])) != {int(swish_mul_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(root_add_out_name), [])) != {int(logistic_idx), int(swish_mul_idx)}:
                continue

            root_add_idx = producers.get(str(root_add_out_name), None)
            if root_add_idx is None:
                continue

            leaf_rewrites: List[Tuple[int, str, str]] = []
            leaf_pre_indices: List[int] = []
            add_indices: List[int] = []
            add_output_names: List[str] = []
            valid_tree = True

            visited_add_outputs: set[str] = set()

            def _collect_tree(tensor_name: str, parent_add_idx: Optional[int]) -> bool:
                nonlocal valid_tree
                prod_idx = producers.get(str(tensor_name), None)
                if prod_idx is None:
                    return False
                prod_op = model_ir.operators[int(prod_idx)]
                prod_type = str(prod_op.op_type)

                if (
                    prod_type == "ADD"
                    and len(prod_op.inputs) == 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(tensor_name)
                    and str(tensor_name) not in model_outputs
                ):
                    if str(tensor_name) in visited_add_outputs:
                        return False
                    visited_add_outputs.add(str(tensor_name))
                    add_indices.append(int(prod_idx))
                    add_output_names.append(str(tensor_name))

                    users = set(int(v) for v in consumers.get(str(tensor_name), []))
                    if parent_add_idx is None:
                        if users != {int(logistic_idx), int(swish_mul_idx)}:
                            return False
                    else:
                        if users != {int(parent_add_idx)}:
                            return False

                    in0 = str(prod_op.inputs[0])
                    in1 = str(prod_op.inputs[1])
                    if not _collect_tree(in0, int(prod_idx)):
                        return False
                    if not _collect_tree(in1, int(prod_idx)):
                        return False
                    return True

                if (
                    prod_type == "MUL"
                    and len(prod_op.inputs) == 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(tensor_name)
                    and str(tensor_name) not in model_outputs
                ):
                    mul_inputs = [str(v) for v in list(prod_op.inputs)]
                    scalar_indices = [
                        int(i)
                        for i, name in enumerate(mul_inputs)
                        if _is_scalar_like_tensor(model_ir, str(name))
                    ]
                    if len(scalar_indices) != 1:
                        return False
                    data_name = str(mul_inputs[1 - int(scalar_indices[0])])
                    pre_idx = producers.get(str(data_name), None)
                    if pre_idx is None:
                        return False
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != str(data_name)
                        or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    ):
                        return False
                    pre_in_name = str(pre_op.inputs[0])
                    pre_out_name = str(pre_op.outputs[0])
                    if pre_in_name in model_outputs or pre_out_name in model_outputs:
                        return False
                    if set(int(v) for v in consumers.get(pre_out_name, [])) != {int(prod_idx)}:
                        return False

                    leaf_rewrites.append((int(prod_idx), str(pre_out_name), str(pre_in_name)))
                    leaf_pre_indices.append(int(pre_idx))
                    return True

                return False

            valid_tree = _collect_tree(str(root_add_out_name), None)
            if not valid_tree:
                continue
            if len(leaf_rewrites) < 3:
                # Keep this pass focused on nested trees; flat 2-leaf pattern is handled elsewhere.
                continue

            swish_out_users = [int(v) for v in consumers.get(swish_out_name, []) if int(v) != int(swish_mul_idx)]
            if len(swish_out_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            for user_idx in swish_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(swish_out_name)
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))
            if len(post_indices) == 0:
                continue

            for mul_idx, pre_out_name, pre_in_name in list(leaf_rewrites):
                mul_op = model_ir.operators[int(mul_idx)]
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                replaced = False
                for idx, name in enumerate(list(mul_inputs)):
                    if str(name) == str(pre_out_name):
                        mul_inputs[idx] = str(pre_in_name)
                        replaced = True
                if not replaced:
                    valid_tree = False
                    break
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(mul_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
            if not valid_tree:
                continue

            for add_output_name in list(dict.fromkeys(add_output_names)):
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(add_output_name), None),
                    perm_nchw_to_nhwc,
                )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(logistic_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(swish_out_name), None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=swish_mul_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            old_swish_out_tensor = model_ir.tensors.get(str(swish_out_name), None)
            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_swish_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_swish_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_swish_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_swish_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_swish_out_tensor.shape_signature)]
                    if old_swish_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_swish_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[swish_out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = set(int(v) for v in post_remove_indices)
            remove_indices.update(int(v) for v in leaf_pre_indices)
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_nested_weighted_add_swish_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_logistic_muladd_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around LOGISTIC + MUL + ADD merge blocks.

    Target:
      skip_nhwc --T(0,3,1,2)--> skip_nchw
      data_nhwc --T(0,3,1,2)--> data_nchw
      gate_nhwc --T(0,3,1,2)--> gate_nchw --LOGISTIC--> sig_nchw
      MUL(data_nchw, sig_nchw) -> mul_nchw
      ADD(mul_nchw, skip_nchw) -> y_nchw
      y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - Bypass three pre-transposes (skip/data/gate) to NHWC.
      - Keep LOGISTIC/MUL/ADD in NHWC.
      - Remove inverse post-transpose by making ADD emit NHWC directly.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            add_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if add_output_name in model_outputs or post_output_name in model_outputs:
                continue

            add_idx = producers.get(add_output_name, None)
            if add_idx is None:
                continue
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            if str(add_op.outputs[0]) != add_output_name:
                continue

            add_inputs = [str(v) for v in list(add_op.inputs)]
            mul_input_side: Optional[int] = None
            mul_idx: Optional[int] = None
            for i, input_name in enumerate(add_inputs):
                producer_idx = producers.get(str(input_name), None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if (
                    str(producer_op.op_type) == "MUL"
                    and len(producer_op.outputs) == 1
                    and str(producer_op.outputs[0]) == str(input_name)
                ):
                    if mul_idx is not None:
                        mul_idx = None
                        break
                    mul_input_side = int(i)
                    mul_idx = int(producer_idx)
            if mul_idx is None or mul_input_side is None:
                continue

            skip_input_name = str(add_inputs[1 - int(mul_input_side)])
            skip_pre_idx = producers.get(skip_input_name, None)
            if skip_pre_idx is None:
                continue
            skip_pre_op = model_ir.operators[int(skip_pre_idx)]
            if (
                str(skip_pre_op.op_type) != "TRANSPOSE"
                or len(skip_pre_op.inputs) < 2
                or len(skip_pre_op.outputs) != 1
                or str(skip_pre_op.outputs[0]) != skip_input_name
                or _read_transpose_perm(model_ir, skip_pre_op) != perm_nhwc_to_nchw
            ):
                continue
            skip_nhwc_name = str(skip_pre_op.inputs[0])
            if skip_nhwc_name in model_outputs or skip_input_name in model_outputs:
                continue

            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_output_name = str(mul_op.outputs[0])
            if mul_output_name in model_outputs:
                continue
            # Optional affine tail:
            #   MUL(data, logistic(gate)) -> MUL(const) -> ADD(skip, ...)
            # Unwrap one scalar-const MUL so this pass can still match and remove
            # surrounding layout transposes.
            core_mul_idx = int(mul_idx)
            core_mul_op = mul_op
            core_mul_output_name = str(mul_output_name)
            if str(core_mul_op.op_type) == "MUL":
                core_mul_inputs = [str(v) for v in list(core_mul_op.inputs)]
                scalar_const_indices = [
                    int(i)
                    for i, input_name in enumerate(core_mul_inputs)
                    if _is_singleton_constant_tensor(model_ir, str(input_name))
                ]
                if len(scalar_const_indices) == 1:
                    main_input_name = str(core_mul_inputs[1 - int(scalar_const_indices[0])])
                    upstream_mul_idx = producers.get(main_input_name, None)
                    if upstream_mul_idx is not None:
                        upstream_mul_op = model_ir.operators[int(upstream_mul_idx)]
                        if (
                            str(upstream_mul_op.op_type) == "MUL"
                            and len(upstream_mul_op.inputs) == 2
                            and len(upstream_mul_op.outputs) == 1
                            and str(upstream_mul_op.outputs[0]) == str(main_input_name)
                        ):
                            core_mul_idx = int(upstream_mul_idx)
                            core_mul_op = upstream_mul_op
                            core_mul_output_name = str(main_input_name)

            pre_data_idx: Optional[int] = None
            pre_data_in_name: Optional[str] = None
            pre_data_out_name: Optional[str] = None
            logistic_idx: Optional[int] = None
            logistic_out_name: Optional[str] = None
            for input_name in [str(v) for v in list(core_mul_op.inputs)]:
                input_producer_idx = producers.get(str(input_name), None)
                if input_producer_idx is None:
                    continue
                input_producer_op = model_ir.operators[int(input_producer_idx)]
                if (
                    str(input_producer_op.op_type) == "TRANSPOSE"
                    and len(input_producer_op.inputs) >= 2
                    and len(input_producer_op.outputs) == 1
                    and str(input_producer_op.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, input_producer_op) == perm_nhwc_to_nchw
                ):
                    if pre_data_idx is not None:
                        pre_data_idx = None
                        break
                    pre_data_idx = int(input_producer_idx)
                    pre_data_in_name = str(input_producer_op.inputs[0])
                    pre_data_out_name = str(input_producer_op.outputs[0])
                    continue
                if (
                    str(input_producer_op.op_type) == "LOGISTIC"
                    and len(input_producer_op.outputs) == 1
                    and str(input_producer_op.outputs[0]) == str(input_name)
                ):
                    if logistic_idx is not None:
                        logistic_idx = None
                        break
                    logistic_idx = int(input_producer_idx)
                    logistic_out_name = str(input_name)
            if (
                pre_data_idx is None
                or pre_data_in_name is None
                or pre_data_out_name is None
                or logistic_idx is None
                or logistic_out_name is None
            ):
                continue

            if pre_data_in_name in model_outputs or pre_data_out_name in model_outputs:
                continue

            logistic_op = model_ir.operators[int(logistic_idx)]
            if str(logistic_op.op_type) != "LOGISTIC" or len(logistic_op.inputs) != 1 or len(logistic_op.outputs) != 1:
                continue
            if str(logistic_op.outputs[0]) != logistic_out_name:
                continue
            gate_pre_out_name = str(logistic_op.inputs[0])
            if gate_pre_out_name in model_outputs:
                continue

            gate_pre_idx = producers.get(gate_pre_out_name, None)
            if gate_pre_idx is None:
                continue
            gate_pre_op = model_ir.operators[int(gate_pre_idx)]
            if (
                str(gate_pre_op.op_type) != "TRANSPOSE"
                or len(gate_pre_op.inputs) < 2
                or len(gate_pre_op.outputs) != 1
                or str(gate_pre_op.outputs[0]) != gate_pre_out_name
                or _read_transpose_perm(model_ir, gate_pre_op) != perm_nhwc_to_nchw
            ):
                continue
            gate_nhwc_name = str(gate_pre_op.inputs[0])
            if gate_nhwc_name in model_outputs or gate_pre_out_name in model_outputs:
                continue

            # The chain user must exist; extra legacy users are allowed and preserved.
            skip_pre_users = [int(v) for v in consumers.get(skip_input_name, [])]
            if int(add_idx) not in set(skip_pre_users):
                continue
            data_pre_users = [int(v) for v in consumers.get(pre_data_out_name, [])]
            if int(core_mul_idx) not in set(data_pre_users):
                continue
            gate_pre_users = [int(v) for v in consumers.get(gate_pre_out_name, [])]
            if int(logistic_idx) not in set(gate_pre_users):
                continue

            # ADD output must have at least one inverse post-transpose. Non-post users
            # are kept via one retained adapter transpose.
            add_out_users = [int(v) for v in consumers.get(add_output_name, [])]
            if len(add_out_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            for user_idx in add_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == add_output_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_indices.append(int(user_idx))
                    post_output_names.append(str(user_op.outputs[0]))
                else:
                    legacy_users.append(int(user_idx))
            if len(post_indices) == 0:
                continue

            # Rewire LOGISTIC/MUL/ADD to NHWC sources.
            _set_operator_inputs(
                model_ir=model_ir,
                op=logistic_op,
                new_inputs=[str(gate_nhwc_name)],
            )

            mul_inputs = [str(v) for v in list(core_mul_op.inputs)]
            mul_replaced = False
            for i, input_name in enumerate(list(mul_inputs)):
                if str(input_name) == str(pre_data_out_name):
                    mul_inputs[i] = str(pre_data_in_name)
                    mul_replaced = True
            if not mul_replaced:
                continue
            _set_operator_inputs(
                model_ir=model_ir,
                op=core_mul_op,
                new_inputs=mul_inputs,
            )

            add_inputs_rewritten = [str(v) for v in list(add_op.inputs)]
            add_replaced = False
            for i, input_name in enumerate(list(add_inputs_rewritten)):
                if str(input_name) == str(skip_input_name):
                    add_inputs_rewritten[i] = str(skip_nhwc_name)
                    add_replaced = True
            if not add_replaced:
                continue
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=add_inputs_rewritten,
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(logistic_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(core_mul_output_name), None),
                perm_nchw_to_nhwc,
            )
            if str(core_mul_output_name) != str(mul_output_name):
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(mul_output_name), None),
                    perm_nchw_to_nhwc,
                )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(add_output_name), None),
                perm_nchw_to_nhwc,
            )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_op,
                new_outputs=[canonical_post_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            old_add_out_tensor = model_ir.tensors.get(add_output_name, None)
            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_add_out_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_add_out_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_add_out_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_add_out_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_add_out_tensor.shape_signature)]
                    if old_add_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_add_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            if len(legacy_users) > 0:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[int(keep_post_idx)]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[add_output_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = set(int(v) for v in post_remove_indices)
            # Remove pre-transposes only when no remaining users need the NCHW view.
            if len([int(v) for v in skip_pre_users if int(v) != int(add_idx)]) == 0:
                remove_indices.add(int(skip_pre_idx))
            if len([int(v) for v in data_pre_users if int(v) != int(core_mul_idx)]) == 0:
                remove_indices.add(int(pre_data_idx))
            if len([int(v) for v in gate_pre_users if int(v) != int(logistic_idx)]) == 0:
                remove_indices.add(int(gate_pre_idx))
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_logistic_muladd_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_osnet_multi_gate_muladd_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips in OSNet-like multi-branch gate blocks.

    Strict target motif (rank-4):
      - Several branch tensors:
          x_i_nhwc --T(0,3,1,2)--> x_i_nchw --RELU--> r_i_nchw
          r_i_nchw --MEAN(axes=[2,3], keepDims=1)--> gap_i_nchw
            -> ... -> RESHAPE([N,C,1,1]) -> LOGISTIC -> g_i_nchw
          MUL(r_i_nchw, g_i_nchw) -> m_i_nchw
      - Nested ADD tree over m_i:
          ADD(...ADD(m0, m1)..., mk) -> y_nchw
      - Output bridge:
          y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - Bypass each branch pre-transpose into RELU.
      - Remap MEAN axes from NCHW to NHWC.
      - Rewrite gate RESHAPE [N,C,1,1] -> [N,1,1,C] (or bypass gate pre-transpose).
      - Keep RELU/LOGISTIC/MUL/ADD in NHWC.
      - Remove all branch pre-transposes and terminal post-transpose bridges.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _is_singleton_spatial_nhwc_to_nchw_reshape(
        *,
        input_name: str,
        output_name: str,
    ) -> bool:
        in_tensor = model_ir.tensors.get(str(input_name), None)
        out_tensor = model_ir.tensors.get(str(output_name), None)
        if in_tensor is None or out_tensor is None:
            return False
        in_shape = [int(v) for v in list(in_tensor.shape)] if in_tensor.shape is not None else []
        out_shape = [int(v) for v in list(out_tensor.shape)] if out_tensor.shape is not None else []
        if len(in_shape) != 4 or len(out_shape) != 4:
            return False
        if any(int(v) < 0 for v in list(in_shape) + list(out_shape)):
            return False
        return (
            int(in_shape[0]) == int(out_shape[0])
            and int(in_shape[1]) == 1
            and int(in_shape[2]) == 1
            and int(in_shape[3]) == int(out_shape[1])
            and int(out_shape[2]) == 1
            and int(out_shape[3]) == 1
        )

    def _clone_const_tensor_for_operator_input(
        *,
        op: OperatorIR,
        input_index: int,
        tag: str,
    ) -> str:
        input_name = str(op.inputs[int(input_index)])
        tensor = model_ir.tensors.get(input_name, None)
        if tensor is None or tensor.data is None:
            return input_name
        cloned_name = f"{input_name}{tag}"
        suffix = 1
        while cloned_name in model_ir.tensors:
            cloned_name = f"{input_name}{tag}_{suffix}"
            suffix += 1
        model_ir.tensors[cloned_name] = TensorIR(
            name=cloned_name,
            dtype=str(tensor.dtype),
            shape=[int(v) for v in list(tensor.shape)] if tensor.shape is not None else [],
            shape_signature=(
                [int(v) for v in list(tensor.shape_signature)]
                if tensor.shape_signature is not None
                else (
                    [int(v) for v in list(tensor.shape)]
                    if tensor.shape is not None
                    else []
                )
            ),
            data=np.array(tensor.data, copy=True),
            is_variable=bool(tensor.is_variable),
            quantization=_clone_quantization(tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=cloned_name,
        )
        return str(cloned_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            add_root_out_name = str(post_op.inputs[0])
            if add_root_out_name in model_outputs:
                continue

            add_root_idx = producers.get(add_root_out_name, None)
            if add_root_idx is None:
                continue
            add_root_op = model_ir.operators[int(add_root_idx)]
            if (
                str(add_root_op.op_type) != "ADD"
                or len(add_root_op.inputs) != 2
                or len(add_root_op.outputs) != 1
                or str(add_root_op.outputs[0]) != add_root_out_name
            ):
                continue

            add_root_users = [int(v) for v in consumers.get(add_root_out_name, [])]
            if len(add_root_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            valid_posts = True
            for user_idx in add_root_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != add_root_out_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    valid_posts = False
                    break
                post_indices.append(int(user_idx))
                post_output_names.append(str(user_op.outputs[0]))
            if not valid_posts or len(post_indices) == 0:
                continue

            # Collect nested ADD tree leaves (MUL outputs).
            visited_add_outputs: set[str] = set()
            add_indices: List[int] = []
            add_output_names: List[str] = []
            leaf_mul_indices: List[int] = []
            tree_valid = True
            stack: List[str] = [str(add_root_out_name)]
            while len(stack) > 0:
                tensor_name = str(stack.pop())
                prod_idx = producers.get(tensor_name, None)
                if prod_idx is None:
                    tree_valid = False
                    break
                prod_op = model_ir.operators[int(prod_idx)]
                prod_type = str(prod_op.op_type)
                if (
                    prod_type == "ADD"
                    and len(prod_op.inputs) == 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == tensor_name
                    and tensor_name not in model_outputs
                ):
                    if tensor_name in visited_add_outputs:
                        tree_valid = False
                        break
                    visited_add_outputs.add(tensor_name)
                    add_indices.append(int(prod_idx))
                    add_output_names.append(str(tensor_name))
                    stack.append(str(prod_op.inputs[0]))
                    stack.append(str(prod_op.inputs[1]))
                elif (
                    prod_type == "MUL"
                    and len(prod_op.inputs) == 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == tensor_name
                    and tensor_name not in model_outputs
                ):
                    leaf_mul_indices.append(int(prod_idx))
                else:
                    tree_valid = False
                    break
            if not tree_valid or len(leaf_mul_indices) < 2:
                continue

            add_index_set = set(int(v) for v in add_indices)
            # Root ADD output may fan out only to inverse post-transposes.
            if set(int(v) for v in consumers.get(add_root_out_name, [])) != set(int(v) for v in post_indices):
                continue
            # Internal ADD outputs must flow only into one ADD parent.
            internal_add_outputs_valid = True
            for add_out_name in add_output_names:
                if str(add_out_name) == str(add_root_out_name):
                    continue
                add_users = [int(v) for v in consumers.get(str(add_out_name), [])]
                if len(add_users) != 1 or int(add_users[0]) not in add_index_set:
                    internal_add_outputs_valid = False
                    break
            if not internal_add_outputs_valid:
                continue
            # Leaf MUL outputs must feed exactly one ADD parent.
            leaf_users_valid = True
            for leaf_mul_idx in leaf_mul_indices:
                leaf_mul_op = model_ir.operators[int(leaf_mul_idx)]
                leaf_out_name = str(leaf_mul_op.outputs[0])
                leaf_users = [int(v) for v in consumers.get(leaf_out_name, [])]
                if len(leaf_users) != 1 or int(leaf_users[0]) not in add_index_set:
                    leaf_users_valid = False
                    break
            if not leaf_users_valid:
                continue

            leaf_infos: List[Dict[str, Any]] = []
            leaf_match_valid = True
            for leaf_mul_idx in leaf_mul_indices:
                leaf_mul_op = model_ir.operators[int(leaf_mul_idx)]
                leaf_inputs = [str(v) for v in list(leaf_mul_op.inputs)]

                matched_info: Optional[Dict[str, Any]] = None
                for relu_input_index in [0, 1]:
                    relu_out_name = str(leaf_inputs[int(relu_input_index)])
                    gate_out_name = str(leaf_inputs[1 - int(relu_input_index)])

                    relu_idx = producers.get(relu_out_name, None)
                    if relu_idx is None:
                        continue
                    relu_op = model_ir.operators[int(relu_idx)]
                    if (
                        str(relu_op.op_type) != "RELU"
                        or len(relu_op.inputs) != 1
                        or len(relu_op.outputs) != 1
                        or str(relu_op.outputs[0]) != relu_out_name
                        or relu_out_name in model_outputs
                    ):
                        continue
                    relu_input_name = str(relu_op.inputs[0])
                    if relu_input_name in model_outputs:
                        continue

                    pre_idx = producers.get(relu_input_name, None)
                    if pre_idx is None:
                        continue
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != relu_input_name
                        or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                        or str(pre_op.inputs[0]) in model_outputs
                    ):
                        continue
                    if set(int(v) for v in consumers.get(relu_input_name, [])) != {int(relu_idx)}:
                        continue

                    relu_users = [int(v) for v in consumers.get(relu_out_name, [])]
                    if int(leaf_mul_idx) not in set(relu_users):
                        continue
                    mean_candidates: List[int] = []
                    for relu_user_idx in relu_users:
                        if int(relu_user_idx) == int(leaf_mul_idx):
                            continue
                        relu_user_op = model_ir.operators[int(relu_user_idx)]
                        if (
                            str(relu_user_op.op_type) == "MEAN"
                            and len(relu_user_op.inputs) >= 2
                            and len(relu_user_op.outputs) == 1
                            and str(relu_user_op.inputs[0]) == relu_out_name
                            and bool(relu_user_op.options.get("keepDims", False))
                            and str(relu_user_op.outputs[0]) not in model_outputs
                        ):
                            mean_candidates.append(int(relu_user_idx))
                    if len(mean_candidates) != 1:
                        continue

                    mean_idx = int(mean_candidates[0])
                    mean_op = model_ir.operators[int(mean_idx)]
                    mean_axes_name = str(mean_op.inputs[1])
                    mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                    mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
                    if mean_axes_vals is None or len(mean_axes_vals) == 0:
                        continue
                    normalized_axes: List[int] = []
                    axes_valid = True
                    for axis in mean_axes_vals:
                        a = int(axis)
                        if a < 0:
                            a += 4
                        if a < 0 or a >= 4:
                            axes_valid = False
                            break
                        normalized_axes.append(int(a))
                    if not axes_valid or sorted(normalized_axes) != [2, 3]:
                        continue
                    mapped_axes = [int(perm_nhwc_to_nchw[int(v)]) for v in normalized_axes]
                    if sorted(mapped_axes) != [1, 2]:
                        continue

                    gate_idx = producers.get(gate_out_name, None)
                    if gate_idx is None:
                        continue
                    gate_op = model_ir.operators[int(gate_idx)]
                    if (
                        str(gate_op.op_type) != "LOGISTIC"
                        or len(gate_op.inputs) != 1
                        or len(gate_op.outputs) != 1
                        or str(gate_op.outputs[0]) != gate_out_name
                        or gate_out_name in model_outputs
                    ):
                        continue
                    if set(int(v) for v in consumers.get(gate_out_name, [])) != {int(leaf_mul_idx)}:
                        continue

                    gate_pre_out_name = str(gate_op.inputs[0])
                    if gate_pre_out_name in model_outputs:
                        continue
                    gate_pre_idx = producers.get(gate_pre_out_name, None)
                    if gate_pre_idx is None:
                        continue
                    gate_pre_op = model_ir.operators[int(gate_pre_idx)]
                    gate_pre_type = str(gate_pre_op.op_type)

                    gate_pre_input_name: Optional[str] = None
                    gate_pre_remove = False
                    gate_shape_tensor: Optional[TensorIR] = None
                    if (
                        gate_pre_type == "TRANSPOSE"
                        and len(gate_pre_op.inputs) >= 2
                        and len(gate_pre_op.outputs) == 1
                        and str(gate_pre_op.outputs[0]) == gate_pre_out_name
                        and _read_transpose_perm(model_ir, gate_pre_op) == perm_nhwc_to_nchw
                        and str(gate_pre_op.inputs[0]) not in model_outputs
                    ):
                        gate_pre_input_name = str(gate_pre_op.inputs[0])
                        gate_pre_remove = True
                    elif (
                        gate_pre_type == "RESHAPE"
                        and len(gate_pre_op.inputs) >= 2
                        and len(gate_pre_op.outputs) == 1
                        and str(gate_pre_op.outputs[0]) == gate_pre_out_name
                        and _is_singleton_spatial_nhwc_to_nchw_reshape(
                            input_name=str(gate_pre_op.inputs[0]),
                            output_name=str(gate_pre_op.outputs[0]),
                        )
                    ):
                        gate_shape_tensor = model_ir.tensors.get(str(gate_pre_op.inputs[1]), None)
                        gate_shape_vals = _read_const_ints_from_tensor(gate_shape_tensor)
                        if gate_shape_vals is None or len(gate_shape_vals) != 4:
                            continue
                        gate_pre_remove = False
                    else:
                        continue
                    if set(int(v) for v in consumers.get(gate_pre_out_name, [])) != {int(gate_idx)}:
                        continue

                    matched_info = {
                        "mul_idx": int(leaf_mul_idx),
                        "mul_op": leaf_mul_op,
                        "mul_out_name": str(leaf_mul_op.outputs[0]),
                        "relu_idx": int(relu_idx),
                        "relu_op": relu_op,
                        "relu_out_name": str(relu_out_name),
                        "pre_idx": int(pre_idx),
                        "pre_input_name": str(pre_op.inputs[0]),
                        "pre_out_name": str(relu_input_name),
                        "mean_idx": int(mean_idx),
                        "mean_op": mean_op,
                        "mean_out_name": str(mean_op.outputs[0]),
                        "mapped_axes": [int(v) for v in mapped_axes],
                        "gate_idx": int(gate_idx),
                        "gate_op": gate_op,
                        "gate_out_name": str(gate_out_name),
                        "gate_pre_idx": int(gate_pre_idx),
                        "gate_pre_type": str(gate_pre_type),
                        "gate_pre_input_name": str(gate_pre_input_name) if gate_pre_input_name is not None else None,
                        "gate_pre_out_name": str(gate_pre_out_name),
                        "gate_pre_remove": bool(gate_pre_remove),
                    }
                    break

                if matched_info is None:
                    leaf_match_valid = False
                    break
                leaf_infos.append(dict(matched_info))
            if not leaf_match_valid:
                continue

            pre_indices = [int(v["pre_idx"]) for v in leaf_infos]
            if len(set(pre_indices)) != len(pre_indices):
                continue

            # Rewrite branch-local ops to NHWC semantics.
            gate_pre_remove_indices: set[int] = set()
            tensor_metadata_targets: set[str] = set()
            for leaf in leaf_infos:
                relu_op = leaf["relu_op"]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=relu_op,
                    new_inputs=[str(leaf["pre_input_name"])],
                )

                mean_op = leaf["mean_op"]
                mean_axes_name = _clone_const_tensor_for_operator_input(
                    op=mean_op,
                    input_index=1,
                    tag="__osnet_nhwc_axes",
                )
                mean_axes_tensor = model_ir.tensors.get(str(mean_axes_name), None)
                mapped_axes = [int(v) for v in list(leaf["mapped_axes"])]
                if not _write_const_ints_to_tensor(mean_axes_tensor, mapped_axes):
                    leaf_match_valid = False
                    break
                if isinstance(mean_op.options, dict):
                    mean_options = dict(mean_op.options)
                    for key in ["axis", "axes", "onnxRawAxes"]:
                        value = mean_options.get(key, None)
                        if isinstance(value, list) and len(value) == len(mapped_axes):
                            mean_options[key] = [int(v) for v in mapped_axes]
                    mean_op.options = mean_options

                gate_op = leaf["gate_op"]
                if bool(leaf["gate_pre_remove"]):
                    gate_pre_input_name = leaf["gate_pre_input_name"]
                    if gate_pre_input_name is None:
                        leaf_match_valid = False
                        break
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=gate_op,
                        new_inputs=[str(gate_pre_input_name)],
                    )
                    gate_pre_remove_indices.add(int(leaf["gate_pre_idx"]))
                else:
                    gate_pre_idx = int(leaf["gate_pre_idx"])
                    gate_pre_op = model_ir.operators[int(gate_pre_idx)]
                    gate_shape_name = _clone_const_tensor_for_operator_input(
                        op=gate_pre_op,
                        input_index=1,
                        tag="__osnet_nhwc_shape",
                    )
                    gate_shape_tensor = model_ir.tensors.get(str(gate_shape_name), None)
                    gate_shape_vals = _read_const_ints_from_tensor(gate_shape_tensor)
                    if gate_shape_vals is None or len(gate_shape_vals) != 4:
                        leaf_match_valid = False
                        break
                    swapped_shape = [
                        int(gate_shape_vals[0]),
                        int(gate_shape_vals[2]),
                        int(gate_shape_vals[3]),
                        int(gate_shape_vals[1]),
                    ]
                    if not _write_const_ints_to_tensor(gate_shape_tensor, swapped_shape):
                        leaf_match_valid = False
                        break

                tensor_metadata_targets.update({
                    str(leaf["relu_out_name"]),
                    str(leaf["mean_out_name"]),
                    str(leaf["gate_pre_out_name"]),
                    str(leaf["gate_out_name"]),
                    str(leaf["mul_out_name"]),
                })
            if not leaf_match_valid:
                continue

            for add_output_name in add_output_names:
                tensor_metadata_targets.add(str(add_output_name))
            tensor_metadata_targets.add(str(add_root_out_name))

            for tensor_name in tensor_metadata_targets:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=add_root_op,
                new_outputs=[canonical_post_output_name],
            )
            _replace_tensor_inputs(model_ir, add_root_out_name, canonical_post_output_name)
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            old_root_tensor = model_ir.tensors.get(add_root_out_name, None)
            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_root_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_root_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_root_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_root_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_root_tensor.shape_signature)]
                    if old_root_tensor.shape_signature is not None
                    else [int(v) for v in list(old_root_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = set(int(v["pre_idx"]) for v in leaf_infos)
            remove_indices.update(int(v) for v in gate_pre_remove_indices)
            remove_indices.update(int(v) for v in post_indices)
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_osnet_multi_gate_muladd_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around a logistic/sub-gated dual-MUL/ADD block feeding two convs.

    Target:
      g_nhwc --T(0,3,1,2)--> g_nchw --LOGISTIC--> sig_nchw --SUB(1, sig)->sub_nchw
      a_nhwc --T(0,3,1,2)--> a_nchw
      b_nhwc --T(0,3,1,2)--> b_nchw
      MUL(sig_nchw, a_nchw) -> m0_nchw
      MUL(sub_nchw, b_nchw) -> m1_nchw
      ADD(b_nchw, m0_nchw)  -> y0_nchw --T(0,2,3,1)--> y0_nhwc -> CONV
      ADD(m1_nchw, a_nchw)  -> y1_nchw --T(0,2,3,1)--> y1_nhwc -> CONV

    Rewrite:
      - Bypass all three pre-transposes to NHWC.
      - Keep LOGISTIC/SUB/MUL/ADD in NHWC.
      - Remove both post-transposes and feed CONV inputs directly.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for gate_pre_idx, gate_pre_op in enumerate(model_ir.operators):
            if str(gate_pre_op.op_type) != "TRANSPOSE" or len(gate_pre_op.inputs) < 2 or len(gate_pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, gate_pre_op) != perm_nhwc_to_nchw:
                continue

            gate_pre_in_name = str(gate_pre_op.inputs[0])
            gate_pre_out_name = str(gate_pre_op.outputs[0])
            if gate_pre_in_name in model_outputs or gate_pre_out_name in model_outputs:
                continue

            gate_users = [int(v) for v in consumers.get(gate_pre_out_name, [])]
            if len(gate_users) != 1:
                continue
            logistic_idx = int(gate_users[0])
            logistic_op = model_ir.operators[int(logistic_idx)]
            if str(logistic_op.op_type) != "LOGISTIC" or len(logistic_op.inputs) != 1 or len(logistic_op.outputs) != 1:
                continue
            if str(logistic_op.inputs[0]) != gate_pre_out_name:
                continue

            logistic_out_name = str(logistic_op.outputs[0])
            if logistic_out_name in model_outputs:
                continue
            logistic_users = [int(v) for v in consumers.get(logistic_out_name, [])]
            if len(logistic_users) != 2:
                continue

            sub_idx: Optional[int] = None
            mul_sig_idx: Optional[int] = None
            for user_idx in logistic_users:
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) == "SUB":
                    if sub_idx is not None:
                        sub_idx = None
                        break
                    sub_idx = int(user_idx)
                elif str(user_op.op_type) == "MUL":
                    if mul_sig_idx is not None:
                        mul_sig_idx = None
                        break
                    mul_sig_idx = int(user_idx)
            if sub_idx is None or mul_sig_idx is None:
                continue

            sub_op = model_ir.operators[int(sub_idx)]
            if len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue
            if str(sub_op.outputs[0]) in model_outputs:
                continue
            if str(logistic_out_name) not in {str(v) for v in list(sub_op.inputs)}:
                continue
            sub_out_name = str(sub_op.outputs[0])
            sub_users = [int(v) for v in consumers.get(sub_out_name, [])]
            if len(sub_users) != 1:
                continue
            mul_sub_idx = int(sub_users[0])
            mul_sub_op = model_ir.operators[int(mul_sub_idx)]
            if str(mul_sub_op.op_type) != "MUL" or len(mul_sub_op.inputs) != 2 or len(mul_sub_op.outputs) != 1:
                continue

            mul_sig_op = model_ir.operators[int(mul_sig_idx)]
            if str(mul_sig_op.op_type) != "MUL" or len(mul_sig_op.inputs) != 2 or len(mul_sig_op.outputs) != 1:
                continue

            def _resolve_pre_for_mul(
                mul_op: OperatorIR,
                gate_tensor_candidates: List[str],
            ) -> Optional[Tuple[int, str, str]]:
                pre_idx_local: Optional[int] = None
                pre_in_name_local: Optional[str] = None
                pre_out_name_local: Optional[str] = None
                for input_name in list(mul_op.inputs):
                    if str(input_name) in set(str(v) for v in gate_tensor_candidates):
                        continue
                    prod_idx = producers.get(str(input_name), None)
                    if prod_idx is None:
                        continue
                    prod_op = model_ir.operators[int(prod_idx)]
                    if (
                        str(prod_op.op_type) == "TRANSPOSE"
                        and len(prod_op.inputs) >= 2
                        and len(prod_op.outputs) == 1
                        and str(prod_op.outputs[0]) == str(input_name)
                        and _read_transpose_perm(model_ir, prod_op) == perm_nhwc_to_nchw
                        and str(prod_op.inputs[0]) not in model_outputs
                        and str(prod_op.outputs[0]) not in model_outputs
                    ):
                        pre_idx_local = int(prod_idx)
                        pre_in_name_local = str(prod_op.inputs[0])
                        pre_out_name_local = str(prod_op.outputs[0])
                        break
                if pre_idx_local is None or pre_in_name_local is None or pre_out_name_local is None:
                    return None
                return int(pre_idx_local), str(pre_in_name_local), str(pre_out_name_local)

            sig_pre_info = _resolve_pre_for_mul(
                mul_sig_op,
                gate_tensor_candidates=[str(logistic_out_name)],
            )
            sub_pre_info = _resolve_pre_for_mul(
                mul_sub_op,
                gate_tensor_candidates=[str(sub_out_name)],
            )
            if sig_pre_info is None or sub_pre_info is None:
                continue

            sig_pre_idx, sig_pre_in_name, sig_pre_out_name = sig_pre_info
            sub_pre_idx, sub_pre_in_name, sub_pre_out_name = sub_pre_info
            if int(sig_pre_idx) == int(sub_pre_idx):
                continue

            mul_sig_out_name = str(mul_sig_op.outputs[0])
            mul_sub_out_name = str(mul_sub_op.outputs[0])
            if mul_sig_out_name in model_outputs or mul_sub_out_name in model_outputs:
                continue
            mul_sig_users = [int(v) for v in consumers.get(mul_sig_out_name, [])]
            mul_sub_users = [int(v) for v in consumers.get(mul_sub_out_name, [])]
            if len(mul_sig_users) != 1 or len(mul_sub_users) != 1:
                continue
            add_sig_idx = int(mul_sig_users[0])
            add_sub_idx = int(mul_sub_users[0])
            if int(add_sig_idx) == int(add_sub_idx):
                continue
            add_sig_op = model_ir.operators[int(add_sig_idx)]
            add_sub_op = model_ir.operators[int(add_sub_idx)]
            if str(add_sig_op.op_type) != "ADD" or len(add_sig_op.inputs) != 2 or len(add_sig_op.outputs) != 1:
                continue
            if str(add_sub_op.op_type) != "ADD" or len(add_sub_op.inputs) != 2 or len(add_sub_op.outputs) != 1:
                continue

            add_sig_inputs = [str(v) for v in list(add_sig_op.inputs)]
            add_sub_inputs = [str(v) for v in list(add_sub_op.inputs)]
            if str(mul_sig_out_name) not in set(add_sig_inputs):
                continue
            if str(mul_sub_out_name) not in set(add_sub_inputs):
                continue
            add_sig_other = (
                add_sig_inputs[1] if add_sig_inputs[0] == str(mul_sig_out_name) else add_sig_inputs[0]
            )
            add_sub_other = (
                add_sub_inputs[1] if add_sub_inputs[0] == str(mul_sub_out_name) else add_sub_inputs[0]
            )
            if set([str(add_sig_other), str(add_sub_other)]) != set([str(sig_pre_out_name), str(sub_pre_out_name)]):
                continue

            add_sig_out_name = str(add_sig_op.outputs[0])
            add_sub_out_name = str(add_sub_op.outputs[0])
            if add_sig_out_name in model_outputs or add_sub_out_name in model_outputs:
                continue

            def _collect_inverse_posts(add_out_name: str) -> Optional[Tuple[List[int], List[str]]]:
                out_users = [int(v) for v in consumers.get(str(add_out_name), [])]
                if len(out_users) == 0:
                    return None
                post_indices_local: List[int] = []
                post_output_names_local: List[str] = []
                for u_idx in out_users:
                    u_op = model_ir.operators[int(u_idx)]
                    if (
                        str(u_op.op_type) == "TRANSPOSE"
                        and len(u_op.inputs) >= 2
                        and len(u_op.outputs) == 1
                        and str(u_op.inputs[0]) == str(add_out_name)
                        and _read_transpose_perm(model_ir, u_op) == perm_nchw_to_nhwc
                        and str(u_op.outputs[0]) not in model_outputs
                    ):
                        post_indices_local.append(int(u_idx))
                        post_output_names_local.append(str(u_op.outputs[0]))
                    else:
                        return None
                return post_indices_local, post_output_names_local

            sig_post_info = _collect_inverse_posts(add_sig_out_name)
            sub_post_info = _collect_inverse_posts(add_sub_out_name)
            if sig_post_info is None or sub_post_info is None:
                continue
            sig_post_indices, sig_post_output_names = sig_post_info
            sub_post_indices, sub_post_output_names = sub_post_info

            # Ensure pre-transpose outputs are local to this chain.
            allowed_users = {
                int(mul_sig_idx),
                int(mul_sub_idx),
                int(add_sig_idx),
                int(add_sub_idx),
            }
            if not set(int(v) for v in consumers.get(str(sig_pre_out_name), [])).issubset(allowed_users):
                continue
            if not set(int(v) for v in consumers.get(str(sub_pre_out_name), [])).issubset(allowed_users):
                continue

            # 1) Bypass gate pre-transpose.
            _set_operator_inputs(
                model_ir=model_ir,
                op=logistic_op,
                new_inputs=[str(gate_pre_in_name)],
            )

            # 2) Bypass the two data pre-transposes in MUL and ADD.
            for mul_op, pre_out_name, pre_in_name in [
                (mul_sig_op, sig_pre_out_name, sig_pre_in_name),
                (mul_sub_op, sub_pre_out_name, sub_pre_in_name),
            ]:
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                mul_inputs = [
                    str(pre_in_name) if str(v) == str(pre_out_name) else str(v)
                    for v in mul_inputs
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )

            for add_op, pre_out_name, pre_in_name in [
                (add_sig_op, sig_pre_out_name, sig_pre_in_name),
                (add_sig_op, sub_pre_out_name, sub_pre_in_name),
                (add_sub_op, sig_pre_out_name, sig_pre_in_name),
                (add_sub_op, sub_pre_out_name, sub_pre_in_name),
            ]:
                add_inputs = [str(v) for v in list(add_op.inputs)]
                rewritten_inputs = [
                    str(pre_in_name) if str(v) == str(pre_out_name) else str(v)
                    for v in add_inputs
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_inputs=rewritten_inputs,
                )

            # 3) Mark rewritten intermediates as NHWC.
            for tensor_name in [
                str(logistic_out_name),
                str(sub_out_name),
                str(mul_sig_out_name),
                str(mul_sub_out_name),
                str(add_sig_out_name),
                str(add_sub_out_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            # 4) Bypass both post-transpose groups.
            for add_op, old_add_out_name, post_output_names in [
                (add_sig_op, add_sig_out_name, sig_post_output_names),
                (add_sub_op, add_sub_out_name, sub_post_output_names),
            ]:
                canonical_post_output_name = str(post_output_names[0])
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_outputs=[canonical_post_output_name],
                )
                for alias_name in list(post_output_names[1:]):
                    _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

                old_add_out_tensor = model_ir.tensors.get(str(old_add_out_name), None)
                canonical_tensor = model_ir.tensors.get(str(canonical_post_output_name), None)
                if old_add_out_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_add_out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_add_out_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_add_out_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_add_out_tensor.shape_signature)]
                        if old_add_out_tensor.shape_signature is not None
                        else [int(v) for v in list(old_add_out_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        canonical_tensor,
                        perm_nchw_to_nhwc,
                    )

            remove_indices = {
                int(gate_pre_idx),
                int(sig_pre_idx),
                int(sub_pre_idx),
                *[int(v) for v in sig_post_indices],
                *[int(v) for v in sub_post_indices],
            }
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains": int(rewritten)}


def _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NDHWC<->NCDHW bridges around 3D LEAKY/ADD + LOGISTIC/MUL/ADD blocks.

    Target:
      base_nhwc --T(0,3,1,2)--> base_nchw --RESHAPE--> base_ncdhw
      skip_nhdwc --T(0,4,1,2,3)--> skip_ncdhw --LEAKY_RELU--> skip_leaky_ncdhw
      ADD(skip_leaky_ncdhw, base_ncdhw) -> add0_ncdhw --T(0,2,3,4,1)--> add0_ndhwc
      gate_ndhwc --T(0,4,1,2,3)--> gate_ncdhw --LOGISTIC--> gate_sig_ncdhw
      MUL(gate_sig_ncdhw, base_ncdhw) -> mul1_ncdhw
      ADD(mul1_ncdhw, skip_leaky_ncdhw) -> add1_ncdhw --T(0,2,3,4,1)--> add1_ndhwc

    Rewrite:
      - Remove the 5D pre/post transposes and keep the chain in NDHWC.
      - Rewrite base RESHAPE to consume base_nhwc directly and emit NDHWC.
      - Remove the 4D base transpose feeding that RESHAPE.

    Safety:
      - Uses a strict single-path motif with local consumers only.
      - Refuses graph/model-output boundary tensors.
    """
    rewritten = 0
    perm_ndhwc_to_ncdhw = [0, 4, 1, 2, 3]
    perm_ncdhw_to_ndhwc = [0, 2, 3, 4, 1]
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    def _copy_tensor_metadata_with_perm(
        *,
        src_name: str,
        dst_name: str,
        perm: List[int],
    ) -> None:
        src_tensor = model_ir.tensors.get(str(src_name), None)
        dst_tensor = model_ir.tensors.get(str(dst_name), None)
        if src_tensor is None or dst_tensor is None:
            return
        dst_tensor.dtype = str(src_tensor.dtype)
        dst_tensor.quantization = _clone_quantization(src_tensor.quantization)
        dst_tensor.shape = [int(v) for v in list(src_tensor.shape)]
        dst_tensor.shape_signature = (
            [int(v) for v in list(src_tensor.shape_signature)]
            if src_tensor.shape_signature is not None
            else [int(v) for v in list(src_tensor.shape)]
        )
        _permute_tensor_metadata_if_rank_matches(
            dst_tensor,
            perm,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post1_idx, post1_op in enumerate(model_ir.operators):
            if str(post1_op.op_type) != "TRANSPOSE" or len(post1_op.inputs) < 2 or len(post1_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post1_op) != perm_ncdhw_to_ndhwc:
                continue

            add1_out_name = str(post1_op.inputs[0])
            add1_post_out_name = str(post1_op.outputs[0])
            if add1_out_name in model_outputs or add1_post_out_name in model_outputs:
                continue

            add1_idx = producers.get(add1_out_name, None)
            if add1_idx is None:
                continue
            add1_op = model_ir.operators[int(add1_idx)]
            if (
                str(add1_op.op_type) != "ADD"
                or len(add1_op.inputs) != 2
                or len(add1_op.outputs) != 1
                or str(add1_op.outputs[0]) != add1_out_name
            ):
                continue
            if set(int(v) for v in consumers.get(add1_out_name, [])) != {int(post1_idx)}:
                continue

            add1_inputs = [str(v) for v in list(add1_op.inputs)]

            skip_name: Optional[str] = None
            base_name: Optional[str] = None
            for lhs_name, rhs_name in [(add1_inputs[0], add1_inputs[1]), (add1_inputs[1], add1_inputs[0])]:
                lhs_prod_idx = producers.get(str(lhs_name), None)
                if lhs_prod_idx is None:
                    continue
                lhs_prod_op = model_ir.operators[int(lhs_prod_idx)]
                if str(lhs_prod_op.op_type) != "LEAKY_RELU":
                    continue
                skip_name = str(lhs_name)
                base_name = str(rhs_name)
                break
            if skip_name is None or base_name is None:
                continue

            # Skip path: NDHWC -> TRANSPOSE -> LEAKY_RELU
            skip_leaky_idx = producers.get(skip_name, None)
            if skip_leaky_idx is None:
                continue
            skip_leaky_op = model_ir.operators[int(skip_leaky_idx)]
            if (
                str(skip_leaky_op.op_type) != "LEAKY_RELU"
                or len(skip_leaky_op.inputs) != 1
                or len(skip_leaky_op.outputs) != 1
                or str(skip_leaky_op.outputs[0]) != skip_name
            ):
                continue
            pre_skip_out_name = str(skip_leaky_op.inputs[0])
            if pre_skip_out_name in model_outputs:
                continue
            pre_skip_idx = producers.get(pre_skip_out_name, None)
            if pre_skip_idx is None:
                continue
            pre_skip_op = model_ir.operators[int(pre_skip_idx)]
            if (
                str(pre_skip_op.op_type) != "TRANSPOSE"
                or len(pre_skip_op.inputs) < 2
                or len(pre_skip_op.outputs) != 1
                or str(pre_skip_op.outputs[0]) != pre_skip_out_name
                or _read_transpose_perm(model_ir, pre_skip_op) != perm_ndhwc_to_ncdhw
            ):
                continue
            skip_ndhwc_name = str(pre_skip_op.inputs[0])
            if skip_ndhwc_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(pre_skip_out_name, [])) != {int(skip_leaky_idx)}:
                continue

            # Find second block:
            #   ADD(mul, skip_name) -> post2 transpose
            add2_candidates = [
                int(v)
                for v in consumers.get(skip_name, [])
                if int(v) != int(add1_idx)
                and str(model_ir.operators[int(v)].op_type) == "ADD"
            ]
            if len(add2_candidates) != 1:
                continue
            add2_idx = int(add2_candidates[0])
            add2_op = model_ir.operators[int(add2_idx)]
            if len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_out_name = str(add2_op.outputs[0])
            if add2_out_name in model_outputs:
                continue
            add2_users = [int(v) for v in consumers.get(add2_out_name, [])]
            if len(add2_users) != 1:
                continue
            post2_idx = int(add2_users[0])
            post2_op = model_ir.operators[int(post2_idx)]
            if (
                str(post2_op.op_type) != "TRANSPOSE"
                or len(post2_op.inputs) < 2
                or len(post2_op.outputs) != 1
                or str(post2_op.inputs[0]) != add2_out_name
                or _read_transpose_perm(model_ir, post2_op) != perm_ncdhw_to_ndhwc
                or str(post2_op.outputs[0]) in model_outputs
            ):
                continue
            add2_post_out_name = str(post2_op.outputs[0])

            # add2 must be ADD(mul2, skip_name)
            add2_inputs = [str(v) for v in list(add2_op.inputs)]
            mul2_idx: Optional[int] = None
            if str(add2_inputs[0]) == str(skip_name):
                cand_idx = producers.get(str(add2_inputs[1]), None)
                if cand_idx is not None and str(model_ir.operators[int(cand_idx)].op_type) == "MUL":
                    mul2_idx = int(cand_idx)
            elif str(add2_inputs[1]) == str(skip_name):
                cand_idx = producers.get(str(add2_inputs[0]), None)
                if cand_idx is not None and str(model_ir.operators[int(cand_idx)].op_type) == "MUL":
                    mul2_idx = int(cand_idx)
            if mul2_idx is None:
                continue
            mul2_op = model_ir.operators[int(mul2_idx)]
            if len(mul2_op.inputs) != 2 or len(mul2_op.outputs) != 1:
                continue
            mul2_out_name = str(mul2_op.outputs[0])
            if mul2_out_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(mul2_out_name, [])) != {int(add2_idx)}:
                continue

            # mul2 must consume LOGISTIC(pre_gate) and base_name.
            mul2_inputs = [str(v) for v in list(mul2_op.inputs)]
            logistic_idx: Optional[int] = None
            if str(mul2_inputs[0]) == str(base_name):
                cand_idx = producers.get(str(mul2_inputs[1]), None)
                if cand_idx is not None and str(model_ir.operators[int(cand_idx)].op_type) == "LOGISTIC":
                    logistic_idx = int(cand_idx)
            elif str(mul2_inputs[1]) == str(base_name):
                cand_idx = producers.get(str(mul2_inputs[0]), None)
                if cand_idx is not None and str(model_ir.operators[int(cand_idx)].op_type) == "LOGISTIC":
                    logistic_idx = int(cand_idx)
            if logistic_idx is None:
                continue
            logistic_op = model_ir.operators[int(logistic_idx)]
            if (
                str(logistic_op.op_type) != "LOGISTIC"
                or len(logistic_op.inputs) != 1
                or len(logistic_op.outputs) != 1
            ):
                continue
            logistic_out_name = str(logistic_op.outputs[0])
            pre_gate_out_name = str(logistic_op.inputs[0])
            if pre_gate_out_name in model_outputs or logistic_out_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(logistic_out_name, [])) != {int(mul2_idx)}:
                continue
            pre_gate_idx = producers.get(pre_gate_out_name, None)
            if pre_gate_idx is None:
                continue
            pre_gate_op = model_ir.operators[int(pre_gate_idx)]
            if (
                str(pre_gate_op.op_type) != "TRANSPOSE"
                or len(pre_gate_op.inputs) < 2
                or len(pre_gate_op.outputs) != 1
                or str(pre_gate_op.outputs[0]) != pre_gate_out_name
                or _read_transpose_perm(model_ir, pre_gate_op) != perm_ndhwc_to_ncdhw
            ):
                continue
            gate_ndhwc_name = str(pre_gate_op.inputs[0])
            if gate_ndhwc_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(pre_gate_out_name, [])) != {int(logistic_idx)}:
                continue

            # base path:
            #   base_nhwc --T(0,3,1,2)--> base_nchw --RESHAPE--> base_name(NCDHW)
            base_reshape_idx = producers.get(base_name, None)
            if base_reshape_idx is None:
                continue
            base_reshape_op = model_ir.operators[int(base_reshape_idx)]
            if (
                str(base_reshape_op.op_type) != "RESHAPE"
                or len(base_reshape_op.inputs) < 2
                or len(base_reshape_op.outputs) != 1
                or str(base_reshape_op.outputs[0]) != str(base_name)
            ):
                continue
            base_shape_name = str(base_reshape_op.inputs[1])
            base_shape_tensor = model_ir.tensors.get(base_shape_name, None)
            base_shape_vals = _read_const_ints_from_tensor(base_shape_tensor)
            if base_shape_vals is None or len(base_shape_vals) != 5:
                continue
            pre_base_out_name = str(base_reshape_op.inputs[0])
            if pre_base_out_name in model_outputs:
                continue
            pre_base_idx = producers.get(pre_base_out_name, None)
            if pre_base_idx is None:
                continue
            pre_base_op = model_ir.operators[int(pre_base_idx)]
            if (
                str(pre_base_op.op_type) != "TRANSPOSE"
                or len(pre_base_op.inputs) < 2
                or len(pre_base_op.outputs) != 1
                or str(pre_base_op.outputs[0]) != pre_base_out_name
                or _read_transpose_perm(model_ir, pre_base_op) != perm_nhwc_to_nchw
            ):
                continue
            base_nhwc_name = str(pre_base_op.inputs[0])
            if base_nhwc_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(pre_base_out_name, [])) != {int(base_reshape_idx)}:
                continue

            # Strict local-consumer checks for safety.
            if set(int(v) for v in consumers.get(str(base_name), [])) != {int(add1_idx), int(mul2_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(skip_name), [])) != {int(add1_idx), int(add2_idx)}:
                continue

            # Rewrite base reshape to NDHWC.
            _set_operator_inputs(
                model_ir=model_ir,
                op=base_reshape_op,
                new_inputs=[str(base_nhwc_name), str(base_shape_name)],
            )
            mapped_base_shape = _permute_shape(
                [int(v) for v in list(base_shape_vals)],
                perm_ncdhw_to_ndhwc,
            )
            if mapped_base_shape is None:
                continue
            _write_const_ints_to_tensor(base_shape_tensor, [int(v) for v in list(mapped_base_shape)])
            if isinstance(base_reshape_op.options, dict):
                reshape_opts = dict(base_reshape_op.options)
                reshape_opts_changed = False
                for key in ["newShape", "onnxRawNewShape"]:
                    value = reshape_opts.get(key, None)
                    if isinstance(value, list) and len(value) == 5:
                        remapped = _permute_shape([int(v) for v in list(value)], perm_ncdhw_to_ndhwc)
                        if remapped is not None and [int(v) for v in remapped] != [int(v) for v in value]:
                            reshape_opts[key] = [int(v) for v in list(remapped)]
                            reshape_opts_changed = True
                if reshape_opts_changed:
                    base_reshape_op.options = reshape_opts
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(base_name), None),
                perm_ncdhw_to_ndhwc,
            )

            # Bypass 5D pre-transposes.
            _set_operator_inputs(
                model_ir=model_ir,
                op=skip_leaky_op,
                new_inputs=[str(skip_ndhwc_name)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(skip_name), None),
                perm_ncdhw_to_ndhwc,
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=logistic_op,
                new_inputs=[str(gate_ndhwc_name)],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(logistic_out_name), None),
                perm_ncdhw_to_ndhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul2_out_name), None),
                perm_ncdhw_to_ndhwc,
            )

            # Bypass post-transpose outputs by letting ADD emit NDHWC tensors directly.
            _set_operator_outputs(
                model_ir=model_ir,
                op=add1_op,
                new_outputs=[str(add1_post_out_name)],
            )
            _copy_tensor_metadata_with_perm(
                src_name=str(add1_out_name),
                dst_name=str(add1_post_out_name),
                perm=perm_ncdhw_to_ndhwc,
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=add2_op,
                new_outputs=[str(add2_post_out_name)],
            )
            _copy_tensor_metadata_with_perm(
                src_name=str(add2_out_name),
                dst_name=str(add2_post_out_name),
                perm=perm_ncdhw_to_ndhwc,
            )

            remove_indices = sorted(
                list(
                    {
                        int(pre_base_idx),
                        int(pre_skip_idx),
                        int(pre_gate_idx),
                        int(post1_idx),
                        int(post2_idx),
                    }
                ),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_3d_leaky_logistic_muladd_ndhwc_chains": int(rewritten)}


def _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCDHW bridge transposes around CONV_3D->LEAKY_RELU gated by Unsqueeze(RESHAPE)+MUL.

    Target:
      conv_ndhwc --T(0,4,1,2,3)--> conv_ncdhw --LEAKY_RELU--> act_ncdhw
      sem_ndhwc  --T(0,4,1,2,3)--> sem_ncdhw --RESHAPE--> gate_ncdhw
      MUL(act_ncdhw, gate_ncdhw) -> mul_ncdhw --T(0,2,3,4,1)--> mul_ndhwc

    Rewrite:
      conv_ndhwc --LEAKY_RELU--> act_ndhwc
      sem_ndhwc  --RESHAPE(remapped shape)--> gate_ndhwc
      MUL(act_ndhwc, gate_ndhwc) -> mul_ndhwc
    """
    rewritten = 0
    perm_ndhwc_to_ncdhw = [0, 4, 1, 2, 3]
    perm_ncdhw_to_ndhwc = [0, 2, 3, 4, 1]
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_ncdhw_to_ndhwc
            ):
                continue

            mul_out_name = str(post_op.inputs[0])
            mul_ndhwc_name = str(post_op.outputs[0])
            if mul_out_name in model_outputs or mul_ndhwc_name in model_outputs:
                continue

            mul_idx = producers.get(mul_out_name, None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if (
                str(mul_op.op_type) != "MUL"
                or len(mul_op.inputs) != 2
                or len(mul_op.outputs) != 1
                or str(mul_op.outputs[0]) != mul_out_name
            ):
                continue
            if set(int(v) for v in consumers.get(mul_out_name, [])) != {int(post_idx)}:
                continue

            mul_in0 = str(mul_op.inputs[0])
            mul_in1 = str(mul_op.inputs[1])
            leaky_idx: Optional[int] = None
            reshape_idx: Optional[int] = None
            if (
                mul_in0 in producers
                and str(model_ir.operators[int(producers[mul_in0])].op_type) == "LEAKY_RELU"
                and mul_in1 in producers
                and str(model_ir.operators[int(producers[mul_in1])].op_type) == "RESHAPE"
            ):
                leaky_idx = int(producers[mul_in0])
                reshape_idx = int(producers[mul_in1])
            elif (
                mul_in1 in producers
                and str(model_ir.operators[int(producers[mul_in1])].op_type) == "LEAKY_RELU"
                and mul_in0 in producers
                and str(model_ir.operators[int(producers[mul_in0])].op_type) == "RESHAPE"
            ):
                leaky_idx = int(producers[mul_in1])
                reshape_idx = int(producers[mul_in0])
            if leaky_idx is None or reshape_idx is None:
                continue

            leaky_op = model_ir.operators[int(leaky_idx)]
            if len(leaky_op.inputs) != 1 or len(leaky_op.outputs) != 1:
                continue
            leaky_out_name = str(leaky_op.outputs[0])
            if set(int(v) for v in consumers.get(leaky_out_name, [])) != {int(mul_idx)}:
                continue
            if leaky_out_name in model_outputs:
                continue

            leaky_pre_name = str(leaky_op.inputs[0])
            leaky_pre_idx = producers.get(leaky_pre_name, None)
            if leaky_pre_idx is None:
                continue
            leaky_pre_op = model_ir.operators[int(leaky_pre_idx)]
            if (
                str(leaky_pre_op.op_type) != "TRANSPOSE"
                or len(leaky_pre_op.inputs) < 2
                or len(leaky_pre_op.outputs) != 1
                or str(leaky_pre_op.outputs[0]) != leaky_pre_name
                or _read_transpose_perm(model_ir, leaky_pre_op) != perm_ndhwc_to_ncdhw
            ):
                continue
            conv_ndhwc_name = str(leaky_pre_op.inputs[0])
            if conv_ndhwc_name in model_outputs or leaky_pre_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(leaky_pre_name, [])) != {int(leaky_idx)}:
                continue

            reshape_op = model_ir.operators[int(reshape_idx)]
            if len(reshape_op.inputs) < 2 or len(reshape_op.outputs) != 1:
                continue
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(reshape_out_name, [])) != {int(mul_idx)}:
                continue

            sem_ncdhw_name = str(reshape_op.inputs[0])
            sem_pre_idx = producers.get(sem_ncdhw_name, None)
            if sem_pre_idx is None:
                continue
            sem_pre_op = model_ir.operators[int(sem_pre_idx)]
            sem_pre_perm = _read_transpose_perm(model_ir, sem_pre_op)
            if (
                str(sem_pre_op.op_type) != "TRANSPOSE"
                or len(sem_pre_op.inputs) < 2
                or len(sem_pre_op.outputs) != 1
                or str(sem_pre_op.outputs[0]) != sem_ncdhw_name
                or (
                    sem_pre_perm != perm_nhwc_to_nchw
                    and sem_pre_perm != perm_ndhwc_to_ncdhw
                )
            ):
                continue
            sem_ndhwc_name = str(sem_pre_op.inputs[0])
            if sem_ndhwc_name in model_outputs or sem_ncdhw_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(sem_ncdhw_name, [])) != {int(reshape_idx)}:
                continue

            reshape_shape_name = str(reshape_op.inputs[1])
            reshape_shape_tensor = model_ir.tensors.get(reshape_shape_name, None)
            reshape_shape_vals = _read_const_ints_from_tensor(reshape_shape_tensor)
            if reshape_shape_vals is None or len(reshape_shape_vals) != 5:
                continue
            mapped_shape = _permute_shape(
                [int(v) for v in list(reshape_shape_vals)],
                perm_ncdhw_to_ndhwc,
            )
            if mapped_shape is None:
                continue

            if not _write_const_ints_to_tensor(
                reshape_shape_tensor,
                [int(v) for v in list(mapped_shape)],
            ):
                continue
            if isinstance(reshape_op.options, dict):
                reshape_opts = dict(reshape_op.options)
                reshape_opts_changed = False
                for key in ["newShape", "onnxRawNewShape"]:
                    value = reshape_opts.get(key, None)
                    if isinstance(value, list) and len(value) == 5:
                        remapped = _permute_shape([int(v) for v in list(value)], perm_ncdhw_to_ndhwc)
                        if remapped is not None and [int(v) for v in list(remapped)] != [int(v) for v in list(value)]:
                            reshape_opts[key] = [int(v) for v in list(remapped)]
                            reshape_opts_changed = True
                if reshape_opts_changed:
                    reshape_op.options = reshape_opts

            _set_operator_inputs(
                model_ir=model_ir,
                op=reshape_op,
                new_inputs=[str(sem_ndhwc_name), str(reshape_shape_name)],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=leaky_op,
                new_inputs=[str(conv_ndhwc_name)],
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(reshape_out_name), None),
                perm_ncdhw_to_ndhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(leaky_out_name), None),
                perm_ncdhw_to_ndhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul_out_name), None),
                perm_ncdhw_to_ndhwc,
            )

            old_mul_out_tensor = model_ir.tensors.get(str(mul_out_name), None)
            mul_ndhwc_tensor = model_ir.tensors.get(str(mul_ndhwc_name), None)
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[str(mul_ndhwc_name)],
            )
            if old_mul_out_tensor is not None and mul_ndhwc_tensor is not None:
                mul_ndhwc_tensor.dtype = str(old_mul_out_tensor.dtype)
                mul_ndhwc_tensor.quantization = _clone_quantization(old_mul_out_tensor.quantization)
                mul_ndhwc_tensor.shape = [int(v) for v in list(old_mul_out_tensor.shape)]
                mul_ndhwc_tensor.shape_signature = (
                    [int(v) for v in list(old_mul_out_tensor.shape_signature)]
                    if old_mul_out_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul_out_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    mul_ndhwc_tensor,
                    perm_ncdhw_to_ndhwc,
                )

            for remove_idx in sorted(
                list({int(sem_pre_idx), int(leaky_pre_idx), int(post_idx)}),
                reverse=True,
            ):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains": int(rewritten)}


def _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove NCHW/NCDHW adapter transposes in cost-volume ScatterND accumulation motifs.

    Target:
      desc0_nhwc --T(0,3,1,2)--> desc0_nchw
      desc1_nhwc --T(0,3,1,2)--> desc1_nchw
      (SLICE/SUM/SQRT/DIV/MEAN/RESHAPE/SCATTER_ND chain in NCHW+NCDHW space)
      vol_ncdhw --T(0,2,3,4,1)--> vol_ndhwc --CONV_3D-->

    Rewrite:
      - Keep rank-4 path in NHWC: rewrite SLICE and reduce axes accordingly.
      - Keep ScatterND result tensors in NDHWC: rewrite ScatterND shape and indices.
      - Remove the two leading NHWC->NCHW transposes and trailing NCDHW->NDHWC transpose.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_ncdhw_to_ndhwc = [0, 2, 3, 4, 1]
    perm_ndhwc_to_ncdhw = [0, 4, 1, 2, 3]
    reduce_ops = {"SUM", "MEAN", "REDUCE_MAX"}
    scatter_layout_ops = {"SCATTER_ND", "ADD", "SUB", "MUL"}
    scatter_const_input_layout_ops = {"ADD", "SUB", "MUL"}
    allowed_ops = {
        "SLICE",
        "MUL",
        "SUM",
        "SQRT",
        "DIV",
        "MEAN",
        "RESHAPE",
        "CAST",
        "SCATTER_ND",
        "SUB",
        "ADD",
        "CONCATENATION",
        "REDUCE_MAX",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _ensure_exclusive_const_input(
        *,
        op: OperatorIR,
        op_idx: int,
        input_index: int,
        consumers: Dict[str, List[int]],
    ) -> Optional[TensorIR]:
        if int(input_index) < 0 or int(input_index) >= len(op.inputs):
            return None
        tensor_name = str(op.inputs[int(input_index)])
        tensor = model_ir.tensors.get(tensor_name, None)
        if tensor is None or tensor.data is None:
            return None
        tensor_consumers = [int(v) for v in consumers.get(tensor_name, [])]
        if len(tensor_consumers) <= 1:
            return tensor
        if set(tensor_consumers) == {int(op_idx)}:
            return tensor
        new_name = _unique_tensor_name(f"{tensor_name}_layout")
        new_data = np.asarray(tensor.data).copy()
        new_shape_signature = (
            [int(v) for v in list(tensor.shape_signature)]
            if tensor.shape_signature is not None
            else [int(v) for v in list(tensor.shape)]
        )
        model_ir.tensors[new_name] = TensorIR(
            name=new_name,
            dtype=str(tensor.dtype),
            shape=[int(v) for v in list(tensor.shape)],
            shape_signature=[int(v) for v in list(new_shape_signature)],
            data=new_data,
            is_variable=bool(tensor.is_variable),
            quantization=_clone_quantization(tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=str(new_name),
        )
        return model_ir.tensors.get(str(new_name), None)

    def _normalize_axis(axis: int, rank: int) -> Optional[int]:
        value = int(axis)
        if value < 0:
            value += int(rank)
        if value < 0 or value >= int(rank):
            return None
        return int(value)

    def _remap_axes_const(
        *,
        op: OperatorIR,
        op_idx: int,
        consumers: Dict[str, List[int]],
        input_index: int,
        rank: int,
        axis_map: List[int],
    ) -> bool:
        axes_tensor = _ensure_exclusive_const_input(
            op=op,
            op_idx=int(op_idx),
            input_index=int(input_index),
            consumers=consumers,
        )
        axes_vals = _read_const_ints_from_tensor(axes_tensor)
        if axes_vals is None or len(axes_vals) == 0:
            return False
        mapped_axes: List[int] = []
        for axis in axes_vals:
            normalized = _normalize_axis(int(axis), int(rank))
            if normalized is None:
                return False
            mapped_axes.append(int(axis_map[int(normalized)]))
        _write_const_ints_to_tensor(axes_tensor, [int(v) for v in mapped_axes])
        return True

    def _remap_slice_consts(
        *,
        op: OperatorIR,
        op_idx: int,
        consumers: Dict[str, List[int]],
    ) -> bool:
        if len(op.inputs) < 3:
            return False
        begin_tensor = _ensure_exclusive_const_input(
            op=op,
            op_idx=int(op_idx),
            input_index=1,
            consumers=consumers,
        )
        size_tensor = _ensure_exclusive_const_input(
            op=op,
            op_idx=int(op_idx),
            input_index=2,
            consumers=consumers,
        )
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None:
            return False
        if len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[int(idx)]) for idx in perm_nchw_to_nhwc]
        new_size = [int(size_vals[int(idx)]) for idx in perm_nchw_to_nhwc]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    def _remap_concat_axis(
        *,
        op: OperatorIR,
        rank: int,
        axis_map: List[int],
    ) -> bool:
        if not isinstance(op.options, dict):
            return False
        raw_axis = int(op.options.get("axis", 1))
        if raw_axis < 0:
            raw_axis += int(rank)
        if raw_axis < 0 or raw_axis >= int(rank):
            return False
        op.options["axis"] = int(axis_map[int(raw_axis)])
        return True

    def _remap_scatter_shape_const(
        *,
        op: OperatorIR,
        op_idx: int,
        consumers: Dict[str, List[int]],
    ) -> Optional[List[int]]:
        if len(op.inputs) < 3:
            return None
        shape_tensor = _ensure_exclusive_const_input(
            op=op,
            op_idx=int(op_idx),
            input_index=2,
            consumers=consumers,
        )
        shape_vals = _read_const_ints_from_tensor(shape_tensor)
        if shape_vals is None or len(shape_vals) != 5:
            return None
        mapped_shape = _permute_shape(
            [int(v) for v in list(shape_vals)],
            perm_ncdhw_to_ndhwc,
        )
        if mapped_shape is None:
            return None
        _write_const_ints_to_tensor(shape_tensor, [int(v) for v in list(mapped_shape)])
        return [int(v) for v in list(mapped_shape)]

    def _remap_scatter_indices_const(
        *,
        op: OperatorIR,
        op_idx: int,
        producers: Dict[str, int],
        consumers: Dict[str, List[int]],
        target_shape: List[int],
    ) -> bool:
        if len(op.inputs) < 1:
            return False
        indices_name = str(op.inputs[0])
        const_owner_op = op
        const_owner_idx = int(op_idx)
        const_input_index = 0

        indices_prod_idx = producers.get(indices_name, None)
        if indices_prod_idx is not None:
            indices_prod_op = model_ir.operators[int(indices_prod_idx)]
            if str(indices_prod_op.op_type) == "CAST" and len(indices_prod_op.inputs) == 1:
                const_owner_op = indices_prod_op
                const_owner_idx = int(indices_prod_idx)
                const_input_index = 0

        const_tensor = _ensure_exclusive_const_input(
            op=const_owner_op,
            op_idx=int(const_owner_idx),
            input_index=int(const_input_index),
            consumers=consumers,
        )
        if const_tensor is None or const_tensor.data is None:
            return False

        coord_array = np.asarray(const_tensor.data)
        if int(coord_array.size) == 0:
            return False
        if int(coord_array.shape[-1]) != 5:
            return False
        flat = coord_array.reshape(-1, 5).copy()

        def _indices_in_bounds(coords: np.ndarray, shape: List[int]) -> bool:
            if int(coords.shape[1]) != int(len(shape)):
                return False
            for axis in range(int(len(shape))):
                dim = int(shape[axis])
                if dim <= 0:
                    return False
                axis_vals = coords[:, int(axis)]
                if np.any(axis_vals < 0):
                    return False
                if np.any(axis_vals >= dim):
                    return False
            return True

        remapped_flat = flat[:, [int(v) for v in perm_ncdhw_to_ndhwc]]
        selected_flat = flat
        if _indices_in_bounds(flat, target_shape):
            selected_flat = flat
        elif _indices_in_bounds(remapped_flat, target_shape):
            selected_flat = remapped_flat
        else:
            return False

        remapped = selected_flat.reshape(coord_array.shape).astype(coord_array.dtype, copy=False)
        if not np.array_equal(remapped, coord_array):
            const_tensor.data = np.asarray(remapped)
            const_tensor.shape = [int(v) for v in list(remapped.shape)]
            const_tensor.shape_signature = [int(v) for v in list(remapped.shape)]
        return True

    def _permute_tensor_data_and_metadata_if_rank_matches(
        tensor: Optional[TensorIR],
        perm: List[int],
    ) -> None:
        if tensor is None:
            return
        if tensor.data is not None:
            try:
                data_array = np.asarray(tensor.data)
                if int(data_array.ndim) == int(len(perm)):
                    transposed = np.transpose(data_array, axes=perm)
                    tensor.data = np.asarray(transposed)
                    tensor.shape = [int(v) for v in list(transposed.shape)]
                    tensor.shape_signature = [int(v) for v in list(transposed.shape)]
                    return
            except Exception:
                pass
        _permute_tensor_metadata_if_rank_matches(tensor, perm)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_ncdhw_to_ndhwc
            ):
                continue

            post_in_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if (
                post_in_name in model_outputs
                or post_out_name in model_outputs
            ):
                continue
            if set(int(v) for v in consumers.get(post_in_name, [])) != {int(post_idx)}:
                continue

            post_users = [int(v) for v in consumers.get(post_out_name, [])]
            if len(post_users) != 1:
                continue
            conv_idx = int(post_users[0])
            conv_op = model_ir.operators[int(conv_idx)]
            if (
                str(conv_op.op_type) != "CONV_3D"
                or len(conv_op.inputs) < 1
                or str(conv_op.inputs[0]) != post_out_name
            ):
                continue

            root_idx = producers.get(post_in_name, None)
            if root_idx is None:
                continue

            candidate_indices: set[int] = set()
            boundary_indices: set[int] = set()
            traversal_ok = True
            stack: List[int] = [int(root_idx)]
            visited: set[int] = set()

            while len(stack) > 0:
                current_idx = int(stack.pop())
                if int(current_idx) in visited:
                    continue
                visited.add(int(current_idx))
                current_op = model_ir.operators[int(current_idx)]

                if str(current_op.op_type) == "TRANSPOSE":
                    current_perm = _read_transpose_perm(model_ir, current_op)
                    if current_perm == perm_nhwc_to_nchw:
                        boundary_indices.add(int(current_idx))
                        continue
                    traversal_ok = False
                    break

                if str(current_op.op_type) not in allowed_ops:
                    traversal_ok = False
                    break

                candidate_indices.add(int(current_idx))
                for input_name in list(current_op.inputs):
                    parent_idx = producers.get(str(input_name), None)
                    if parent_idx is not None:
                        stack.append(int(parent_idx))

            if not traversal_ok:
                continue
            if len(candidate_indices) == 0:
                continue
            if len(boundary_indices) != 2:
                continue

            boundary_plans: List[Tuple[int, str, str]] = []
            for boundary_idx in sorted(list(boundary_indices)):
                boundary_op = model_ir.operators[int(boundary_idx)]
                if (
                    str(boundary_op.op_type) != "TRANSPOSE"
                    or len(boundary_op.inputs) < 2
                    or len(boundary_op.outputs) != 1
                    or _read_transpose_perm(model_ir, boundary_op) != perm_nhwc_to_nchw
                ):
                    traversal_ok = False
                    break
                boundary_in = str(boundary_op.inputs[0])
                boundary_out = str(boundary_op.outputs[0])
                if boundary_in in model_outputs or boundary_out in model_outputs:
                    traversal_ok = False
                    break
                boundary_users = set(int(v) for v in consumers.get(boundary_out, []))
                if len(boundary_users) == 0:
                    traversal_ok = False
                    break
                if any(int(user_idx) not in candidate_indices for user_idx in boundary_users):
                    traversal_ok = False
                    break
                boundary_plans.append((int(boundary_idx), str(boundary_in), str(boundary_out)))
            if not traversal_ok:
                continue

            apply_ok = True
            for op_idx in sorted(list(candidate_indices)):
                op = model_ir.operators[int(op_idx)]
                op_type = str(op.op_type)
                if op_type == "SLICE":
                    if not _remap_slice_consts(
                        op=op,
                        op_idx=int(op_idx),
                        consumers=consumers,
                    ):
                        apply_ok = False
                        break
                elif op_type in reduce_ops:
                    if not _remap_axes_const(
                        op=op,
                        op_idx=int(op_idx),
                        consumers=consumers,
                        input_index=1,
                        rank=4,
                        axis_map=perm_nhwc_to_nchw,
                    ):
                        apply_ok = False
                        break
                elif op_type == "CONCATENATION":
                    rank_hint = 4
                    if len(op.inputs) > 0:
                        in_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
                        if in_tensor is not None and in_tensor.shape is not None:
                            rank_hint = int(len(list(in_tensor.shape)))
                    if int(rank_hint) == 4:
                        if not _remap_concat_axis(
                            op=op,
                            rank=4,
                            axis_map=perm_nhwc_to_nchw,
                        ):
                            apply_ok = False
                            break
                elif op_type == "SCATTER_ND":
                    mapped_shape = _remap_scatter_shape_const(
                        op=op,
                        op_idx=int(op_idx),
                        consumers=consumers,
                    )
                    if mapped_shape is None:
                        apply_ok = False
                        break
                    if not _remap_scatter_indices_const(
                        op=op,
                        op_idx=int(op_idx),
                        producers=producers,
                        consumers=consumers,
                        target_shape=[int(v) for v in list(mapped_shape)],
                    ):
                        apply_ok = False
                        break

            if not apply_ok:
                continue

            for _, boundary_in, boundary_out in boundary_plans:
                _replace_tensor_inputs(
                    model_ir=model_ir,
                    src_name=str(boundary_out),
                    dst_name=str(boundary_in),
                )

            conv_inputs = [
                str(post_in_name) if str(v) == str(post_out_name) else str(v)
                for v in list(conv_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=conv_op,
                new_inputs=conv_inputs,
            )

            for op_idx in sorted(list(candidate_indices)):
                op = model_ir.operators[int(op_idx)]
                if str(op.op_type) not in scatter_const_input_layout_ops:
                    continue
                for input_name in list(op.inputs):
                    input_tensor = model_ir.tensors.get(str(input_name), None)
                    if input_tensor is None:
                        continue
                    rank = (
                        len(list(input_tensor.shape))
                        if input_tensor.shape is not None
                        else (
                            len(list(input_tensor.shape_signature))
                            if input_tensor.shape_signature is not None
                            else 0
                        )
                    )
                    if int(rank) != 5:
                        continue
                    producer_idx = producers.get(str(input_name), None)
                    if producer_idx is not None:
                        continue
                    _permute_tensor_data_and_metadata_if_rank_matches(
                        input_tensor,
                        perm_ncdhw_to_ndhwc,
                    )

            for op_idx in sorted(list(candidate_indices)):
                op = model_ir.operators[int(op_idx)]
                for output_name in list(op.outputs):
                    output_tensor = model_ir.tensors.get(str(output_name), None)
                    if output_tensor is None:
                        continue
                    rank = (
                        len(list(output_tensor.shape))
                        if output_tensor.shape is not None
                        else (
                            len(list(output_tensor.shape_signature))
                            if output_tensor.shape_signature is not None
                            else 0
                        )
                    )
                    if int(rank) == 4:
                        _permute_tensor_metadata_if_rank_matches(
                            output_tensor,
                            perm_nchw_to_nhwc,
                        )
                    elif int(rank) == 5 and str(op.op_type) in scatter_layout_ops:
                        _permute_tensor_metadata_if_rank_matches(
                            output_tensor,
                            perm_ncdhw_to_ndhwc,
                        )

            remove_indices = sorted(
                [int(post_idx)] + [int(v[0]) for v in list(boundary_plans)],
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_cost_volume_scatter_ndhwc_chains": int(rewritten)}


def _optimize_transpose_conv_attention_nhwc_propagation_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Propagate NHWC layout through Conv-attention motifs and remove bridge transposes.

    Target motif (rank-4):
      x_nhwc --T(0,3,1,2)--> x_nchw --ACT--> a_nchw
      a_nchw --MEAN(keepDims=True, axes in NCHW)--> m_nchw --T(0,2,3,1)--> m_nhwc --CONV--CONV--> c_nhwc
      c_nhwc --T(0,3,1,2)--> c_nchw --GATE--> s_nchw
      MUL(a_nchw, s_nchw) --> y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - bypass both NCHW pre-transposes and all post-transposes in the motif
      - remap MEAN axes from NCHW to NHWC
      - keep ACT/MEAN/GATE/MUL in NHWC

    Supported GATE:
      - LOGISTIC
      - HardSigmoid expansion: MUL(const)->ADD(const)->RELU_0_TO_1
      - HardSigmoid expansion: MUL(const)->ADD(const)->MAXIMUM(const)->MINIMUM(const)
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw_const_name = "__nhwc_to_nchw_perm_rank4__"
    activation_ops = {"RELU", "LEAKY_RELU", "RELU6"}
    legacy_consumer_safe_ops = {
        "ADD",
        "SUB",
        "MUL",
        "DIV",
        "MAXIMUM",
        "MINIMUM",
        "RELU",
        "RELU6",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "NEG",
        "ABS",
        "SQRT",
        "EXP",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    if perm_nhwc_to_nchw_const_name not in model_ir.tensors:
        model_ir.tensors[perm_nhwc_to_nchw_const_name] = TensorIR(
            name=perm_nhwc_to_nchw_const_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        def _match_hardsigmoid_gate_from_output(gate_output_name: str) -> Optional[Dict[str, Any]]:
            gate_idx = producers.get(str(gate_output_name), None)
            if gate_idx is None:
                return None
            gate_op = model_ir.operators[int(gate_idx)]
            gate_type = str(gate_op.op_type)
            add_output_name: Optional[str] = None
            chain_indices: List[int] = []
            metadata_names: List[str] = []

            if gate_type == "RELU_0_TO_1":
                if len(gate_op.inputs) != 1 or len(gate_op.outputs) != 1:
                    return None
                add_output_name = str(gate_op.inputs[0])
                if set(int(v) for v in consumers.get(add_output_name, [])) != {int(gate_idx)}:
                    return None
                chain_indices.append(int(gate_idx))
                metadata_names.append(str(gate_output_name))
            elif gate_type == "MINIMUM":
                if len(gate_op.inputs) != 2 or len(gate_op.outputs) != 1:
                    return None
                min_inputs = [str(v) for v in list(gate_op.inputs)]
                max_output_name: Optional[str] = None
                min_side_name: Optional[str] = None
                for candidate_name, other_name in [(min_inputs[0], min_inputs[1]), (min_inputs[1], min_inputs[0])]:
                    producer_idx = producers.get(candidate_name, None)
                    if producer_idx is None:
                        continue
                    producer_op = model_ir.operators[int(producer_idx)]
                    if str(producer_op.op_type) != "MAXIMUM":
                        continue
                    max_output_name = str(candidate_name)
                    min_side_name = str(other_name)
                    break
                if max_output_name is None or min_side_name is None:
                    return None
                if not _is_singleton_constant_tensor(model_ir, min_side_name):
                    return None
                if set(int(v) for v in consumers.get(max_output_name, [])) != {int(gate_idx)}:
                    return None

                max_idx = producers.get(max_output_name, None)
                if max_idx is None:
                    return None
                max_op = model_ir.operators[int(max_idx)]
                if str(max_op.op_type) != "MAXIMUM" or len(max_op.inputs) != 2 or len(max_op.outputs) != 1:
                    return None
                max_inputs = [str(v) for v in list(max_op.inputs)]
                for candidate_name, other_name in [(max_inputs[0], max_inputs[1]), (max_inputs[1], max_inputs[0])]:
                    producer_idx = producers.get(candidate_name, None)
                    if producer_idx is None:
                        continue
                    producer_op = model_ir.operators[int(producer_idx)]
                    if str(producer_op.op_type) != "ADD":
                        continue
                    add_output_name = str(candidate_name)
                    if not _is_singleton_constant_tensor(model_ir, str(other_name)):
                        return None
                    break
                if add_output_name is None:
                    return None
                if set(int(v) for v in consumers.get(add_output_name, [])) != {int(max_idx)}:
                    return None
                chain_indices.extend([int(max_idx), int(gate_idx)])
                metadata_names.extend([max_output_name, str(gate_output_name)])
            else:
                return None

            add_idx = producers.get(str(add_output_name), None)
            if add_idx is None:
                return None
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                return None
            add_inputs = [str(v) for v in list(add_op.inputs)]
            mul_output_name: Optional[str] = None
            for candidate_name, other_name in [(add_inputs[0], add_inputs[1]), (add_inputs[1], add_inputs[0])]:
                producer_idx = producers.get(candidate_name, None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if str(producer_op.op_type) != "MUL":
                    continue
                if not _is_singleton_constant_tensor(model_ir, str(other_name)):
                    return None
                mul_output_name = str(candidate_name)
                break
            if mul_output_name is None:
                return None
            if set(int(v) for v in consumers.get(mul_output_name, [])) != {int(add_idx)}:
                return None
            chain_indices.append(int(add_idx))
            metadata_names.insert(0, str(add_output_name))

            mul_idx = producers.get(mul_output_name, None)
            if mul_idx is None:
                return None
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                return None
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            main_input_name: Optional[str] = None
            main_input_index: Optional[int] = None
            for input_index, input_name in enumerate(mul_inputs):
                side_name = str(mul_inputs[1 - input_index])
                if _is_singleton_constant_tensor(model_ir, side_name):
                    main_input_name = str(input_name)
                    main_input_index = int(input_index)
                    break
            if main_input_name is None or main_input_index is None:
                return None

            chain_indices.append(int(mul_idx))
            metadata_names.insert(0, str(mul_output_name))

            return {
                "head_idx": int(mul_idx),
                "head_input_name": str(main_input_name),
                "head_input_index": int(main_input_index),
                "output_name": str(gate_output_name),
                "chain_indices": [int(v) for v in chain_indices],
                "metadata_names": [str(v) for v in metadata_names],
            }

        def _match_hardswish_activation_from_output(activation_output_name: str) -> Optional[Dict[str, Any]]:
            activation_idx = producers.get(str(activation_output_name), None)
            if activation_idx is None:
                return None
            activation_op = model_ir.operators[int(activation_idx)]
            if (
                str(activation_op.op_type) != "MUL"
                or len(activation_op.inputs) != 2
                or len(activation_op.outputs) != 1
                or str(activation_op.outputs[0]) != str(activation_output_name)
            ):
                return None

            activation_inputs = [str(v) for v in list(activation_op.inputs)]
            for data_input_index in [0, 1]:
                data_input_name = str(activation_inputs[int(data_input_index)])
                gate_output_name = str(activation_inputs[1 - int(data_input_index)])

                gate_match = _match_hardsigmoid_gate_from_output(gate_output_name)
                if gate_match is None:
                    continue
                if str(gate_match["head_input_name"]) != str(data_input_name):
                    continue

                gate_users = set(int(v) for v in consumers.get(str(gate_output_name), []))
                if gate_users != {int(activation_idx)}:
                    continue

                expected_data_users = {int(activation_idx), int(gate_match["head_idx"])}
                data_users = set(int(v) for v in consumers.get(str(data_input_name), []))
                if data_users != expected_data_users:
                    continue

                activation_metadata_names = [str(v) for v in list(gate_match.get("metadata_names", []))]
                activation_metadata_names.append(str(activation_output_name))
                rewire_specs = [
                    (int(gate_match["head_idx"]), int(gate_match["head_input_index"])),
                    (int(activation_idx), int(data_input_index)),
                ]

                return {
                    "activation_idx": int(activation_idx),
                    "activation_input_name": str(data_input_name),
                    "activation_input_expected_users": sorted(list(expected_data_users)),
                    "activation_output_name": str(activation_output_name),
                    "activation_metadata_names": activation_metadata_names,
                    "rewire_specs": rewire_specs,
                }
            return None

        def _match_self_hardswish_from_source(source_nchw_name: str) -> Optional[Dict[str, Any]]:
            source_users = [int(v) for v in consumers.get(str(source_nchw_name), [])]
            if len(source_users) == 0:
                return None

            add_idx: Optional[int] = None
            mul_idx: Optional[int] = None
            add_out_name: Optional[str] = None
            gate_name: Optional[str] = None

            for user_idx in source_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if user_type == "ADD":
                    if len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                        return None
                    add_inputs = [str(v) for v in list(user_op.inputs)]
                    if source_nchw_name == add_inputs[0]:
                        add_side_name = str(add_inputs[1])
                    elif source_nchw_name == add_inputs[1]:
                        add_side_name = str(add_inputs[0])
                    else:
                        return None
                    if not _is_singleton_constant_tensor(model_ir, add_side_name):
                        return None
                    if add_idx is not None:
                        return None
                    add_idx = int(user_idx)
                    add_out_name = str(user_op.outputs[0])
                    continue
                if user_type == "MUL":
                    if len(user_op.inputs) != 2 or len(user_op.outputs) != 1:
                        return None
                    mul_inputs = [str(v) for v in list(user_op.inputs)]
                    if source_nchw_name == mul_inputs[0]:
                        other_name = str(mul_inputs[1])
                    elif source_nchw_name == mul_inputs[1]:
                        other_name = str(mul_inputs[0])
                    else:
                        return None
                    if mul_idx is not None:
                        return None
                    mul_idx = int(user_idx)
                    gate_name = str(other_name)
                    continue
                return None

            if add_idx is None or mul_idx is None or add_out_name is None or gate_name is None:
                return None
            if set(int(v) for v in source_users) != {int(add_idx), int(mul_idx)}:
                return None

            scale_idx = producers.get(str(gate_name), None)
            if scale_idx is None:
                return None
            scale_op = model_ir.operators[int(scale_idx)]
            if (
                str(scale_op.op_type) not in {"MUL", "DIV"}
                or len(scale_op.inputs) != 2
                or len(scale_op.outputs) != 1
                or str(scale_op.outputs[0]) != str(gate_name)
            ):
                return None
            scale_inputs = [str(v) for v in list(scale_op.inputs)]
            if str(add_out_name) == scale_inputs[0]:
                scale_side_name = str(scale_inputs[1])
            elif str(add_out_name) == scale_inputs[1]:
                scale_side_name = str(scale_inputs[0])
            else:
                return None
            if not _is_singleton_constant_tensor(model_ir, scale_side_name):
                return None
            if set(int(v) for v in consumers.get(str(add_out_name), [])) != {int(scale_idx)}:
                return None
            if set(int(v) for v in consumers.get(str(gate_name), [])) != {int(mul_idx)}:
                return None

            mul_op = model_ir.operators[int(mul_idx)]
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            if not (
                (str(source_nchw_name) == mul_inputs[0] and str(gate_name) == mul_inputs[1])
                or (str(source_nchw_name) == mul_inputs[1] and str(gate_name) == mul_inputs[0])
            ):
                return None

            return {
                "add_idx": int(add_idx),
                "scale_idx": int(scale_idx),
                "mul_idx": int(mul_idx),
                "add_out_name": str(add_out_name),
                "scale_out_name": str(gate_name),
                "mul_out_name": str(mul_op.outputs[0]),
            }

        for candidate_post_idx, candidate_post_op in enumerate(model_ir.operators):
            if (
                str(candidate_post_op.op_type) != "TRANSPOSE"
                or len(candidate_post_op.inputs) < 2
                or len(candidate_post_op.outputs) != 1
                or _read_transpose_perm(model_ir, candidate_post_op) != perm_nchw_to_nhwc
            ):
                continue

            mul_output_name = str(candidate_post_op.inputs[0])
            if mul_output_name in model_outputs:
                continue
            mul_idx = producers.get(mul_output_name, None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(mul_op.outputs[0]) != mul_output_name:
                continue

            mul_output_users = [int(v) for v in consumers.get(mul_output_name, [])]
            if len(mul_output_users) == 0:
                continue
            post_pairs: List[Tuple[int, OperatorIR]] = []
            legacy_slots: List[Tuple[int, int]] = []
            valid_posts = True
            for user_idx in mul_output_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != mul_output_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    if str(user_op.op_type) not in legacy_consumer_safe_ops:
                        valid_posts = False
                        break
                    input_indices = [
                        int(input_idx)
                        for input_idx, input_name in enumerate(list(user_op.inputs))
                        if str(input_name) == str(mul_output_name)
                    ]
                    if len(input_indices) == 0:
                        valid_posts = False
                        break
                    for input_idx in input_indices:
                        legacy_slots.append((int(id(user_op)), int(input_idx)))
                    continue
                post_pairs.append((int(user_idx), user_op))
            if not valid_posts or len(post_pairs) == 0:
                continue
            post_pairs = sorted(post_pairs, key=lambda v: int(v[0]))
            post_indices = [int(v[0]) for v in post_pairs]
            post_output_names = [str(v[1].outputs[0]) for v in post_pairs]

            input0_name = str(mul_op.inputs[0])
            input1_name = str(mul_op.inputs[1])
            activation_idx: Optional[int] = None
            activation_input_name: Optional[str] = None
            activation_input_expected_users: Optional[set[int]] = None
            activation_output_name: Optional[str] = None
            activation_rewire_specs: List[Tuple[int, int]] = []
            activation_metadata_names: List[str] = []
            gate_kind: Optional[str] = None
            gate_head_idx: Optional[int] = None
            gate_head_input_name: Optional[str] = None
            gate_head_input_index: Optional[int] = None
            gate_output_name: Optional[str] = None
            gate_metadata_names: List[str] = []
            for lhs_name, rhs_name in [(input0_name, input1_name), (input1_name, input0_name)]:
                lhs_prod_idx = producers.get(lhs_name, None)
                if lhs_prod_idx is None:
                    continue
                lhs_op = model_ir.operators[int(lhs_prod_idx)]
                lhs_activation_match: Optional[Dict[str, Any]] = None
                if str(lhs_op.op_type) in activation_ops:
                    if len(lhs_op.inputs) == 1 and len(lhs_op.outputs) == 1 and str(lhs_op.outputs[0]) == str(lhs_name):
                        lhs_activation_match = {
                            "activation_idx": int(lhs_prod_idx),
                            "activation_input_name": str(lhs_op.inputs[0]),
                            "activation_input_expected_users": [int(lhs_prod_idx)],
                            "activation_output_name": str(lhs_name),
                            "activation_metadata_names": [str(lhs_name)],
                            "rewire_specs": [(int(lhs_prod_idx), 0)],
                        }
                else:
                    lhs_activation_match = _match_hardswish_activation_from_output(str(lhs_name))
                if lhs_activation_match is None:
                    continue

                rhs_prod_idx = producers.get(rhs_name, None)
                if rhs_prod_idx is None:
                    continue
                rhs_op = model_ir.operators[int(rhs_prod_idx)]

                if str(rhs_op.op_type) == "LOGISTIC":
                    if len(rhs_op.inputs) != 1 or len(rhs_op.outputs) != 1:
                        continue
                    if str(rhs_op.outputs[0]) != rhs_name:
                        continue
                    activation_idx = int(lhs_activation_match["activation_idx"])
                    activation_input_name = str(lhs_activation_match["activation_input_name"])
                    activation_input_expected_users = set(
                        int(v) for v in list(lhs_activation_match["activation_input_expected_users"])
                    )
                    activation_output_name = str(lhs_activation_match["activation_output_name"])
                    activation_rewire_specs = [
                        (int(v[0]), int(v[1])) for v in list(lhs_activation_match["rewire_specs"])
                    ]
                    activation_metadata_names = [
                        str(v) for v in list(lhs_activation_match.get("activation_metadata_names", []))
                    ]
                    gate_kind = "LOGISTIC"
                    gate_head_idx = int(rhs_prod_idx)
                    gate_head_input_name = str(rhs_op.inputs[0])
                    gate_head_input_index = 0
                    gate_output_name = str(rhs_name)
                    gate_metadata_names = [str(rhs_name)]
                    break

                hsig_match = _match_hardsigmoid_gate_from_output(str(rhs_name))
                if hsig_match is None:
                    continue
                activation_idx = int(lhs_activation_match["activation_idx"])
                activation_input_name = str(lhs_activation_match["activation_input_name"])
                activation_input_expected_users = set(
                    int(v) for v in list(lhs_activation_match["activation_input_expected_users"])
                )
                activation_output_name = str(lhs_activation_match["activation_output_name"])
                activation_rewire_specs = [
                    (int(v[0]), int(v[1])) for v in list(lhs_activation_match["rewire_specs"])
                ]
                activation_metadata_names = [
                    str(v) for v in list(lhs_activation_match.get("activation_metadata_names", []))
                ]
                gate_kind = "HARDSIGMOID"
                gate_head_idx = int(hsig_match["head_idx"])
                gate_head_input_name = str(hsig_match["head_input_name"])
                gate_head_input_index = int(hsig_match["head_input_index"])
                gate_output_name = str(hsig_match["output_name"])
                gate_metadata_names = [str(v) for v in list(hsig_match.get("metadata_names", []))]
                break
            if (
                activation_idx is None
                or activation_input_name is None
                or activation_input_expected_users is None
                or activation_output_name is None
                or gate_kind is None
                or gate_head_idx is None
                or gate_head_input_name is None
                or gate_head_input_index is None
                or gate_output_name is None
            ):
                continue

            gate_head_op = model_ir.operators[int(gate_head_idx)]
            if (
                activation_input_name in model_outputs
                or activation_output_name in model_outputs
                or gate_head_input_name in model_outputs
                or gate_output_name in model_outputs
            ):
                continue

            pre_activation_idx = producers.get(activation_input_name, None)
            if pre_activation_idx is None:
                continue
            pre_activation_op = model_ir.operators[int(pre_activation_idx)]
            if (
                str(pre_activation_op.op_type) != "TRANSPOSE"
                or len(pre_activation_op.inputs) < 2
                or len(pre_activation_op.outputs) != 1
                or str(pre_activation_op.outputs[0]) != activation_input_name
                or _read_transpose_perm(model_ir, pre_activation_op) != perm_nhwc_to_nchw
            ):
                continue
            source_nhwc_name = str(pre_activation_op.inputs[0])
            if source_nhwc_name in model_outputs:
                continue

            mean_idx_candidates = [
                int(v)
                for v in consumers.get(activation_output_name, [])
                if int(v) != int(mul_idx)
            ]
            if len(mean_idx_candidates) != 1:
                continue
            mean_idx = int(mean_idx_candidates[0])
            mean_op = model_ir.operators[int(mean_idx)]
            if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                continue
            if str(mean_op.inputs[0]) != activation_output_name:
                continue
            if not bool(mean_op.options.get("keepDims", False)):
                continue
            mean_output_name = str(mean_op.outputs[0])
            if mean_output_name in model_outputs:
                continue

            mean_post_users = [int(v) for v in consumers.get(mean_output_name, [])]
            if len(mean_post_users) != 1:
                continue
            mean_post_idx = int(mean_post_users[0])
            mean_post_op = model_ir.operators[int(mean_post_idx)]
            if (
                str(mean_post_op.op_type) != "TRANSPOSE"
                or len(mean_post_op.inputs) < 2
                or len(mean_post_op.outputs) != 1
                or str(mean_post_op.inputs[0]) != mean_output_name
                or _read_transpose_perm(model_ir, mean_post_op) != perm_nchw_to_nhwc
            ):
                continue
            conv1_input_name = str(mean_post_op.outputs[0])
            if conv1_input_name in model_outputs:
                continue

            conv1_users = [int(v) for v in consumers.get(conv1_input_name, [])]
            if len(conv1_users) != 1:
                continue
            conv1_idx = int(conv1_users[0])
            conv1_op = model_ir.operators[int(conv1_idx)]
            if str(conv1_op.op_type) != "CONV_2D" or len(conv1_op.inputs) < 1 or len(conv1_op.outputs) != 1:
                continue
            if str(conv1_op.inputs[0]) != conv1_input_name:
                continue
            conv1_output_name = str(conv1_op.outputs[0])
            if conv1_output_name in model_outputs:
                continue
            conv1_output_users = [int(v) for v in consumers.get(conv1_output_name, [])]
            if len(conv1_output_users) != 1:
                continue
            conv2_idx = int(conv1_output_users[0])
            conv2_op = model_ir.operators[int(conv2_idx)]
            if str(conv2_op.op_type) != "CONV_2D" or len(conv2_op.inputs) < 1 or len(conv2_op.outputs) != 1:
                continue
            if str(conv2_op.inputs[0]) != conv1_output_name:
                continue
            conv2_output_name = str(conv2_op.outputs[0])
            if conv2_output_name in model_outputs:
                continue

            pre_gate_idx = producers.get(gate_head_input_name, None)
            if pre_gate_idx is None:
                continue
            pre_gate_op = model_ir.operators[int(pre_gate_idx)]
            if (
                str(pre_gate_op.op_type) != "TRANSPOSE"
                or len(pre_gate_op.inputs) < 2
                or len(pre_gate_op.outputs) != 1
                or str(pre_gate_op.outputs[0]) != gate_head_input_name
                or _read_transpose_perm(model_ir, pre_gate_op) != perm_nhwc_to_nchw
                or str(pre_gate_op.inputs[0]) != conv2_output_name
            ):
                continue

            # Locality guards: this rewrite assumes a closed motif.
            if set(int(v) for v in consumers.get(activation_input_name, [])) != set(int(v) for v in activation_input_expected_users):
                continue
            if set(int(v) for v in consumers.get(activation_output_name, [])) != {int(mean_idx), int(mul_idx)}:
                continue
            if set(int(v) for v in consumers.get(mean_output_name, [])) != {int(mean_post_idx)}:
                continue
            if set(int(v) for v in consumers.get(conv1_input_name, [])) != {int(conv1_idx)}:
                continue
            if set(int(v) for v in consumers.get(conv1_output_name, [])) != {int(conv2_idx)}:
                continue
            if set(int(v) for v in consumers.get(conv2_output_name, [])) != {int(pre_gate_idx)}:
                continue
            if set(int(v) for v in consumers.get(gate_head_input_name, [])) != {int(gate_head_idx)}:
                continue
            if set(int(v) for v in consumers.get(gate_output_name, [])) != {int(mul_idx)}:
                continue

            # Remap MEAN axes from NCHW to NHWC.
            mean_axes_tensor = model_ir.tensors.get(str(mean_op.inputs[1]), None)
            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_vals is None or len(mean_axes_vals) == 0:
                continue
            source_tensor = model_ir.tensors.get(source_nhwc_name, None)
            rank = 4
            if source_tensor is not None and source_tensor.shape is not None and len(list(source_tensor.shape)) > 0:
                rank = int(len(list(source_tensor.shape)))
            if rank != 4:
                continue
            normalized_axes: List[int] = []
            valid_axes = True
            for axis in mean_axes_vals:
                a = int(axis)
                if a < 0:
                    a += int(rank)
                if a < 0 or a >= int(rank):
                    valid_axes = False
                    break
                normalized_axes.append(int(a))
            if not valid_axes:
                continue
            mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]
            if not _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes]):
                continue

            # Rewire to NHWC.
            for rewire_op_idx, rewire_input_idx in activation_rewire_specs:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=model_ir.operators[int(rewire_op_idx)],
                    input_index=int(rewire_input_idx),
                    new_input_name=source_nhwc_name,
                )
            conv1_inputs = [str(v) for v in list(conv1_op.inputs)]
            conv1_inputs[0] = str(mean_output_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=conv1_op,
                new_inputs=conv1_inputs,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=gate_head_op,
                input_index=int(gate_head_input_index),
                new_input_name=conv2_output_name,
            )

            # Update metadata of tensors that switch from NCHW to NHWC.
            old_mul_tensor = model_ir.tensors.get(mul_output_name, None)
            legacy_mul_shape = (
                [int(v) for v in list(old_mul_tensor.shape)]
                if old_mul_tensor is not None and old_mul_tensor.shape is not None
                else None
            )
            legacy_mul_shape_signature = None
            if old_mul_tensor is not None:
                legacy_mul_shape_signature = (
                    [int(v) for v in list(old_mul_tensor.shape_signature)]
                    if old_mul_tensor.shape_signature is not None
                    else (
                        [int(v) for v in list(old_mul_tensor.shape)]
                        if old_mul_tensor.shape is not None
                        else None
                    )
                )
            metadata_targets = [str(v) for v in list(activation_metadata_names)]
            metadata_targets.append(str(mean_output_name))
            metadata_targets.extend([str(v) for v in list(gate_metadata_names)])
            metadata_targets.append(str(mul_output_name))
            for tensor_name in list(dict.fromkeys(metadata_targets)):
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[canonical_post_output_name],
            )
            _replace_tensor_inputs(model_ir, mul_output_name, canonical_post_output_name)
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if old_mul_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_mul_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_mul_tensor.shape_signature)]
                    if old_mul_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            pending_legacy_slots = list(legacy_slots)

            remove_indices = {
                int(pre_activation_idx),
                int(mean_post_idx),
                int(pre_gate_idx),
                *[int(v) for v in post_indices],
            }
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            if len(pending_legacy_slots) > 0:
                op_index_by_id = {int(id(op)): int(op_idx) for op_idx, op in enumerate(model_ir.operators)}
                valid_legacy_slots: List[Tuple[int, int]] = []
                for op_id, input_idx in pending_legacy_slots:
                    op_idx = op_index_by_id.get(int(op_id), None)
                    if op_idx is None:
                        continue
                    if int(op_idx) < 0 or int(op_idx) >= len(model_ir.operators):
                        continue
                    legacy_op = model_ir.operators[int(op_idx)]
                    if int(input_idx) < 0 or int(input_idx) >= len(legacy_op.inputs):
                        continue
                    if str(legacy_op.inputs[int(input_idx)]) != str(canonical_post_output_name):
                        continue
                    valid_legacy_slots.append((int(op_idx), int(input_idx)))

                if len(valid_legacy_slots) > 0:
                    adapter_name = _unique_tensor_name(f"{mul_output_name}_nchw_adapter")
                    adapter_dtype = str(old_mul_tensor.dtype) if old_mul_tensor is not None else "FLOAT32"
                    adapter_quant = (
                        _clone_quantization(old_mul_tensor.quantization)
                        if old_mul_tensor is not None
                        else None
                    )
                    adapter_shape = (
                        [int(v) for v in list(legacy_mul_shape)]
                        if legacy_mul_shape is not None
                        else [1]
                    )
                    adapter_shape_signature = (
                        [int(v) for v in list(legacy_mul_shape_signature)]
                        if legacy_mul_shape_signature is not None
                        else [int(v) for v in list(adapter_shape)]
                    )
                    model_ir.tensors[adapter_name] = TensorIR(
                        name=adapter_name,
                        dtype=str(adapter_dtype),
                        shape=[int(v) for v in list(adapter_shape)],
                        shape_signature=[int(v) for v in list(adapter_shape_signature)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(adapter_quant),
                    )
                    adapter_op = OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(canonical_post_output_name), perm_nhwc_to_nchw_const_name],
                        outputs=[str(adapter_name)],
                    )
                    for op_idx, input_idx in valid_legacy_slots:
                        legacy_op = model_ir.operators[int(op_idx)]
                        new_inputs = [str(v) for v in list(legacy_op.inputs)]
                        new_inputs[int(input_idx)] = str(adapter_name)
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=legacy_op,
                            new_inputs=new_inputs,
                        )
                    insert_index = int(min(v[0] for v in valid_legacy_slots))
                    model_ir.operators.insert(int(insert_index), adapter_op)

            rewritten += 1
            changed = True
            break

        if not changed:
            for pre_idx, pre_op in enumerate(model_ir.operators):
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                ):
                    continue
                source_nhwc_name = str(pre_op.inputs[0])
                source_nchw_name = str(pre_op.outputs[0])
                if source_nhwc_name in model_outputs or source_nchw_name in model_outputs:
                    continue

                hs0_match = _match_self_hardswish_from_source(source_nchw_name)
                if hs0_match is None:
                    continue
                hsw0_out_name = str(hs0_match["mul_out_name"])
                if hsw0_out_name in model_outputs:
                    continue

                hsw0_users = [int(v) for v in consumers.get(hsw0_out_name, [])]
                if len(hsw0_users) != 1:
                    continue
                mean0_idx = int(hsw0_users[0])
                mean0_op = model_ir.operators[int(mean0_idx)]
                if (
                    str(mean0_op.op_type) != "MEAN"
                    or len(mean0_op.inputs) < 2
                    or len(mean0_op.outputs) != 1
                    or str(mean0_op.inputs[0]) != str(hsw0_out_name)
                    or not bool(mean0_op.options.get("keepDims", False))
                ):
                    continue
                mean0_out_name = str(mean0_op.outputs[0])
                if mean0_out_name in model_outputs:
                    continue
                mean0_axes_tensor = model_ir.tensors.get(str(mean0_op.inputs[1]), None)
                mean0_axes_vals = _read_const_ints_from_tensor(mean0_axes_tensor)
                if mean0_axes_vals is None or len(mean0_axes_vals) == 0:
                    continue

                mean0_users = [int(v) for v in consumers.get(mean0_out_name, [])]
                if len(mean0_users) != 1:
                    continue
                mean0_post_idx = int(mean0_users[0])
                mean0_post_op = model_ir.operators[int(mean0_post_idx)]
                if (
                    str(mean0_post_op.op_type) != "TRANSPOSE"
                    or len(mean0_post_op.inputs) < 2
                    or len(mean0_post_op.outputs) != 1
                    or str(mean0_post_op.inputs[0]) != str(mean0_out_name)
                    or _read_transpose_perm(model_ir, mean0_post_op) != perm_nchw_to_nhwc
                ):
                    continue
                conv_input_name = str(mean0_post_op.outputs[0])
                if conv_input_name in model_outputs:
                    continue

                conv_users = [int(v) for v in consumers.get(conv_input_name, [])]
                if len(conv_users) != 1:
                    continue
                conv_idx = int(conv_users[0])
                conv_op = model_ir.operators[int(conv_idx)]
                if (
                    str(conv_op.op_type) not in {"CONV_2D", "DEPTHWISE_CONV_2D"}
                    or len(conv_op.inputs) < 1
                    or len(conv_op.outputs) != 1
                    or str(conv_op.inputs[0]) != str(conv_input_name)
                ):
                    continue
                conv_out_name = str(conv_op.outputs[0])
                if conv_out_name in model_outputs:
                    continue

                conv_out_users = [int(v) for v in consumers.get(conv_out_name, [])]
                if len(conv_out_users) != 1:
                    continue
                pre1_idx = int(conv_out_users[0])
                pre1_op = model_ir.operators[int(pre1_idx)]
                if (
                    str(pre1_op.op_type) != "TRANSPOSE"
                    or len(pre1_op.inputs) < 2
                    or len(pre1_op.outputs) != 1
                    or str(pre1_op.inputs[0]) != str(conv_out_name)
                    or _read_transpose_perm(model_ir, pre1_op) != perm_nhwc_to_nchw
                ):
                    continue
                gate_nchw_name = str(pre1_op.outputs[0])
                if gate_nchw_name in model_outputs:
                    continue

                hs1_match = _match_self_hardswish_from_source(gate_nchw_name)
                if hs1_match is None:
                    continue
                hsw1_out_name = str(hs1_match["mul_out_name"])
                if hsw1_out_name in model_outputs:
                    continue

                hsw1_users = [int(v) for v in consumers.get(hsw1_out_name, [])]
                if len(hsw1_users) != 1:
                    continue
                mean1_idx = int(hsw1_users[0])
                mean1_op = model_ir.operators[int(mean1_idx)]
                if (
                    str(mean1_op.op_type) != "MEAN"
                    or len(mean1_op.inputs) < 2
                    or len(mean1_op.outputs) != 1
                    or str(mean1_op.inputs[0]) != str(hsw1_out_name)
                    or not bool(mean1_op.options.get("keepDims", False))
                ):
                    continue
                mean1_out_name = str(mean1_op.outputs[0])
                if mean1_out_name in model_outputs:
                    continue
                mean1_axes_tensor = model_ir.tensors.get(str(mean1_op.inputs[1]), None)
                mean1_axes_vals = _read_const_ints_from_tensor(mean1_axes_tensor)
                if mean1_axes_vals is None or len(mean1_axes_vals) == 0:
                    continue
                mean1_users = [int(v) for v in consumers.get(mean1_out_name, [])]
                if len(mean1_users) == 0:
                    continue
                # Keep this rewrite conservative: terminal mean should not feed post-transpose bridges.
                if any(str(model_ir.operators[int(v)].op_type) == "TRANSPOSE" for v in mean1_users):
                    continue

                source_tensor = model_ir.tensors.get(source_nhwc_name, None)
                rank = 4
                if source_tensor is not None and source_tensor.shape is not None and len(list(source_tensor.shape)) > 0:
                    rank = int(len(list(source_tensor.shape)))
                if rank != 4:
                    continue

                def _map_axes_to_nhwc_side(axes_vals: List[int]) -> Optional[List[int]]:
                    normalized_axes: List[int] = []
                    for axis in axes_vals:
                        a = int(axis)
                        if a < 0:
                            a += int(rank)
                        if a < 0 or a >= int(rank):
                            return None
                        normalized_axes.append(int(a))
                    return [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]

                mapped_axes0 = _map_axes_to_nhwc_side([int(v) for v in list(mean0_axes_vals)])
                mapped_axes1 = _map_axes_to_nhwc_side([int(v) for v in list(mean1_axes_vals)])
                if mapped_axes0 is None or mapped_axes1 is None:
                    continue
                if not (
                    _write_const_ints_to_tensor(mean0_axes_tensor, [int(v) for v in mapped_axes0])
                    or _read_const_ints_from_tensor(mean0_axes_tensor) == [int(v) for v in mapped_axes0]
                ):
                    continue
                if not (
                    _write_const_ints_to_tensor(mean1_axes_tensor, [int(v) for v in mapped_axes1])
                    or _read_const_ints_from_tensor(mean1_axes_tensor) == [int(v) for v in mapped_axes1]
                ):
                    continue

                hs0_add_op = model_ir.operators[int(hs0_match["add_idx"])]
                hs0_add_inputs = [
                    str(source_nhwc_name) if str(v) == str(source_nchw_name) else str(v)
                    for v in list(hs0_add_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=hs0_add_op,
                    new_inputs=hs0_add_inputs,
                )
                hs0_mul_op = model_ir.operators[int(hs0_match["mul_idx"])]
                hs0_mul_inputs = [
                    str(source_nhwc_name) if str(v) == str(source_nchw_name) else str(v)
                    for v in list(hs0_mul_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=hs0_mul_op,
                    new_inputs=hs0_mul_inputs,
                )

                conv_inputs = [str(v) for v in list(conv_op.inputs)]
                conv_inputs[0] = str(mean0_out_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=conv_op,
                    new_inputs=conv_inputs,
                )

                hs1_add_op = model_ir.operators[int(hs1_match["add_idx"])]
                hs1_add_inputs = [
                    str(conv_out_name) if str(v) == str(gate_nchw_name) else str(v)
                    for v in list(hs1_add_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=hs1_add_op,
                    new_inputs=hs1_add_inputs,
                )
                hs1_mul_op = model_ir.operators[int(hs1_match["mul_idx"])]
                hs1_mul_inputs = [
                    str(conv_out_name) if str(v) == str(gate_nchw_name) else str(v)
                    for v in list(hs1_mul_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=hs1_mul_op,
                    new_inputs=hs1_mul_inputs,
                )

                for tensor_name in [
                    str(hs0_match["add_out_name"]),
                    str(hs0_match["scale_out_name"]),
                    str(hs0_match["mul_out_name"]),
                    str(mean0_out_name),
                    str(hs1_match["add_out_name"]),
                    str(hs1_match["scale_out_name"]),
                    str(hs1_match["mul_out_name"]),
                    str(mean1_out_name),
                ]:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(tensor_name), None),
                        perm_nchw_to_nhwc,
                    )

                remove_indices = {int(mean0_post_idx), int(pre1_idx)}
                pre_remaining_users = [
                    int(v)
                    for v in consumers.get(source_nchw_name, [])
                    if int(v) not in {int(hs0_match["add_idx"]), int(hs0_match["mul_idx"])}
                ]
                if len(pre_remaining_users) == 0:
                    remove_indices.add(int(pre_idx))
                for remove_idx in sorted(remove_indices, reverse=True):
                    del model_ir.operators[int(remove_idx)]

                rewritten += 1
                changed = True
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_conv_attention_nhwc_propagation_chains": int(rewritten)}


def _optimize_batchmatmul_affine_transpose_input_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold NCHW bridge transposes around affine branches feeding BATCH_MATMUL.

    Target:
      lhs_nhwc --T(0,3,1,2)--> lhs_nchw --MUL(const)--ADD(const)--RESHAPE([N,C,S])--T(0,2,1)--> lhs
      rhs_nhwc --T(0,3,1,2)--> rhs_nchw --MUL(const)--ADD(const)--RESHAPE([N,C,M])--> rhs
      BATCH_MATMUL(lhs, rhs, adjY=False)

    Rewrite:
      lhs_nhwc --MUL(const_nhwc)--ADD(const_nhwc)--RESHAPE([N,S,C])--> lhs
      rhs_nhwc --MUL(const_nhwc)--ADD(const_nhwc)--RESHAPE([N,M,C])--> rhs
      BATCH_MATMUL(lhs, rhs, adjY=True)
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_swap_last2_rank3 = [0, 2, 1]

    def _is_nchw_channelwise_or_singleton_const(tensor_name: str) -> bool:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return False
        try:
            array = np.asarray(tensor.data)
        except Exception:
            return False
        if int(array.size) == 1:
            return True
        if array.ndim != 4:
            return False
        return (
            int(array.shape[0]) == 1
            and int(array.shape[2]) == 1
            and int(array.shape[3]) == 1
            and int(array.shape[1]) > 0
        )

    def _permute_const_nchw_to_nhwc(tensor_name: str) -> bool:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return False
        try:
            array = np.asarray(tensor.data)
        except Exception:
            return False
        if array.ndim != 4:
            return False
        try:
            transposed = np.transpose(array, axes=perm_nchw_to_nhwc)
        except Exception:
            return False
        tensor.data = np.asarray(transposed)
        tensor.shape = [int(v) for v in list(transposed.shape)]
        tensor.shape_signature = [int(v) for v in list(transposed.shape)]
        return True

    def _match_affine_branch_from_reshape_output(
        *,
        reshape_output_name: str,
        terminal_consumer_idx: int,
        consumers: Dict[str, List[int]],
        producers: Dict[str, int],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        reshape_idx = producers.get(str(reshape_output_name), None)
        if reshape_idx is None:
            return None
        reshape_op = model_ir.operators[int(reshape_idx)]
        if (
            str(reshape_op.op_type) != "RESHAPE"
            or len(reshape_op.inputs) < 2
            or len(reshape_op.outputs) != 1
            or str(reshape_op.outputs[0]) != str(reshape_output_name)
        ):
            return None
        if set(int(v) for v in consumers.get(str(reshape_output_name), [])) != {int(terminal_consumer_idx)}:
            return None

        add_out_name = str(reshape_op.inputs[0])
        add_idx = producers.get(add_out_name, None)
        if add_idx is None:
            return None
        add_op = model_ir.operators[int(add_idx)]
        if (
            str(add_op.op_type) != "ADD"
            or len(add_op.inputs) != 2
            or len(add_op.outputs) != 1
            or str(add_op.outputs[0]) != str(add_out_name)
        ):
            return None
        if set(int(v) for v in consumers.get(add_out_name, [])) != {int(reshape_idx)}:
            return None

        add_inputs = [str(v) for v in list(add_op.inputs)]
        mul_out_name: Optional[str] = None
        add_side_name: Optional[str] = None
        for candidate_name, side_name in [(add_inputs[0], add_inputs[1]), (add_inputs[1], add_inputs[0])]:
            candidate_idx = producers.get(str(candidate_name), None)
            if candidate_idx is None:
                continue
            candidate_op = model_ir.operators[int(candidate_idx)]
            if str(candidate_op.op_type) != "MUL":
                continue
            if not _is_nchw_channelwise_or_singleton_const(str(side_name)):
                return None
            mul_out_name = str(candidate_name)
            add_side_name = str(side_name)
            break
        if mul_out_name is None or add_side_name is None:
            return None

        mul_idx = producers.get(mul_out_name, None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if (
            str(mul_op.op_type) != "MUL"
            or len(mul_op.inputs) != 2
            or len(mul_op.outputs) != 1
            or str(mul_op.outputs[0]) != str(mul_out_name)
        ):
            return None
        if set(int(v) for v in consumers.get(mul_out_name, [])) != {int(add_idx)}:
            return None

        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        pre_output_name: Optional[str] = None
        mul_const_name: Optional[str] = None
        mul_data_input_index: Optional[int] = None
        pre_idx: Optional[int] = None
        pre_op: Optional[OperatorIR] = None
        for input_index, input_name in enumerate(mul_inputs):
            side_name = str(mul_inputs[1 - int(input_index)])
            candidate_idx = producers.get(str(input_name), None)
            if candidate_idx is None:
                continue
            candidate_op = model_ir.operators[int(candidate_idx)]
            if (
                str(candidate_op.op_type) != "TRANSPOSE"
                or len(candidate_op.inputs) < 2
                or len(candidate_op.outputs) != 1
                or str(candidate_op.outputs[0]) != str(input_name)
                or _read_transpose_perm(model_ir, candidate_op) != perm_nhwc_to_nchw
            ):
                continue
            if not _is_nchw_channelwise_or_singleton_const(side_name):
                return None
            pre_output_name = str(input_name)
            mul_const_name = str(side_name)
            mul_data_input_index = int(input_index)
            pre_idx = int(candidate_idx)
            pre_op = candidate_op
            break
        if (
            pre_idx is None
            or pre_op is None
            or pre_output_name is None
            or mul_const_name is None
            or mul_data_input_index is None
        ):
            return None
        if pre_output_name in model_outputs:
            return None
        if set(int(v) for v in consumers.get(pre_output_name, [])) != {int(mul_idx)}:
            return None

        source_nhwc_name = str(pre_op.inputs[0])
        if source_nhwc_name in model_outputs:
            return None

        return {
            "pre_idx": int(pre_idx),
            "mul_idx": int(mul_idx),
            "add_idx": int(add_idx),
            "reshape_idx": int(reshape_idx),
            "source_nhwc_name": str(source_nhwc_name),
            "mul_out_name": str(mul_out_name),
            "add_out_name": str(add_out_name),
            "reshape_out_name": str(reshape_output_name),
            "mul_const_name": str(mul_const_name),
            "add_const_name": str(add_side_name),
            "mul_data_input_index": int(mul_data_input_index),
            "reshape_shape_name": str(reshape_op.inputs[1]),
        }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for bmm_idx, bmm_op in enumerate(model_ir.operators):
            if str(bmm_op.op_type) != "BATCH_MATMUL" or len(bmm_op.inputs) != 2:
                continue
            bmm_options = dict(bmm_op.options) if isinstance(bmm_op.options, dict) else {}
            if bool(bmm_options.get("adjY", False)):
                continue

            lhs_input_name = str(bmm_op.inputs[0])
            rhs_input_name = str(bmm_op.inputs[1])

            lhs_post_idx = producers.get(lhs_input_name, None)
            if lhs_post_idx is None:
                continue
            lhs_post_op = model_ir.operators[int(lhs_post_idx)]
            if (
                str(lhs_post_op.op_type) != "TRANSPOSE"
                or len(lhs_post_op.inputs) < 2
                or len(lhs_post_op.outputs) != 1
                or str(lhs_post_op.outputs[0]) != lhs_input_name
                or _read_transpose_perm(model_ir, lhs_post_op) != perm_swap_last2_rank3
            ):
                continue
            if set(int(v) for v in consumers.get(lhs_input_name, [])) != {int(bmm_idx)}:
                continue
            lhs_reshape_out_name = str(lhs_post_op.inputs[0])

            lhs_branch = _match_affine_branch_from_reshape_output(
                reshape_output_name=lhs_reshape_out_name,
                terminal_consumer_idx=int(lhs_post_idx),
                consumers=consumers,
                producers=producers,
                model_outputs=model_outputs,
            )
            if lhs_branch is None:
                continue

            rhs_branch = _match_affine_branch_from_reshape_output(
                reshape_output_name=rhs_input_name,
                terminal_consumer_idx=int(bmm_idx),
                consumers=consumers,
                producers=producers,
                model_outputs=model_outputs,
            )
            if rhs_branch is None:
                continue

            shape_rewrite_ok = True
            for branch in [lhs_branch, rhs_branch]:
                mul_op = model_ir.operators[int(branch["mul_idx"])]
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=mul_op,
                    input_index=int(branch["mul_data_input_index"]),
                    new_input_name=str(branch["source_nhwc_name"]),
                )

                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(branch["mul_out_name"]), None),
                    perm_nchw_to_nhwc,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(branch["add_out_name"]), None),
                    perm_nchw_to_nhwc,
                )

                _permute_const_nchw_to_nhwc(str(branch["mul_const_name"]))
                _permute_const_nchw_to_nhwc(str(branch["add_const_name"]))

                reshape_op = model_ir.operators[int(branch["reshape_idx"])]
                shape_tensor = model_ir.tensors.get(str(branch["reshape_shape_name"]), None)
                shape_vals = _read_const_ints_from_tensor(shape_tensor)
                if shape_vals is None or len(shape_vals) != 3:
                    shape_rewrite_ok = False
                    break
                new_shape = [int(shape_vals[0]), int(shape_vals[2]), int(shape_vals[1])]
                _write_const_ints_to_tensor(shape_tensor, [int(v) for v in new_shape])

                if isinstance(reshape_op.options, dict):
                    reshape_opts = dict(reshape_op.options)
                    for key in ["newShape", "onnxRawNewShape"]:
                        value = reshape_opts.get(key, None)
                        if isinstance(value, list) and len(value) == 3:
                            reshape_opts[key] = [int(new_shape[0]), int(new_shape[1]), int(new_shape[2])]
                    reshape_op.options = reshape_opts

                reshape_out_tensor = model_ir.tensors.get(str(branch["reshape_out_name"]), None)
                if reshape_out_tensor is not None:
                    reshape_out_tensor.shape = [int(v) for v in new_shape]
                    reshape_out_tensor.shape_signature = [int(v) for v in new_shape]

            if not shape_rewrite_ok:
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=bmm_op,
                input_index=0,
                new_input_name=str(lhs_branch["reshape_out_name"]),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=bmm_op,
                input_index=1,
                new_input_name=str(rhs_branch["reshape_out_name"]),
            )
            bmm_options["adjY"] = True
            bmm_op.options = bmm_options

            remove_indices = {
                int(lhs_post_idx),
                int(lhs_branch["pre_idx"]),
                int(rhs_branch["pre_idx"]),
            }
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_batchmatmul_affine_transpose_input_chains": int(rewritten)}


def _optimize_batchmatmul_reshape_se_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove SE-like transpose bridges after BATCH_MATMUL -> RESHAPE tails.

    Target:
      x_nchw = RESHAPE(batch_matmul_out, [N,C,H,W])
      x_nchw --MEAN([2,3],keepDims=1)--> m_nchw --T(0,2,3,1)--> m_nhwc --CONV--CONV--> g_nhwc
      g_nhwc --T(0,3,1,2)--> g_nchw --LOGISTIC--> s_nchw
      MUL(x_nchw, s_nchw) -> y_nchw --T(0,2,3,1)--> y_nhwc

    Rewrite:
      - rewrite RESHAPE output to NHWC: [N,H,W,C]
      - remap MEAN axes to NHWC and bypass all three transposes
      - keep LOGISTIC/MUL in NHWC and preserve post-transpose output tensor name
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_out_idx, post_out_op in enumerate(model_ir.operators):
            if (
                str(post_out_op.op_type) != "TRANSPOSE"
                or len(post_out_op.inputs) < 2
                or len(post_out_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_out_op) != perm_nchw_to_nhwc
            ):
                continue

            mul_out_name = str(post_out_op.inputs[0])
            if mul_out_name in model_outputs:
                continue
            mul_idx = producers.get(mul_out_name, None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if (
                str(mul_op.op_type) != "MUL"
                or len(mul_op.inputs) != 2
                or len(mul_op.outputs) != 1
                or str(mul_op.outputs[0]) != mul_out_name
            ):
                continue

            mul_out_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_out_users) == 0:
                continue
            post_indices: List[int] = []
            post_output_names: List[str] = []
            valid_posts = True
            for user_idx in mul_out_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != mul_out_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    valid_posts = False
                    break
                post_indices.append(int(user_idx))
                post_output_names.append(str(user_op.outputs[0]))
            if not valid_posts or len(post_indices) == 0:
                continue

            source_name: Optional[str] = None
            gate_name: Optional[str] = None
            for lhs_name, rhs_name in [
                (str(mul_op.inputs[0]), str(mul_op.inputs[1])),
                (str(mul_op.inputs[1]), str(mul_op.inputs[0])),
            ]:
                lhs_idx = producers.get(lhs_name, None)
                if lhs_idx is None:
                    continue
                lhs_op = model_ir.operators[int(lhs_idx)]
                if (
                    str(lhs_op.op_type) == "RESHAPE"
                    and len(lhs_op.inputs) >= 2
                    and len(lhs_op.outputs) == 1
                    and str(lhs_op.outputs[0]) == lhs_name
                    and lhs_name not in model_outputs
                ):
                    source_name = str(lhs_name)
                    gate_name = str(rhs_name)
                    break
            if source_name is None or gate_name is None:
                continue

            source_idx = producers.get(source_name, None)
            if source_idx is None:
                continue
            source_op = model_ir.operators[int(source_idx)]
            if str(source_op.op_type) != "RESHAPE" or len(source_op.inputs) < 2:
                continue
            source_users = set(int(v) for v in consumers.get(source_name, []))
            if int(mul_idx) not in source_users:
                continue
            mean_candidates = [int(v) for v in source_users if int(v) != int(mul_idx)]
            if len(mean_candidates) != 1:
                continue

            mean_idx = int(mean_candidates[0])
            mean_op = model_ir.operators[int(mean_idx)]
            if (
                str(mean_op.op_type) != "MEAN"
                or len(mean_op.inputs) < 2
                or len(mean_op.outputs) != 1
                or str(mean_op.inputs[0]) != source_name
                or not bool(mean_op.options.get("keepDims", False))
            ):
                continue
            mean_out_name = str(mean_op.outputs[0])
            if mean_out_name in model_outputs:
                continue

            mean_post_users = [int(v) for v in consumers.get(mean_out_name, [])]
            if len(mean_post_users) != 1:
                continue
            mean_post_idx = int(mean_post_users[0])
            mean_post_op = model_ir.operators[int(mean_post_idx)]
            if (
                str(mean_post_op.op_type) != "TRANSPOSE"
                or len(mean_post_op.inputs) < 2
                or len(mean_post_op.outputs) != 1
                or str(mean_post_op.inputs[0]) != mean_out_name
                or _read_transpose_perm(model_ir, mean_post_op) != perm_nchw_to_nhwc
            ):
                continue

            conv1_input_name = str(mean_post_op.outputs[0])
            conv1_users = [int(v) for v in consumers.get(conv1_input_name, [])]
            if len(conv1_users) != 1:
                continue
            conv1_idx = int(conv1_users[0])
            conv1_op = model_ir.operators[int(conv1_idx)]
            if (
                str(conv1_op.op_type) != "CONV_2D"
                or len(conv1_op.inputs) < 1
                or len(conv1_op.outputs) != 1
                or str(conv1_op.inputs[0]) != conv1_input_name
            ):
                continue
            conv1_out_name = str(conv1_op.outputs[0])
            conv1_out_users = [int(v) for v in consumers.get(conv1_out_name, [])]
            if len(conv1_out_users) != 1:
                continue
            conv2_idx = int(conv1_out_users[0])
            conv2_op = model_ir.operators[int(conv2_idx)]
            if (
                str(conv2_op.op_type) != "CONV_2D"
                or len(conv2_op.inputs) < 1
                or len(conv2_op.outputs) != 1
                or str(conv2_op.inputs[0]) != conv1_out_name
            ):
                continue
            conv2_out_name = str(conv2_op.outputs[0])

            gate_idx = producers.get(gate_name, None)
            if gate_idx is None:
                continue
            gate_op = model_ir.operators[int(gate_idx)]
            if (
                str(gate_op.op_type) != "LOGISTIC"
                or len(gate_op.inputs) != 1
                or len(gate_op.outputs) != 1
                or str(gate_op.outputs[0]) != gate_name
            ):
                continue
            if set(int(v) for v in consumers.get(gate_name, [])) != {int(mul_idx)}:
                continue

            pre_gate_name = str(gate_op.inputs[0])
            pre_gate_idx = producers.get(pre_gate_name, None)
            if pre_gate_idx is None:
                continue
            pre_gate_op = model_ir.operators[int(pre_gate_idx)]
            if (
                str(pre_gate_op.op_type) != "TRANSPOSE"
                or len(pre_gate_op.inputs) < 2
                or len(pre_gate_op.outputs) != 1
                or str(pre_gate_op.outputs[0]) != pre_gate_name
                or _read_transpose_perm(model_ir, pre_gate_op) != perm_nhwc_to_nchw
                or str(pre_gate_op.inputs[0]) != conv2_out_name
            ):
                continue
            if set(int(v) for v in consumers.get(pre_gate_name, [])) != {int(gate_idx)}:
                continue

            # Keep semantics when converting source from NCHW to NHWC by
            # transposing BATCH_MATMUL output layout (without inserting transpose ops).
            source_bmm_out_name = str(source_op.inputs[0])
            source_bmm_idx = producers.get(source_bmm_out_name, None)
            if source_bmm_idx is None:
                continue
            source_bmm_op = model_ir.operators[int(source_bmm_idx)]
            if (
                str(source_bmm_op.op_type) != "BATCH_MATMUL"
                or len(source_bmm_op.inputs) != 2
                or len(source_bmm_op.outputs) != 1
                or str(source_bmm_op.outputs[0]) != source_bmm_out_name
            ):
                continue
            if set(int(v) for v in consumers.get(source_bmm_out_name, [])) != {int(source_idx)}:
                continue

            source_bmm_options = dict(source_bmm_op.options) if isinstance(source_bmm_op.options, dict) else {}
            # This rewrite is tuned for the affine-input optimized attention path:
            #   adjX=False, adjY=True.
            if bool(source_bmm_options.get("adjX", False)):
                continue
            if not bool(source_bmm_options.get("adjY", False)):
                continue

            old_bmm_inputs = [str(v) for v in list(source_bmm_op.inputs)]
            _set_operator_inputs(
                model_ir=model_ir,
                op=source_bmm_op,
                new_inputs=[str(old_bmm_inputs[1]), str(old_bmm_inputs[0])],
            )
            old_adj_x = bool(source_bmm_options.get("adjX", False))
            old_adj_y = bool(source_bmm_options.get("adjY", False))
            source_bmm_options["adjX"] = bool(not old_adj_y)
            source_bmm_options["adjY"] = bool(not old_adj_x)
            source_bmm_op.options = source_bmm_options

            source_bmm_tensor = model_ir.tensors.get(source_bmm_out_name, None)
            if source_bmm_tensor is not None:
                bmm_shape = (
                    [int(v) for v in list(source_bmm_tensor.shape)]
                    if source_bmm_tensor.shape is not None
                    else None
                )
                if bmm_shape is not None and len(bmm_shape) >= 2:
                    bmm_shape[-1], bmm_shape[-2] = int(bmm_shape[-2]), int(bmm_shape[-1])
                    source_bmm_tensor.shape = [int(v) for v in bmm_shape]
                bmm_signature = (
                    [int(v) for v in list(source_bmm_tensor.shape_signature)]
                    if source_bmm_tensor.shape_signature is not None
                    else (
                        [int(v) for v in list(source_bmm_tensor.shape)]
                        if source_bmm_tensor.shape is not None
                        else None
                    )
                )
                if bmm_signature is not None and len(bmm_signature) >= 2:
                    bmm_signature[-1], bmm_signature[-2] = int(bmm_signature[-2]), int(bmm_signature[-1])
                    source_bmm_tensor.shape_signature = [int(v) for v in bmm_signature]

            # Rewrite RESHAPE output layout from NCHW to NHWC.
            source_shape_name = str(source_op.inputs[1])
            source_shape_tensor = model_ir.tensors.get(source_shape_name, None)
            source_shape_vals = _read_const_ints_from_tensor(source_shape_tensor)
            if source_shape_vals is None or len(source_shape_vals) != 4:
                continue
            source_new_shape = [
                int(source_shape_vals[0]),
                int(source_shape_vals[2]),
                int(source_shape_vals[3]),
                int(source_shape_vals[1]),
            ]
            _write_const_ints_to_tensor(source_shape_tensor, [int(v) for v in source_new_shape])

            if isinstance(source_op.options, dict):
                source_opts = dict(source_op.options)
                for key in ["newShape", "onnxRawNewShape"]:
                    value = source_opts.get(key, None)
                    if isinstance(value, list) and len(value) == 4:
                        source_opts[key] = [int(v) for v in source_new_shape]
                source_op.options = source_opts

            source_tensor = model_ir.tensors.get(source_name, None)
            if source_tensor is not None:
                source_tensor.shape = [int(v) for v in source_new_shape]
                source_tensor.shape_signature = [int(v) for v in source_new_shape]

            # Remap MEAN axes from NCHW to NHWC.
            mean_axes_tensor = model_ir.tensors.get(str(mean_op.inputs[1]), None)
            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_vals is None or len(mean_axes_vals) == 0:
                continue
            mapped_axes = [int(perm_nhwc_to_nchw[int(v)]) for v in list(mean_axes_vals)]
            _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes])
            if isinstance(mean_op.options, dict):
                mean_opts = dict(mean_op.options)
                for key in ["axis", "axes", "onnxRawAxes"]:
                    value = mean_opts.get(key, None)
                    if isinstance(value, list) and len(value) == len(mapped_axes):
                        mean_opts[key] = [int(v) for v in mapped_axes]
                mean_op.options = mean_opts

            conv1_inputs = [str(v) for v in list(conv1_op.inputs)]
            conv1_inputs[0] = str(mean_out_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=conv1_op,
                new_inputs=conv1_inputs,
            )

            _replace_operator_input_at(
                model_ir=model_ir,
                op=gate_op,
                input_index=0,
                new_input_name=conv2_out_name,
            )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mean_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(gate_name, None),
                perm_nchw_to_nhwc,
            )

            representative_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[representative_output_name],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_name, representative_output_name)

            old_mul_tensor = model_ir.tensors.get(mul_out_name, None)
            representative_tensor = model_ir.tensors.get(representative_output_name, None)
            if old_mul_tensor is not None and representative_tensor is not None:
                representative_tensor.dtype = str(old_mul_tensor.dtype)
                representative_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                representative_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                representative_tensor.shape_signature = (
                    [int(v) for v in list(old_mul_tensor.shape_signature)]
                    if old_mul_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    representative_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {
                int(mean_post_idx),
                int(pre_gate_idx),
                *[int(v) for v in post_indices],
            }
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_batchmatmul_reshape_se_nhwc_chains": int(rewritten)}


def _optimize_batchmatmul_transpose_input_to_adj_flags(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate rank-3 transpose adapters directly feeding BATCH_MATMUL inputs.

    Pattern A:
      x --TRANSPOSE([0,2,1])--> x_t --BATCH_MATMUL(...)
      -> remove transpose and toggle adjX/adjY.

    Pattern B (singleton-preserving pre-reshape):
      x[s,1,c] --TRANSPOSE([1,2,0])--> x_t[1,c,s] --BATCH_MATMUL(...)
      -> rewrite TRANSPOSE to RESHAPE([1,s,c]) and toggle adjX/adjY.
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _is_memory_order_equivalent_by_singletons(signature: List[int], perm: List[int]) -> bool:
        if len(signature) != 3 or len(perm) != 3:
            return False
        if any(int(v) <= 0 for v in signature):
            return False
        non_singleton_input_axes = [int(i) for i, dim in enumerate(signature) if int(dim) != 1]
        non_singleton_permuted_axes = [int(i) for i in perm if int(signature[int(i)]) != 1]
        return non_singleton_input_axes == non_singleton_permuted_axes

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for bmm_idx, bmm_op in enumerate(model_ir.operators):
            if str(bmm_op.op_type) != "BATCH_MATMUL" or len(bmm_op.inputs) != 2:
                continue

            bmm_options = dict(bmm_op.options) if isinstance(bmm_op.options, dict) else {}

            rewritten_this_bmm = False
            for input_index in [0, 1]:
                bmm_input_name = str(bmm_op.inputs[int(input_index)])
                transpose_idx = producers.get(str(bmm_input_name), None)
                if transpose_idx is None:
                    continue
                transpose_op = model_ir.operators[int(transpose_idx)]
                if (
                    str(transpose_op.op_type) != "TRANSPOSE"
                    or len(transpose_op.inputs) < 2
                    or len(transpose_op.outputs) != 1
                    or str(transpose_op.outputs[0]) != str(bmm_input_name)
                    or str(bmm_input_name) in model_outputs
                ):
                    continue
                if set(int(v) for v in consumers.get(str(bmm_input_name), [])) != {int(bmm_idx)}:
                    continue

                perm = _read_transpose_perm(model_ir, transpose_op)
                if perm is None or len(list(perm)) != 3:
                    continue

                transpose_input_name = str(transpose_op.inputs[0])
                transpose_input_tensor = model_ir.tensors.get(str(transpose_input_name), None)
                transpose_output_tensor = model_ir.tensors.get(str(bmm_input_name), None)
                if (
                    transpose_input_tensor is None
                    or transpose_output_tensor is None
                    or not _is_fully_known_positive_shape(list(transpose_input_tensor.shape))
                    or not _is_fully_known_positive_shape(list(transpose_output_tensor.shape))
                ):
                    continue
                input_shape = [int(v) for v in list(transpose_input_tensor.shape)]
                output_shape = [int(v) for v in list(transpose_output_tensor.shape)]
                if len(input_shape) != 3 or len(output_shape) != 3:
                    continue
                expected_output_shape = _permute_shape(input_shape, [int(v) for v in list(perm)])
                if expected_output_shape is None or [int(v) for v in list(expected_output_shape)] != output_shape:
                    continue

                removed_transpose = False
                if [int(v) for v in list(perm)] == [0, 2, 1]:
                    _replace_operator_input_at(
                        model_ir=model_ir,
                        op=bmm_op,
                        input_index=int(input_index),
                        new_input_name=str(transpose_input_name),
                    )
                    del model_ir.operators[int(transpose_idx)]
                    removed_transpose = True
                else:
                    reshape_axes = [int(perm[0]), int(perm[2]), int(perm[1])]
                    if sorted(reshape_axes) != [0, 1, 2]:
                        continue
                    input_signature = (
                        [int(v) for v in list(transpose_input_tensor.shape_signature)]
                        if transpose_input_tensor.shape_signature is not None
                        else [int(v) for v in list(input_shape)]
                    )
                    if not _is_memory_order_equivalent_by_singletons(input_signature, reshape_axes):
                        continue
                    reshape_shape = [int(input_shape[int(axis)]) for axis in reshape_axes]
                    if [int(reshape_shape[0]), int(reshape_shape[2]), int(reshape_shape[1])] != output_shape:
                        continue

                    reshape_shape_name = _unique_tensor_name(f"{bmm_input_name}_reshape_shape")
                    model_ir.tensors[reshape_shape_name] = TensorIR(
                        name=reshape_shape_name,
                        dtype="INT32",
                        shape=[3],
                        shape_signature=[3],
                        data=np.asarray([int(v) for v in list(reshape_shape)], dtype=np.int32),
                        is_variable=False,
                    )
                    transpose_op.op_type = "RESHAPE"
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=transpose_op,
                        new_inputs=[str(transpose_input_name), reshape_shape_name],
                    )
                    transpose_op.options = {"newShape": [int(v) for v in list(reshape_shape)]}
                    transpose_output_tensor.shape = [int(v) for v in list(reshape_shape)]
                    transpose_output_tensor.shape_signature = [int(v) for v in list(reshape_shape)]

                flag_name = "adjX" if int(input_index) == 0 else "adjY"
                bmm_options[flag_name] = bool(not bool(bmm_options.get(flag_name, False)))
                bmm_op.options = bmm_options

                rewritten += 1
                changed = True
                rewritten_this_bmm = True
                break

            if rewritten_this_bmm:
                break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_batchmatmul_transpose_input_to_adj_flags": int(rewritten)}


def _optimize_attention_qkv_gather_reshape_transpose_hoist_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Hoist shared QKV/KV layout transforms before per-branch GATHERs.

    Target branches (`N in {2,3}`, scalar gather indices `0..N-1`):
      src[N,T,1,D] --GATHER(i)--> g_i[T,1,D]
                    --RESHAPE([T,H,C])--> r_i[T,H,C]
                    --TRANSPOSE([1,0,2])--> t_i[H,T,C]
                    --RESHAPE([1,H,T,C])--> z_i

    Rewrite:
      src --RESHAPE([N,T,H,C])--> src_qkv
          --TRANSPOSE([0,2,1,3])--> src_qkv_ht
      src_qkv_ht --GATHER(i)--> g_i'[H,T,C]
      g_i' --RESHAPE([1,H,T,C])--> z_i
    """
    rewritten = 0
    candidate_branch_counts = [3, 2]
    perm_hcw_to_hwc = [0, 2, 1]

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    def _gather_axis(op: OperatorIR) -> int:
        opts = dict(op.options) if isinstance(op.options, dict) else {}
        axis_value = opts.get("axis", 0)
        if isinstance(axis_value, list):
            if len(axis_value) != 1:
                return 0
            axis_value = axis_value[0]
        try:
            return int(axis_value)
        except Exception:
            return 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for source_name, user_indices in list(consumers.items()):
            source_shape = _shape_list(str(source_name))
            if source_shape is None or len(source_shape) != 4:
                continue
            branch_count: Optional[int] = None
            gather_branches: Dict[int, Dict[str, Any]] = {}
            gather_indices: set[int] = set()

            for candidate_branch_count in list(candidate_branch_counts):
                if not _dims_compatible(int(source_shape[0]), int(candidate_branch_count)):
                    continue
                expected_indices = set(range(int(candidate_branch_count)))
                candidate_branches: Dict[int, Dict[str, Any]] = {}

                for user_idx in [int(v) for v in user_indices]:
                    gather_op = model_ir.operators[int(user_idx)]
                    if (
                        str(gather_op.op_type) != "GATHER"
                        or len(gather_op.inputs) < 2
                        or len(gather_op.outputs) != 1
                        or str(gather_op.inputs[0]) != str(source_name)
                        or _gather_axis(gather_op) != 0
                    ):
                        continue
                    index_vals = _read_const_ints_from_tensor(
                        model_ir.tensors.get(str(gather_op.inputs[1]), None)
                    )
                    if index_vals is None or len(index_vals) != 1:
                        continue
                    gather_index = int(index_vals[0])
                    if gather_index not in expected_indices or gather_index in candidate_branches:
                        continue

                    gather_out_name = str(gather_op.outputs[0])
                    if gather_out_name in model_outputs:
                        continue
                    gather_out_users = [int(v) for v in consumers.get(gather_out_name, [])]
                    if len(gather_out_users) != 1:
                        continue
                    reshape1_idx = int(gather_out_users[0])
                    reshape1_op = model_ir.operators[int(reshape1_idx)]
                    if (
                        str(reshape1_op.op_type) != "RESHAPE"
                        or len(reshape1_op.inputs) < 1
                        or len(reshape1_op.outputs) != 1
                        or str(reshape1_op.inputs[0]) != gather_out_name
                    ):
                        continue
                    reshape1_out_name = str(reshape1_op.outputs[0])
                    if reshape1_out_name in model_outputs:
                        continue
                    reshape1_out_users = [int(v) for v in consumers.get(reshape1_out_name, [])]
                    if len(reshape1_out_users) != 1:
                        continue
                    transpose_idx = int(reshape1_out_users[0])
                    transpose_op = model_ir.operators[int(transpose_idx)]
                    if (
                        str(transpose_op.op_type) != "TRANSPOSE"
                        or len(transpose_op.inputs) < 2
                        or len(transpose_op.outputs) != 1
                        or str(transpose_op.inputs[0]) != reshape1_out_name
                        or _read_transpose_perm(model_ir, transpose_op) != [1, 0, 2]
                    ):
                        continue
                    transpose_out_name = str(transpose_op.outputs[0])
                    if transpose_out_name in model_outputs:
                        continue
                    transpose_out_users = [int(v) for v in consumers.get(transpose_out_name, [])]
                    if len(transpose_out_users) != 1:
                        continue
                    reshape2_idx = int(transpose_out_users[0])
                    reshape2_op = model_ir.operators[int(reshape2_idx)]
                    if (
                        str(reshape2_op.op_type) != "RESHAPE"
                        or len(reshape2_op.inputs) < 1
                        or len(reshape2_op.outputs) != 1
                        or str(reshape2_op.inputs[0]) != transpose_out_name
                    ):
                        continue

                    gather_shape = _shape_list(gather_out_name)
                    reshape1_shape = _shape_list(reshape1_out_name)
                    transpose_shape = _shape_list(transpose_out_name)
                    reshape2_shape = _shape_list(str(reshape2_op.outputs[0]))
                    if (
                        gather_shape is None
                        or reshape1_shape is None
                        or transpose_shape is None
                        or reshape2_shape is None
                        or len(gather_shape) != 3
                        or len(reshape1_shape) != 3
                        or len(transpose_shape) != 3
                        or len(reshape2_shape) != 4
                    ):
                        continue

                    candidate_branches[gather_index] = {
                        "gather_idx": int(user_idx),
                        "gather_op": gather_op,
                        "gather_out_name": gather_out_name,
                        "reshape1_idx": int(reshape1_idx),
                        "reshape1_op": reshape1_op,
                        "reshape1_out_name": reshape1_out_name,
                        "transpose_idx": int(transpose_idx),
                        "transpose_op": transpose_op,
                        "transpose_out_name": transpose_out_name,
                        "reshape2_idx": int(reshape2_idx),
                        "reshape2_op": reshape2_op,
                        "gather_shape": [int(v) for v in gather_shape],
                        "reshape1_shape": [int(v) for v in reshape1_shape],
                        "transpose_shape": [int(v) for v in transpose_shape],
                        "reshape2_shape": [int(v) for v in reshape2_shape],
                    }

                if set(candidate_branches.keys()) == expected_indices:
                    branch_count = int(candidate_branch_count)
                    gather_branches = candidate_branches
                    gather_indices = set(expected_indices)
                    break

            if branch_count is None:
                continue

            branch0 = gather_branches[0]
            t = int(branch0["reshape1_shape"][0])
            h = int(branch0["reshape1_shape"][1])
            c = int(branch0["reshape1_shape"][2])
            if not _shape_compatible(branch0["transpose_shape"], [h, t, c]):
                continue
            if not _shape_compatible(branch0["reshape2_shape"], [1, h, t, c]):
                continue

            is_consistent = True
            for gather_index in sorted(list(gather_indices)):
                branch = gather_branches[gather_index]
                if not _shape_compatible(branch["gather_shape"], [t, 1, h * c]):
                    is_consistent = False
                    break
                if not _shape_compatible(branch["reshape1_shape"], [t, h, c]):
                    is_consistent = False
                    break
                if not _shape_compatible(branch["transpose_shape"], [h, t, c]):
                    is_consistent = False
                    break
                if not _shape_compatible(branch["reshape2_shape"], [1, h, t, c]):
                    is_consistent = False
                    break
            if not is_consistent:
                continue

            n = int(source_shape[0])
            src_t = int(source_shape[1])
            src_one = int(source_shape[2])
            src_d = int(source_shape[3])
            if not _dims_compatible(n, int(branch_count)):
                continue
            if not _dims_compatible(src_t, t) or not _dims_compatible(src_one, 1):
                continue
            if int(src_d) > 0 and int(h) > 0 and int(c) > 0 and int(src_d) != int(h) * int(c):
                continue

            shared_shape_name = _unique_tensor_name(f"{source_name}_qkv_reshape_shape")
            shared_reshape_out_name = _unique_tensor_name(f"{source_name}_qkv_thc")
            shared_perm_name = _unique_tensor_name(f"{source_name}_qkv_perm")
            shared_transpose_out_name = _unique_tensor_name(f"{source_name}_qkv_htc")

            model_ir.tensors[shared_shape_name] = TensorIR(
                name=shared_shape_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray([int(branch_count), int(t), int(h), int(c)], dtype=np.int32),
                is_variable=False,
            )
            source_tensor = model_ir.tensors.get(str(source_name), None)
            model_ir.tensors[shared_reshape_out_name] = TensorIR(
                name=shared_reshape_out_name,
                dtype=str(source_tensor.dtype if source_tensor is not None else "FLOAT32"),
                shape=[int(branch_count), int(t), int(h), int(c)],
                shape_signature=[int(branch_count), int(t), int(h), int(c)],
            )
            model_ir.tensors[shared_perm_name] = TensorIR(
                name=shared_perm_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray([0, 2, 1, 3], dtype=np.int32),
                is_variable=False,
            )
            model_ir.tensors[shared_transpose_out_name] = TensorIR(
                name=shared_transpose_out_name,
                dtype=str(source_tensor.dtype if source_tensor is not None else "FLOAT32"),
                shape=[int(branch_count), int(h), int(t), int(c)],
                shape_signature=[int(branch_count), int(h), int(t), int(c)],
            )

            shared_reshape_op = OperatorIR(
                op_type="RESHAPE",
                inputs=[str(source_name), str(shared_shape_name)],
                outputs=[str(shared_reshape_out_name)],
                options={"newShape": [int(branch_count), int(t), int(h), int(c)]},
            )
            shared_transpose_op = OperatorIR(
                op_type="TRANSPOSE",
                inputs=[str(shared_reshape_out_name), str(shared_perm_name)],
                outputs=[str(shared_transpose_out_name)],
                options={},
            )
            insert_base_idx = min(int(gather_branches[idx]["gather_idx"]) for idx in sorted(list(gather_indices)))
            model_ir.operators.insert(int(insert_base_idx), shared_reshape_op)
            model_ir.operators.insert(int(insert_base_idx) + 1, shared_transpose_op)

            remove_indices: set[int] = set()
            for gather_index in sorted(list(gather_indices)):
                branch = gather_branches[gather_index]
                gather_op = branch["gather_op"]
                gather_out_name = str(branch["gather_out_name"])
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=gather_op,
                    input_index=0,
                    new_input_name=str(shared_transpose_out_name),
                )
                gather_tensor = model_ir.tensors.get(gather_out_name, None)
                if gather_tensor is not None:
                    gather_tensor.shape = [int(h), int(t), int(c)]
                    gather_tensor.shape_signature = [int(h), int(t), int(c)]

                _replace_tensor_inputs(
                    model_ir,
                    str(branch["transpose_out_name"]),
                    gather_out_name,
                )
                reshape1_remove_idx = next(
                    (idx for idx, op in enumerate(model_ir.operators) if op is branch["reshape1_op"]),
                    None,
                )
                transpose_remove_idx = next(
                    (idx for idx, op in enumerate(model_ir.operators) if op is branch["transpose_op"]),
                    None,
                )
                if reshape1_remove_idx is not None:
                    remove_indices.add(int(reshape1_remove_idx))
                if transpose_remove_idx is not None:
                    remove_indices.add(int(transpose_remove_idx))

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_qkv_gather_reshape_transpose_hoist_chains": int(rewritten)}


def _optimize_attention_qkv_slice_replace_gather_reshape_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Replace attention QKV/KV branch `GATHER -> RESHAPE` with direct `SLICE`.

    Target (post hoist):
      qkv_htc[N,H,T,C] --GATHER(i, axis=0)--> g_i[H,T,C]
                       --RESHAPE([1,H,T,C])--> z_i      (N in {2,3})

    Rewrite:
      qkv_htc --SLICE(begin=[i,0,0,0], size=[1,H,T,C])--> z_i
    """
    rewritten = 0
    candidate_branch_counts = [3, 2]

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    def _gather_axis(op: OperatorIR) -> int:
        opts = dict(op.options) if isinstance(op.options, dict) else {}
        axis_value = opts.get("axis", 0)
        if isinstance(axis_value, list):
            if len(axis_value) != 1:
                return 0
            axis_value = axis_value[0]
        try:
            return int(axis_value)
        except Exception:
            return 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for source_name, user_indices in list(consumers.items()):
            source_shape = _shape_list(str(source_name))
            if source_shape is None or len(source_shape) != 4:
                continue
            n0, h, t, c = [int(v) for v in list(source_shape)]
            if not _is_fully_known_positive_shape([h, t, c]):
                continue

            gather_indices: set[int] = set()
            branch_by_index: Dict[int, Dict[str, Any]] = {}
            for candidate_branch_count in list(candidate_branch_counts):
                if not _dims_compatible(n0, int(candidate_branch_count)):
                    continue
                expected_indices = set(range(int(candidate_branch_count)))
                candidate_branch_by_index: Dict[int, Dict[str, Any]] = {}

                for user_idx in [int(v) for v in user_indices]:
                    gather_op = model_ir.operators[int(user_idx)]
                    if (
                        str(gather_op.op_type) != "GATHER"
                        or len(gather_op.inputs) < 2
                        or len(gather_op.outputs) != 1
                        or str(gather_op.inputs[0]) != str(source_name)
                        or _gather_axis(gather_op) != 0
                    ):
                        continue
                    index_vals = _read_const_ints_from_tensor(
                        model_ir.tensors.get(str(gather_op.inputs[1]), None)
                    )
                    if index_vals is None or len(index_vals) != 1:
                        continue
                    gather_index = int(index_vals[0])
                    if gather_index not in expected_indices or gather_index in candidate_branch_by_index:
                        continue

                    gather_out_name = str(gather_op.outputs[0])
                    if gather_out_name in model_outputs:
                        continue
                    gather_out_shape = _shape_list(gather_out_name)
                    if gather_out_shape is None or len(gather_out_shape) != 3:
                        continue
                    if not _shape_compatible(gather_out_shape, [h, t, c]):
                        continue

                    gather_out_users = [int(v) for v in consumers.get(gather_out_name, [])]
                    if len(gather_out_users) != 1:
                        continue
                    reshape_idx = int(gather_out_users[0])
                    reshape_op = model_ir.operators[int(reshape_idx)]
                    if (
                        str(reshape_op.op_type) != "RESHAPE"
                        or len(reshape_op.inputs) < 1
                        or len(reshape_op.outputs) != 1
                        or str(reshape_op.inputs[0]) != gather_out_name
                    ):
                        continue
                    reshape_out_name = str(reshape_op.outputs[0])
                    reshape_out_shape = _shape_list(reshape_out_name)
                    if reshape_out_shape is None or len(reshape_out_shape) != 4:
                        continue
                    if not _shape_compatible(reshape_out_shape, [1, h, t, c]):
                        continue

                    candidate_branch_by_index[gather_index] = {
                        "gather_idx": int(user_idx),
                        "gather_op": gather_op,
                        "reshape_idx": int(reshape_idx),
                        "reshape_op": reshape_op,
                        "reshape_out_name": str(reshape_out_name),
                    }

                if set(candidate_branch_by_index.keys()) == expected_indices:
                    gather_indices = set(expected_indices)
                    branch_by_index = candidate_branch_by_index
                    break

            if len(gather_indices) <= 0:
                continue

            remove_indices: set[int] = set()
            for gather_index in sorted(list(gather_indices)):
                branch = branch_by_index[gather_index]
                reshape_op = branch["reshape_op"]
                reshape_out_name = str(branch["reshape_out_name"])

                begin_name = _unique_tensor_name(f"{reshape_out_name}_slice_begin")
                size_name = _unique_tensor_name(f"{reshape_out_name}_slice_size")
                model_ir.tensors[begin_name] = TensorIR(
                    name=begin_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray([int(gather_index), 0, 0, 0], dtype=np.int32),
                    is_variable=False,
                )
                model_ir.tensors[size_name] = TensorIR(
                    name=size_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray([1, int(h), int(t), int(c)], dtype=np.int32),
                    is_variable=False,
                )

                reshape_op.op_type = "SLICE"
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=reshape_op,
                    new_inputs=[str(source_name), str(begin_name), str(size_name)],
                )
                reshape_op.options = {}
                out_tensor = model_ir.tensors.get(reshape_out_name, None)
                if out_tensor is not None:
                    out_tensor.shape = [1, int(h), int(t), int(c)]
                    out_tensor.shape_signature = [1, int(h), int(t), int(c)]

                gather_remove_idx = next(
                    (idx for idx, op in enumerate(model_ir.operators) if op is branch["gather_op"]),
                    None,
                )
                if gather_remove_idx is not None:
                    remove_indices.add(int(gather_remove_idx))

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_qkv_slice_replace_gather_reshape_chains": int(rewritten)}


def _optimize_attention_qkv_slice_to_split_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Replace attention QKV/KV sibling `SLICE` groups with a single `SPLIT`.

    Target (generalized):
      src[...]
        --SLICE(begin,size)--> out_i  (i in 0..N-1, N in {2,3})
      where all branches slice along one axis with equal chunk size and
      non-overlapping contiguous starts.

    Rewrite:
      SPLIT(numSplits=N, axis=split_axis): src -> [out_0, ..., out_{N-1}]
    """
    rewritten = 0
    candidate_branch_counts = [3, 2]

    def _shape_list(name: str) -> Optional[List[int]]:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or tensor.shape is None:
            return None
        return [int(v) for v in list(tensor.shape)]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) < 0 or int(b) < 0:
            return True
        return int(a) == int(b)

    def _shape_compatible(a: List[int], b: List[int]) -> bool:
        if len(a) != len(b):
            return False
        return all(_dims_compatible(int(x), int(y)) for x, y in zip(a, b))

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for source_name, user_indices in list(consumers.items()):
            source_shape = _shape_list(str(source_name))
            if source_shape is None or len(source_shape) <= 0:
                continue
            if not _is_fully_known_positive_shape([int(v) for v in list(source_shape)]):
                continue
            rank = int(len(source_shape))

            branch_count: Optional[int] = None
            split_axis: Optional[int] = None
            branch_by_index: Dict[int, Dict[str, Any]] = {}
            for candidate_branch_count in list(candidate_branch_counts):
                expected_indices = set(range(int(candidate_branch_count)))
                for axis in range(rank):
                    axis_dim = int(source_shape[axis])
                    if int(axis_dim) <= 0 or int(axis_dim) % int(candidate_branch_count) != 0:
                        continue
                    chunk = int(axis_dim // int(candidate_branch_count))
                    if int(chunk) <= 0:
                        continue

                    candidate_branch_by_index: Dict[int, Dict[str, Any]] = {}
                    for user_idx in [int(v) for v in user_indices]:
                        slice_op = model_ir.operators[int(user_idx)]
                        if (
                            str(slice_op.op_type) != "SLICE"
                            or len(slice_op.inputs) < 3
                            or len(slice_op.outputs) != 1
                            or str(slice_op.inputs[0]) != str(source_name)
                        ):
                            continue

                        begin_vals = _read_const_ints_from_tensor(
                            model_ir.tensors.get(str(slice_op.inputs[1]), None)
                        )
                        size_vals = _read_const_ints_from_tensor(
                            model_ir.tensors.get(str(slice_op.inputs[2]), None)
                        )
                        if (
                            begin_vals is None
                            or size_vals is None
                            or len(begin_vals) != rank
                            or len(size_vals) != rank
                        ):
                            continue
                        begin = [int(v) for v in list(begin_vals)]
                        size = [int(v) for v in list(size_vals)]

                        valid = True
                        for dim_idx in range(rank):
                            dim = int(source_shape[dim_idx])
                            b = int(begin[dim_idx])
                            s = int(size[dim_idx])
                            if dim_idx == axis:
                                if int(b) < 0 or int(b) % int(chunk) != 0:
                                    valid = False
                                    break
                                if int(s) != int(chunk):
                                    valid = False
                                    break
                            else:
                                if int(b) != 0 or int(s) != int(dim):
                                    valid = False
                                    break
                        if not valid:
                            continue

                        slice_index = int(begin[axis] // chunk)
                        if slice_index not in expected_indices or slice_index in candidate_branch_by_index:
                            continue

                        out_name = str(slice_op.outputs[0])
                        out_shape = _shape_list(out_name)
                        expected_out_shape = [int(v) for v in list(source_shape)]
                        expected_out_shape[axis] = int(chunk)
                        if out_shape is None or len(out_shape) != rank:
                            continue
                        if not _shape_compatible(out_shape, expected_out_shape):
                            continue

                        candidate_branch_by_index[slice_index] = {
                            "slice_idx": int(user_idx),
                            "slice_op": slice_op,
                            "out_name": out_name,
                        }

                    if set(candidate_branch_by_index.keys()) == expected_indices:
                        branch_count = int(candidate_branch_count)
                        split_axis = int(axis)
                        branch_by_index = candidate_branch_by_index
                        break
                if branch_count is not None:
                    break

            if branch_count is None:
                continue
            if split_axis is None:
                continue

            axis_name = _unique_tensor_name(f"{source_name}_qkv_split_axis")
            model_ir.tensors[axis_name] = TensorIR(
                name=axis_name,
                dtype="INT32",
                shape=[1],
                shape_signature=[1],
                data=np.asarray([int(split_axis)], dtype=np.int32),
                is_variable=False,
            )

            split_outputs = [str(branch_by_index[idx]["out_name"]) for idx in sorted(list(branch_by_index.keys()))]
            split_op = OperatorIR(
                op_type="SPLIT",
                inputs=[str(axis_name), str(source_name)],
                outputs=split_outputs,
                options={"numSplits": int(branch_count)},
            )
            insert_base_idx = min(int(branch_by_index[idx]["slice_idx"]) for idx in sorted(list(branch_by_index.keys())))
            model_ir.operators.insert(int(insert_base_idx), split_op)

            remove_indices: set[int] = set()
            chunk = int(source_shape[int(split_axis)] // int(branch_count))
            for idx in sorted(list(branch_by_index.keys())):
                slice_op = branch_by_index[idx]["slice_op"]
                slice_remove_idx = next(
                    (op_idx for op_idx, op in enumerate(model_ir.operators) if op is slice_op),
                    None,
                )
                if slice_remove_idx is not None:
                    remove_indices.add(int(slice_remove_idx))

                out_name = str(branch_by_index[idx]["out_name"])
                out_tensor = model_ir.tensors.get(out_name, None)
                if out_tensor is not None:
                    out_shape = [int(v) for v in list(source_shape)]
                    out_shape[int(split_axis)] = int(chunk)
                    out_tensor.shape = [int(v) for v in list(out_shape)]
                    out_tensor.shape_signature = [int(v) for v in list(out_shape)]

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_qkv_slice_to_split_chains": int(rewritten)}


def _optimize_attention_split_post_reshape_collapse_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse `SPLIT -> RESHAPE xN` into `RESHAPE -> SPLIT` when safe.

    Target:
      src --SPLIT(axis=a, numSplits=N)--> s_i
      s_i --RESHAPE(shape=tgt)--> r_i      (for all i in 0..N-1)

    Rewrite:
      src --RESHAPE(shape=src_pre)--> src_pre
      src_pre --SPLIT(axis=a', numSplits=N)--> r_i

    Safety:
    - `SPLIT` outputs must each feed exactly one `RESHAPE`.
    - All reshape targets must be identical, static, and rank-preserving.
    - Reshape must be memory-order-equivalent by singleton axes only.
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        serial = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{serial}"
            serial += 1
        return candidate

    def _is_memory_order_equivalent_perm(
        *,
        src_shape: List[int],
        perm: List[int],
    ) -> bool:
        non_singleton_src_axes = [int(idx) for idx, dim in enumerate(src_shape) if int(dim) != 1]
        non_singleton_perm_axes = [int(axis) for axis in list(perm) if int(src_shape[int(axis)]) != 1]
        return non_singleton_src_axes == non_singleton_perm_axes

    def _find_singleton_preserving_perm(
        *,
        src_shape: List[int],
        dst_shape: List[int],
    ) -> Optional[List[int]]:
        rank = len(src_shape)
        if int(rank) <= 0 or len(dst_shape) != int(rank) or int(rank) > 6:
            return None
        src = [int(v) for v in list(src_shape)]
        dst = [int(v) for v in list(dst_shape)]
        for perm_tuple in permutations(range(int(rank))):
            perm = [int(v) for v in list(perm_tuple)]
            candidate = _permute_shape(src, perm)
            if candidate is None:
                continue
            if [int(v) for v in list(candidate)] != dst:
                continue
            if not _is_memory_order_equivalent_perm(src_shape=src, perm=perm):
                continue
            return [int(v) for v in list(perm)]
        return None

    def _assign_axis_const(
        *,
        split_op: OperatorIR,
        split_idx: int,
        new_axis: int,
        consumers: Dict[str, List[int]],
    ) -> None:
        if len(split_op.inputs) < 1:
            return
        axis_name = str(split_op.inputs[0])
        axis_tensor = model_ir.tensors.get(axis_name, None)
        if axis_tensor is None:
            return
        axis_users = set(int(v) for v in consumers.get(str(axis_name), []))
        if axis_users == {int(split_idx)}:
            _write_const_ints_to_tensor(axis_tensor, [int(new_axis)])
            return

        axis_const_name = _unique_tensor_name(f"{axis_name}_split_axis")
        np_dtype = np.int32
        if axis_tensor.data is not None:
            try:
                np_dtype = np.asarray(axis_tensor.data).dtype
            except Exception:
                np_dtype = np.int32
        model_ir.tensors[axis_const_name] = TensorIR(
            name=axis_const_name,
            dtype=str(axis_tensor.dtype),
            shape=[1],
            shape_signature=[1],
            data=np.asarray([int(new_axis)], dtype=np_dtype),
            is_variable=False,
            quantization=_clone_quantization(axis_tensor.quantization),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=split_op,
            input_index=0,
            new_input_name=axis_const_name,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for split_idx, split_op in enumerate(model_ir.operators):
            if str(split_op.op_type) != "SPLIT" or len(split_op.inputs) < 2 or len(split_op.outputs) < 2:
                continue

            split_input_name = str(split_op.inputs[1])
            split_input_tensor = model_ir.tensors.get(split_input_name, None)
            if split_input_tensor is None or not _is_fully_known_positive_shape(split_input_tensor.shape):
                continue
            split_input_shape = [int(v) for v in list(split_input_tensor.shape)]
            split_rank = int(len(split_input_shape))
            if int(split_rank) <= 0:
                continue

            axis_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(split_op.inputs[0]), None))
            if axis_vals is None or len(axis_vals) <= 0:
                continue
            split_axis = int(axis_vals[0])
            if split_axis < 0:
                split_axis += int(split_rank)
            if split_axis < 0 or split_axis >= int(split_rank):
                continue

            num_splits = int(split_op.options.get("numSplits", len(split_op.outputs)))
            if int(num_splits) <= 1 or int(num_splits) != len(split_op.outputs):
                continue
            if int(split_input_shape[int(split_axis)]) % int(num_splits) != 0:
                continue
            split_chunk = int(split_input_shape[int(split_axis)] // int(num_splits))
            if int(split_chunk) <= 0:
                continue

            split_output_shape = [int(v) for v in list(split_input_shape)]
            split_output_shape[int(split_axis)] = int(split_chunk)

            reshape_indices: List[int] = []
            reshape_output_names: List[str] = []
            reshape_target_shape: Optional[List[int]] = None
            valid_chain = True

            for split_out_name in [str(v) for v in list(split_op.outputs)]:
                split_out_users = [int(v) for v in consumers.get(str(split_out_name), [])]
                if len(split_out_users) != 1:
                    valid_chain = False
                    break
                reshape_idx = int(split_out_users[0])
                reshape_op = model_ir.operators[int(reshape_idx)]
                if (
                    str(reshape_op.op_type) != "RESHAPE"
                    or len(reshape_op.inputs) < 2
                    or len(reshape_op.outputs) != 1
                    or str(reshape_op.inputs[0]) != str(split_out_name)
                ):
                    valid_chain = False
                    break

                split_out_tensor = model_ir.tensors.get(str(split_out_name), None)
                if (
                    split_out_tensor is None
                    or not _is_fully_known_positive_shape(split_out_tensor.shape)
                    or [int(v) for v in list(split_out_tensor.shape)] != [int(v) for v in list(split_output_shape)]
                ):
                    valid_chain = False
                    break

                target_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(reshape_op.inputs[1]), None))
                if (
                    target_vals is None
                    or len(target_vals) != int(split_rank)
                    or not _is_fully_known_positive_shape([int(v) for v in list(target_vals)])
                ):
                    valid_chain = False
                    break
                target_shape = [int(v) for v in list(target_vals)]
                if int(np.prod(np.asarray(target_shape, dtype=np.int64))) != int(
                    np.prod(np.asarray(split_output_shape, dtype=np.int64))
                ):
                    valid_chain = False
                    break
                if reshape_target_shape is None:
                    reshape_target_shape = [int(v) for v in list(target_shape)]
                elif [int(v) for v in list(reshape_target_shape)] != [int(v) for v in list(target_shape)]:
                    valid_chain = False
                    break

                reshape_indices.append(int(reshape_idx))
                reshape_output_names.append(str(reshape_op.outputs[0]))

            if not valid_chain or reshape_target_shape is None:
                continue

            perm = _find_singleton_preserving_perm(
                src_shape=[int(v) for v in list(split_output_shape)],
                dst_shape=[int(v) for v in list(reshape_target_shape)],
            )
            if perm is None:
                continue
            if int(split_axis) not in set(int(v) for v in list(perm)):
                continue

            pre_reshape_shape = _permute_shape(
                [int(v) for v in list(split_input_shape)],
                [int(v) for v in list(perm)],
            )
            if pre_reshape_shape is None:
                continue
            pre_reshape_shape = [int(v) for v in list(pre_reshape_shape)]
            new_split_axis = int([int(i) for i, axis in enumerate(list(perm)) if int(axis) == int(split_axis)][0])
            if int(pre_reshape_shape[int(new_split_axis)]) % int(num_splits) != 0:
                continue
            if int(pre_reshape_shape[int(new_split_axis)] // int(num_splits)) != int(
                reshape_target_shape[int(new_split_axis)]
            ):
                continue

            pre_reshape_shape_name = _unique_tensor_name(f"{split_input_name}_pre_split_shape")
            model_ir.tensors[pre_reshape_shape_name] = TensorIR(
                name=pre_reshape_shape_name,
                dtype="INT32",
                shape=[int(split_rank)],
                shape_signature=[int(split_rank)],
                data=np.asarray([int(v) for v in list(pre_reshape_shape)], dtype=np.int32),
                is_variable=False,
            )
            pre_reshape_out_name = _unique_tensor_name(f"{split_input_name}_pre_split_reshape")
            input_signature = (
                [int(v) for v in list(split_input_tensor.shape_signature)]
                if split_input_tensor.shape_signature is not None
                else [int(v) for v in list(split_input_shape)]
            )
            pre_signature = _permute_shape([int(v) for v in list(input_signature)], [int(v) for v in list(perm)])
            if pre_signature is None:
                pre_signature = [int(v) for v in list(pre_reshape_shape)]
            model_ir.tensors[pre_reshape_out_name] = TensorIR(
                name=pre_reshape_out_name,
                dtype=str(split_input_tensor.dtype),
                shape=[int(v) for v in list(pre_reshape_shape)],
                shape_signature=[int(v) for v in list(pre_signature)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(split_input_tensor.quantization),
            )
            pre_reshape_op = OperatorIR(
                op_type="RESHAPE",
                inputs=[str(split_input_name), str(pre_reshape_shape_name)],
                outputs=[str(pre_reshape_out_name)],
                options={"newShape": [int(v) for v in list(pre_reshape_shape)]},
            )
            model_ir.operators.insert(int(split_idx), pre_reshape_op)

            shifted_split_idx = int(split_idx + 1)
            shifted_split_op = model_ir.operators[int(shifted_split_idx)]
            _replace_operator_input_at(
                model_ir=model_ir,
                op=shifted_split_op,
                input_index=1,
                new_input_name=str(pre_reshape_out_name),
            )
            _assign_axis_const(
                split_op=shifted_split_op,
                split_idx=int(shifted_split_idx),
                new_axis=int(new_split_axis),
                consumers=consumers,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=shifted_split_op,
                new_outputs=[str(v) for v in list(reshape_output_names)],
            )

            shifted_reshape_indices = sorted([int(v) + 1 for v in list(reshape_indices)], reverse=True)
            for remove_idx in shifted_reshape_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_attention_split_post_reshape_collapse_chains": int(rewritten)}


def _optimize_attention_qkv_shared_pretranspose_slice_nchw_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Hoist sibling NHWC->NCHW transposes on Q/K/V branches to a single
    pre-slice transpose on the shared source tensor.

    Target (rank-4):
      src_nhwc --SLICE(c-range for q)--> q_nhwc --RESHAPE--> q_nchw_like --SOFTMAX
      src_nhwc --SLICE(c-range for k)--> k_nhwc --T(0,3,1,2)--> k_nchw
      src_nhwc --SLICE(c-range for v)--> v_nhwc --RELU--> v_relu_nhwc --T(0,3,1,2)--> v_nchw

    Rewrite:
      src_nhwc --T(0,3,1,2)--> src_nchw
      SLICE begin/size are remapped to NCHW indexing for q/k/v slices.
      Drop per-branch transposes on k and v-relu branches.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        serial = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{serial}"
            serial += 1
        return candidate

    def _permute_begin_or_size_nhwc_to_nchw(values: List[int]) -> List[int]:
        # [N,H,W,C] -> [N,C,H,W]
        return [int(values[0]), int(values[3]), int(values[1]), int(values[2])]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for k_transpose_idx, k_transpose_op in enumerate(model_ir.operators):
            if (
                str(k_transpose_op.op_type) != "TRANSPOSE"
                or len(k_transpose_op.inputs) < 2
                or len(k_transpose_op.outputs) != 1
                or _read_transpose_perm(model_ir, k_transpose_op) != perm_nhwc_to_nchw
            ):
                continue

            k_slice_out_name = str(k_transpose_op.inputs[0])
            k_transpose_out_name = str(k_transpose_op.outputs[0])
            shared_perm_name = str(k_transpose_op.inputs[1])
            if k_slice_out_name in model_outputs or k_transpose_out_name in model_outputs:
                continue

            k_slice_idx = producers.get(k_slice_out_name, None)
            if k_slice_idx is None:
                continue
            k_slice_op = model_ir.operators[int(k_slice_idx)]
            if (
                str(k_slice_op.op_type) != "SLICE"
                or len(k_slice_op.inputs) < 3
                or len(k_slice_op.outputs) != 1
                or str(k_slice_op.outputs[0]) != k_slice_out_name
            ):
                continue

            src_name = str(k_slice_op.inputs[0])
            if src_name in model_outputs:
                continue
            src_users = [int(v) for v in consumers.get(src_name, [])]
            if len(src_users) < 3:
                continue

            slice_user_indices = []
            for user_idx in src_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "SLICE"
                    and len(user_op.inputs) >= 3
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == src_name
                ):
                    slice_user_indices.append(int(user_idx))
            if len(slice_user_indices) != 3 or set(int(v) for v in src_users) != set(slice_user_indices):
                continue

            slice_by_out_name: Dict[str, Dict[str, Any]] = {}
            valid_slices = True
            for slice_idx in slice_user_indices:
                slice_op = model_ir.operators[int(slice_idx)]
                begin_name = str(slice_op.inputs[1])
                size_name = str(slice_op.inputs[2])
                begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(begin_name, None))
                size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(size_name, None))
                if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
                    valid_slices = False
                    break
                out_name = str(slice_op.outputs[0])
                slice_by_out_name[out_name] = {
                    "idx": int(slice_idx),
                    "begin_name": begin_name,
                    "size_name": size_name,
                    "begin_vals": [int(v) for v in list(begin_vals)],
                    "size_vals": [int(v) for v in list(size_vals)],
                }
            if not valid_slices:
                continue

            if k_slice_out_name not in slice_by_out_name:
                continue
            if set(int(v) for v in consumers.get(k_slice_out_name, [])) != {int(k_transpose_idx)}:
                continue

            v_relu_idx: Optional[int] = None
            v_transpose_idx: Optional[int] = None
            v_transpose_out_name = ""
            v_relu_out_name = ""
            v_slice_out_name = ""
            q_slice_out_name = ""
            for out_name, meta in list(slice_by_out_name.items()):
                if out_name == k_slice_out_name:
                    continue
                user_indices = [int(v) for v in consumers.get(str(out_name), [])]
                if len(user_indices) == 1:
                    only_user_idx = int(user_indices[0])
                    only_user_op = model_ir.operators[int(only_user_idx)]
                    if (
                        str(only_user_op.op_type) == "RESHAPE"
                        and len(only_user_op.outputs) == 1
                    ):
                        reshape_out_name = str(only_user_op.outputs[0])
                        softmax_users = [int(v) for v in consumers.get(reshape_out_name, [])]
                        if len(softmax_users) == 1:
                            softmax_op = model_ir.operators[int(softmax_users[0])]
                            if str(softmax_op.op_type) == "SOFTMAX":
                                q_slice_out_name = str(out_name)
                                continue
                if len(user_indices) == 1:
                    relu_candidate_idx = int(user_indices[0])
                    relu_candidate_op = model_ir.operators[int(relu_candidate_idx)]
                    if (
                        str(relu_candidate_op.op_type) == "RELU"
                        and len(relu_candidate_op.outputs) == 1
                    ):
                        relu_out_name = str(relu_candidate_op.outputs[0])
                        relu_users = [int(v) for v in consumers.get(relu_out_name, [])]
                        if len(relu_users) == 1:
                            t_idx = int(relu_users[0])
                            t_op = model_ir.operators[int(t_idx)]
                            if (
                                str(t_op.op_type) == "TRANSPOSE"
                                and len(t_op.inputs) >= 2
                                and len(t_op.outputs) == 1
                                and str(t_op.inputs[0]) == relu_out_name
                                and _read_transpose_perm(model_ir, t_op) == perm_nhwc_to_nchw
                                and str(t_op.outputs[0]) not in model_outputs
                            ):
                                v_slice_out_name = str(out_name)
                                v_relu_idx = int(relu_candidate_idx)
                                v_relu_out_name = relu_out_name
                                v_transpose_idx = int(t_idx)
                                v_transpose_out_name = str(t_op.outputs[0])
                                continue

            if (
                q_slice_out_name == ""
                or v_slice_out_name == ""
                or v_relu_idx is None
                or v_transpose_idx is None
                or v_relu_out_name == ""
                or v_transpose_out_name == ""
            ):
                continue

            src_tensor = model_ir.tensors.get(src_name, None)
            if src_tensor is None or src_tensor.shape is None or len(list(src_tensor.shape)) != 4:
                continue

            src_shape_nhwc = [int(v) for v in list(src_tensor.shape)]
            src_shape_nchw = _permute_shape(src_shape_nhwc, perm_nhwc_to_nchw)
            if src_shape_nchw is None:
                continue
            src_sig_nhwc = (
                [int(v) for v in list(src_tensor.shape_signature)]
                if src_tensor.shape_signature is not None
                else [int(v) for v in list(src_shape_nhwc)]
            )
            src_sig_nchw = _permute_shape(src_sig_nhwc, perm_nhwc_to_nchw)
            if src_sig_nchw is None:
                src_sig_nchw = [int(v) for v in list(src_shape_nchw)]

            shared_nchw_name = _unique_tensor_name(f"{src_name}_qkv_nchw")
            model_ir.tensors[shared_nchw_name] = TensorIR(
                name=shared_nchw_name,
                dtype=str(src_tensor.dtype),
                shape=[int(v) for v in list(src_shape_nchw)],
                shape_signature=[int(v) for v in list(src_sig_nchw)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(src_tensor.quantization),
            )

            shared_transpose_op = OperatorIR(
                op_type="TRANSPOSE",
                inputs=[str(src_name), str(shared_perm_name)],
                outputs=[str(shared_nchw_name)],
            )

            for _, branch_meta in list(slice_by_out_name.items()):
                branch_slice_op = model_ir.operators[int(branch_meta["idx"])]
                begin_name = str(branch_meta["begin_name"])
                size_name = str(branch_meta["size_name"])
                begin_vals = [int(v) for v in list(branch_meta["begin_vals"])]
                size_vals = [int(v) for v in list(branch_meta["size_vals"])]
                mapped_begin = _permute_begin_or_size_nhwc_to_nchw(begin_vals)
                mapped_size = _permute_begin_or_size_nhwc_to_nchw(size_vals)
                if not (
                    _write_const_ints_to_tensor(model_ir.tensors.get(begin_name, None), mapped_begin)
                    or _read_const_ints_from_tensor(model_ir.tensors.get(begin_name, None)) == mapped_begin
                ):
                    valid_slices = False
                    break
                if not (
                    _write_const_ints_to_tensor(model_ir.tensors.get(size_name, None), mapped_size)
                    or _read_const_ints_from_tensor(model_ir.tensors.get(size_name, None)) == mapped_size
                ):
                    valid_slices = False
                    break
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=branch_slice_op,
                    new_inputs=[str(shared_nchw_name), begin_name, size_name],
                )
                out_name = str(branch_slice_op.outputs[0])
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_nhwc_to_nchw,
                )
            if not valid_slices:
                continue

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(v_relu_out_name, None),
                perm_nhwc_to_nchw,
            )

            _replace_tensor_inputs(model_ir, k_transpose_out_name, k_slice_out_name)
            _replace_tensor_inputs(model_ir, v_transpose_out_name, v_relu_out_name)

            remove_indices = sorted([int(k_transpose_idx), int(v_transpose_idx)], reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            insert_index = min(int(v) for v in list(slice_user_indices))
            removed_before_insert = sum(1 for v in remove_indices if int(v) < int(insert_index))
            adjusted_insert_index = int(insert_index) - int(removed_before_insert)
            model_ir.operators.insert(int(adjusted_insert_index), shared_transpose_op)

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_qkv_shared_pretranspose_slice_nchw_chains": int(rewritten)}


def _optimize_attention_qkv_weighted_sum_bridge_to_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Convert NCHW-side weighted-sum attention bridge to NHWC and remove
    redundant branch transposes.

    Target:
      k_nhwc --T(0,3,1,2)--> k_nchw
      soft_nchw --(from SOFTMAX)-->
      MUL(k_nchw, soft_nchw) -> m_nchw
      SUM(m_nchw, axis=[3], keepDims=True) -> s_nchw
      MUL(s_nchw, ones_nchw_const) -> e_nchw
      e_nchw --T(0,2,3,1)--> e_nhwc
      MUL(relu_nhwc, e_nhwc) -> out_nhwc

    Rewrite:
      soft_nchw --T(0,2,3,1)--> soft_nhwc
      MUL(k_nhwc, soft_nhwc) -> m_nhwc
      SUM(axis=[2], keepDims=True) -> s_nhwc
      MUL(s_nhwc, ones_nhwc_const) -> e_nhwc
      MUL(relu_nhwc, e_nhwc) -> out_nhwc
      Remove k/e bridge transposes.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        serial = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{serial}"
            serial += 1
        return candidate

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_t_idx, post_t_op in enumerate(model_ir.operators):
            if (
                str(post_t_op.op_type) != "TRANSPOSE"
                or len(post_t_op.inputs) < 2
                or len(post_t_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_t_op) != perm_nchw_to_nhwc
            ):
                continue

            expand_nchw_name = str(post_t_op.inputs[0])
            expand_nhwc_name = str(post_t_op.outputs[0])
            perm_post_name = str(post_t_op.inputs[1])
            if expand_nchw_name in model_outputs or expand_nhwc_name in model_outputs:
                continue

            relu_mul_users = [int(v) for v in consumers.get(expand_nhwc_name, [])]
            if len(relu_mul_users) != 1:
                continue
            relu_mul_idx = int(relu_mul_users[0])
            relu_mul_op = model_ir.operators[int(relu_mul_idx)]
            if str(relu_mul_op.op_type) != "MUL" or len(relu_mul_op.inputs) != 2 or len(relu_mul_op.outputs) != 1:
                continue
            if str(relu_mul_op.outputs[0]) in model_outputs:
                continue

            relu_mul_inputs = [str(v) for v in list(relu_mul_op.inputs)]
            if relu_mul_inputs[0] == expand_nhwc_name:
                relu_input_name = str(relu_mul_inputs[1])
                relu_mul_expand_input_index = 0
            elif relu_mul_inputs[1] == expand_nhwc_name:
                relu_input_name = str(relu_mul_inputs[0])
                relu_mul_expand_input_index = 1
            else:
                continue
            relu_prod_idx = producers.get(relu_input_name, None)
            if relu_prod_idx is None or str(model_ir.operators[int(relu_prod_idx)].op_type) != "RELU":
                continue

            expand_mul_idx = producers.get(expand_nchw_name, None)
            if expand_mul_idx is None:
                continue
            expand_mul_op = model_ir.operators[int(expand_mul_idx)]
            if (
                str(expand_mul_op.op_type) != "MUL"
                or len(expand_mul_op.inputs) != 2
                or len(expand_mul_op.outputs) != 1
                or str(expand_mul_op.outputs[0]) != expand_nchw_name
            ):
                continue

            expand_mul_inputs = [str(v) for v in list(expand_mul_op.inputs)]
            sum_out_name = ""
            expand_const_name = ""
            for input_name, other_name in [
                (str(expand_mul_inputs[0]), str(expand_mul_inputs[1])),
                (str(expand_mul_inputs[1]), str(expand_mul_inputs[0])),
            ]:
                other_tensor = model_ir.tensors.get(other_name, None)
                if other_tensor is not None and other_tensor.data is not None:
                    sum_out_name = str(input_name)
                    expand_const_name = str(other_name)
                    break
            if sum_out_name == "" or expand_const_name == "":
                continue
            if set(int(v) for v in consumers.get(sum_out_name, [])) != {int(expand_mul_idx)}:
                continue

            sum_idx = producers.get(sum_out_name, None)
            if sum_idx is None:
                continue
            sum_op = model_ir.operators[int(sum_idx)]
            if (
                str(sum_op.op_type) != "SUM"
                or len(sum_op.inputs) < 2
                or len(sum_op.outputs) != 1
                or str(sum_op.outputs[0]) != sum_out_name
                or not bool(sum_op.options.get("keepDims", False))
            ):
                continue
            sum_axes_name = str(sum_op.inputs[1])
            sum_axes_vals = _read_const_ints_from_tensor(model_ir.tensors.get(sum_axes_name, None))
            if sum_axes_vals is None or [int(v) for v in list(sum_axes_vals)] != [3]:
                continue

            kmul_out_name = str(sum_op.inputs[0])
            kmul_idx = producers.get(kmul_out_name, None)
            if kmul_idx is None:
                continue
            kmul_op = model_ir.operators[int(kmul_idx)]
            if (
                str(kmul_op.op_type) != "MUL"
                or len(kmul_op.inputs) != 2
                or len(kmul_op.outputs) != 1
                or str(kmul_op.outputs[0]) != kmul_out_name
            ):
                continue

            kmul_inputs = [str(v) for v in list(kmul_op.inputs)]
            k_nchw_name = ""
            soft_nchw_name = ""
            k_input_index = None
            soft_input_index = None
            for idx, input_name in enumerate(kmul_inputs):
                prod_idx = producers.get(str(input_name), None)
                if prod_idx is None:
                    continue
                prod_op = model_ir.operators[int(prod_idx)]
                if (
                    str(prod_op.op_type) == "TRANSPOSE"
                    and len(prod_op.inputs) >= 2
                    and len(prod_op.outputs) == 1
                    and str(prod_op.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, prod_op) == perm_nhwc_to_nchw
                ):
                    k_nchw_name = str(input_name)
                    k_input_index = int(idx)
                    soft_nchw_name = str(kmul_inputs[1 - idx])
                    soft_input_index = int(1 - idx)
                    break
            if (
                k_nchw_name == ""
                or soft_nchw_name == ""
                or k_input_index is None
                or soft_input_index is None
            ):
                continue

            k_pre_t_idx = producers.get(k_nchw_name, None)
            if k_pre_t_idx is None:
                continue
            k_pre_t_op = model_ir.operators[int(k_pre_t_idx)]
            if (
                str(k_pre_t_op.op_type) != "TRANSPOSE"
                or len(k_pre_t_op.inputs) < 2
                or len(k_pre_t_op.outputs) != 1
                or str(k_pre_t_op.outputs[0]) != k_nchw_name
                or _read_transpose_perm(model_ir, k_pre_t_op) != perm_nhwc_to_nchw
                or str(k_pre_t_op.outputs[0]) in model_outputs
            ):
                continue
            k_nhwc_name = str(k_pre_t_op.inputs[0])
            if k_nhwc_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(k_nchw_name, [])) != {int(kmul_idx)}:
                continue

            soft_prod_idx = producers.get(soft_nchw_name, None)
            if soft_prod_idx is None:
                continue
            soft_prod_op = model_ir.operators[int(soft_prod_idx)]
            if str(soft_prod_op.op_type) != "SOFTMAX":
                continue
            if set(int(v) for v in consumers.get(soft_nchw_name, [])) != {int(kmul_idx)}:
                continue

            if set(int(v) for v in consumers.get(expand_nchw_name, [])) != {int(post_t_idx)}:
                continue

            expand_const_tensor = model_ir.tensors.get(expand_const_name, None)
            if expand_const_tensor is None or expand_const_tensor.data is None:
                continue
            if set(int(v) for v in consumers.get(expand_const_name, [])) != {int(expand_mul_idx)}:
                continue
            expand_const_data = np.asarray(expand_const_tensor.data)
            if int(expand_const_data.ndim) != 4:
                continue
            rotated_const = np.transpose(expand_const_data, axes=perm_nchw_to_nhwc).astype(
                expand_const_data.dtype,
                copy=False,
            )

            kmul_out_tensor = model_ir.tensors.get(kmul_out_name, None)
            sum_out_tensor = model_ir.tensors.get(sum_out_name, None)
            expand_out_tensor = model_ir.tensors.get(expand_nchw_name, None)
            if kmul_out_tensor is None or sum_out_tensor is None or expand_out_tensor is None:
                continue
            if (
                kmul_out_tensor.shape is None
                or sum_out_tensor.shape is None
                or expand_out_tensor.shape is None
                or len(list(kmul_out_tensor.shape)) != 4
                or len(list(sum_out_tensor.shape)) != 4
                or len(list(expand_out_tensor.shape)) != 4
            ):
                continue

            soft_nhwc_name = _unique_tensor_name(f"{soft_nchw_name}_to_nhwc")
            soft_nhwc_shape = _permute_shape(
                [int(v) for v in list(model_ir.tensors[soft_nchw_name].shape)],
                perm_nchw_to_nhwc,
            )
            if soft_nhwc_shape is None:
                continue
            soft_nhwc_sig = _permute_shape(
                (
                    [int(v) for v in list(model_ir.tensors[soft_nchw_name].shape_signature)]
                    if model_ir.tensors[soft_nchw_name].shape_signature is not None
                    else [int(v) for v in list(model_ir.tensors[soft_nchw_name].shape)]
                ),
                perm_nchw_to_nhwc,
            )
            if soft_nhwc_sig is None:
                soft_nhwc_sig = [int(v) for v in list(soft_nhwc_shape)]
            model_ir.tensors[soft_nhwc_name] = TensorIR(
                name=soft_nhwc_name,
                dtype=str(model_ir.tensors[soft_nchw_name].dtype),
                shape=[int(v) for v in list(soft_nhwc_shape)],
                shape_signature=[int(v) for v in list(soft_nhwc_sig)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(model_ir.tensors[soft_nchw_name].quantization),
            )
            soft_to_nhwc_op = OperatorIR(
                op_type="TRANSPOSE",
                inputs=[str(soft_nchw_name), str(perm_post_name)],
                outputs=[str(soft_nhwc_name)],
            )

            kmul_new_inputs = [str(v) for v in list(kmul_inputs)]
            kmul_new_inputs[int(k_input_index)] = str(k_nhwc_name)
            kmul_new_inputs[int(soft_input_index)] = str(soft_nhwc_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=kmul_op,
                new_inputs=kmul_new_inputs,
            )

            if not (
                _write_const_ints_to_tensor(model_ir.tensors.get(sum_axes_name, None), [2])
                or _read_const_ints_from_tensor(model_ir.tensors.get(sum_axes_name, None)) == [2]
            ):
                continue

            expand_const_tensor.data = np.asarray(rotated_const)
            expand_const_tensor.shape = [int(v) for v in list(rotated_const.shape)]
            expand_const_tensor.shape_signature = [int(v) for v in list(rotated_const.shape)]

            _permute_tensor_metadata_if_rank_matches(kmul_out_tensor, perm_nchw_to_nhwc)
            _permute_tensor_metadata_if_rank_matches(sum_out_tensor, perm_nchw_to_nhwc)
            _permute_tensor_metadata_if_rank_matches(expand_out_tensor, perm_nchw_to_nhwc)

            _replace_operator_input_at(
                model_ir=model_ir,
                op=relu_mul_op,
                input_index=int(relu_mul_expand_input_index),
                new_input_name=str(expand_nchw_name),
            )

            remove_indices = sorted([int(k_pre_t_idx), int(post_t_idx)], reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            insert_index = int(kmul_idx)
            removed_before_insert = sum(1 for v in remove_indices if int(v) < int(insert_index))
            adjusted_insert_index = int(insert_index) - int(removed_before_insert)
            model_ir.operators.insert(int(adjusted_insert_index), soft_to_nhwc_op)

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_attention_qkv_weighted_sum_bridge_to_nhwc_chains": int(rewritten)}


def _optimize_split_conv_concat_transpose_bridge_to_single_post_nchw(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Collapse a NCHW bridge around split/conv/concat into NHWC with a single
    post-concat transpose.

    Target motif:
      src_nhwc --T(0,3,1,2)--> src_nchw
      SPLIT(axis=1): src_nchw -> [keep_nchw, branch_nchw]
      branch_nchw --T(0,2,3,1)--> branch_nhwc -> conv_subgraph -> y*_nhwc
      y*_nhwc --T(0,3,1,2)--> y*_nchw
      CONCAT(axis=1): [keep_nchw, y*_nchw] -> z_nchw

    Rewrite:
      SPLIT(axis=3): src_nhwc -> [keep_nhwc, branch_nhwc]
      conv_subgraph stays NHWC
      CONCAT(axis=3): [keep_nhwc, y*_nhwc] -> z_nhwc
      z_nhwc --T(0,3,1,2)--> z_nchw
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        serial = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{serial}"
            serial += 1
        return candidate

    while True:
        changed = False
        producers = _build_tensor_producer_map(model_ir)
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for split_idx, split_op in enumerate(model_ir.operators):
            if str(split_op.op_type) != "SPLIT" or len(split_op.inputs) < 2:
                continue
            split_axis_name = str(split_op.inputs[0])
            split_data_name = str(split_op.inputs[1])
            split_axis_vals = _read_const_ints_from_tensor(
                model_ir.tensors.get(split_axis_name, None)
            )
            if split_axis_vals is None or len(split_axis_vals) <= 0:
                continue
            split_axis = int(split_axis_vals[0])
            # Operate only on channel split in NCHW.
            if split_axis not in [1, -3]:
                continue
            if split_data_name in model_outputs:
                continue

            pre_t_idx = producers.get(split_data_name, None)
            if pre_t_idx is None:
                continue
            pre_t_op = model_ir.operators[int(pre_t_idx)]
            if (
                str(pre_t_op.op_type) != "TRANSPOSE"
                or len(pre_t_op.inputs) < 2
                or len(pre_t_op.outputs) != 1
                or str(pre_t_op.outputs[0]) != split_data_name
                or _read_transpose_perm(model_ir, pre_t_op) != perm_nhwc_to_nchw
            ):
                continue
            src_nhwc_name = str(pre_t_op.inputs[0])
            if src_nhwc_name in model_outputs:
                continue

            split_outputs = [str(v) for v in list(split_op.outputs)]
            if len(split_outputs) < 2:
                continue

            branch_t_idx: Optional[int] = None
            branch_split_out_name = ""
            branch_nhwc_name = ""
            for out_name in list(split_outputs):
                user_indices = [int(v) for v in consumers.get(out_name, [])]
                if len(user_indices) != 1:
                    continue
                candidate_idx = int(user_indices[0])
                candidate_op = model_ir.operators[int(candidate_idx)]
                if (
                    str(candidate_op.op_type) == "TRANSPOSE"
                    and len(candidate_op.inputs) >= 2
                    and len(candidate_op.outputs) == 1
                    and str(candidate_op.inputs[0]) == out_name
                    and _read_transpose_perm(model_ir, candidate_op) == perm_nchw_to_nhwc
                ):
                    branch_t_idx = int(candidate_idx)
                    branch_split_out_name = str(out_name)
                    branch_nhwc_name = str(candidate_op.outputs[0])
                    break
            if branch_t_idx is None or branch_split_out_name == "" or branch_nhwc_name == "":
                continue
            if branch_nhwc_name in model_outputs:
                continue

            # Identify concat that consumes one or more post NHWC->NCHW transposes
            # and at least one direct split output.
            candidate_concat_indices: List[int] = []
            for out_name in list(split_outputs):
                if out_name == branch_split_out_name:
                    continue
                for user_idx in consumers.get(out_name, []):
                    user_op = model_ir.operators[int(user_idx)]
                    if str(user_op.op_type) == "CONCATENATION":
                        candidate_concat_indices.append(int(user_idx))
            if len(candidate_concat_indices) <= 0:
                continue

            selected_concat_idx: Optional[int] = None
            selected_concat_op: Optional[OperatorIR] = None
            selected_post_t_outputs: List[str] = []
            selected_post_t_inputs: List[str] = []
            for concat_idx in list(dict.fromkeys(candidate_concat_indices)):
                concat_op = model_ir.operators[int(concat_idx)]
                if len(concat_op.outputs) != 1:
                    continue
                concat_output_name = str(concat_op.outputs[0])
                if concat_output_name in model_outputs:
                    continue
                concat_axis = int(concat_op.options.get("axis", 0))
                if concat_axis not in [1, -3]:
                    continue
                concat_inputs = [str(v) for v in list(concat_op.inputs)]
                if branch_split_out_name in concat_inputs:
                    # branch split output should be consumed via branch transpose path.
                    continue
                direct_split_inputs = [
                    v for v in concat_inputs if v in split_outputs and v != branch_split_out_name
                ]
                if len(direct_split_inputs) <= 0:
                    continue
                post_t_indices: List[int] = []
                post_t_outputs: List[str] = []
                post_t_inputs: List[str] = []
                valid = True
                for input_name in list(concat_inputs):
                    prod_idx = producers.get(input_name, None)
                    if prod_idx is None:
                        continue
                    prod_op = model_ir.operators[int(prod_idx)]
                    if (
                        str(prod_op.op_type) == "TRANSPOSE"
                        and len(prod_op.inputs) >= 2
                        and len(prod_op.outputs) == 1
                        and str(prod_op.outputs[0]) == input_name
                        and _read_transpose_perm(model_ir, prod_op) == perm_nhwc_to_nchw
                    ):
                        if str(prod_op.outputs[0]) in model_outputs:
                            valid = False
                            break
                        if set(int(v) for v in consumers.get(str(prod_op.outputs[0]), [])) != {int(concat_idx)}:
                            valid = False
                            break
                        post_t_indices.append(int(prod_idx))
                        post_t_outputs.append(str(prod_op.outputs[0]))
                        post_t_inputs.append(str(prod_op.inputs[0]))
                if not valid or len(post_t_indices) <= 0:
                    continue
                selected_concat_idx = int(concat_idx)
                selected_concat_op = concat_op
                selected_post_t_outputs = list(post_t_outputs)
                selected_post_t_inputs = list(post_t_inputs)
                break

            if selected_concat_idx is None or selected_concat_op is None:
                continue

            # 1) Bypass pre-branch transpose.
            _replace_tensor_inputs(model_ir, branch_nhwc_name, branch_split_out_name)
            branch_tensor = model_ir.tensors.get(branch_split_out_name, None)
            if branch_tensor is not None:
                _permute_tensor_metadata_if_rank_matches(
                    branch_tensor,
                    perm_nchw_to_nhwc,
                )

            # 2) Rewire split to consume NHWC source and split channel axis=3.
            new_split_axis_name = _unique_tensor_name(f"{split_axis_name}_nhwc")
            model_ir.tensors[new_split_axis_name] = TensorIR(
                name=new_split_axis_name,
                dtype="INT32",
                shape=[1],
                shape_signature=[1],
                data=np.asarray([3], dtype=np.int32),
                is_variable=False,
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=split_op,
                new_inputs=[new_split_axis_name, src_nhwc_name],
            )
            for split_out_name in list(split_outputs):
                split_out_tensor = model_ir.tensors.get(split_out_name, None)
                if split_out_tensor is None:
                    continue
                _permute_tensor_metadata_if_rank_matches(
                    split_out_tensor,
                    perm_nchw_to_nhwc,
                )

            # 3) Replace concat post-transpose inputs with NHWC tensors and switch axis.
            concat_inputs = [str(v) for v in list(selected_concat_op.inputs)]
            replaced_concat_inputs: List[str] = []
            post_out_to_in = {
                str(out_name): str(in_name)
                for out_name, in_name in zip(
                    selected_post_t_outputs,
                    selected_post_t_inputs,
                )
            }
            for input_name in list(concat_inputs):
                replaced_concat_inputs.append(post_out_to_in.get(str(input_name), str(input_name)))
            _set_operator_inputs(
                model_ir=model_ir,
                op=selected_concat_op,
                new_inputs=replaced_concat_inputs,
            )
            selected_concat_op.options["axis"] = int(3)

            # 4) Keep original concat output contract by appending one NHWC->NCHW transpose.
            concat_output_nchw_name = str(selected_concat_op.outputs[0])
            concat_output_nhwc_name = _unique_tensor_name(f"{concat_output_nchw_name}__nhwc")
            concat_output_nchw_tensor = model_ir.tensors.get(concat_output_nchw_name, None)
            if concat_output_nchw_tensor is None:
                continue
            nhwc_shape = _permute_shape(
                [int(v) for v in list(concat_output_nchw_tensor.shape)],
                perm_nchw_to_nhwc,
            )
            if nhwc_shape is None:
                continue
            nhwc_signature = _permute_shape(
                (
                    [int(v) for v in list(concat_output_nchw_tensor.shape_signature)]
                    if concat_output_nchw_tensor.shape_signature is not None
                    else [int(v) for v in list(concat_output_nchw_tensor.shape)]
                ),
                perm_nchw_to_nhwc,
            )
            if nhwc_signature is None:
                nhwc_signature = [int(v) for v in list(nhwc_shape)]
            model_ir.tensors[concat_output_nhwc_name] = TensorIR(
                name=concat_output_nhwc_name,
                dtype=str(concat_output_nchw_tensor.dtype),
                shape=[int(v) for v in list(nhwc_shape)],
                shape_signature=[int(v) for v in list(nhwc_signature)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(concat_output_nchw_tensor.quantization),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=selected_concat_op,
                new_outputs=[concat_output_nhwc_name],
            )
            bridge_perm_name = _unique_tensor_name(f"{concat_output_nchw_name}_bridge_perm")
            model_ir.tensors[bridge_perm_name] = TensorIR(
                name=bridge_perm_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                is_variable=False,
            )
            model_ir.operators.insert(
                int(selected_concat_idx + 1),
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[concat_output_nhwc_name, bridge_perm_name],
                    outputs=[concat_output_nchw_name],
                ),
            )

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_split_conv_concat_transpose_bridge_to_single_post_nchw": int(rewritten)}


def _optimize_transpose_csp_attention_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate RTMDet-like CSP attention NCHW/NHWC bridge chains.

    Strict target motif (rank-4):
      short_nhwc --T(0,3,1,2)--> short_nchw --LOGISTIC--> s_sig --MUL(short_nchw, s_sig)--> short_branch
      main_nhwc  --T(0,3,1,2)--> main_nchw
      point_nhwc --T(0,3,1,2)--> point_nchw --LOGISTIC--> p_sig --MUL(point_nchw, p_sig)--> point_branch
      ADD(main_nchw, point_branch) -> add_nchw
      CONCAT(axis=1, [add_nchw, short_branch]) -> feat_nchw
      MEAN(keepDims=True, axes=[2,3]) -> gap_nchw --T(0,2,3,1)--> gap_nhwc --CONV_2D--> gate_nhwc
      gate_nhwc --T(0,3,1,2)--> gate_in_nchw --(LOGISTIC|HardSigmoid expansion)--> gate_nchw
      MUL(feat_nchw, gate_nchw) -> out_nchw --T(0,2,3,1)--> out_nhwc

    Rewrite:
      - keep all branch tensors in NHWC
      - remap CONCAT axis to NHWC channel axis and MEAN axes to NHWC
      - remove bridge transposes in the motif
      - preserve post-transpose output tensor names
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _is_singleton_spatial_nchw_to_nhwc_reshape(
        *,
        input_name: str,
        output_name: str,
    ) -> bool:
        in_tensor = model_ir.tensors.get(str(input_name), None)
        out_tensor = model_ir.tensors.get(str(output_name), None)
        if in_tensor is None or out_tensor is None:
            return False
        in_shape = [int(v) for v in list(in_tensor.shape)]
        out_shape = [int(v) for v in list(out_tensor.shape)]
        if len(in_shape) != 4 or len(out_shape) != 4:
            return False
        if any(int(v) < 0 for v in list(in_shape) + list(out_shape)):
            return False
        return (
            int(in_shape[0]) == int(out_shape[0])
            and int(in_shape[1]) == int(out_shape[3])
            and int(in_shape[2]) == 1
            and int(in_shape[3]) == 1
            and int(out_shape[1]) == 1
            and int(out_shape[2]) == 1
        )

    def _is_singleton_spatial_nhwc_to_nchw_reshape(
        *,
        input_name: str,
        output_name: str,
    ) -> bool:
        in_tensor = model_ir.tensors.get(str(input_name), None)
        out_tensor = model_ir.tensors.get(str(output_name), None)
        if in_tensor is None or out_tensor is None:
            return False
        in_shape = [int(v) for v in list(in_tensor.shape)]
        out_shape = [int(v) for v in list(out_tensor.shape)]
        if len(in_shape) != 4 or len(out_shape) != 4:
            return False
        if any(int(v) < 0 for v in list(in_shape) + list(out_shape)):
            return False
        return (
            int(in_shape[0]) == int(out_shape[0])
            and int(in_shape[1]) == 1
            and int(in_shape[2]) == 1
            and int(in_shape[3]) == int(out_shape[1])
            and int(out_shape[2]) == 1
            and int(out_shape[3]) == 1
        )

    def _match_sigmoid_self_mul_from_transpose_output(
        branch_output_name: str,
        consumers: Dict[str, List[int]],
        producers: Dict[str, int],
    ) -> Optional[Dict[str, Any]]:
        mul_idx = producers.get(str(branch_output_name), None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if (
            str(mul_op.op_type) != "MUL"
            or len(mul_op.inputs) != 2
            or len(mul_op.outputs) != 1
            or str(mul_op.outputs[0]) != str(branch_output_name)
        ):
            return None

        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        for candidate_input_idx in [0, 1]:
            transpose_output_name = str(mul_inputs[int(candidate_input_idx)])
            sigmoid_output_name = str(mul_inputs[1 - int(candidate_input_idx)])

            pre_idx = producers.get(transpose_output_name, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != transpose_output_name
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue

            sig_idx = producers.get(sigmoid_output_name, None)
            if sig_idx is None:
                continue
            sig_op = model_ir.operators[int(sig_idx)]
            if (
                str(sig_op.op_type) != "LOGISTIC"
                or len(sig_op.inputs) != 1
                or len(sig_op.outputs) != 1
                or str(sig_op.inputs[0]) != transpose_output_name
                or str(sig_op.outputs[0]) != sigmoid_output_name
            ):
                continue

            if set(int(v) for v in consumers.get(sigmoid_output_name, [])) != {int(mul_idx)}:
                continue
            if set(int(v) for v in consumers.get(transpose_output_name, [])) != {int(sig_idx), int(mul_idx)}:
                continue

            return {
                "transpose_idx": int(pre_idx),
                "transpose_output_name": str(transpose_output_name),
                "source_nhwc_name": str(pre_op.inputs[0]),
                "sigmoid_idx": int(sig_idx),
                "sigmoid_output_name": str(sigmoid_output_name),
                "mul_idx": int(mul_idx),
                "branch_output_name": str(branch_output_name),
            }
        return None

    def _match_hardsigmoid_gate_from_output(
        gate_output_name: str,
        consumers: Dict[str, List[int]],
        producers: Dict[str, int],
    ) -> Optional[Dict[str, Any]]:
        gate_idx = producers.get(str(gate_output_name), None)
        if gate_idx is None:
            return None
        gate_op = model_ir.operators[int(gate_idx)]
        gate_type = str(gate_op.op_type)
        add_output_name: Optional[str] = None
        chain_indices: List[int] = []
        metadata_names: List[str] = []

        if gate_type == "LOGISTIC":
            if len(gate_op.inputs) != 1 or len(gate_op.outputs) != 1:
                return None
            return {
                "head_idx": int(gate_idx),
                "head_input_name": str(gate_op.inputs[0]),
                "head_input_index": 0,
                "output_name": str(gate_output_name),
                "chain_indices": [int(gate_idx)],
                "metadata_names": [str(gate_output_name)],
            }

        if gate_type == "RELU_0_TO_1":
            if len(gate_op.inputs) != 1 or len(gate_op.outputs) != 1:
                return None
            add_output_name = str(gate_op.inputs[0])
            if set(int(v) for v in consumers.get(add_output_name, [])) != {int(gate_idx)}:
                return None
            chain_indices.append(int(gate_idx))
            metadata_names.append(str(gate_output_name))
        elif gate_type == "MINIMUM":
            if len(gate_op.inputs) != 2 or len(gate_op.outputs) != 1:
                return None
            min_inputs = [str(v) for v in list(gate_op.inputs)]
            max_output_name: Optional[str] = None
            min_side_name: Optional[str] = None
            for candidate_name, other_name in [(min_inputs[0], min_inputs[1]), (min_inputs[1], min_inputs[0])]:
                producer_idx = producers.get(candidate_name, None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if str(producer_op.op_type) != "MAXIMUM":
                    continue
                max_output_name = str(candidate_name)
                min_side_name = str(other_name)
                break
            if max_output_name is None or min_side_name is None:
                return None
            if not _is_singleton_constant_tensor(model_ir, min_side_name):
                return None
            if set(int(v) for v in consumers.get(max_output_name, [])) != {int(gate_idx)}:
                return None

            max_idx = producers.get(max_output_name, None)
            if max_idx is None:
                return None
            max_op = model_ir.operators[int(max_idx)]
            if str(max_op.op_type) != "MAXIMUM" or len(max_op.inputs) != 2 or len(max_op.outputs) != 1:
                return None
            max_inputs = [str(v) for v in list(max_op.inputs)]
            for candidate_name, other_name in [(max_inputs[0], max_inputs[1]), (max_inputs[1], max_inputs[0])]:
                producer_idx = producers.get(candidate_name, None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if str(producer_op.op_type) != "ADD":
                    continue
                add_output_name = str(candidate_name)
                if not _is_singleton_constant_tensor(model_ir, str(other_name)):
                    return None
                break
            if add_output_name is None:
                return None
            if set(int(v) for v in consumers.get(add_output_name, [])) != {int(max_idx)}:
                return None
            chain_indices.extend([int(max_idx), int(gate_idx)])
            metadata_names.extend([str(max_output_name), str(gate_output_name)])
        else:
            return None

        add_idx = producers.get(str(add_output_name), None)
        if add_idx is None:
            return None
        add_op = model_ir.operators[int(add_idx)]
        if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
            return None
        add_inputs = [str(v) for v in list(add_op.inputs)]
        mul_output_name: Optional[str] = None
        for candidate_name, other_name in [(add_inputs[0], add_inputs[1]), (add_inputs[1], add_inputs[0])]:
            producer_idx = producers.get(candidate_name, None)
            if producer_idx is None:
                continue
            producer_op = model_ir.operators[int(producer_idx)]
            if str(producer_op.op_type) != "MUL":
                continue
            if not _is_singleton_constant_tensor(model_ir, str(other_name)):
                return None
            mul_output_name = str(candidate_name)
            break
        if mul_output_name is None:
            return None
        if set(int(v) for v in consumers.get(mul_output_name, [])) != {int(add_idx)}:
            return None
        chain_indices.append(int(add_idx))
        metadata_names.insert(0, str(add_output_name))

        mul_idx = producers.get(mul_output_name, None)
        if mul_idx is None:
            return None
        mul_op = model_ir.operators[int(mul_idx)]
        if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
            return None
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        main_input_name: Optional[str] = None
        main_input_index: Optional[int] = None
        for input_index, input_name in enumerate(mul_inputs):
            side_name = str(mul_inputs[1 - input_index])
            if _is_singleton_constant_tensor(model_ir, side_name):
                main_input_name = str(input_name)
                main_input_index = int(input_index)
                break
        if main_input_name is None or main_input_index is None:
            return None

        chain_indices.append(int(mul_idx))
        metadata_names.insert(0, str(mul_output_name))
        return {
            "head_idx": int(mul_idx),
            "head_input_name": str(main_input_name),
            "head_input_index": int(main_input_index),
            "output_name": str(gate_output_name),
            "chain_indices": [int(v) for v in chain_indices],
            "metadata_names": [str(v) for v in metadata_names],
        }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for candidate_post_idx, candidate_post_op in enumerate(model_ir.operators):
            if (
                str(candidate_post_op.op_type) != "TRANSPOSE"
                or len(candidate_post_op.inputs) < 2
                or len(candidate_post_op.outputs) != 1
                or _read_transpose_perm(model_ir, candidate_post_op) != perm_nchw_to_nhwc
            ):
                continue

            mul_output_name = str(candidate_post_op.inputs[0])
            if mul_output_name in model_outputs:
                continue
            mul_idx = producers.get(mul_output_name, None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(mul_op.outputs[0]) != mul_output_name:
                continue

            mul_output_users = [int(v) for v in consumers.get(mul_output_name, [])]
            post_pairs: List[Tuple[int, OperatorIR]] = []
            valid_posts = len(mul_output_users) > 0
            for user_idx in mul_output_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) != "TRANSPOSE"
                    or len(user_op.inputs) < 2
                    or len(user_op.outputs) != 1
                    or str(user_op.inputs[0]) != mul_output_name
                    or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                    or str(user_op.outputs[0]) in model_outputs
                ):
                    valid_posts = False
                    break
                post_pairs.append((int(user_idx), user_op))
            if not valid_posts or len(post_pairs) == 0:
                continue
            post_pairs = sorted(post_pairs, key=lambda v: int(v[0]))
            post_indices = [int(v[0]) for v in post_pairs]
            post_output_names = [str(v[1].outputs[0]) for v in post_pairs]

            concat_name: Optional[str] = None
            gate_output_name: Optional[str] = None
            input0_name = str(mul_op.inputs[0])
            input1_name = str(mul_op.inputs[1])
            for lhs_name, rhs_name in [(input0_name, input1_name), (input1_name, input0_name)]:
                lhs_prod_idx = producers.get(lhs_name, None)
                if lhs_prod_idx is None:
                    continue
                lhs_prod_op = model_ir.operators[int(lhs_prod_idx)]
                if str(lhs_prod_op.op_type) != "CONCATENATION":
                    continue
                concat_name = str(lhs_name)
                gate_output_name = str(rhs_name)
                break
            if concat_name is None or gate_output_name is None:
                continue

            concat_idx = producers.get(concat_name, None)
            if concat_idx is None:
                continue
            concat_op = model_ir.operators[int(concat_idx)]
            if (
                str(concat_op.op_type) != "CONCATENATION"
                or len(concat_op.inputs) != 2
                or len(concat_op.outputs) != 1
                or str(concat_op.outputs[0]) != concat_name
                or int(concat_op.options.get("axis", 1)) != 1
                or concat_name in model_outputs
            ):
                continue

            gate_match = _match_hardsigmoid_gate_from_output(
                gate_output_name=gate_output_name,
                consumers=consumers,
                producers=producers,
            )
            if gate_match is None:
                continue
            gate_head_idx = int(gate_match["head_idx"])
            gate_head_input_name = str(gate_match["head_input_name"])
            gate_head_input_index = int(gate_match["head_input_index"])
            gate_metadata_names = [str(v) for v in list(gate_match.get("metadata_names", []))]
            gate_head_op = model_ir.operators[int(gate_head_idx)]
            if gate_head_input_name in model_outputs or gate_output_name in model_outputs:
                continue

            pre_gate_idx = producers.get(gate_head_input_name, None)
            if pre_gate_idx is None:
                continue
            pre_gate_op = model_ir.operators[int(pre_gate_idx)]
            gate_conv_output_name: Optional[str] = None
            if (
                str(pre_gate_op.op_type) == "TRANSPOSE"
                and len(pre_gate_op.inputs) >= 2
                and len(pre_gate_op.outputs) == 1
                and str(pre_gate_op.outputs[0]) == gate_head_input_name
                and _read_transpose_perm(model_ir, pre_gate_op) == perm_nhwc_to_nchw
            ):
                gate_conv_output_name = str(pre_gate_op.inputs[0])
            elif (
                str(pre_gate_op.op_type) == "RESHAPE"
                and len(pre_gate_op.inputs) >= 2
                and len(pre_gate_op.outputs) == 1
                and str(pre_gate_op.outputs[0]) == gate_head_input_name
                and _is_singleton_spatial_nhwc_to_nchw_reshape(
                    input_name=str(pre_gate_op.inputs[0]),
                    output_name=str(pre_gate_op.outputs[0]),
                )
            ):
                gate_conv_output_name = str(pre_gate_op.inputs[0])
            if gate_conv_output_name is None:
                continue
            if gate_conv_output_name in model_outputs:
                continue

            concat_users = set(int(v) for v in consumers.get(concat_name, []))
            if concat_users != {int(mul_idx)} and len(concat_users) != 2:
                continue
            mean_idx_candidates = [int(v) for v in concat_users if int(v) != int(mul_idx)]
            if len(mean_idx_candidates) != 1:
                continue
            mean_idx = int(mean_idx_candidates[0])
            mean_op = model_ir.operators[int(mean_idx)]
            if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                continue
            if str(mean_op.inputs[0]) != concat_name or not bool(mean_op.options.get("keepDims", False)):
                continue
            mean_output_name = str(mean_op.outputs[0])
            if mean_output_name in model_outputs:
                continue

            mean_post_users = [int(v) for v in consumers.get(mean_output_name, [])]
            if len(mean_post_users) != 1:
                continue
            mean_post_idx = int(mean_post_users[0])
            mean_post_op = model_ir.operators[int(mean_post_idx)]
            conv_input_name: Optional[str] = None
            if (
                str(mean_post_op.op_type) == "TRANSPOSE"
                and len(mean_post_op.inputs) >= 2
                and len(mean_post_op.outputs) == 1
                and str(mean_post_op.inputs[0]) == mean_output_name
                and _read_transpose_perm(model_ir, mean_post_op) == perm_nchw_to_nhwc
            ):
                conv_input_name = str(mean_post_op.outputs[0])
            elif (
                str(mean_post_op.op_type) == "RESHAPE"
                and len(mean_post_op.inputs) >= 2
                and len(mean_post_op.outputs) == 1
                and str(mean_post_op.inputs[0]) == mean_output_name
                and _is_singleton_spatial_nchw_to_nhwc_reshape(
                    input_name=str(mean_post_op.inputs[0]),
                    output_name=str(mean_post_op.outputs[0]),
                )
            ):
                conv_input_name = str(mean_post_op.outputs[0])
            if conv_input_name is None:
                continue
            if conv_input_name in model_outputs:
                continue

            conv_idx_candidates = [int(v) for v in consumers.get(conv_input_name, [])]
            if len(conv_idx_candidates) != 1:
                continue
            conv_idx = int(conv_idx_candidates[0])
            conv_op = model_ir.operators[int(conv_idx)]
            if (
                str(conv_op.op_type) != "CONV_2D"
                or len(conv_op.inputs) < 1
                or len(conv_op.outputs) != 1
                or str(conv_op.inputs[0]) != conv_input_name
            ):
                continue
            if str(conv_op.outputs[0]) != gate_conv_output_name:
                continue

            concat_input_names = [str(v) for v in list(concat_op.inputs)]
            use_add_main_path = False
            add_idx: Optional[int] = None
            add_op: Optional[OperatorIR] = None
            add_output_name: Optional[str] = None
            main_pre_idx: Optional[int] = None
            main_transpose_output_name: Optional[str] = None
            main_source_nhwc_name: Optional[str] = None
            point_branch_output_name: Optional[str] = None
            short_branch_output_name: Optional[str] = None

            # Variant A: CONCAT(ADD(main_nchw, point_branch_nchw), short_branch_nchw)
            for a_name, b_name in [(concat_input_names[0], concat_input_names[1]), (concat_input_names[1], concat_input_names[0])]:
                add_candidate_idx = producers.get(a_name, None)
                if add_candidate_idx is None:
                    continue
                add_candidate_op = model_ir.operators[int(add_candidate_idx)]
                if str(add_candidate_op.op_type) != "ADD":
                    continue
                add_input_names = [str(v) for v in list(add_candidate_op.inputs)]
                for m_name, p_name in [(add_input_names[0], add_input_names[1]), (add_input_names[1], add_input_names[0])]:
                    pre_idx = producers.get(m_name, None)
                    if pre_idx is None:
                        continue
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != m_name
                        or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    ):
                        continue
                    add_idx = int(add_candidate_idx)
                    add_op = add_candidate_op
                    add_output_name = str(a_name)
                    short_branch_output_name = str(b_name)
                    main_pre_idx = int(pre_idx)
                    main_transpose_output_name = str(m_name)
                    main_source_nhwc_name = str(pre_op.inputs[0])
                    point_branch_output_name = str(p_name)
                    use_add_main_path = True
                    break
                if use_add_main_path:
                    break

            # Variant B: CONCAT(point_branch_nchw, short_branch_nchw)
            if not use_add_main_path:
                point_branch_output_name = str(concat_input_names[0])
                short_branch_output_name = str(concat_input_names[1])

            if point_branch_output_name is None or short_branch_output_name is None:
                continue

            if use_add_main_path:
                if (
                    add_idx is None
                    or add_op is None
                    or add_output_name is None
                    or main_pre_idx is None
                    or main_transpose_output_name is None
                    or main_source_nhwc_name is None
                ):
                    continue
                if (
                    str(add_op.op_type) != "ADD"
                    or len(add_op.inputs) != 2
                    or len(add_op.outputs) != 1
                    or str(add_op.outputs[0]) != add_output_name
                    or add_output_name in model_outputs
                    or main_source_nhwc_name in model_outputs
                ):
                    continue

            point_match = _match_sigmoid_self_mul_from_transpose_output(
                branch_output_name=str(point_branch_output_name),
                consumers=consumers,
                producers=producers,
            )
            short_match = _match_sigmoid_self_mul_from_transpose_output(
                branch_output_name=str(short_branch_output_name),
                consumers=consumers,
                producers=producers,
            )
            if point_match is None or short_match is None:
                # Try swapped interpretation for variant-B.
                if not use_add_main_path:
                    point_match = _match_sigmoid_self_mul_from_transpose_output(
                        branch_output_name=str(short_branch_output_name),
                        consumers=consumers,
                        producers=producers,
                    )
                    short_match = _match_sigmoid_self_mul_from_transpose_output(
                        branch_output_name=str(point_branch_output_name),
                        consumers=consumers,
                        producers=producers,
                    )
                    if point_match is not None and short_match is not None:
                        point_branch_output_name, short_branch_output_name = (
                            str(short_branch_output_name),
                            str(point_branch_output_name),
                        )
                if point_match is None or short_match is None:
                    continue

            if int(point_match["transpose_idx"]) == int(short_match["transpose_idx"]):
                continue
            if use_add_main_path and (
                int(point_match["transpose_idx"]) == int(main_pre_idx)
                or int(short_match["transpose_idx"]) == int(main_pre_idx)
            ):
                continue

            if use_add_main_path:
                if set(int(v) for v in consumers.get(str(main_transpose_output_name), [])) != {int(add_idx)}:
                    continue
                if set(int(v) for v in consumers.get(str(point_branch_output_name), [])) != {int(add_idx)}:
                    continue
                if set(int(v) for v in consumers.get(str(short_branch_output_name), [])) != {int(concat_idx)}:
                    continue
            else:
                if set(int(v) for v in consumers.get(str(point_branch_output_name), [])) != {int(concat_idx)}:
                    continue
                if set(int(v) for v in consumers.get(str(short_branch_output_name), [])) != {int(concat_idx)}:
                    continue
            if set(int(v) for v in consumers.get(mean_output_name, [])) != {int(mean_post_idx)}:
                continue
            if set(int(v) for v in consumers.get(conv_input_name, [])) != {int(conv_idx)}:
                continue
            if set(int(v) for v in consumers.get(gate_conv_output_name, [])) != {int(pre_gate_idx)}:
                continue
            if set(int(v) for v in consumers.get(gate_head_input_name, [])) != {int(gate_head_idx)}:
                continue
            if set(int(v) for v in consumers.get(gate_output_name, [])) != {int(mul_idx)}:
                continue

            mean_axes_tensor = model_ir.tensors.get(str(mean_op.inputs[1]), None)
            mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_vals is None or len(mean_axes_vals) == 0:
                continue
            normalized_axes: List[int] = []
            valid_axes = True
            for axis in mean_axes_vals:
                a = int(axis)
                if a < 0:
                    a += 4
                if a < 0 or a >= 4:
                    valid_axes = False
                    break
                normalized_axes.append(int(a))
            if not valid_axes:
                continue
            mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]
            if not _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes]):
                continue

            short_sig_op = model_ir.operators[int(short_match["sigmoid_idx"])]
            short_mul_op = model_ir.operators[int(short_match["mul_idx"])]
            _set_operator_inputs(
                model_ir=model_ir,
                op=short_sig_op,
                new_inputs=[str(short_match["source_nhwc_name"])],
            )
            short_mul_inputs = [
                str(short_match["source_nhwc_name"])
                if str(v) == str(short_match["transpose_output_name"])
                else str(v)
                for v in list(short_mul_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=short_mul_op,
                new_inputs=short_mul_inputs,
            )

            point_sig_op = model_ir.operators[int(point_match["sigmoid_idx"])]
            point_mul_op = model_ir.operators[int(point_match["mul_idx"])]
            _set_operator_inputs(
                model_ir=model_ir,
                op=point_sig_op,
                new_inputs=[str(point_match["source_nhwc_name"])],
            )
            point_mul_inputs = [
                str(point_match["source_nhwc_name"])
                if str(v) == str(point_match["transpose_output_name"])
                else str(v)
                for v in list(point_mul_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=point_mul_op,
                new_inputs=point_mul_inputs,
            )

            if use_add_main_path and add_op is not None:
                add_inputs = [
                    str(main_source_nhwc_name)
                    if str(v) == str(main_transpose_output_name)
                    else str(v)
                    for v in list(add_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_inputs=add_inputs,
                )

            concat_op.options = dict(concat_op.options) if isinstance(concat_op.options, dict) else {}
            concat_op.options["axis"] = 3

            conv_inputs = [str(v) for v in list(conv_op.inputs)]
            conv_inputs[0] = str(mean_output_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=conv_op,
                new_inputs=conv_inputs,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=gate_head_op,
                input_index=int(gate_head_input_index),
                new_input_name=str(gate_conv_output_name),
            )

            old_mul_tensor = model_ir.tensors.get(mul_output_name, None)
            old_mul_shape = (
                [int(v) for v in list(old_mul_tensor.shape)]
                if old_mul_tensor is not None and old_mul_tensor.shape is not None
                else None
            )
            old_mul_signature = (
                [int(v) for v in list(old_mul_tensor.shape_signature)]
                if old_mul_tensor is not None and old_mul_tensor.shape_signature is not None
                else (
                    [int(v) for v in list(old_mul_tensor.shape)]
                    if old_mul_tensor is not None and old_mul_tensor.shape is not None
                    else None
                )
            )

            metadata_targets = [
                str(short_match["sigmoid_output_name"]),
                str(short_match["branch_output_name"]),
                str(point_match["sigmoid_output_name"]),
                str(point_match["branch_output_name"]),
                str(concat_name),
                str(mean_output_name),
                *[str(v) for v in gate_metadata_names],
                str(mul_output_name),
            ]
            if use_add_main_path and add_output_name is not None:
                metadata_targets.insert(4, str(add_output_name))
            for tensor_name in list(dict.fromkeys(metadata_targets)):
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )

            canonical_post_output_name = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[canonical_post_output_name],
            )
            _replace_tensor_inputs(model_ir, mul_output_name, canonical_post_output_name)
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

            canonical_tensor = model_ir.tensors.get(canonical_post_output_name, None)
            if canonical_tensor is not None and old_mul_tensor is not None:
                canonical_tensor.dtype = str(old_mul_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                if old_mul_shape is not None:
                    canonical_tensor.shape = [int(v) for v in list(old_mul_shape)]
                if old_mul_signature is not None:
                    canonical_tensor.shape_signature = [int(v) for v in list(old_mul_signature)]
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = {
                int(short_match["transpose_idx"]),
                int(point_match["transpose_idx"]),
                int(mean_post_idx),
                int(pre_gate_idx),
                *[int(v) for v in post_indices],
            }
            if use_add_main_path and main_pre_idx is not None:
                remove_indices.add(int(main_pre_idx))
            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_csp_attention_nhwc_chains": int(rewritten)}


def _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Propagate NHWC across SiNet-like SA/PA attention branches with MIRROR_PAD.

    This pass targets the strict motif:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw -> (MEAN, REDUCE_MAX, RESHAPE[unsqueeze0])
      MEAN/REDUCE_MAX -> CONCAT(axis=1) -> MIRROR_PAD -> T(0,2,3,1) -> CONV -> RESHAPE
      CA branch output_nhwc -> T(0,3,1,2)
      ADD(sa_nchw, ca_nchw) -> RESHAPE[unsqueeze1] -> CONCAT(axis=2) -> RESHAPE -> MIRROR_PAD
      -> T(0,2,3,1) -> CONV -> LOGISTIC -> MUL(x_nhwc, gate)

    Rewrite:
      - keep SA/PA branches in NHWC (rewrite axes/pads/reshape-shapes/concat-axis)
      - remove bridge transposes and redundant SA reshape
      - keep output tensor names stable for downstream users
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    # For PA unsqueeze/concat branch: [N,C,K,H,W] -> [N,H,W,C,K]
    # This preserves channel-interleaving order when reshaping [C,2] -> [2C].
    perm5_nckhw_to_nhwck = [0, 3, 4, 1, 2]

    def _normalize_axis(axis: int, rank: int) -> Optional[int]:
        value = int(axis)
        if value < 0:
            value += int(rank)
        if value < 0 or value >= int(rank):
            return None
        return int(value)

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 0
        while candidate in model_ir.tensors:
            suffix += 1
            candidate = f"{base}_{suffix}"
        return str(candidate)

    def _rewrite_reduce_axes_nchw_to_nhwc(op: OperatorIR) -> bool:
        if len(op.inputs) < 2:
            return False
        axes_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
        axes_vals = _read_const_ints_from_tensor(axes_tensor)
        if axes_vals is None or len(axes_vals) == 0:
            return False
        mapped_axes: List[int] = []
        for axis in axes_vals:
            normalized = _normalize_axis(int(axis), 4)
            if normalized is None:
                return False
            mapped_axes.append(int(perm_nhwc_to_nchw[int(normalized)]))
        _write_const_ints_to_tensor(axes_tensor, [int(v) for v in mapped_axes])
        return True

    def _rewrite_pad_pairs_nchw_to_nhwc(op: OperatorIR) -> bool:
        if len(op.inputs) < 2:
            return False
        pads_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
        if pads_tensor is None or pads_tensor.data is None:
            return False
        try:
            pads_array = np.asarray(pads_tensor.data)
            pads_pairs = np.asarray(pads_array).reshape(4, 2)
        except Exception:
            return False
        if int(pads_pairs.size) != 8:
            return False
        # NCHW [N,C,H,W] -> NHWC [N,H,W,C]
        pads_nhwc = np.asarray(
            [
                pads_pairs[0],
                pads_pairs[2],
                pads_pairs[3],
                pads_pairs[1],
            ],
            dtype=pads_array.dtype,
        )
        pads_tensor.data = np.asarray(pads_nhwc)
        pads_tensor.shape = [4, 2]
        pads_tensor.shape_signature = [4, 2]
        return True

    def _rewrite_reshape_shape_by_perm(
        *,
        reshape_op: OperatorIR,
        index_perm: List[int],
    ) -> bool:
        if len(reshape_op.inputs) < 2:
            return False
        shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
        shape_vals = _read_const_ints_from_tensor(shape_tensor)
        if shape_vals is None or len(shape_vals) != len(index_perm):
            return False
        mapped_shape = [int(shape_vals[int(idx)]) for idx in index_perm]
        _write_const_ints_to_tensor(shape_tensor, mapped_shape)
        if isinstance(reshape_op.options, dict):
            opts = dict(reshape_op.options)
            opts_changed = False
            for key in ["newShape", "onnxRawNewShape"]:
                value = opts.get(key, None)
                if not isinstance(value, list) or len(value) != len(index_perm):
                    continue
                mapped_value = [int(value[int(idx)]) for idx in index_perm]
                if mapped_value != [int(v) for v in value]:
                    opts[key] = [int(v) for v in mapped_value]
                    opts_changed = True
            if opts_changed:
                reshape_op.options = opts
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            source_nhwc_name = str(pre_op.inputs[0])
            source_nchw_name = str(pre_op.outputs[0])
            if source_nchw_name in model_outputs:
                continue

            source_users = [int(v) for v in consumers.get(source_nchw_name, [])]
            if len(source_users) < 3 or len(source_users) > 4:
                continue

            mean_idx: Optional[int] = None
            max_idx: Optional[int] = None
            unsq0_idx: Optional[int] = None
            optional_source_mul_idx: Optional[int] = None
            invalid_source_fanout = False
            for user_idx in source_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == source_nchw_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean_idx = int(user_idx)
                    continue
                if (
                    str(user_op.op_type) == "REDUCE_MAX"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == source_nchw_name
                    and bool(user_op.options.get("keepDims", False))
                ):
                    max_idx = int(user_idx)
                    continue
                if (
                    str(user_op.op_type) == "RESHAPE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == source_nchw_name
                ):
                    unsq0_idx = int(user_idx)
                    continue
                if optional_source_mul_idx is not None:
                    invalid_source_fanout = True
                    break
                optional_source_mul_idx = int(user_idx)
            if invalid_source_fanout:
                continue
            if mean_idx is None or max_idx is None or unsq0_idx is None:
                continue

            mean_op = model_ir.operators[int(mean_idx)]
            max_op = model_ir.operators[int(max_idx)]
            unsq0_op = model_ir.operators[int(unsq0_idx)]

            mean_out_name = str(mean_op.outputs[0])
            max_out_name = str(max_op.outputs[0])
            unsq0_out_name = str(unsq0_op.outputs[0])
            if (
                mean_out_name in model_outputs
                or max_out_name in model_outputs
                or unsq0_out_name in model_outputs
            ):
                continue

            if len(consumers.get(mean_out_name, [])) != 1 or len(consumers.get(max_out_name, [])) != 1:
                continue
            concat_sa_idx = int(consumers[mean_out_name][0])
            if int(consumers[max_out_name][0]) != int(concat_sa_idx):
                continue
            concat_sa_op = model_ir.operators[int(concat_sa_idx)]
            if (
                str(concat_sa_op.op_type) != "CONCATENATION"
                or len(concat_sa_op.inputs) != 2
                or len(concat_sa_op.outputs) != 1
            ):
                continue
            concat_sa_axis = int(concat_sa_op.options.get("axis", 1))
            if concat_sa_axis < 0:
                concat_sa_axis += 4
            if concat_sa_axis != 1:
                continue

            concat_sa_out_name = str(concat_sa_op.outputs[0])
            concat_sa_users = [int(v) for v in consumers.get(concat_sa_out_name, [])]
            if len(concat_sa_users) != 1:
                continue
            mirror_sa_idx = int(concat_sa_users[0])
            mirror_sa_op = model_ir.operators[int(mirror_sa_idx)]
            if (
                str(mirror_sa_op.op_type) != "MIRROR_PAD"
                or len(mirror_sa_op.inputs) < 2
                or len(mirror_sa_op.outputs) != 1
                or str(mirror_sa_op.inputs[0]) != concat_sa_out_name
            ):
                continue

            mirror_sa_out_name = str(mirror_sa_op.outputs[0])
            mirror_sa_users = [int(v) for v in consumers.get(mirror_sa_out_name, [])]
            if len(mirror_sa_users) != 1:
                continue
            trans_sa_idx = int(mirror_sa_users[0])
            trans_sa_op = model_ir.operators[int(trans_sa_idx)]
            if (
                str(trans_sa_op.op_type) != "TRANSPOSE"
                or len(trans_sa_op.inputs) < 2
                or len(trans_sa_op.outputs) != 1
                or str(trans_sa_op.inputs[0]) != mirror_sa_out_name
                or _read_transpose_perm(model_ir, trans_sa_op) != perm_nchw_to_nhwc
            ):
                continue

            sa_conv_in_name = str(trans_sa_op.outputs[0])
            sa_conv_users = [int(v) for v in consumers.get(sa_conv_in_name, [])]
            if len(sa_conv_users) != 1:
                continue
            sa_conv_idx = int(sa_conv_users[0])
            sa_conv_op = model_ir.operators[int(sa_conv_idx)]
            if (
                str(sa_conv_op.op_type) != "CONV_2D"
                or len(sa_conv_op.inputs) < 1
                or len(sa_conv_op.outputs) != 1
                or str(sa_conv_op.inputs[0]) != sa_conv_in_name
            ):
                continue

            sa_conv_out_name = str(sa_conv_op.outputs[0])
            sa_conv_out_users = [int(v) for v in consumers.get(sa_conv_out_name, [])]
            if len(sa_conv_out_users) != 1:
                continue
            sa_reshape_idx = int(sa_conv_out_users[0])
            sa_reshape_op = model_ir.operators[int(sa_reshape_idx)]
            if (
                str(sa_reshape_op.op_type) != "RESHAPE"
                or len(sa_reshape_op.inputs) < 1
                or len(sa_reshape_op.outputs) != 1
                or str(sa_reshape_op.inputs[0]) != sa_conv_out_name
            ):
                continue
            sa_nchw_name = str(sa_reshape_op.outputs[0])
            if sa_nchw_name in model_outputs:
                continue

            sa_add_users = [int(v) for v in consumers.get(sa_nchw_name, [])]
            if len(sa_add_users) != 1:
                continue
            add_idx = int(sa_add_users[0])
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_inputs = [str(v) for v in list(add_op.inputs)]
            if sa_nchw_name not in add_inputs:
                continue
            ca_nchw_name = (
                str(add_inputs[1]) if str(add_inputs[0]) == sa_nchw_name else str(add_inputs[0])
            )

            trans_ca_idx = producers.get(ca_nchw_name, None)
            if trans_ca_idx is None:
                continue
            trans_ca_op = model_ir.operators[int(trans_ca_idx)]
            if (
                str(trans_ca_op.op_type) != "TRANSPOSE"
                or len(trans_ca_op.inputs) < 2
                or len(trans_ca_op.outputs) != 1
                or str(trans_ca_op.outputs[0]) != ca_nchw_name
                or _read_transpose_perm(model_ir, trans_ca_op) != perm_nhwc_to_nchw
            ):
                continue
            ca_nhwc_name = str(trans_ca_op.inputs[0])
            if ca_nhwc_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(ca_nchw_name, [])) != {int(add_idx)}:
                continue

            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue
            add_out_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_out_users) != 1:
                continue
            unsq1_idx = int(add_out_users[0])
            unsq1_op = model_ir.operators[int(unsq1_idx)]
            if (
                str(unsq1_op.op_type) != "RESHAPE"
                or len(unsq1_op.inputs) < 2
                or len(unsq1_op.outputs) != 1
                or str(unsq1_op.inputs[0]) != add_out_name
            ):
                continue

            unsq1_out_name = str(unsq1_op.outputs[0])
            if unsq1_out_name in model_outputs:
                continue
            if len(consumers.get(unsq0_out_name, [])) != 1 or len(consumers.get(unsq1_out_name, [])) != 1:
                continue
            concat_pa_idx = int(consumers[unsq0_out_name][0])
            if int(consumers[unsq1_out_name][0]) != int(concat_pa_idx):
                continue
            concat_pa_op = model_ir.operators[int(concat_pa_idx)]
            if (
                str(concat_pa_op.op_type) != "CONCATENATION"
                or len(concat_pa_op.inputs) != 2
                or len(concat_pa_op.outputs) != 1
            ):
                continue
            concat_pa_axis = int(concat_pa_op.options.get("axis", 2))
            if concat_pa_axis < 0:
                concat_pa_axis += 5
            if concat_pa_axis != 2:
                continue
            concat_pa_inputs = [str(v) for v in list(concat_pa_op.inputs)]
            if set(concat_pa_inputs) != {str(unsq0_out_name), str(unsq1_out_name)}:
                continue

            concat_pa_out_name = str(concat_pa_op.outputs[0])
            if concat_pa_out_name in model_outputs:
                continue
            concat_pa_users = [int(v) for v in consumers.get(concat_pa_out_name, [])]
            if len(concat_pa_users) != 1:
                continue
            reshape_pa_idx = int(concat_pa_users[0])
            reshape_pa_op = model_ir.operators[int(reshape_pa_idx)]
            if (
                str(reshape_pa_op.op_type) != "RESHAPE"
                or len(reshape_pa_op.inputs) < 2
                or len(reshape_pa_op.outputs) != 1
                or str(reshape_pa_op.inputs[0]) != concat_pa_out_name
            ):
                continue

            reshape_pa_out_name = str(reshape_pa_op.outputs[0])
            if reshape_pa_out_name in model_outputs:
                continue
            reshape_pa_users = [int(v) for v in consumers.get(reshape_pa_out_name, [])]
            if len(reshape_pa_users) != 1:
                continue
            mirror_pa_idx = int(reshape_pa_users[0])
            mirror_pa_op = model_ir.operators[int(mirror_pa_idx)]
            if (
                str(mirror_pa_op.op_type) != "MIRROR_PAD"
                or len(mirror_pa_op.inputs) < 2
                or len(mirror_pa_op.outputs) != 1
                or str(mirror_pa_op.inputs[0]) != reshape_pa_out_name
            ):
                continue

            mirror_pa_out_name = str(mirror_pa_op.outputs[0])
            mirror_pa_users = [int(v) for v in consumers.get(mirror_pa_out_name, [])]
            if len(mirror_pa_users) != 1:
                continue
            trans_pa_idx = int(mirror_pa_users[0])
            trans_pa_op = model_ir.operators[int(trans_pa_idx)]
            if (
                str(trans_pa_op.op_type) != "TRANSPOSE"
                or len(trans_pa_op.inputs) < 2
                or len(trans_pa_op.outputs) != 1
                or str(trans_pa_op.inputs[0]) != mirror_pa_out_name
                or _read_transpose_perm(model_ir, trans_pa_op) != perm_nchw_to_nhwc
            ):
                continue

            pa_conv_in_name = str(trans_pa_op.outputs[0])
            pa_conv_users = [int(v) for v in consumers.get(pa_conv_in_name, [])]
            if len(pa_conv_users) != 1:
                continue
            pa_conv_idx = int(pa_conv_users[0])
            pa_conv_op = model_ir.operators[int(pa_conv_idx)]
            if (
                str(pa_conv_op.op_type) != "CONV_2D"
                or len(pa_conv_op.inputs) < 1
                or len(pa_conv_op.outputs) != 1
                or str(pa_conv_op.inputs[0]) != pa_conv_in_name
            ):
                continue

            pa_conv_out_name = str(pa_conv_op.outputs[0])
            pa_conv_out_users = [int(v) for v in consumers.get(pa_conv_out_name, [])]
            if len(pa_conv_out_users) != 1:
                continue
            gate_pre_idx: Optional[int] = None
            gate_idx = int(pa_conv_out_users[0])
            gate_op = model_ir.operators[int(gate_idx)]
            gate_input_name = str(pa_conv_out_name)
            if str(gate_op.op_type) == "TRANSPOSE":
                if (
                    len(gate_op.inputs) < 2
                    or len(gate_op.outputs) != 1
                    or str(gate_op.inputs[0]) != pa_conv_out_name
                    or _read_transpose_perm(model_ir, gate_op) != perm_nhwc_to_nchw
                ):
                    continue
                gate_pre_idx = int(gate_idx)
                gate_input_name = str(gate_op.outputs[0])
                if gate_input_name in model_outputs:
                    continue
                gate_pre_users = [int(v) for v in consumers.get(gate_input_name, [])]
                if len(gate_pre_users) != 1:
                    continue
                gate_idx = int(gate_pre_users[0])
                gate_op = model_ir.operators[int(gate_idx)]
            if (
                str(gate_op.op_type) != "LOGISTIC"
                or len(gate_op.inputs) != 1
                or len(gate_op.outputs) != 1
                or str(gate_op.inputs[0]) != gate_input_name
            ):
                continue

            gate_out_name = str(gate_op.outputs[0])
            gate_out_users = [int(v) for v in consumers.get(gate_out_name, [])]
            if len(gate_out_users) != 1:
                continue
            mul_idx = int(gate_out_users[0])
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            if str(gate_out_name) not in mul_inputs:
                continue
            mul_uses_source_nhwc = str(source_nhwc_name) in mul_inputs
            mul_uses_source_nchw = str(source_nchw_name) in mul_inputs
            if (not mul_uses_source_nhwc and not mul_uses_source_nchw) or (
                mul_uses_source_nhwc and mul_uses_source_nchw
            ):
                continue
            if optional_source_mul_idx is None:
                if mul_uses_source_nchw:
                    continue
            else:
                if int(optional_source_mul_idx) != int(mul_idx) or not mul_uses_source_nchw:
                    continue

            if source_nchw_name in model_outputs:
                continue

            # Strict locality guards on branch-only outputs.
            if set(int(v) for v in consumers.get(sa_conv_out_name, [])) != {int(sa_reshape_idx)}:
                continue
            if set(int(v) for v in consumers.get(sa_nchw_name, [])) != {int(add_idx)}:
                continue
            if set(int(v) for v in consumers.get(add_out_name, [])) != {int(unsq1_idx)}:
                continue
            if set(int(v) for v in consumers.get(concat_sa_out_name, [])) != {int(mirror_sa_idx)}:
                continue
            if set(int(v) for v in consumers.get(mirror_sa_out_name, [])) != {int(trans_sa_idx)}:
                continue
            if set(int(v) for v in consumers.get(pa_conv_in_name, [])) != {int(pa_conv_idx)}:
                continue
            if gate_pre_idx is None:
                if set(int(v) for v in consumers.get(pa_conv_out_name, [])) != {int(gate_idx)}:
                    continue
            else:
                if set(int(v) for v in consumers.get(pa_conv_out_name, [])) != {int(gate_pre_idx)}:
                    continue
                if set(int(v) for v in consumers.get(gate_input_name, [])) != {int(gate_idx)}:
                    continue
            if set(int(v) for v in consumers.get(gate_out_name, [])) != {int(mul_idx)}:
                continue

            if not _rewrite_reduce_axes_nchw_to_nhwc(mean_op):
                continue
            if not _rewrite_reduce_axes_nchw_to_nhwc(max_op):
                continue
            if not _rewrite_pad_pairs_nchw_to_nhwc(mirror_sa_op):
                continue
            if not _rewrite_reshape_shape_by_perm(
                reshape_op=unsq0_op,
                index_perm=perm5_nckhw_to_nhwck,
            ):
                continue
            if not _rewrite_reshape_shape_by_perm(
                reshape_op=unsq1_op,
                index_perm=perm5_nckhw_to_nhwck,
            ):
                continue
            if not _rewrite_reshape_shape_by_perm(
                reshape_op=reshape_pa_op,
                index_perm=perm_nchw_to_nhwc,
            ):
                continue
            if not _rewrite_pad_pairs_nchw_to_nhwc(mirror_pa_op):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean_op,
                input_index=0,
                new_input_name=str(source_nhwc_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=max_op,
                input_index=0,
                new_input_name=str(source_nhwc_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=unsq0_op,
                input_index=0,
                new_input_name=str(source_nhwc_name),
            )
            concat_sa_op.options["axis"] = 3
            _replace_operator_input_at(
                model_ir=model_ir,
                op=sa_conv_op,
                input_index=0,
                new_input_name=str(mirror_sa_out_name),
            )

            add_new_inputs = [
                str(sa_conv_out_name) if str(v) == str(sa_nchw_name) else (
                    str(ca_nhwc_name) if str(v) == str(ca_nchw_name) else str(v)
                )
                for v in list(add_inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=add_new_inputs,
            )

            concat_pa_op.options["axis"] = 4
            _replace_operator_input_at(
                model_ir=model_ir,
                op=pa_conv_op,
                input_index=0,
                new_input_name=str(mirror_pa_out_name),
            )
            if gate_pre_idx is not None:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=gate_op,
                    input_index=0,
                    new_input_name=str(pa_conv_out_name),
                )
            if mul_uses_source_nchw:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=[
                        str(source_nhwc_name) if str(v) == str(source_nchw_name) else str(v)
                        for v in list(mul_inputs)
                    ],
                )

            for tensor_name in [
                str(mean_out_name),
                str(max_out_name),
                str(concat_sa_out_name),
                str(mirror_sa_out_name),
                str(add_out_name),
                str(reshape_pa_out_name),
                str(mirror_pa_out_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )
            if gate_pre_idx is not None:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(gate_out_name), None),
                    perm_nchw_to_nhwc,
                )
            for tensor_name in [
                str(unsq0_out_name),
                str(unsq1_out_name),
                str(concat_pa_out_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm5_nckhw_to_nhwck,
                )

            if mul_uses_source_nchw:
                legacy_mul_output_name = str(mul_op.outputs[0])
                legacy_mul_tensor = model_ir.tensors.get(legacy_mul_output_name, None)
                mul_nhwc_output_name = _unique_tensor_name(f"{legacy_mul_output_name}_nhwc")
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_outputs=[str(mul_nhwc_output_name)],
                )
                if legacy_mul_tensor is not None:
                    model_ir.tensors[str(mul_nhwc_output_name)] = TensorIR(
                        name=str(mul_nhwc_output_name),
                        dtype=str(legacy_mul_tensor.dtype),
                        shape=[int(v) for v in list(legacy_mul_tensor.shape)],
                        shape_signature=(
                            [int(v) for v in list(legacy_mul_tensor.shape_signature)]
                            if legacy_mul_tensor.shape_signature is not None
                            else [int(v) for v in list(legacy_mul_tensor.shape)]
                        ),
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(legacy_mul_tensor.quantization),
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(mul_nhwc_output_name), None),
                        perm_nchw_to_nhwc,
                    )
                else:
                    model_ir.tensors[str(mul_nhwc_output_name)] = TensorIR(
                        name=str(mul_nhwc_output_name),
                        dtype="FLOAT32",
                        shape=[1],
                        shape_signature=[1],
                        data=None,
                        is_variable=False,
                        quantization=None,
                    )
                nhwc_to_nchw_perm_name = _unique_tensor_name(
                    f"{legacy_mul_output_name}_perm_nhwc_to_nchw"
                )
                model_ir.tensors[str(nhwc_to_nchw_perm_name)] = TensorIR(
                    name=str(nhwc_to_nchw_perm_name),
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                model_ir.operators.insert(
                    int(mul_idx) + 1,
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(mul_nhwc_output_name), str(nhwc_to_nchw_perm_name)],
                        outputs=[str(legacy_mul_output_name)],
                    ),
                )

            remove_indices = sorted(
                list(
                    {
                        int(pre_idx),
                        int(trans_ca_idx),
                        int(trans_sa_idx),
                        int(sa_reshape_idx),
                        int(trans_pa_idx),
                        *([int(gate_pre_idx)] if gate_pre_idx is not None else []),
                    }
                ),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_sa_pa_mirrorpad_nhwc_propagation_chains": int(rewritten)}


def _optimize_sinet_mix_attention_double_logistic_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW bridge transposes around SiNet mix attention tails (mix1-like).

    Strict target sketch:
      src_nchw = ADD(branch_nchw, residual_nchw)
      (SA/PA attention branches on src_nchw with MIRROR_PAD + TRANSPOSE wrappers)
      gate_nchw = LOGISTIC(LOGISTIC(T(pa_conv_nhwc)))
      mix_mul_nchw = MUL(gate_nchw, branch_nchw)
      mix_add2_nchw = ADD(src_nchw, mix_mul_nchw)
      mix_mul1_nchw = MUL(SUB(const, gate_nchw), residual_nchw)
      mix_add3_nchw = ADD(mix_add2_nchw, mix_mul1_nchw)
      mix_add3_nchw --T(0,2,3,1)--> mix_conv_in_nhwc

      Upstream (strict for this rewrite):
        branch_nchw = T(branch_nhwc)
        residual_nchw = ADD(T(residual_a_nhwc), T(residual_b_nhwc))

    Rewrite:
      - Keep source/attention/tail composition in NHWC.
      - Remove bridge transposes in SA/PA and tail.
      - Remove upstream branch/residual adapters when safely local.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm5_nckhw_to_nhwck = [0, 3, 4, 1, 2]

    def _normalize_axis(axis: int, rank: int) -> Optional[int]:
        value = int(axis)
        if value < 0:
            value += int(rank)
        if value < 0 or value >= int(rank):
            return None
        return int(value)

    def _rewrite_reduce_axes_nchw_to_nhwc(op: OperatorIR) -> bool:
        if len(op.inputs) < 2:
            return False
        axes_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
        axes_vals = _read_const_ints_from_tensor(axes_tensor)
        if axes_vals is None or len(axes_vals) == 0:
            return False
        mapped_axes: List[int] = []
        for axis in axes_vals:
            normalized = _normalize_axis(int(axis), 4)
            if normalized is None:
                return False
            mapped_axes.append(int(perm_nhwc_to_nchw[int(normalized)]))
        return bool(_write_const_ints_to_tensor(axes_tensor, [int(v) for v in mapped_axes]))

    def _rewrite_pad_pairs_nchw_to_nhwc(op: OperatorIR) -> bool:
        if len(op.inputs) < 2:
            return False
        pads_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
        if pads_tensor is None or pads_tensor.data is None:
            return False
        try:
            pads_array = np.asarray(pads_tensor.data)
            pads_pairs = np.asarray(pads_array).reshape(4, 2)
        except Exception:
            return False
        if int(pads_pairs.size) != 8:
            return False
        pads_nhwc = np.asarray(
            [pads_pairs[0], pads_pairs[2], pads_pairs[3], pads_pairs[1]],
            dtype=pads_array.dtype,
        )
        pads_tensor.data = np.asarray(pads_nhwc)
        pads_tensor.shape = [4, 2]
        pads_tensor.shape_signature = [4, 2]
        return True

    def _rewrite_reshape_shape_by_perm(
        *,
        reshape_op: OperatorIR,
        index_perm: List[int],
    ) -> bool:
        if len(reshape_op.inputs) < 2:
            return False
        shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
        shape_vals = _read_const_ints_from_tensor(shape_tensor)
        if shape_vals is None or len(shape_vals) != len(index_perm):
            return False
        mapped_shape = [int(shape_vals[int(idx)]) for idx in index_perm]
        if not _write_const_ints_to_tensor(shape_tensor, mapped_shape):
            return False
        if isinstance(reshape_op.options, dict):
            opts = dict(reshape_op.options)
            opts_changed = False
            for key in ["newShape", "onnxRawNewShape"]:
                value = opts.get(key, None)
                if not isinstance(value, list) or len(value) != len(index_perm):
                    continue
                mapped_value = [int(value[int(idx)]) for idx in index_perm]
                if mapped_value != [int(v) for v in value]:
                    opts[key] = [int(v) for v in mapped_value]
                    opts_changed = True
            if opts_changed:
                reshape_op.options = opts
        return True

    def _match_input_from_producer_type(
        *,
        target_op: OperatorIR,
        producers: Dict[str, int],
        required_type: str,
    ) -> Optional[Tuple[int, int, OperatorIR]]:
        for input_idx, input_name in enumerate(list(target_op.inputs)):
            prod_idx = producers.get(str(input_name), None)
            if prod_idx is None:
                continue
            prod_op = model_ir.operators[int(prod_idx)]
            if str(prod_op.op_type) == str(required_type):
                return int(input_idx), int(prod_idx), prod_op
        return None

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_in_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if post_in_name in model_outputs:
                continue

            post_users = [int(v) for v in consumers.get(post_out_name, [])]
            if len(post_users) != 1:
                continue
            post_conv_idx = int(post_users[0])
            post_conv_op = model_ir.operators[int(post_conv_idx)]
            if (
                str(post_conv_op.op_type) != "CONV_2D"
                or len(post_conv_op.inputs) < 1
                or str(post_conv_op.inputs[0]) != post_out_name
            ):
                continue

            add3_idx = producers.get(post_in_name, None)
            if add3_idx is None:
                continue
            add3_op = model_ir.operators[int(add3_idx)]
            if str(add3_op.op_type) != "ADD" or len(add3_op.inputs) != 2 or len(add3_op.outputs) != 1:
                continue

            add3_left_name = str(add3_op.inputs[0])
            add3_right_name = str(add3_op.inputs[1])
            add2_idx = producers.get(add3_left_name, None)
            mul1_idx = producers.get(add3_right_name, None)
            if add2_idx is None or mul1_idx is None:
                add2_idx = producers.get(add3_right_name, None)
                mul1_idx = producers.get(add3_left_name, None)
                if add2_idx is None or mul1_idx is None:
                    continue
                add3_left_name, add3_right_name = add3_right_name, add3_left_name
            add2_op = model_ir.operators[int(add2_idx)]
            mul1_op = model_ir.operators[int(mul1_idx)]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            if str(mul1_op.op_type) != "MUL" or len(mul1_op.inputs) != 2 or len(mul1_op.outputs) != 1:
                continue
            if str(add2_op.outputs[0]) != str(add3_left_name) or str(mul1_op.outputs[0]) != str(add3_right_name):
                continue

            # add2 = ADD(source, mul0)
            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])
            add2_prod0 = producers.get(add2_in0, None)
            add2_prod1 = producers.get(add2_in1, None)
            if add2_prod0 is None or add2_prod1 is None:
                continue
            if str(model_ir.operators[int(add2_prod0)].op_type) == "ADD" and str(model_ir.operators[int(add2_prod1)].op_type) == "MUL":
                source_idx = int(add2_prod0)
                mul0_idx = int(add2_prod1)
                source_name = str(add2_in0)
                mul0_name = str(add2_in1)
            elif str(model_ir.operators[int(add2_prod1)].op_type) == "ADD" and str(model_ir.operators[int(add2_prod0)].op_type) == "MUL":
                source_idx = int(add2_prod1)
                mul0_idx = int(add2_prod0)
                source_name = str(add2_in1)
                mul0_name = str(add2_in0)
            else:
                continue
            source_op = model_ir.operators[int(source_idx)]
            mul0_op = model_ir.operators[int(mul0_idx)]
            if len(source_op.inputs) != 2 or len(source_op.outputs) != 1:
                continue
            if str(source_op.outputs[0]) != str(source_name):
                continue
            if len(mul0_op.inputs) != 2 or len(mul0_op.outputs) != 1:
                continue
            if str(mul0_op.outputs[0]) != str(mul0_name):
                continue

            source_users = set(int(v) for v in consumers.get(source_name, []))
            if int(add2_idx) not in source_users:
                continue

            # source inputs
            source_in0 = str(source_op.inputs[0])
            source_in1 = str(source_op.inputs[1])
            if source_in0 in model_outputs or source_in1 in model_outputs:
                continue

            # mul1 = MUL(sub, residual)
            mul1_in0 = str(mul1_op.inputs[0])
            mul1_in1 = str(mul1_op.inputs[1])
            sub_idx = producers.get(mul1_in0, None)
            residual_name = str(mul1_in1)
            if sub_idx is None or str(model_ir.operators[int(sub_idx)].op_type) != "SUB":
                sub_idx = producers.get(mul1_in1, None)
                residual_name = str(mul1_in0)
                if sub_idx is None or str(model_ir.operators[int(sub_idx)].op_type) != "SUB":
                    continue
            sub_op = model_ir.operators[int(sub_idx)]
            if len(sub_op.inputs) != 2 or len(sub_op.outputs) != 1:
                continue
            sub_out_name = str(sub_op.outputs[0])
            if sub_out_name not in {mul1_in0, mul1_in1}:
                continue

            # residual must be one of source inputs.
            if residual_name not in {source_in0, source_in1}:
                continue
            branch_nchw_name = str(source_in1) if str(source_in0) == str(residual_name) else str(source_in0)

            # mul0 must use branch tensor and gate2 tensor.
            mul0_inputs = [str(v) for v in list(mul0_op.inputs)]
            if branch_nchw_name not in mul0_inputs:
                continue
            gate2_name = str(mul0_inputs[1]) if str(mul0_inputs[0]) == str(branch_nchw_name) else str(mul0_inputs[0])
            gate2_idx = producers.get(gate2_name, None)
            if gate2_idx is None:
                continue
            gate2_op = model_ir.operators[int(gate2_idx)]
            if str(gate2_op.op_type) != "LOGISTIC" or len(gate2_op.inputs) != 1 or len(gate2_op.outputs) != 1:
                continue
            if str(gate2_op.outputs[0]) != str(gate2_name):
                continue
            gate1_name = str(gate2_op.inputs[0])
            gate1_idx = producers.get(gate1_name, None)
            if gate1_idx is None:
                continue
            gate1_op = model_ir.operators[int(gate1_idx)]
            if str(gate1_op.op_type) != "LOGISTIC" or len(gate1_op.inputs) != 1 or len(gate1_op.outputs) != 1:
                continue
            if str(gate1_op.outputs[0]) != str(gate1_name):
                continue

            # sub must consume gate2
            if str(gate2_name) not in [str(v) for v in list(sub_op.inputs)]:
                continue

            # branch_nchw = TRANSPOSE(branch_nhwc, NHWC->NCHW)
            trans_branch_idx = producers.get(branch_nchw_name, None)
            if trans_branch_idx is None:
                continue
            trans_branch_op = model_ir.operators[int(trans_branch_idx)]
            if (
                str(trans_branch_op.op_type) != "TRANSPOSE"
                or len(trans_branch_op.inputs) < 2
                or len(trans_branch_op.outputs) != 1
                or str(trans_branch_op.outputs[0]) != branch_nchw_name
                or _read_transpose_perm(model_ir, trans_branch_op) != perm_nhwc_to_nchw
            ):
                continue
            branch_nhwc_name = str(trans_branch_op.inputs[0])

            # residual_nchw adapter:
            #   A) residual_nchw = ADD(T(a_nhwc), T(b_nhwc))
            #   B) residual_nchw = T(residual_nhwc)
            residual_idx = producers.get(residual_name, None)
            if residual_idx is None:
                continue
            residual_op = model_ir.operators[int(residual_idx)]
            residual_add_mode = False
            residual_a_nhwc_name: Optional[str] = None
            residual_b_nhwc_name: Optional[str] = None
            residual_a_nchw_name: Optional[str] = None
            residual_b_nchw_name: Optional[str] = None
            trans_res_a_idx: Optional[int] = None
            trans_res_b_idx: Optional[int] = None
            trans_res_direct_idx: Optional[int] = None
            residual_direct_nhwc_name: Optional[str] = None
            if str(residual_op.op_type) == "ADD" and len(residual_op.inputs) == 2 and len(residual_op.outputs) == 1:
                if str(residual_op.outputs[0]) != str(residual_name):
                    continue
                trans_res_a_idx = producers.get(str(residual_op.inputs[0]), None)
                trans_res_b_idx = producers.get(str(residual_op.inputs[1]), None)
                if trans_res_a_idx is None or trans_res_b_idx is None:
                    continue
                trans_res_a_op = model_ir.operators[int(trans_res_a_idx)]
                trans_res_b_op = model_ir.operators[int(trans_res_b_idx)]
                if (
                    str(trans_res_a_op.op_type) != "TRANSPOSE"
                    or str(trans_res_b_op.op_type) != "TRANSPOSE"
                    or len(trans_res_a_op.inputs) < 2
                    or len(trans_res_b_op.inputs) < 2
                    or len(trans_res_a_op.outputs) != 1
                    or len(trans_res_b_op.outputs) != 1
                    or _read_transpose_perm(model_ir, trans_res_a_op) != perm_nhwc_to_nchw
                    or _read_transpose_perm(model_ir, trans_res_b_op) != perm_nhwc_to_nchw
                ):
                    continue
                residual_add_mode = True
                residual_a_nhwc_name = str(trans_res_a_op.inputs[0])
                residual_b_nhwc_name = str(trans_res_b_op.inputs[0])
                residual_a_nchw_name = str(trans_res_a_op.outputs[0])
                residual_b_nchw_name = str(trans_res_b_op.outputs[0])
            elif (
                str(residual_op.op_type) == "TRANSPOSE"
                and len(residual_op.inputs) >= 2
                and len(residual_op.outputs) == 1
                and str(residual_op.outputs[0]) == str(residual_name)
                and _read_transpose_perm(model_ir, residual_op) == perm_nhwc_to_nchw
            ):
                trans_res_direct_idx = int(residual_idx)
                residual_direct_nhwc_name = str(residual_op.inputs[0])
            else:
                continue

            # Locality guards for upstream adapters.
            if set(int(v) for v in consumers.get(branch_nchw_name, [])) != {int(source_idx), int(mul0_idx)}:
                continue
            if set(int(v) for v in consumers.get(residual_name, [])) != {int(source_idx), int(mul1_idx)}:
                continue
            if residual_add_mode:
                if residual_a_nchw_name is None or residual_b_nchw_name is None:
                    continue
                if set(int(v) for v in consumers.get(residual_a_nchw_name, [])) != {int(residual_idx)}:
                    continue
                if set(int(v) for v in consumers.get(residual_b_nchw_name, [])) != {int(residual_idx)}:
                    continue

            # Parse SA/PA attention branches on source tensor.
            mean_from_source: List[int] = []
            max_idx: Optional[int] = None
            unsq0_idx: Optional[int] = None
            for user_idx in list(source_users):
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "MEAN"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(source_name)
                    and bool(user_op.options.get("keepDims", False))
                ):
                    mean_from_source.append(int(user_idx))
                    continue
                if (
                    str(user_op.op_type) == "REDUCE_MAX"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(source_name)
                    and bool(user_op.options.get("keepDims", False))
                ):
                    max_idx = int(user_idx)
                    continue
                if (
                    str(user_op.op_type) == "RESHAPE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(source_name)
                ):
                    unsq0_idx = int(user_idx)
                    continue
            if len(mean_from_source) != 2 or max_idx is None or unsq0_idx is None:
                continue

            # mean_ca path: MEAN -> T(NCHW->NHWC) -> CONV -> CONV -> T(NHWC->NCHW)
            mean_ca_idx: Optional[int] = None
            mean_sa_idx: Optional[int] = None
            trans_ca0_idx: Optional[int] = None
            conv_ca0_idx: Optional[int] = None
            conv_ca2_idx: Optional[int] = None
            trans_ca2_idx: Optional[int] = None
            ca_nchw_name: Optional[str] = None
            ca_conv2_out_nhwc_name: Optional[str] = None
            for mean_idx in list(mean_from_source):
                mean_op = model_ir.operators[int(mean_idx)]
                mean_out = str(mean_op.outputs[0])
                mean_users = [int(v) for v in consumers.get(mean_out, [])]
                if len(mean_users) == 1:
                    t_idx = int(mean_users[0])
                    t_op = model_ir.operators[int(t_idx)]
                    if (
                        str(t_op.op_type) == "TRANSPOSE"
                        and len(t_op.inputs) >= 2
                        and len(t_op.outputs) == 1
                        and str(t_op.inputs[0]) == str(mean_out)
                        and _read_transpose_perm(model_ir, t_op) == perm_nchw_to_nhwc
                    ):
                        ca0_in = str(t_op.outputs[0])
                        ca0_users = [int(v) for v in consumers.get(ca0_in, [])]
                        if len(ca0_users) != 1:
                            continue
                        c0_idx = int(ca0_users[0])
                        c0_op = model_ir.operators[int(c0_idx)]
                        if (
                            str(c0_op.op_type) != "CONV_2D"
                            or len(c0_op.inputs) < 1
                            or len(c0_op.outputs) != 1
                            or str(c0_op.inputs[0]) != str(ca0_in)
                        ):
                            continue
                        c0_out = str(c0_op.outputs[0])
                        c0_users = [int(v) for v in consumers.get(c0_out, [])]
                        if len(c0_users) != 1:
                            continue
                        c2_idx = int(c0_users[0])
                        c2_op = model_ir.operators[int(c2_idx)]
                        if (
                            str(c2_op.op_type) != "CONV_2D"
                            or len(c2_op.inputs) < 1
                            or len(c2_op.outputs) != 1
                            or str(c2_op.inputs[0]) != str(c0_out)
                        ):
                            continue
                        c2_out = str(c2_op.outputs[0])
                        c2_users = [int(v) for v in consumers.get(c2_out, [])]
                        if len(c2_users) != 1:
                            continue
                        t2_idx = int(c2_users[0])
                        t2_op = model_ir.operators[int(t2_idx)]
                        if (
                            str(t2_op.op_type) != "TRANSPOSE"
                            or len(t2_op.inputs) < 2
                            or len(t2_op.outputs) != 1
                            or str(t2_op.inputs[0]) != str(c2_out)
                            or _read_transpose_perm(model_ir, t2_op) != perm_nhwc_to_nchw
                        ):
                            continue
                        mean_ca_idx = int(mean_idx)
                        trans_ca0_idx = int(t_idx)
                        conv_ca0_idx = int(c0_idx)
                        conv_ca2_idx = int(c2_idx)
                        trans_ca2_idx = int(t2_idx)
                        ca_nchw_name = str(t2_op.outputs[0])
                        ca_conv2_out_nhwc_name = str(c2_out)
                    else:
                        mean_sa_idx = int(mean_idx)
                else:
                    mean_sa_idx = int(mean_idx)
            if (
                mean_ca_idx is None
                or mean_sa_idx is None
                or trans_ca0_idx is None
                or conv_ca0_idx is None
                or conv_ca2_idx is None
                or trans_ca2_idx is None
                or ca_nchw_name is None
                or ca_conv2_out_nhwc_name is None
            ):
                continue

            mean_sa_op = model_ir.operators[int(mean_sa_idx)]
            max_op = model_ir.operators[int(max_idx)]
            unsq0_op = model_ir.operators[int(unsq0_idx)]
            mean_sa_out_name = str(mean_sa_op.outputs[0])
            max_out_name = str(max_op.outputs[0])
            unsq0_out_name = str(unsq0_op.outputs[0])

            # SA concat
            if len(consumers.get(mean_sa_out_name, [])) != 1 or len(consumers.get(max_out_name, [])) != 1:
                continue
            concat_sa_idx = int(consumers[mean_sa_out_name][0])
            if int(consumers[max_out_name][0]) != int(concat_sa_idx):
                continue
            concat_sa_op = model_ir.operators[int(concat_sa_idx)]
            if str(concat_sa_op.op_type) != "CONCATENATION" or len(concat_sa_op.inputs) != 2 or len(concat_sa_op.outputs) != 1:
                continue
            concat_sa_axis = int(concat_sa_op.options.get("axis", 1))
            if concat_sa_axis < 0:
                concat_sa_axis += 4
            if concat_sa_axis != 1:
                continue
            concat_sa_out_name = str(concat_sa_op.outputs[0])
            if len(consumers.get(concat_sa_out_name, [])) != 1:
                continue
            mirror_sa_idx = int(consumers[concat_sa_out_name][0])
            mirror_sa_op = model_ir.operators[int(mirror_sa_idx)]
            if str(mirror_sa_op.op_type) != "MIRROR_PAD" or len(mirror_sa_op.inputs) < 2 or len(mirror_sa_op.outputs) != 1:
                continue
            mirror_sa_out_name = str(mirror_sa_op.outputs[0])
            if str(mirror_sa_op.inputs[0]) != str(concat_sa_out_name):
                continue
            if len(consumers.get(mirror_sa_out_name, [])) != 1:
                continue
            trans_sa_idx = int(consumers[mirror_sa_out_name][0])
            trans_sa_op = model_ir.operators[int(trans_sa_idx)]
            if (
                str(trans_sa_op.op_type) != "TRANSPOSE"
                or len(trans_sa_op.inputs) < 2
                or len(trans_sa_op.outputs) != 1
                or str(trans_sa_op.inputs[0]) != str(mirror_sa_out_name)
                or _read_transpose_perm(model_ir, trans_sa_op) != perm_nchw_to_nhwc
            ):
                continue
            sa_conv_in_name = str(trans_sa_op.outputs[0])
            if len(consumers.get(sa_conv_in_name, [])) != 1:
                continue
            sa_conv_idx = int(consumers[sa_conv_in_name][0])
            sa_conv_op = model_ir.operators[int(sa_conv_idx)]
            if str(sa_conv_op.op_type) != "CONV_2D" or len(sa_conv_op.inputs) < 1 or len(sa_conv_op.outputs) != 1:
                continue
            if str(sa_conv_op.inputs[0]) != str(sa_conv_in_name):
                continue
            sa_conv_out_name = str(sa_conv_op.outputs[0])
            if len(consumers.get(sa_conv_out_name, [])) != 1:
                continue
            sa_reshape_idx = int(consumers[sa_conv_out_name][0])
            sa_reshape_op = model_ir.operators[int(sa_reshape_idx)]
            if str(sa_reshape_op.op_type) != "RESHAPE" or len(sa_reshape_op.inputs) < 1 or len(sa_reshape_op.outputs) != 1:
                continue
            if str(sa_reshape_op.inputs[0]) != str(sa_conv_out_name):
                continue
            sa_nchw_name = str(sa_reshape_op.outputs[0])

            # add_attn
            if len(consumers.get(sa_nchw_name, [])) != 1:
                continue
            add_attn_idx = int(consumers[sa_nchw_name][0])
            add_attn_op = model_ir.operators[int(add_attn_idx)]
            if str(add_attn_op.op_type) != "ADD" or len(add_attn_op.inputs) != 2 or len(add_attn_op.outputs) != 1:
                continue
            add_attn_inputs = [str(v) for v in list(add_attn_op.inputs)]
            if str(ca_nchw_name) not in add_attn_inputs or str(sa_nchw_name) not in add_attn_inputs:
                continue
            add_attn_out_name = str(add_attn_op.outputs[0])

            if len(consumers.get(add_attn_out_name, [])) != 1:
                continue
            unsq1_idx = int(consumers[add_attn_out_name][0])
            unsq1_op = model_ir.operators[int(unsq1_idx)]
            if str(unsq1_op.op_type) != "RESHAPE" or len(unsq1_op.inputs) < 2 or len(unsq1_op.outputs) != 1:
                continue
            if str(unsq1_op.inputs[0]) != str(add_attn_out_name):
                continue
            unsq1_out_name = str(unsq1_op.outputs[0])

            if len(consumers.get(unsq0_out_name, [])) != 1 or len(consumers.get(unsq1_out_name, [])) != 1:
                continue
            concat_pa_idx = int(consumers[unsq0_out_name][0])
            if int(consumers[unsq1_out_name][0]) != int(concat_pa_idx):
                continue
            concat_pa_op = model_ir.operators[int(concat_pa_idx)]
            if str(concat_pa_op.op_type) != "CONCATENATION" or len(concat_pa_op.inputs) != 2 or len(concat_pa_op.outputs) != 1:
                continue
            concat_pa_axis = int(concat_pa_op.options.get("axis", 2))
            if concat_pa_axis < 0:
                concat_pa_axis += 5
            if concat_pa_axis != 2:
                continue
            concat_pa_out_name = str(concat_pa_op.outputs[0])
            if len(consumers.get(concat_pa_out_name, [])) != 1:
                continue
            reshape_pa_idx = int(consumers[concat_pa_out_name][0])
            reshape_pa_op = model_ir.operators[int(reshape_pa_idx)]
            if str(reshape_pa_op.op_type) != "RESHAPE" or len(reshape_pa_op.inputs) < 2 or len(reshape_pa_op.outputs) != 1:
                continue
            if str(reshape_pa_op.inputs[0]) != str(concat_pa_out_name):
                continue
            reshape_pa_out_name = str(reshape_pa_op.outputs[0])
            if len(consumers.get(reshape_pa_out_name, [])) != 1:
                continue
            mirror_pa_idx = int(consumers[reshape_pa_out_name][0])
            mirror_pa_op = model_ir.operators[int(mirror_pa_idx)]
            if str(mirror_pa_op.op_type) != "MIRROR_PAD" or len(mirror_pa_op.inputs) < 2 or len(mirror_pa_op.outputs) != 1:
                continue
            if str(mirror_pa_op.inputs[0]) != str(reshape_pa_out_name):
                continue
            mirror_pa_out_name = str(mirror_pa_op.outputs[0])
            if len(consumers.get(mirror_pa_out_name, [])) != 1:
                continue
            trans_pa_idx = int(consumers[mirror_pa_out_name][0])
            trans_pa_op = model_ir.operators[int(trans_pa_idx)]
            if (
                str(trans_pa_op.op_type) != "TRANSPOSE"
                or len(trans_pa_op.inputs) < 2
                or len(trans_pa_op.outputs) != 1
                or str(trans_pa_op.inputs[0]) != str(mirror_pa_out_name)
                or _read_transpose_perm(model_ir, trans_pa_op) != perm_nchw_to_nhwc
            ):
                continue
            pa_conv_in_name = str(trans_pa_op.outputs[0])
            pa_conv_users = [int(v) for v in consumers.get(pa_conv_in_name, [])]
            if len(pa_conv_users) != 1:
                continue
            pa_conv_idx = int(pa_conv_users[0])
            pa_conv_op = model_ir.operators[int(pa_conv_idx)]
            if str(pa_conv_op.op_type) != "CONV_2D" or len(pa_conv_op.inputs) < 1 or len(pa_conv_op.outputs) != 1:
                continue
            if str(pa_conv_op.inputs[0]) != str(pa_conv_in_name):
                continue
            pa_conv_out_nhwc_name = str(pa_conv_op.outputs[0])
            pa_conv_out_nchw_name = str(gate1_op.inputs[0])
            trans_pa_out_idx = producers.get(pa_conv_out_nchw_name, None)
            if trans_pa_out_idx is None or int(trans_pa_out_idx) != int(producers.get(pa_conv_out_nchw_name, -1)):
                continue
            trans_pa_out_op = model_ir.operators[int(trans_pa_out_idx)]
            if (
                str(trans_pa_out_op.op_type) != "TRANSPOSE"
                or len(trans_pa_out_op.inputs) < 2
                or len(trans_pa_out_op.outputs) != 1
                or str(trans_pa_out_op.inputs[0]) != str(pa_conv_out_nhwc_name)
                or str(trans_pa_out_op.outputs[0]) != str(pa_conv_out_nchw_name)
                or _read_transpose_perm(model_ir, trans_pa_out_op) != perm_nhwc_to_nchw
            ):
                continue

            # Locality and output guards.
            if any(
                t in model_outputs
                for t in [
                    str(source_name),
                    str(residual_name),
                    str(branch_nchw_name),
                    str(add_attn_out_name),
                    str(add3_op.outputs[0]),
                    str(post_in_name),
                ]
            ):
                continue
            if set(int(v) for v in consumers.get(post_out_name, [])) != {int(post_conv_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(ca_nchw_name), [])) != {int(add_attn_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(sa_nchw_name), [])) != {int(add_attn_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(pa_conv_out_nchw_name), [])) != {int(gate1_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(gate1_name), [])) != {int(gate2_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(gate2_name), [])) != {int(mul0_idx), int(sub_idx)}:
                continue

            # Rewrite constants.
            mean_ca_op = model_ir.operators[int(mean_ca_idx)]
            if not _rewrite_reduce_axes_nchw_to_nhwc(mean_ca_op):
                continue
            if not _rewrite_reduce_axes_nchw_to_nhwc(mean_sa_op):
                continue
            if not _rewrite_reduce_axes_nchw_to_nhwc(max_op):
                continue
            if not _rewrite_pad_pairs_nchw_to_nhwc(mirror_sa_op):
                continue
            if not _rewrite_reshape_shape_by_perm(reshape_op=unsq0_op, index_perm=perm5_nckhw_to_nhwck):
                continue
            if not _rewrite_reshape_shape_by_perm(reshape_op=unsq1_op, index_perm=perm5_nckhw_to_nhwck):
                continue
            if not _rewrite_reshape_shape_by_perm(reshape_op=reshape_pa_op, index_perm=perm_nchw_to_nhwc):
                continue
            if not _rewrite_pad_pairs_nchw_to_nhwc(mirror_pa_op):
                continue

            # Rewire upstream residual/branch adapters.
            residual_nhwc_name = str(residual_name)
            if residual_add_mode:
                if residual_a_nhwc_name is None or residual_b_nhwc_name is None:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=residual_op,
                    new_inputs=[str(residual_a_nhwc_name), str(residual_b_nhwc_name)],
                )
            else:
                if residual_direct_nhwc_name is None:
                    continue
                residual_nhwc_name = str(residual_direct_nhwc_name)
            source_new_inputs = [
                str(branch_nhwc_name) if str(v) == str(branch_nchw_name) else (
                    str(residual_nhwc_name) if str(v) == str(residual_name) else str(v)
                )
                for v in list(source_op.inputs)
            ]
            _set_operator_inputs(model_ir=model_ir, op=source_op, new_inputs=source_new_inputs)

            mul0_new_inputs = [
                str(branch_nhwc_name) if str(v) == str(branch_nchw_name) else str(v)
                for v in list(mul0_op.inputs)
            ]
            _set_operator_inputs(model_ir=model_ir, op=mul0_op, new_inputs=mul0_new_inputs)
            if str(residual_nhwc_name) != str(residual_name):
                mul1_new_inputs = [
                    str(residual_nhwc_name) if str(v) == str(residual_name) else str(v)
                    for v in list(mul1_op.inputs)
                ]
                _set_operator_inputs(model_ir=model_ir, op=mul1_op, new_inputs=mul1_new_inputs)

            # Rewire CA/SA/PA branches to NHWC.
            _replace_operator_input_at(
                model_ir=model_ir,
                op=model_ir.operators[int(conv_ca0_idx)],
                input_index=0,
                new_input_name=str(mean_ca_op.outputs[0]),
            )
            concat_sa_op.options["axis"] = 3
            _replace_operator_input_at(
                model_ir=model_ir,
                op=sa_conv_op,
                input_index=0,
                new_input_name=str(mirror_sa_out_name),
            )
            add_attn_new_inputs = [
                str(sa_conv_out_name) if str(v) == str(sa_nchw_name) else (
                    str(ca_conv2_out_nhwc_name) if str(v) == str(ca_nchw_name) else str(v)
                )
                for v in list(add_attn_inputs)
            ]
            _set_operator_inputs(model_ir=model_ir, op=add_attn_op, new_inputs=add_attn_new_inputs)
            concat_pa_op.options["axis"] = 4
            _replace_operator_input_at(
                model_ir=model_ir,
                op=pa_conv_op,
                input_index=0,
                new_input_name=str(mirror_pa_out_name),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=gate1_op,
                input_index=0,
                new_input_name=str(pa_conv_out_nhwc_name),
            )

            # Remove terminal mix transpose.
            _replace_operator_input_at(
                model_ir=model_ir,
                op=post_conv_op,
                input_index=0,
                new_input_name=str(post_in_name),
            )

            # Permute metadata for tensors switched from NCHW to NHWC.
            for tensor_name in [
                str(residual_name),
                str(source_name),
                str(mean_ca_op.outputs[0]),
                str(mean_sa_out_name),
                str(max_out_name),
                str(concat_sa_out_name),
                str(mirror_sa_out_name),
                str(add_attn_out_name),
                str(reshape_pa_out_name),
                str(mirror_pa_out_name),
                str(gate1_name),
                str(gate2_name),
                str(mul0_name),
                str(sub_out_name),
                str(add2_op.outputs[0]),
                str(mul1_op.outputs[0]),
                str(add3_op.outputs[0]),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm_nchw_to_nhwc,
                )
            for tensor_name in [
                str(unsq0_out_name),
                str(unsq1_out_name),
                str(concat_pa_out_name),
            ]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(tensor_name), None),
                    perm5_nckhw_to_nhwck,
                )

            remove_index_set = {
                int(trans_branch_idx),
                int(trans_ca0_idx),
                int(trans_ca2_idx),
                int(trans_sa_idx),
                int(sa_reshape_idx),
                int(trans_pa_idx),
                int(trans_pa_out_idx),
                int(post_idx),
            }
            if residual_add_mode:
                if trans_res_a_idx is None or trans_res_b_idx is None:
                    continue
                remove_index_set.add(int(trans_res_a_idx))
                remove_index_set.add(int(trans_res_b_idx))
            else:
                if trans_res_direct_idx is None:
                    continue
                remove_index_set.add(int(trans_res_direct_idx))
            remove_indices = sorted(list(remove_index_set), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_sinet_mix_attention_double_logistic_nhwc_chains": int(optimized)}


def _optimize_transpose_add_concat_const_suffix_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around ADD fan-in concat blocks with const suffixes.

    Target:
      (x_i_nhwc --T(0,3,1,2)--> x_i_nchw)
      (b_nhwc  --T(0,3,1,2)--> b_nchw; shared)
      ADD(x_i_nchw, b_nchw) -> a_i_nchw
      CONCAT(axis=1, [a_i_nchw...]) -> c_nchw
      c_nchw --(MUL|ADD with const)*--> z_nchw
      z_nchw --T(0,2,3,1)--> z_nhwc

    Rewrite:
      ADD(x_i_nhwc, b_nhwc) -> a_i_nhwc
      CONCAT(axis=3, [a_i_nhwc...]) -> c_nhwc
      c_nhwc --(MUL|ADD with const_nhwc)*--> z_nhwc
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue
            concat_axis = int(concat_op.options.get("axis", 1))
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue

            add_rewrites: List[Tuple[int, List[str]]] = []
            branch_pre_remove_indices: List[int] = []
            base_pre_idx: Optional[int] = None
            base_pre_input_name: Optional[str] = None
            rewritable = True

            for add_out_name in concat_inputs:
                add_idx = producers.get(str(add_out_name), None)
                if add_idx is None:
                    rewritable = False
                    break
                add_op = model_ir.operators[int(add_idx)]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    rewritable = False
                    break
                if str(add_op.outputs[0]) != str(add_out_name):
                    rewritable = False
                    break
                if set(int(v) for v in consumers.get(str(add_out_name), [])) != {int(concat_idx)}:
                    rewritable = False
                    break

                add_inputs = [str(v) for v in list(add_op.inputs)]
                transpose_inputs: List[Tuple[int, int, str, bool]] = []
                for input_index, add_input_name in enumerate(add_inputs):
                    pre_idx = producers.get(str(add_input_name), None)
                    if pre_idx is None:
                        continue
                    pre_op = model_ir.operators[int(pre_idx)]
                    if (
                        str(pre_op.op_type) != "TRANSPOSE"
                        or len(pre_op.inputs) < 2
                        or len(pre_op.outputs) != 1
                        or str(pre_op.outputs[0]) != str(add_input_name)
                        or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                        or str(add_input_name) in model_outputs
                    ):
                        continue
                    direct_only = set(int(v) for v in consumers.get(str(add_input_name), [])) == {int(add_idx)}
                    transpose_inputs.append((int(input_index), int(pre_idx), str(pre_op.inputs[0]), bool(direct_only)))

                if len(transpose_inputs) < 2:
                    rewritable = False
                    break

                branch_plan = None
                base_plan = None
                for plan in transpose_inputs:
                    if bool(plan[3]):
                        branch_plan = plan
                        break
                if branch_plan is None:
                    rewritable = False
                    break
                for plan in transpose_inputs:
                    if int(plan[1]) != int(branch_plan[1]):
                        base_plan = plan
                        break
                if base_plan is None:
                    rewritable = False
                    break

                if base_pre_idx is None:
                    base_pre_idx = int(base_plan[1])
                    base_pre_input_name = str(base_plan[2])
                elif int(base_pre_idx) != int(base_plan[1]):
                    rewritable = False
                    break

                new_add_inputs = list(add_inputs)
                new_add_inputs[int(branch_plan[0])] = str(branch_plan[2])
                new_add_inputs[int(base_plan[0])] = str(base_plan[2])
                add_rewrites.append((int(add_idx), [str(v) for v in new_add_inputs]))
                branch_pre_remove_indices.append(int(branch_plan[1]))

            if not rewritable or base_pre_idx is None or base_pre_input_name is None:
                continue

            # Strict suffix: CONCAT -> MUL(const) -> ADD(const) -> TRANSPOSE(NCHW->NHWC)
            concat_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(concat_users) != 1:
                continue
            mul_idx = int(concat_users[0])
            mul_op = model_ir.operators[mul_idx]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_in0 = str(mul_op.inputs[0])
            mul_in1 = str(mul_op.inputs[1])
            if mul_in0 == concat_out_name:
                mul_side_index = 1
                mul_side_name = mul_in1
            elif mul_in1 == concat_out_name:
                mul_side_index = 0
                mul_side_name = mul_in0
            else:
                continue
            mul_side_tensor = model_ir.tensors.get(mul_side_name, None)
            if mul_side_tensor is None or mul_side_tensor.data is None:
                continue
            mul_out_name = str(mul_op.outputs[0])
            if mul_out_name in model_outputs:
                continue

            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_users) != 1:
                continue
            add2_idx = int(mul_users[0])
            add2_op = model_ir.operators[add2_idx]
            if str(add2_op.op_type) != "ADD" or len(add2_op.inputs) != 2 or len(add2_op.outputs) != 1:
                continue
            add2_in0 = str(add2_op.inputs[0])
            add2_in1 = str(add2_op.inputs[1])
            if add2_in0 == mul_out_name:
                add2_side_index = 1
                add2_side_name = add2_in1
            elif add2_in1 == mul_out_name:
                add2_side_index = 0
                add2_side_name = add2_in0
            else:
                continue
            add2_side_tensor = model_ir.tensors.get(add2_side_name, None)
            if add2_side_tensor is None or add2_side_tensor.data is None:
                continue
            add2_out_name = str(add2_op.outputs[0])
            if add2_out_name in model_outputs:
                continue

            add2_users = [int(v) for v in consumers.get(add2_out_name, [])]
            if len(add2_users) != 1:
                continue
            post_idx = int(add2_users[0])
            post_op = model_ir.operators[post_idx]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.inputs[0]) != add2_out_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            post_out_name = str(post_op.outputs[0])
            if post_out_name in model_outputs:
                continue

            # Apply rewrites atomically after full validation.
            for add_idx, new_inputs in add_rewrites:
                add_op = model_ir.operators[int(add_idx)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_inputs=[str(v) for v in new_inputs],
                )
                if len(add_op.outputs) == 1:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(add_op.outputs[0]), None),
                        perm_nchw_to_nhwc,
                    )

            if np.asarray(mul_side_tensor.data).ndim == 4:
                mul_nhwc = np.transpose(np.asarray(mul_side_tensor.data), perm_nchw_to_nhwc).astype(
                    np.asarray(mul_side_tensor.data).dtype, copy=False
                )
                mul_side_tensor.data = np.asarray(mul_nhwc)
                mul_side_tensor.shape = [int(v) for v in list(mul_nhwc.shape)]
                mul_side_tensor.shape_signature = [int(v) for v in list(mul_nhwc.shape)]
            if np.asarray(add2_side_tensor.data).ndim == 4:
                add2_nhwc = np.transpose(np.asarray(add2_side_tensor.data), perm_nchw_to_nhwc).astype(
                    np.asarray(add2_side_tensor.data).dtype, copy=False
                )
                add2_side_tensor.data = np.asarray(add2_nhwc)
                add2_side_tensor.shape = [int(v) for v in list(add2_nhwc.shape)]
                add2_side_tensor.shape_signature = [int(v) for v in list(add2_nhwc.shape)]

            concat_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(concat_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mul_out_name, None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(add2_out_name, None),
                perm_nchw_to_nhwc,
            )

            _set_operator_outputs(
                model_ir=model_ir,
                op=add2_op,
                new_outputs=[post_out_name],
            )

            old_add2_tensor = model_ir.tensors.get(add2_out_name, None)
            post_tensor = model_ir.tensors.get(post_out_name, None)
            if old_add2_tensor is not None and post_tensor is not None:
                post_tensor.dtype = str(old_add2_tensor.dtype)
                post_tensor.quantization = _clone_quantization(old_add2_tensor.quantization)
                post_tensor.shape = [int(v) for v in list(old_add2_tensor.shape)]
                post_tensor.shape_signature = (
                    [int(v) for v in list(old_add2_tensor.shape_signature)]
                    if old_add2_tensor.shape_signature is not None
                    else [int(v) for v in list(old_add2_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = set(int(v) for v in branch_pre_remove_indices)
            remove_indices.add(int(post_idx))
            # Remove base transpose only when it became dead.
            base_nchw_name = str(model_ir.operators[int(base_pre_idx)].outputs[0])
            remaining_users = [int(v) for v in consumers.get(base_nchw_name, [])]
            if set(int(v) for v in remaining_users).issubset(set(int(idx) for idx, _ in add_rewrites)):
                remove_indices.add(int(base_pre_idx))

            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_add_concat_const_suffix_nhwc_chains": int(optimized)}


def _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around dual-MUL concat blocks.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      MUL(x_nchw, c0_nchw) -> m0_nchw
      MUL(x_nchw, c1_nchw) -> m1_nchw
      CONCAT(axis=1, [m0_nchw, m1_nchw]) -> c_nchw
      c_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      MUL(x_nhwc, c0_nhwc) -> m0_nhwc
      MUL(x_nhwc, c1_nhwc) -> m1_nhwc
      CONCAT(axis=3, [m0_nhwc, m1_nhwc]) -> y_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _clone_or_rewrite_mul_const_to_nhwc(
        *,
        mul_idx: int,
        side_input_index: int,
        side_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        chain_indices: set[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        side_tensor = model_ir.tensors.get(str(side_input_name), None)
        if side_tensor is None or side_tensor.data is None:
            return False
        side_data = np.asarray(side_tensor.data)
        if int(side_data.size) == 1:
            return True
        if side_data.ndim != 4:
            return False

        target_shape = (
            [int(v) for v in list(target_shape_nhwc)]
            if _is_fully_known_positive_shape(target_shape_nhwc)
            else None
        )
        nhwc_data: Optional[np.ndarray] = None
        if target_shape is not None:
            side_shape = [int(v) for v in list(side_data.shape)]
            if _broadcast_static_shapes(target_shape, side_shape) is not None:
                nhwc_data = np.asarray(side_data)
            else:
                rotated = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                rotated_shape = [int(v) for v in list(rotated.shape)]
                if _broadcast_static_shapes(target_shape, rotated_shape) is None:
                    return False
                nhwc_data = np.asarray(rotated)
        else:
            side_shape = [int(v) for v in list(side_data.shape)]
            if (
                len(side_shape) == 4
                and int(side_shape[0]) == 1
                and int(side_shape[1]) > 0
                and int(side_shape[2]) == 1
                and int(side_shape[3]) == 1
            ):
                nhwc_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
            else:
                nhwc_data = np.asarray(side_data)
        if nhwc_data is None:
            return False

        side_users = [int(v) for v in consumers.get(str(side_input_name), [])]
        shared_outside_chain = any(int(u) not in chain_indices for u in side_users)
        rewritten_side_name = str(side_input_name)
        if shared_outside_chain:
            rewritten_side_name = _unique_tensor_name(f"{side_input_name}_nhwc")
            model_ir.tensors[rewritten_side_name] = TensorIR(
                name=rewritten_side_name,
                dtype=str(side_tensor.dtype),
                shape=[int(v) for v in list(nhwc_data.shape)],
                shape_signature=[int(v) for v in list(nhwc_data.shape)],
                data=np.asarray(nhwc_data),
                is_variable=False,
                quantization=_clone_quantization(side_tensor.quantization),
            )
        else:
            side_tensor.data = np.asarray(nhwc_data)
            side_tensor.shape = [int(v) for v in list(nhwc_data.shape)]
            side_tensor.shape_signature = [int(v) for v in list(nhwc_data.shape)]

        _replace_operator_input_at(
            model_ir=model_ir,
            op=model_ir.operators[int(mul_idx)],
            input_index=int(side_input_index),
            new_input_name=str(rewritten_side_name),
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue
            concat_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if concat_out_name in model_outputs or post_out_name in model_outputs:
                continue

            concat_idx = producers.get(str(concat_out_name), None)
            if concat_idx is None:
                continue
            concat_op = model_ir.operators[int(concat_idx)]
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.inputs) != 2 or len(concat_op.outputs) != 1:
                continue
            axis = int(concat_op.options.get("axis", 1))
            if axis < 0:
                axis += 4
            if axis != 1:
                continue
            if set(int(v) for v in consumers.get(str(concat_out_name), [])) != {int(post_idx)}:
                continue

            mul_indices: List[int] = []
            mul_plans: List[Tuple[int, int, str, str]] = []
            shared_data_name: Optional[str] = None
            rewritable = True
            for concat_input_name in [str(v) for v in list(concat_op.inputs)]:
                mul_idx = producers.get(str(concat_input_name), None)
                if mul_idx is None:
                    rewritable = False
                    break
                mul_op = model_ir.operators[int(mul_idx)]
                if (
                    str(mul_op.op_type) != "MUL"
                    or len(mul_op.inputs) != 2
                    or len(mul_op.outputs) != 1
                    or str(mul_op.outputs[0]) != str(concat_input_name)
                ):
                    rewritable = False
                    break
                if set(int(v) for v in consumers.get(str(concat_input_name), [])) != {int(concat_idx)}:
                    rewritable = False
                    break
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                data_input_index = None
                side_input_index = None
                if shared_data_name is None:
                    for candidate_data_index in [0, 1]:
                        candidate_side_index = 1 - int(candidate_data_index)
                        candidate_data_name = str(mul_inputs[int(candidate_data_index)])
                        candidate_prod_idx = producers.get(candidate_data_name, None)
                        if candidate_prod_idx is None:
                            continue
                        candidate_prod_op = model_ir.operators[int(candidate_prod_idx)]
                        if (
                            str(candidate_prod_op.op_type) == "TRANSPOSE"
                            and len(candidate_prod_op.inputs) >= 2
                            and len(candidate_prod_op.outputs) == 1
                            and str(candidate_prod_op.outputs[0]) == candidate_data_name
                            and _read_transpose_perm(model_ir, candidate_prod_op) == perm_nhwc_to_nchw
                        ):
                            shared_data_name = str(candidate_data_name)
                            data_input_index = int(candidate_data_index)
                            side_input_index = int(candidate_side_index)
                            break
                else:
                    if mul_inputs[0] == str(shared_data_name):
                        data_input_index = 0
                        side_input_index = 1
                    elif mul_inputs[1] == str(shared_data_name):
                        data_input_index = 1
                        side_input_index = 0
                if data_input_index is None or side_input_index is None:
                    rewritable = False
                    break
                side_input_name = str(mul_inputs[int(side_input_index)])
                side_tensor = model_ir.tensors.get(str(side_input_name), None)
                if side_tensor is None or side_tensor.data is None:
                    rewritable = False
                    break

                mul_indices.append(int(mul_idx))
                mul_plans.append(
                    (
                        int(mul_idx),
                        int(side_input_index),
                        str(side_input_name),
                        str(mul_inputs[int(data_input_index)]),
                    )
                )

            if not rewritable or shared_data_name is None:
                continue

            pre_idx = producers.get(str(shared_data_name), None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if (
                str(pre_op.op_type) != "TRANSPOSE"
                or len(pre_op.inputs) < 2
                or len(pre_op.outputs) != 1
                or str(pre_op.outputs[0]) != str(shared_data_name)
                or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
            ):
                continue
            if shared_data_name in model_outputs or str(pre_op.inputs[0]) in model_outputs:
                continue
            if set(int(v) for v in consumers.get(str(shared_data_name), [])) != set(int(v) for v in mul_indices):
                continue

            chain_indices = {int(pre_idx), int(concat_idx), int(post_idx)} | set(int(v) for v in mul_indices)
            target_shape_nhwc = (
                [int(v) for v in list(model_ir.tensors[str(pre_op.inputs[0])].shape)]
                if str(pre_op.inputs[0]) in model_ir.tensors
                else None
            )
            for mul_idx, side_input_index, side_input_name, _ in mul_plans:
                if not _clone_or_rewrite_mul_const_to_nhwc(
                    mul_idx=int(mul_idx),
                    side_input_index=int(side_input_index),
                    side_input_name=str(side_input_name),
                    target_shape_nhwc=target_shape_nhwc,
                    chain_indices=chain_indices,
                    consumers=consumers,
                ):
                    rewritable = False
                    break
            if not rewritable:
                continue

            for mul_idx in mul_indices:
                mul_op = model_ir.operators[int(mul_idx)]
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                for input_index, input_name in enumerate(list(mul_inputs)):
                    if str(input_name) == str(shared_data_name):
                        mul_inputs[int(input_index)] = str(pre_op.inputs[0])
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(mul_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )

            concat_op.options["axis"] = 3
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[post_out_name],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat_out_name), None),
                perm_nchw_to_nhwc,
            )

            old_concat_tensor = model_ir.tensors.get(str(concat_out_name), None)
            post_out_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_concat_tensor is not None and post_out_tensor is not None:
                post_out_tensor.dtype = str(old_concat_tensor.dtype)
                post_out_tensor.quantization = _clone_quantization(old_concat_tensor.quantization)
                post_out_tensor.shape = [int(v) for v in list(old_concat_tensor.shape)]
                post_out_tensor.shape_signature = (
                    [int(v) for v in list(old_concat_tensor.shape_signature)]
                    if old_concat_tensor.shape_signature is not None
                    else [int(v) for v in list(old_concat_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_out_tensor,
                    perm_nchw_to_nhwc,
                )

            for remove_idx in sorted([int(post_idx), int(pre_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_dual_mul_concat_prepost_nhwc_chains": int(rewritten)}


def _optimize_transpose_axis3_const_concat_bridge_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Collapse NCHW adapters around axis=3 CONCAT with const suffixes to a single NHWC->NCHW bridge.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      const_i_nchw --------------------^
      CONCAT(axis=3, [const..., x_nchw, ...]) -> y_nchw
      y_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc   (at least one branch)
      y_nchw may also have legacy NCHW consumers.

    Rewrite:
      const_i_nchw -> const_i_nhwc (constant data permuted)
      CONCAT(axis=2, [const..., x_nhwc, ...]) -> y_nhwc
      remove post TRANSPOSE branches to y_nhwc
      if legacy NCHW consumers exist:
        y_nhwc --TRANSPOSE(0,3,1,2)--> y_nchw_bridge
        legacy consumers read y_nchw_bridge
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _dims_compatible(a: int, b: int) -> bool:
        if int(a) <= 0 or int(b) <= 0:
            return True
        return int(a) == int(b)

    def _shape_compatible_except_axis2(shape_a: List[int], shape_b: List[int]) -> bool:
        if len(shape_a) != 4 or len(shape_b) != 4:
            return False
        for dim_idx in [0, 1, 3]:
            if not _dims_compatible(int(shape_a[dim_idx]), int(shape_b[dim_idx])):
                return False
        return True

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _op_index(op_ref: OperatorIR) -> Optional[int]:
        for idx, op in enumerate(model_ir.operators):
            if int(id(op)) == int(id(op_ref)):
                return int(idx)
        return None

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue

            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            if (
                concat_out_tensor is None
                or concat_out_tensor.shape is None
                or len(list(concat_out_tensor.shape)) != 4
            ):
                continue

            concat_axis = int(concat_op.options.get("axis", 3))
            if concat_axis < 0:
                concat_axis += 4
            if int(concat_axis) != 3:
                continue

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue

            adapter_input_index: Optional[int] = None
            adapter_pre_idx: Optional[int] = None
            adapter_pre_op: Optional[OperatorIR] = None
            adapter_source_name: Optional[str] = None
            adapter_output_name: Optional[str] = None
            adapter_perm_name: Optional[str] = None

            for input_index, input_name in enumerate(concat_inputs):
                producer_idx = producers.get(str(input_name), None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if (
                    str(producer_op.op_type) == "TRANSPOSE"
                    and len(producer_op.inputs) >= 2
                    and len(producer_op.outputs) == 1
                    and str(producer_op.outputs[0]) == str(input_name)
                    and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                ):
                    if adapter_pre_idx is not None:
                        adapter_pre_idx = None
                        break
                    adapter_input_index = int(input_index)
                    adapter_pre_idx = int(producer_idx)
                    adapter_pre_op = producer_op
                    adapter_source_name = str(producer_op.inputs[0])
                    adapter_output_name = str(producer_op.outputs[0])
                    adapter_perm_name = str(producer_op.inputs[1])

            if (
                adapter_input_index is None
                or adapter_pre_idx is None
                or adapter_pre_op is None
                or adapter_source_name is None
                or adapter_output_name is None
                or adapter_perm_name is None
            ):
                continue

            adapter_source_tensor = model_ir.tensors.get(str(adapter_source_name), None)
            if (
                adapter_source_tensor is None
                or adapter_source_tensor.shape is None
                or len(list(adapter_source_tensor.shape)) != 4
            ):
                continue
            adapter_shape_nhwc = [int(v) for v in list(adapter_source_tensor.shape)]

            const_rewrites: Dict[str, np.ndarray] = {}
            new_concat_inputs: List[str] = list(concat_inputs)
            new_concat_inputs[int(adapter_input_index)] = str(adapter_source_name)
            valid_inputs = True

            for input_index, input_name in enumerate(concat_inputs):
                if int(input_index) == int(adapter_input_index):
                    continue
                tensor = model_ir.tensors.get(str(input_name), None)
                if (
                    tensor is None
                    or tensor.data is None
                    or tensor.shape is None
                    or len(list(tensor.shape)) != 4
                ):
                    valid_inputs = False
                    break
                # Avoid mutating shared constants used by non-concat consumers.
                if set(int(v) for v in consumers.get(str(input_name), [])) != {int(concat_idx)}:
                    valid_inputs = False
                    break
                data = np.asarray(tensor.data)
                if int(data.ndim) != 4:
                    valid_inputs = False
                    break
                converted = np.transpose(data, axes=perm_nchw_to_nhwc).astype(data.dtype, copy=False)
                if not _shape_compatible_except_axis2(
                    adapter_shape_nhwc,
                    [int(v) for v in list(converted.shape)],
                ):
                    valid_inputs = False
                    break
                const_rewrites[str(input_name)] = np.asarray(converted)
            if not valid_inputs:
                continue

            concat_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(concat_users) == 0:
                continue
            post_ops: List[OperatorIR] = []
            post_output_names: List[str] = []
            legacy_ops: List[OperatorIR] = []
            valid_consumers = True
            for user_idx in concat_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                ):
                    post_output_name = str(user_op.outputs[0])
                    if post_output_name in model_outputs:
                        valid_consumers = False
                        break
                    post_ops.append(user_op)
                    post_output_names.append(post_output_name)
                else:
                    legacy_ops.append(user_op)
            if not valid_consumers or len(post_ops) == 0:
                continue

            # Apply constants rewrite and concat axis/layout conversion.
            for input_name, converted in const_rewrites.items():
                tensor = model_ir.tensors.get(str(input_name), None)
                if tensor is None:
                    valid_inputs = False
                    break
                tensor.data = np.asarray(converted)
                tensor.shape = [int(v) for v in list(converted.shape)]
                tensor.shape_signature = [int(v) for v in list(converted.shape)]
            if not valid_inputs:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in new_concat_inputs],
            )
            concat_op.options["axis"] = 2
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(concat_out_name, None),
                perm_nchw_to_nhwc,
            )

            # Bypass all post inverse transposes.
            for post_output_name in post_output_names:
                _replace_tensor_inputs(model_ir, str(post_output_name), concat_out_name)

            # Bridge legacy consumers (if any) back to NCHW once.
            bridge_name: Optional[str] = None
            bridge_insert_before: Optional[int] = None
            if len(legacy_ops) > 0:
                concat_out_tensor_after = model_ir.tensors.get(concat_out_name, None)
                if (
                    concat_out_tensor_after is None
                    or concat_out_tensor_after.shape is None
                    or len(list(concat_out_tensor_after.shape)) != 4
                ):
                    continue
                bridge_shape = _permute_shape(
                    [int(v) for v in list(concat_out_tensor_after.shape)],
                    perm_nhwc_to_nchw,
                )
                bridge_sig_src = (
                    list(concat_out_tensor_after.shape_signature)
                    if concat_out_tensor_after.shape_signature is not None
                    else list(concat_out_tensor_after.shape)
                )
                bridge_signature = _permute_shape(
                    [int(v) for v in list(bridge_sig_src)],
                    perm_nhwc_to_nchw,
                )
                if bridge_shape is None or bridge_signature is None:
                    continue
                bridge_name = _unique_tensor_name(f"{concat_out_name}_nchw_bridge")
                model_ir.tensors[str(bridge_name)] = TensorIR(
                    name=str(bridge_name),
                    dtype=str(concat_out_tensor_after.dtype),
                    shape=[int(v) for v in list(bridge_shape)],
                    shape_signature=[int(v) for v in list(bridge_signature)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(concat_out_tensor_after.quantization),
                )
                for legacy_op in legacy_ops:
                    legacy_inputs = [
                        str(bridge_name) if str(inp) == str(concat_out_name) else str(inp)
                        for inp in list(legacy_op.inputs)
                    ]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=legacy_op,
                        new_inputs=legacy_inputs,
                    )

            remove_ops: List[OperatorIR] = list(post_ops)
            adapter_users = [int(v) for v in consumers.get(str(adapter_output_name), [])]
            if set(int(v) for v in adapter_users) == {int(concat_idx)}:
                remove_ops.append(adapter_pre_op)

            remove_indices: List[int] = []
            for remove_op in remove_ops:
                remove_index = _op_index(remove_op)
                if remove_index is not None:
                    remove_indices.append(int(remove_index))
            for remove_index in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_index)]

            if bridge_name is not None:
                legacy_indices_after = [
                    _op_index(legacy_op)
                    for legacy_op in legacy_ops
                ]
                legacy_indices_after = [
                    int(v)
                    for v in legacy_indices_after
                    if v is not None
                ]
                if len(legacy_indices_after) == 0:
                    continue
                bridge_insert_before = int(min(legacy_indices_after))
                model_ir.operators.insert(
                    int(bridge_insert_before),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(concat_out_name), str(adapter_perm_name)],
                        outputs=[str(bridge_name)],
                    ),
                )

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_axis3_const_concat_bridge_nhwc_chains": int(optimized)}


def _optimize_transposeconv_output_nhwc_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Propagate NHWC through transpose wrappers around Conv-family outputs.

    Target:
      Y_nhwc(from CONV_2D|DEPTHWISE_CONV_2D|TRANSPOSE_CONV)
        --TRANSPOSE(0,3,1,2)--> Y_nchw
      Y_nchw --(QUANTIZE|DEQUANTIZE|RELU|RELU6|ADD|SUB|MUL|DIV|MAXIMUM|MINIMUM)*--> Z_nchw
      Z_nchw --TRANSPOSE(0,2,3,1)--> Z_nhwc

    Rewrite:
      Y_nhwc --(same chain with NHWC metadata/constants)*--> Z_nhwc

    Safety:
    - Leading transpose input must be produced by CONV_2D, DEPTHWISE_CONV_2D, or TRANSPOSE_CONV.
    - Chain must be strictly linear and end with exactly one inverse transpose.
    - Binary side inputs must be constants and broadcast-safe after NCHW->NHWC remap.
    - Activation quantization must remain per-tensor.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_passthrough_ops = {
        "QUANTIZE",
        "DEQUANTIZE",
        "CAST",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
    }
    binary_passthrough_ops = {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM"}

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_outputs or pre_output_name in model_outputs:
                continue

            pre_input_prod_idx = producers.get(pre_input_name, None)
            if pre_input_prod_idx is None:
                continue
            pre_input_prod_op = model_ir.operators[int(pre_input_prod_idx)]
            if str(pre_input_prod_op.op_type) not in {
                "CONV_2D",
                "DEPTHWISE_CONV_2D",
                "TRANSPOSE_CONV",
            }:
                continue

            chain_indices: List[int] = []
            chain_ops: List[OperatorIR] = []
            chain_output_names: List[str] = []
            chain_binary_specs: List[Tuple[int, int, str, str]] = []
            current_tensor = pre_output_name
            chain_build_valid = True

            while True:
                current_users = [int(v) for v in consumers.get(current_tensor, [])]
                if len(current_users) != 1:
                    break
                op_idx = int(current_users[0])
                op = model_ir.operators[int(op_idx)]
                op_type = str(op.op_type)

                if op_type in unary_passthrough_ops:
                    if len(op.inputs) != 1 or len(op.outputs) != 1:
                        break
                    if str(op.inputs[0]) != current_tensor:
                        break
                    out_name = str(op.outputs[0])
                    if out_name in model_outputs:
                        chain_build_valid = False
                        break
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        chain_build_valid = False
                        break
                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    current_tensor = out_name
                    continue

                if op_type in binary_passthrough_ops:
                    if len(op.inputs) != 2 or len(op.outputs) != 1:
                        break
                    input_0 = str(op.inputs[0])
                    input_1 = str(op.inputs[1])
                    if input_0 == current_tensor:
                        side_input_index = 1
                        side_input_name = input_1
                    elif input_1 == current_tensor:
                        side_input_index = 0
                        side_input_name = input_0
                    else:
                        break

                    side_tensor = model_ir.tensors.get(side_input_name, None)
                    if side_tensor is None or side_tensor.data is None:
                        break
                    if not _is_per_tensor_quantization(side_tensor.quantization):
                        break

                    side_data = np.asarray(side_tensor.data)
                    if int(side_data.size) != 1:
                        if side_data.ndim != 4:
                            break
                        current_tensor_ir = model_ir.tensors.get(current_tensor, None)
                        if (
                            current_tensor_ir is None
                            or not _is_fully_known_positive_shape(current_tensor_ir.shape)
                            or len(list(current_tensor_ir.shape)) != 4
                        ):
                            break
                        nchw_shape = [int(v) for v in list(current_tensor_ir.shape)]
                        side_shape = [int(v) for v in list(side_data.shape)]
                        if _broadcast_static_shapes(nchw_shape, side_shape) is None:
                            break
                        side_shape_nhwc = _permute_shape(side_shape, perm_nchw_to_nhwc)
                        nchw_to_nhwc_shape = _permute_shape(nchw_shape, perm_nchw_to_nhwc)
                        if side_shape_nhwc is None or nchw_to_nhwc_shape is None:
                            break
                        if _broadcast_static_shapes(nchw_to_nhwc_shape, side_shape_nhwc) is None:
                            break

                    out_name = str(op.outputs[0])
                    if out_name in model_outputs:
                        chain_build_valid = False
                        break
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        chain_build_valid = False
                        break

                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    chain_binary_specs.append((int(op_idx), int(side_input_index), side_input_name, current_tensor))
                    current_tensor = out_name
                    continue

                break

            if len(chain_ops) == 0 or not chain_build_valid:
                continue

            tail_users = [int(v) for v in consumers.get(current_tensor, [])]
            if len(tail_users) != 1:
                continue
            post_idx = int(tail_users[0])
            if post_idx in set(chain_indices):
                continue
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != current_tensor:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            post_output_name = str(post_op.outputs[0])
            if post_output_name in model_outputs:
                continue

            linear_ok = True
            previous_tensor_name = pre_output_name
            chain_index_to_pos = {int(idx): pos for pos, idx in enumerate(chain_indices)}
            for op_idx, op in zip(chain_indices, chain_ops):
                if previous_tensor_name not in set(str(v) for v in op.inputs):
                    linear_ok = False
                    break
                out_name = str(op.outputs[0]) if len(op.outputs) > 0 else ""
                if out_name == "":
                    linear_ok = False
                    break
                expected_users = [int(v) for v in consumers.get(out_name, [])]
                pos = int(chain_index_to_pos[int(op_idx)])
                if pos < len(chain_indices) - 1:
                    if expected_users != [int(chain_indices[pos + 1])]:
                        linear_ok = False
                        break
                else:
                    if expected_users != [int(post_idx)]:
                        linear_ok = False
                        break
                previous_tensor_name = out_name
            if not linear_ok:
                continue

            tensors_for_quant_validation: List[Optional[TensorIR]] = [
                model_ir.tensors.get(pre_input_name, None),
                model_ir.tensors.get(pre_output_name, None),
            ] + [model_ir.tensors.get(name, None) for name in chain_output_names]
            if not _all_per_tensor_quantized(tensors_for_quant_validation):
                continue

            chain_index_set = set(int(v) for v in chain_indices)
            const_rewrite_ok = True
            for op_idx, side_input_index, side_input_name, _ in chain_binary_specs:
                side_tensor = model_ir.tensors.get(side_input_name, None)
                if side_tensor is None or side_tensor.data is None:
                    const_rewrite_ok = False
                    break
                side_data = np.asarray(side_tensor.data)
                if int(side_data.size) == 1:
                    continue
                if side_data.ndim != 4:
                    const_rewrite_ok = False
                    break

                permuted_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                side_users = [int(v) for v in consumers.get(side_input_name, [])]
                shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)

                target_const_name = str(side_input_name)
                if shared_outside_chain:
                    target_const_name = _unique_tensor_name(f"{side_input_name}_nhwc")
                    model_ir.tensors[target_const_name] = TensorIR(
                        name=target_const_name,
                        dtype=str(side_tensor.dtype),
                        shape=[int(v) for v in list(permuted_data.shape)],
                        shape_signature=[int(v) for v in list(permuted_data.shape)],
                        data=np.asarray(permuted_data),
                        is_variable=False,
                        quantization=_clone_quantization(side_tensor.quantization),
                    )
                else:
                    side_tensor.data = np.asarray(permuted_data)
                    side_tensor.shape = [int(v) for v in list(permuted_data.shape)]
                    side_tensor.shape_signature = [int(v) for v in list(permuted_data.shape)]

                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=model_ir.operators[int(op_idx)],
                    input_index=int(side_input_index),
                    new_input_name=target_const_name,
                )

            if not const_rewrite_ok:
                continue

            old_last_name = str(chain_ops[-1].outputs[0])
            old_last_tensor = model_ir.tensors.get(old_last_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)

            first_op = chain_ops[0]
            first_input_names = [str(v) for v in first_op.inputs]
            if first_input_names[0] == pre_output_name:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=[pre_input_name] + first_input_names[1:],
                )
            elif len(first_input_names) > 1 and first_input_names[1] == pre_output_name:
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=[first_input_names[0], pre_input_name],
                )
            else:
                continue

            _set_operator_outputs(
                model_ir=model_ir,
                op=chain_ops[-1],
                new_outputs=[post_output_name],
            )

            for out_name in chain_output_names[:-1]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_nchw_to_nhwc,
                )

            if post_output_tensor is not None:
                if old_last_tensor is not None:
                    post_output_tensor.dtype = str(old_last_tensor.dtype)
                    post_output_tensor.quantization = _clone_quantization(old_last_tensor.quantization)
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_nchw_to_nhwc,
                )

            remove_indices = sorted(list({int(pre_idx), int(post_idx)}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transposeconv_output_nhwc_passthrough_chains": int(rewritten)}


def _optimize_transposeconv_output_channel1_terminal_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove terminal NCHW adapter transposes after TRANSPOSE_CONV when channel dim is singleton.

    Target:
      Y_nhwc(from TRANSPOSE_CONV) --TRANSPOSE(0,3,1,2)--> Y_nchw
      Y_nchw --(ELEMENTWISE|SQUEEZE)*--> Z_out(graph output)

    Rewrite:
      Y_nhwc --(same chain with NHWC metadata/constants and remapped squeeze axes)*--> Z_out

    Safety:
    - Leading transpose input must be produced by TRANSPOSE_CONV.
    - TRANSPOSE_CONV output channel dim must be statically 1.
    - Chain must be strictly linear and terminate at a graph output tensor.
    - Chain must include at least one SQUEEZE with remappable singleton axes.
    - Binary side inputs must be constants and broadcast-safe after NCHW->NHWC remap.
    - Activation quantization must remain per-tensor.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    nchw_axis_to_nhwc_axis = {0: 0, 1: 3, 2: 1, 3: 2}
    unary_passthrough_ops = {
        "QUANTIZE",
        "DEQUANTIZE",
        "CAST",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "LOGISTIC",
        "TANH",
        "GELU",
        "HARD_SWISH",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
    }
    binary_passthrough_ops = {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM"}

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_outputs or pre_output_name in model_outputs:
                continue

            pre_input_prod_idx = producers.get(pre_input_name, None)
            if pre_input_prod_idx is None:
                continue
            pre_input_prod_op = model_ir.operators[int(pre_input_prod_idx)]
            if str(pre_input_prod_op.op_type) != "TRANSPOSE_CONV":
                continue

            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            if (
                pre_input_tensor is None
                or not _is_fully_known_positive_shape(pre_input_tensor.shape)
                or len(list(pre_input_tensor.shape)) != 4
            ):
                continue
            pre_input_shape = [int(v) for v in list(pre_input_tensor.shape)]
            if int(pre_input_shape[3]) != 1:
                continue

            chain_indices: List[int] = []
            chain_ops: List[OperatorIR] = []
            chain_output_names: List[str] = []
            chain_binary_specs: List[Tuple[int, int, str]] = []
            chain_main_input_index: Dict[int, int] = {}
            squeeze_axis_rewrites: List[Tuple[int, List[int]]] = []
            current_tensor = pre_output_name
            chain_build_valid = True
            seen_squeeze = False

            while True:
                current_users = [int(v) for v in consumers.get(current_tensor, [])]
                if len(current_users) != 1:
                    break
                op_idx = int(current_users[0])
                op = model_ir.operators[int(op_idx)]
                op_type = str(op.op_type)

                if op_type in unary_passthrough_ops:
                    if len(op.inputs) != 1 or len(op.outputs) != 1:
                        break
                    if str(op.inputs[0]) != current_tensor:
                        break
                    out_name = str(op.outputs[0])
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        chain_build_valid = False
                        break
                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    current_tensor = out_name
                    continue

                if op_type in binary_passthrough_ops:
                    if len(op.inputs) != 2 or len(op.outputs) != 1:
                        break
                    input_0 = str(op.inputs[0])
                    input_1 = str(op.inputs[1])
                    if input_0 == current_tensor:
                        side_input_index = 1
                        side_input_name = input_1
                        main_input_index = 0
                    elif input_1 == current_tensor:
                        side_input_index = 0
                        side_input_name = input_0
                        main_input_index = 1
                    else:
                        break

                    side_tensor = model_ir.tensors.get(side_input_name, None)
                    if side_tensor is None or side_tensor.data is None:
                        break
                    if not _is_per_tensor_quantization(side_tensor.quantization):
                        break

                    side_data = np.asarray(side_tensor.data)
                    if int(side_data.size) != 1:
                        if side_data.ndim != 4:
                            break
                        current_tensor_ir = model_ir.tensors.get(current_tensor, None)
                        if (
                            current_tensor_ir is None
                            or not _is_fully_known_positive_shape(current_tensor_ir.shape)
                            or len(list(current_tensor_ir.shape)) != 4
                        ):
                            break
                        nchw_shape = [int(v) for v in list(current_tensor_ir.shape)]
                        side_shape = [int(v) for v in list(side_data.shape)]
                        if _broadcast_static_shapes(nchw_shape, side_shape) is None:
                            break
                        side_shape_nhwc = _permute_shape(side_shape, perm_nchw_to_nhwc)
                        nhwc_shape = _permute_shape(nchw_shape, perm_nchw_to_nhwc)
                        if side_shape_nhwc is None or nhwc_shape is None:
                            break
                        if _broadcast_static_shapes(nhwc_shape, side_shape_nhwc) is None:
                            break

                    out_name = str(op.outputs[0])
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        chain_build_valid = False
                        break

                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    chain_binary_specs.append((int(op_idx), int(side_input_index), side_input_name))
                    chain_main_input_index[int(op_idx)] = int(main_input_index)
                    current_tensor = out_name
                    continue

                if op_type == "SQUEEZE":
                    if seen_squeeze:
                        break
                    if len(op.inputs) != 1 or len(op.outputs) != 1:
                        break
                    if str(op.inputs[0]) != current_tensor:
                        break
                    current_tensor_ir = model_ir.tensors.get(current_tensor, None)
                    if (
                        current_tensor_ir is None
                        or not _is_fully_known_positive_shape(current_tensor_ir.shape)
                        or len(list(current_tensor_ir.shape)) != 4
                    ):
                        chain_build_valid = False
                        break
                    nchw_shape = [int(v) for v in list(current_tensor_ir.shape)]
                    nhwc_shape = _permute_shape(nchw_shape, perm_nchw_to_nhwc)
                    if nhwc_shape is None:
                        chain_build_valid = False
                        break

                    raw_axes = op.options.get("squeezeDims", [])
                    if raw_axes is None:
                        raw_axes_list: List[int] = []
                    elif isinstance(raw_axes, (list, tuple, np.ndarray)):
                        raw_axes_list = [int(v) for v in list(raw_axes)]
                    else:
                        raw_axes_list = [int(raw_axes)]
                    normalized_axes: List[int] = []
                    if len(raw_axes_list) == 0:
                        normalized_axes = [int(idx) for idx, dim in enumerate(nchw_shape) if int(dim) == 1]
                    else:
                        for axis in raw_axes_list:
                            axis_norm = int(axis)
                            if axis_norm < 0:
                                axis_norm += 4
                            if axis_norm < 0 or axis_norm >= 4:
                                chain_build_valid = False
                                break
                            if axis_norm not in normalized_axes:
                                normalized_axes.append(int(axis_norm))
                        if not chain_build_valid:
                            break

                    if len(normalized_axes) == 0:
                        chain_build_valid = False
                        break

                    mapped_axes: List[int] = []
                    for axis in normalized_axes:
                        if int(nchw_shape[int(axis)]) != 1:
                            chain_build_valid = False
                            break
                        mapped_axis = int(nchw_axis_to_nhwc_axis[int(axis)])
                        if int(nhwc_shape[int(mapped_axis)]) != 1:
                            chain_build_valid = False
                            break
                        if mapped_axis not in mapped_axes:
                            mapped_axes.append(int(mapped_axis))
                    if not chain_build_valid:
                        break

                    out_name = str(op.outputs[0])
                    out_tensor = model_ir.tensors.get(out_name, None)
                    if out_tensor is not None and not _is_per_tensor_quantization(out_tensor.quantization):
                        chain_build_valid = False
                        break

                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    squeeze_axis_rewrites.append((int(op_idx), [int(v) for v in mapped_axes]))
                    current_tensor = out_name
                    seen_squeeze = True
                    continue

                break

            if len(chain_ops) == 0 or not chain_build_valid or not seen_squeeze:
                continue
            if current_tensor not in model_outputs:
                continue
            if len(consumers.get(current_tensor, [])) != 0:
                continue

            # Verify strict linearity on the main path.
            linear_ok = True
            previous_tensor_name = pre_output_name
            chain_index_to_pos = {int(idx): pos for pos, idx in enumerate(chain_indices)}
            for op_idx, op in zip(chain_indices, chain_ops):
                if previous_tensor_name not in set(str(v) for v in op.inputs):
                    linear_ok = False
                    break
                out_name = str(op.outputs[0]) if len(op.outputs) > 0 else ""
                if out_name == "":
                    linear_ok = False
                    break
                expected_users = [int(v) for v in consumers.get(out_name, [])]
                pos = int(chain_index_to_pos[int(op_idx)])
                if pos < len(chain_indices) - 1:
                    if expected_users != [int(chain_indices[pos + 1])]:
                        linear_ok = False
                        break
                else:
                    if len(expected_users) != 0:
                        linear_ok = False
                        break
                previous_tensor_name = out_name
            if not linear_ok:
                continue

            tensors_for_quant_validation: List[Optional[TensorIR]] = [
                model_ir.tensors.get(pre_input_name, None),
                model_ir.tensors.get(pre_output_name, None),
            ] + [model_ir.tensors.get(name, None) for name in chain_output_names]
            if not _all_per_tensor_quantized(tensors_for_quant_validation):
                continue

            chain_index_set = set(int(v) for v in chain_indices)
            const_rewrite_ok = True
            for op_idx, side_input_index, side_input_name in chain_binary_specs:
                side_tensor = model_ir.tensors.get(side_input_name, None)
                if side_tensor is None or side_tensor.data is None:
                    const_rewrite_ok = False
                    break
                side_data = np.asarray(side_tensor.data)
                if int(side_data.size) == 1:
                    continue
                if side_data.ndim != 4:
                    const_rewrite_ok = False
                    break

                permuted_data = np.transpose(side_data, perm_nchw_to_nhwc).astype(side_data.dtype, copy=False)
                side_users = [int(v) for v in consumers.get(side_input_name, [])]
                shared_outside_chain = any(int(u) not in chain_index_set for u in side_users)

                target_const_name = str(side_input_name)
                if shared_outside_chain:
                    target_const_name = _unique_tensor_name(f"{side_input_name}_nhwc")
                    model_ir.tensors[target_const_name] = TensorIR(
                        name=target_const_name,
                        dtype=str(side_tensor.dtype),
                        shape=[int(v) for v in list(permuted_data.shape)],
                        shape_signature=[int(v) for v in list(permuted_data.shape)],
                        data=np.asarray(permuted_data),
                        is_variable=False,
                        quantization=_clone_quantization(side_tensor.quantization),
                    )
                else:
                    side_tensor.data = np.asarray(permuted_data)
                    side_tensor.shape = [int(v) for v in list(permuted_data.shape)]
                    side_tensor.shape_signature = [int(v) for v in list(permuted_data.shape)]

                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=model_ir.operators[int(op_idx)],
                    input_index=int(side_input_index),
                    new_input_name=target_const_name,
                )

            if not const_rewrite_ok:
                continue

            first_op = chain_ops[0]
            first_input_names = [str(v) for v in first_op.inputs]
            if len(first_input_names) == 1:
                if first_input_names[0] != pre_output_name:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=[pre_input_name],
                )
            elif len(first_input_names) == 2:
                first_op_idx = int(chain_indices[0])
                main_input_index = int(chain_main_input_index.get(first_op_idx, -1))
                if main_input_index not in [0, 1]:
                    if first_input_names[0] == pre_output_name:
                        main_input_index = 0
                    elif first_input_names[1] == pre_output_name:
                        main_input_index = 1
                    else:
                        continue
                first_input_names[main_input_index] = str(pre_input_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=first_input_names,
                )
            else:
                continue

            for squeeze_op_idx, mapped_axes in squeeze_axis_rewrites:
                squeeze_op = model_ir.operators[int(squeeze_op_idx)]
                squeeze_options = dict(squeeze_op.options) if isinstance(squeeze_op.options, dict) else {}
                squeeze_options["squeezeDims"] = [int(v) for v in mapped_axes]
                squeeze_op.options = squeeze_options

            for out_name in chain_output_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_nchw_to_nhwc,
                )

            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transposeconv_output_channel1_terminal_transpose_chains": int(rewritten)}


def _optimize_singleton_channel_layout_transpose_to_reshape(model_ir: ModelIR) -> Dict[str, int]:
    """
    Replace TRANSPOSE with RESHAPE when permutation is memory-order
    equivalent due to singleton dimensions.

    Safe cases:
    - NHWC -> NCHW with input shape [N,H,W,1]
    - NCHW -> NHWC with input shape [N,1,H,W]
    - NHWC -> NCHW with input shape [N,1,1,C]
    - NCHW -> NHWC with input shape [N,C,1,1]
    In these rank-4 layout cases, moved axes are singleton and the transpose
    is an exact reshape.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    consumers = _build_tensor_consumer_map(model_ir)

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _is_memory_order_equivalent_by_singletons(
        *,
        input_signature: List[int],
        perm: List[int],
    ) -> bool:
        rank = len(input_signature)
        if rank < 2 or len(perm) != rank:
            return False
        if any(int(v) <= 0 for v in input_signature):
            return False
        non_singleton_input_axes = [
            int(axis_idx)
            for axis_idx, dim in enumerate(input_signature)
            if int(dim) != 1
        ]
        non_singleton_permuted_axes = [
            int(axis_idx)
            for axis_idx in list(perm)
            if int(input_signature[int(axis_idx)]) != 1
        ]
        return non_singleton_input_axes == non_singleton_permuted_axes

    def _feeds_quantized_logistic_concat_axis1_chain(transpose_out_name: str) -> bool:
        """
        Keep transpose form for detection-head style paths:
          TRANSPOSE -> DQ -> LOGISTIC -> Q -> DQ -> CONCAT(axis=1)
        where concat has any non-singleton channel sibling input.
        """
        lv1_users = [int(v) for v in consumers.get(str(transpose_out_name), [])]
        if len(lv1_users) != 1:
            return False
        dq1_op = model_ir.operators[int(lv1_users[0])]
        if (
            str(dq1_op.op_type) != "DEQUANTIZE"
            or len(dq1_op.inputs) != 1
            or len(dq1_op.outputs) != 1
            or str(dq1_op.inputs[0]) != str(transpose_out_name)
        ):
            return False
        dq1_out = str(dq1_op.outputs[0])

        lv2_users = [int(v) for v in consumers.get(dq1_out, [])]
        if len(lv2_users) != 1:
            return False
        logistic_op = model_ir.operators[int(lv2_users[0])]
        if (
            str(logistic_op.op_type) != "LOGISTIC"
            or len(logistic_op.inputs) != 1
            or len(logistic_op.outputs) != 1
            or str(logistic_op.inputs[0]) != dq1_out
        ):
            return False
        logistic_out = str(logistic_op.outputs[0])

        lv3_users = [int(v) for v in consumers.get(logistic_out, [])]
        if len(lv3_users) != 1:
            return False
        q_op = model_ir.operators[int(lv3_users[0])]
        if (
            str(q_op.op_type) != "QUANTIZE"
            or len(q_op.inputs) != 1
            or len(q_op.outputs) != 1
            or str(q_op.inputs[0]) != logistic_out
        ):
            return False
        q_out = str(q_op.outputs[0])

        lv4_users = [int(v) for v in consumers.get(q_out, [])]
        if len(lv4_users) != 1:
            return False
        dq2_op = model_ir.operators[int(lv4_users[0])]
        if (
            str(dq2_op.op_type) != "DEQUANTIZE"
            or len(dq2_op.inputs) != 1
            or len(dq2_op.outputs) != 1
            or str(dq2_op.inputs[0]) != q_out
        ):
            return False
        dq2_out = str(dq2_op.outputs[0])

        lv5_users = [int(v) for v in consumers.get(dq2_out, [])]
        if len(lv5_users) != 1:
            return False
        concat_op = model_ir.operators[int(lv5_users[0])]
        if (
            str(concat_op.op_type) != "CONCATENATION"
            or len(concat_op.inputs) < 2
            or str(dq2_out) not in set(str(v) for v in list(concat_op.inputs))
        ):
            return False
        concat_axis = int(concat_op.options.get("axis", 1))
        if concat_axis < 0:
            concat_axis += 4
        if concat_axis != 1:
            return False

        # If siblings include non-singleton channel inputs, this branch is part
        # of an NCHW concat contract and downstream shape passes can optimize it.
        for input_name in [str(v) for v in list(concat_op.inputs)]:
            if str(input_name) == str(dq2_out):
                continue
            input_tensor = model_ir.tensors.get(str(input_name), None)
            if input_tensor is None or len(list(input_tensor.shape)) != 4:
                continue
            if int(input_tensor.shape[1]) > 1:
                return True
        return False

    for op in model_ir.operators:
        if str(op.op_type) != "TRANSPOSE" or len(op.inputs) < 2 or len(op.outputs) != 1:
            continue

        perm = _read_transpose_perm(model_ir, op)
        if perm is None:
            continue

        input_name = str(op.inputs[0])
        output_name = str(op.outputs[0])
        input_tensor = model_ir.tensors.get(input_name, None)
        output_tensor = model_ir.tensors.get(output_name, None)
        if (
            input_tensor is None
            or output_tensor is None
            or not _is_fully_known_positive_shape(input_tensor.shape)
            or not _is_fully_known_positive_shape(output_tensor.shape)
        ):
            continue

        input_shape = [int(v) for v in list(input_tensor.shape)]
        output_shape = [int(v) for v in list(output_tensor.shape)]
        rank = len(input_shape)
        # This pass is intended only for NHWC<->NCHW style rank-4 layout bridges.
        # Rank-3 permutations (e.g. [0,2,1]) can be axis-semantic transforms and
        # rewriting them to RESHAPE may break downstream ops (e.g. LogSoftmax/Add).
        if rank != 4:
            continue
        input_signature = (
            [int(v) for v in list(input_tensor.shape_signature)]
            if input_tensor.shape_signature is not None
            else [int(v) for v in list(input_shape)]
        )
        output_signature = (
            [int(v) for v in list(output_tensor.shape_signature)]
            if output_tensor.shape_signature is not None
            else [int(v) for v in list(output_shape)]
        )
        if len(output_shape) != rank or len(list(perm)) != rank:
            continue
        # Dynamic dimensions can be represented as placeholder static 1s in `shape`.
        # Replacing transpose with reshape is only safe when the layout is fully
        # static in signatures as well.
        if (
            len(input_signature) != rank
            or len(output_signature) != rank
            or any(int(v) < 0 for v in input_signature)
            or any(int(v) < 0 for v in output_signature)
        ):
            continue

        canonical_singleton_safe = False
        if perm == perm_nhwc_to_nchw:
            channel_singleton = int(input_shape[3]) == 1 and int(output_shape[1]) == 1
            spatial_singleton = (
                int(input_shape[1]) == 1
                and int(input_shape[2]) == 1
                and int(output_shape[2]) == 1
                and int(output_shape[3]) == 1
            )
            canonical_singleton_safe = bool(channel_singleton or spatial_singleton)
        elif perm == perm_nchw_to_nhwc:
            channel_singleton = int(input_shape[1]) == 1 and int(output_shape[3]) == 1
            spatial_singleton = (
                int(input_shape[2]) == 1
                and int(input_shape[3]) == 1
                and int(output_shape[1]) == 1
                and int(output_shape[2]) == 1
            )
            canonical_singleton_safe = bool(channel_singleton or spatial_singleton)

        expected_output_shape = _permute_shape(input_shape, perm)
        if expected_output_shape is None or [int(v) for v in list(expected_output_shape)] != output_shape:
            continue
        expected_output_signature = _permute_shape(input_signature, perm)
        if (
            expected_output_signature is None
            or [int(v) for v in list(expected_output_signature)] != output_signature
        ):
            continue

        if not canonical_singleton_safe and not _is_memory_order_equivalent_by_singletons(
            input_signature=[int(v) for v in list(input_signature)],
            perm=[int(v) for v in list(perm)],
        ):
            continue

        # Preserve transpose (do not downcast to reshape) for quantized
        # detection-head singleton gates that feed NCHW concat chains.
        if (
            rank == 4
            and [int(v) for v in list(perm)] == perm_nhwc_to_nchw
            and int(input_shape[3]) == 1
            and _feeds_quantized_logistic_concat_axis1_chain(str(output_name))
        ):
            continue

        reshape_shape_name = _unique_tensor_name(f"{output_name}_reshape_shape")
        model_ir.tensors[reshape_shape_name] = TensorIR(
            name=reshape_shape_name,
            dtype="INT32",
            shape=[len(output_shape)],
            shape_signature=[len(output_shape)],
            data=np.asarray(output_shape, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )

        op.op_type = "RESHAPE"
        op.inputs = [input_name, reshape_shape_name]
        op.options = {"newShape": [int(v) for v in output_shape]}
        rewritten += 1

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_singleton_channel_layout_transpose_to_reshape": int(rewritten)}


def _optimize_singleton_layout_reshape_unary_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove singleton-layout RESHAPE wrappers around layout-agnostic unary ops.

    Target:
      x --RESHAPE(layout wrapper)--> a --UNARY--> b --RESHAPE(inverse wrapper)--> y

    Rewrite:
      x --UNARY--> y

    Safety:
    - Both RESHAPEs must be rank-4 NHWC<->NCHW permutations.
    - Permutations must be inverse of each other.
    - The permutation must be memory-order equivalent due singleton dimensions.
    - Chain must be strict linear: pre-reshape output and unary output have one consumer.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "FLOOR",
        "CEIL",
        "ROUND",
    }

    def _is_static_rank4_shape(tensor: Optional[TensorIR]) -> bool:
        if tensor is None:
            return False
        if tensor.shape is None or len(list(tensor.shape)) != 4:
            return False
        if any(int(v) <= 0 for v in list(tensor.shape)):
            return False
        signature = (
            list(tensor.shape_signature)
            if tensor.shape_signature is not None
            else list(tensor.shape)
        )
        if len(signature) != 4:
            return False
        if any(int(v) < 0 for v in signature):
            return False
        return True

    def _detect_layout_perm(src_shape: List[int], dst_shape: List[int]) -> Optional[List[int]]:
        for perm in [perm_nhwc_to_nchw, perm_nchw_to_nhwc]:
            expected = _permute_shape(src_shape, perm)
            if expected is not None and [int(v) for v in list(expected)] == [int(v) for v in list(dst_shape)]:
                return [int(v) for v in list(perm)]
        return None

    def _is_singleton_layout_permute_safe(src_shape: List[int], perm: List[int]) -> bool:
        if len(src_shape) != 4:
            return False
        if [int(v) for v in list(perm)] == perm_nhwc_to_nchw:
            return int(src_shape[3]) == 1 or (int(src_shape[1]) == 1 and int(src_shape[2]) == 1)
        if [int(v) for v in list(perm)] == perm_nchw_to_nhwc:
            return int(src_shape[1]) == 1 or (int(src_shape[2]) == 1 and int(src_shape[3]) == 1)
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "RESHAPE" or len(pre_op.inputs) < 1 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue

            unary_idx = int(pre_users[0])
            unary_op = model_ir.operators[int(unary_idx)]
            if str(unary_op.op_type) not in unary_passthrough_ops:
                continue
            if len(unary_op.inputs) != 1 or len(unary_op.outputs) != 1:
                continue
            if str(unary_op.inputs[0]) != pre_output_name:
                continue

            unary_output_name = str(unary_op.outputs[0])
            if unary_output_name in model_outputs:
                continue
            unary_users = [int(v) for v in consumers.get(unary_output_name, [])]
            if len(unary_users) != 1:
                continue

            post_idx = int(unary_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "RESHAPE" or len(post_op.inputs) < 1 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != unary_output_name:
                continue

            post_output_name = str(post_op.outputs[0])

            pre_input_tensor = model_ir.tensors.get(pre_input_name, None)
            pre_output_tensor = model_ir.tensors.get(pre_output_name, None)
            unary_output_tensor = model_ir.tensors.get(unary_output_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if (
                not _is_static_rank4_shape(pre_input_tensor)
                or not _is_static_rank4_shape(pre_output_tensor)
                or not _is_static_rank4_shape(unary_output_tensor)
                or not _is_static_rank4_shape(post_output_tensor)
            ):
                continue

            pre_input_shape = [int(v) for v in list(pre_input_tensor.shape)]
            pre_output_shape = [int(v) for v in list(pre_output_tensor.shape)]
            unary_output_shape = [int(v) for v in list(unary_output_tensor.shape)]
            post_output_shape = [int(v) for v in list(post_output_tensor.shape)]
            if pre_output_shape != unary_output_shape:
                continue

            perm_pre = _detect_layout_perm(pre_input_shape, pre_output_shape)
            if perm_pre is None:
                continue
            if not _is_singleton_layout_permute_safe(pre_input_shape, perm_pre):
                continue
            perm_post = _detect_layout_perm(pre_output_shape, post_output_shape)
            if perm_post is None:
                continue
            if perm_post != _invert_perm(perm_pre):
                continue
            if not _is_singleton_layout_permute_safe(pre_output_shape, perm_post):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[pre_input_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=unary_op,
                new_outputs=[post_output_name],
            )

            if unary_output_tensor is not None and post_output_tensor is not None:
                post_output_tensor.dtype = str(unary_output_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(unary_output_tensor.quantization)

            remove_indices: List[int] = []
            pre_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is pre_op), None)
            if pre_remove_idx is not None:
                remove_indices.append(int(pre_remove_idx))
            post_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is post_op), None)
            if post_remove_idx is not None:
                remove_indices.append(int(post_remove_idx))
            for remove_idx in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_singleton_layout_reshape_unary_passthrough_chains": int(rewritten)}


def _optimize_consecutive_inverse_singleton_layout_reshapes(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove consecutive inverse NHWC<->NCHW layout RESHAPE pairs.

    Target:
      x --RESHAPE(NHWC->NCHW)--> a --RESHAPE(NCHW->NHWC)--> y
      (or reverse direction)

    Rewrite:
      x --> y

    Safety:
    - Both RESHAPEs must be rank-4 and form inverse canonical layout permutations.
    - Memory-order equivalence must hold due singleton channel or singleton spatial dims.
    - Intermediate tensor must be single-consumer.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _is_static_rank4_shape(tensor: Optional[TensorIR]) -> bool:
        if tensor is None:
            return False
        if tensor.shape is None or len(list(tensor.shape)) != 4:
            return False
        if any(int(v) <= 0 for v in list(tensor.shape)):
            return False
        signature = (
            list(tensor.shape_signature)
            if tensor.shape_signature is not None
            else list(tensor.shape)
        )
        if len(signature) != 4:
            return False
        if any(int(v) < 0 for v in signature):
            return False
        return True

    def _detect_layout_perm(src_shape: List[int], dst_shape: List[int]) -> Optional[List[int]]:
        for perm in [perm_nhwc_to_nchw, perm_nchw_to_nhwc]:
            expected = _permute_shape(src_shape, perm)
            if expected is not None and [int(v) for v in list(expected)] == [int(v) for v in list(dst_shape)]:
                return [int(v) for v in list(perm)]
        return None

    def _is_singleton_layout_permute_safe(src_shape: List[int], perm: List[int]) -> bool:
        if len(src_shape) != 4:
            return False
        if [int(v) for v in list(perm)] == perm_nhwc_to_nchw:
            return int(src_shape[3]) == 1 or (int(src_shape[1]) == 1 and int(src_shape[2]) == 1)
        if [int(v) for v in list(perm)] == perm_nchw_to_nhwc:
            return int(src_shape[1]) == 1 or (int(src_shape[2]) == 1 and int(src_shape[3]) == 1)
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for first_idx, first_op in enumerate(model_ir.operators):
            if str(first_op.op_type) != "RESHAPE" or len(first_op.inputs) < 1 or len(first_op.outputs) != 1:
                continue

            src_name = str(first_op.inputs[0])
            mid_name = str(first_op.outputs[0])
            mid_users = [int(v) for v in consumers.get(mid_name, [])]
            if len(mid_users) != 1:
                continue

            second_idx = int(mid_users[0])
            second_op = model_ir.operators[int(second_idx)]
            if str(second_op.op_type) != "RESHAPE" or len(second_op.inputs) < 1 or len(second_op.outputs) != 1:
                continue
            if str(second_op.inputs[0]) != mid_name:
                continue

            dst_name = str(second_op.outputs[0])

            src_tensor = model_ir.tensors.get(src_name, None)
            mid_tensor = model_ir.tensors.get(mid_name, None)
            dst_tensor = model_ir.tensors.get(dst_name, None)
            if (
                not _is_static_rank4_shape(src_tensor)
                or not _is_static_rank4_shape(mid_tensor)
                or not _is_static_rank4_shape(dst_tensor)
            ):
                continue

            src_shape = [int(v) for v in list(src_tensor.shape)]
            mid_shape = [int(v) for v in list(mid_tensor.shape)]
            dst_shape = [int(v) for v in list(dst_tensor.shape)]

            perm_1 = _detect_layout_perm(src_shape, mid_shape)
            perm_2 = _detect_layout_perm(mid_shape, dst_shape)
            if perm_1 is None or perm_2 is None:
                continue
            if perm_2 != _invert_perm(perm_1):
                continue
            if not _is_singleton_layout_permute_safe(src_shape, perm_1):
                continue
            if not _is_singleton_layout_permute_safe(mid_shape, perm_2):
                continue

            # Keep graph-output tensor names stable when possible.
            if dst_name in model_outputs:
                src_users = [int(v) for v in consumers.get(src_name, [])]
                if (
                    src_name in model_ir.inputs
                    or src_name in model_outputs
                    or set(src_users) != {int(first_idx)}
                ):
                    continue
                _rename_tensor_globally(
                    model_ir=model_ir,
                    old_name=str(src_name),
                    new_name=str(dst_name),
                )
            else:
                _replace_tensor_inputs(
                    model_ir=model_ir,
                    src_name=str(dst_name),
                    dst_name=str(src_name),
                )

            for remove_idx in sorted([int(first_idx), int(second_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_consecutive_inverse_singleton_layout_reshapes": int(rewritten)}


def _optimize_singleton_layout_reshape_maxpool_binary_cast_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant singleton layout RESHAPE bridges around consecutive MAX_POOL blocks.

    Target:
      a_nchw --RESHAPE--> a_nhwc --MAX_POOL--> b_nhwc --RESHAPE--> b_nchw
      BIN(a_nchw, b_nchw) --(optional CAST)--> c_nchw
      c_nchw --RESHAPE--> c_nhwc --MAX_POOL--> d_nhwc --RESHAPE--> d_nchw

    Rewrite:
      a_nchw --RESHAPE--> a_nhwc --MAX_POOL--> b_nhwc
      BIN(a_nhwc, b_nhwc) --(optional CAST)--> c_nhwc
      c_nhwc --MAX_POOL--> d_nhwc --RESHAPE--> d_nchw

    This removes the middle `RESHAPE->...->RESHAPE` layout roundtrip while
    preserving external NCHW contracts.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    binary_ops = {
        "EQUAL",
        "NOT_EQUAL",
        "GREATER",
        "GREATER_EQUAL",
        "LESS",
        "LESS_EQUAL",
        "MAXIMUM",
        "MINIMUM",
        "ADD",
        "SUB",
        "MUL",
        "DIV",
    }

    def _read_reshape_target_shape(op: OperatorIR) -> Optional[List[int]]:
        if str(op.op_type) != "RESHAPE" or len(op.inputs) < 2:
            return None
        shape_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        if shape_vals is None or len(shape_vals) != 4:
            return None
        return [int(v) for v in list(shape_vals)]

    def _set_tensor_shape_metadata(tensor_name: str, shape: List[int]) -> None:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            return
        normalized = [int(v) for v in list(shape)]
        tensor.shape = [int(v) for v in list(normalized)]
        tensor.shape_signature = [int(v) for v in list(normalized)]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post1_idx, post1_op in enumerate(model_ir.operators):
            if str(post1_op.op_type) != "RESHAPE" or len(post1_op.inputs) < 1 or len(post1_op.outputs) != 1:
                continue
            post1_in_name = str(post1_op.inputs[0])
            post1_out_name = str(post1_op.outputs[0])
            if post1_out_name in model_outputs:
                continue

            pool1_idx = producers.get(post1_in_name, None)
            if pool1_idx is None:
                continue
            pool1_op = model_ir.operators[int(pool1_idx)]
            if str(pool1_op.op_type) != "MAX_POOL_2D" or len(pool1_op.inputs) != 1 or len(pool1_op.outputs) != 1:
                continue
            if str(pool1_op.outputs[0]) != post1_in_name:
                continue

            pre1_out_name = str(pool1_op.inputs[0])
            pre1_idx = producers.get(pre1_out_name, None)
            if pre1_idx is None:
                continue
            pre1_op = model_ir.operators[int(pre1_idx)]
            if str(pre1_op.op_type) != "RESHAPE" or len(pre1_op.inputs) < 1 or len(pre1_op.outputs) != 1:
                continue
            if str(pre1_op.outputs[0]) != pre1_out_name:
                continue

            a_nchw_name = str(pre1_op.inputs[0])
            pre1_target_shape = _read_reshape_target_shape(pre1_op)
            post1_target_shape = _read_reshape_target_shape(post1_op)
            if pre1_target_shape is None or post1_target_shape is None:
                continue
            if pre1_target_shape[3] != 1 or post1_target_shape[1] != 1:
                continue
            expected_post1 = _permute_shape(pre1_target_shape, perm_nhwc_to_nchw)
            if expected_post1 is None or [int(v) for v in list(expected_post1)] != post1_target_shape:
                continue

            post1_users = [int(v) for v in consumers.get(post1_out_name, [])]
            if len(post1_users) != 1:
                continue
            bin_idx = int(post1_users[0])
            bin_op = model_ir.operators[int(bin_idx)]
            if str(bin_op.op_type) not in binary_ops or len(bin_op.inputs) != 2 or len(bin_op.outputs) != 1:
                continue
            if a_nchw_name not in set(str(v) for v in bin_op.inputs):
                continue
            if post1_out_name not in set(str(v) for v in bin_op.inputs):
                continue
            bin_out_name = str(bin_op.outputs[0])
            if bin_out_name in model_outputs:
                continue

            cast_path: Optional[Dict[str, Any]] = None
            non_cast_user_indices: List[int] = []
            invalid_fanout = False
            bin_users = [int(v) for v in consumers.get(str(bin_out_name), [])]
            for bin_user_idx in list(bin_users):
                bin_user_op = model_ir.operators[int(bin_user_idx)]
                if str(bin_user_op.op_type) != "CAST":
                    non_cast_user_indices.append(int(bin_user_idx))
                    continue
                if len(bin_user_op.inputs) != 1 or len(bin_user_op.outputs) != 1:
                    invalid_fanout = True
                    break
                if str(bin_user_op.inputs[0]) != str(bin_out_name):
                    invalid_fanout = True
                    break
                cast_out_name = str(bin_user_op.outputs[0])
                if cast_out_name in model_outputs:
                    invalid_fanout = True
                    break

                pre2_users = [int(v) for v in consumers.get(cast_out_name, [])]
                if len(pre2_users) != 1:
                    invalid_fanout = True
                    break
                pre2_idx = int(pre2_users[0])
                pre2_op = model_ir.operators[int(pre2_idx)]
                if str(pre2_op.op_type) != "RESHAPE" or len(pre2_op.inputs) < 1 or len(pre2_op.outputs) != 1:
                    invalid_fanout = True
                    break
                if str(pre2_op.inputs[0]) != cast_out_name:
                    invalid_fanout = True
                    break
                pre2_out_name = str(pre2_op.outputs[0])
                if pre2_out_name in model_outputs:
                    invalid_fanout = True
                    break

                pool2_users = [int(v) for v in consumers.get(pre2_out_name, [])]
                if len(pool2_users) != 1:
                    invalid_fanout = True
                    break
                pool2_idx = int(pool2_users[0])
                pool2_op = model_ir.operators[int(pool2_idx)]
                if (
                    str(pool2_op.op_type) != "MAX_POOL_2D"
                    or len(pool2_op.inputs) != 1
                    or len(pool2_op.outputs) != 1
                    or str(pool2_op.inputs[0]) != pre2_out_name
                ):
                    invalid_fanout = True
                    break
                pool2_out_name = str(pool2_op.outputs[0])
                if pool2_out_name in model_outputs:
                    invalid_fanout = True
                    break

                post2_users = [int(v) for v in consumers.get(pool2_out_name, [])]
                if len(post2_users) != 1:
                    invalid_fanout = True
                    break
                post2_idx = int(post2_users[0])
                post2_op = model_ir.operators[int(post2_idx)]
                if (
                    str(post2_op.op_type) != "RESHAPE"
                    or len(post2_op.inputs) < 1
                    or len(post2_op.outputs) != 1
                    or str(post2_op.inputs[0]) != pool2_out_name
                ):
                    invalid_fanout = True
                    break

                pre2_target_shape = _read_reshape_target_shape(pre2_op)
                post2_target_shape = _read_reshape_target_shape(post2_op)
                if pre2_target_shape is None or post2_target_shape is None:
                    invalid_fanout = True
                    break
                if pre2_target_shape[3] != 1 or post2_target_shape[1] != 1:
                    invalid_fanout = True
                    break
                expected_post2 = _permute_shape(pre2_target_shape, perm_nhwc_to_nchw)
                if expected_post2 is None or [int(v) for v in list(expected_post2)] != post2_target_shape:
                    invalid_fanout = True
                    break

                if cast_path is not None:
                    invalid_fanout = True
                    break
                cast_path = {
                    "cast_idx": int(bin_user_idx),
                    "cast_op": bin_user_op,
                    "cast_out_name": str(cast_out_name),
                    "pre2_idx": int(pre2_idx),
                    "pre2_op": pre2_op,
                    "pool2_op": pool2_op,
                    "pre2_target_shape": [int(v) for v in list(pre2_target_shape)],
                    "post2_target_shape": [int(v) for v in list(post2_target_shape)],
                }

            if invalid_fanout or cast_path is None:
                continue

            # Rewire binary inputs to NHWC path.
            bin_inputs = [str(v) for v in list(bin_op.inputs)]
            replaced = False
            for input_index, input_name in enumerate(bin_inputs):
                if input_name == post1_out_name:
                    bin_inputs[input_index] = str(post1_in_name)
                    replaced = True
                elif input_name == a_nchw_name:
                    bin_inputs[input_index] = str(pre1_out_name)
                    replaced = True
            if not replaced:
                continue
            _set_operator_inputs(
                model_ir=model_ir,
                op=bin_op,
                new_inputs=bin_inputs,
            )

            pre2_target_shape = [int(v) for v in list(cast_path["pre2_target_shape"])]
            _set_tensor_shape_metadata(str(bin_out_name), pre2_target_shape)

            if len(non_cast_user_indices) > 0:
                bin_out_tensor = model_ir.tensors.get(str(bin_out_name), None)
                adapter_shape_name = _unique_tensor_name(f"{bin_out_name}_nchw_shape")
                adapter_out_name = _unique_tensor_name(f"{bin_out_name}_nchw")
                model_ir.tensors[adapter_shape_name] = TensorIR(
                    name=adapter_shape_name,
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(post1_target_shape, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                model_ir.tensors[adapter_out_name] = TensorIR(
                    name=adapter_out_name,
                    dtype=(
                        str(bin_out_tensor.dtype)
                        if bin_out_tensor is not None
                        else "BOOL"
                    ),
                    shape=[int(v) for v in list(post1_target_shape)],
                    shape_signature=[int(v) for v in list(post1_target_shape)],
                    data=None,
                    is_variable=False,
                    quantization=(
                        _clone_quantization(bin_out_tensor.quantization)
                        if bin_out_tensor is not None
                        else None
                    ),
                )

                for user_idx in list(non_cast_user_indices):
                    user_op = model_ir.operators[int(user_idx)]
                    user_inputs = [str(v) for v in list(user_op.inputs)]
                    if str(bin_out_name) not in set(user_inputs):
                        continue
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=user_op,
                        new_inputs=[
                            str(adapter_out_name) if str(v) == str(bin_out_name) else str(v)
                            for v in list(user_inputs)
                        ],
                    )

                bin_current_idx = next(
                    (int(idx) for idx, op_ref in enumerate(model_ir.operators) if op_ref is bin_op),
                    None,
                )
                if bin_current_idx is None:
                    continue
                model_ir.operators.insert(
                    int(bin_current_idx) + 1,
                    OperatorIR(
                        op_type="RESHAPE",
                        inputs=[str(bin_out_name), str(adapter_shape_name)],
                        outputs=[str(adapter_out_name)],
                        options={"newShape": [int(v) for v in list(post1_target_shape)]},
                    ),
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=cast_path["pool2_op"],
                new_inputs=[str(cast_path["cast_out_name"])],
            )
            _set_tensor_shape_metadata(str(cast_path["cast_out_name"]), pre2_target_shape)

            remove_indices: List[int] = []
            for op_idx, op_ref in enumerate(model_ir.operators):
                if op_ref is post1_op or op_ref is cast_path["pre2_op"]:
                    remove_indices.append(int(op_idx))
            for remove_idx in sorted(list(set(remove_indices)), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_singleton_layout_reshape_maxpool_binary_cast_chains": int(rewritten)}


def _optimize_constant_input_cast_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold CAST when its input tensor already has embedded constant data.

    Rewrite:
      CONST(in) -> CAST -> out
    Into:
      CONST(out_casted_data) and remove CAST op.
    """
    rewritten = 0

    while True:
        changed = False
        for cast_idx, cast_op in enumerate(model_ir.operators):
            if str(cast_op.op_type) != "CAST" or len(cast_op.inputs) != 1 or len(cast_op.outputs) != 1:
                continue

            in_name = str(cast_op.inputs[0])
            out_name = str(cast_op.outputs[0])
            if in_name == "" or out_name == "":
                continue
            if in_name == out_name:
                continue

            in_tensor = model_ir.tensors.get(in_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if in_tensor is None or out_tensor is None:
                continue
            if bool(in_tensor.is_variable) or in_tensor.data is None:
                continue

            out_dtype = str(cast_op.options.get("outDataType", "")).upper()
            if out_dtype == "":
                out_dtype = str(out_tensor.dtype).upper()
            if out_dtype == "":
                continue

            casted = _cast_const_array_to_tflite_dtype(in_tensor.data, out_dtype)
            if casted is None:
                continue

            out_tensor.data = casted
            out_tensor.dtype = str(out_dtype)
            _append_tensor_lineage_event(
                model_ir=model_ir,
                event={
                    "kind": "fold_constant_cast",
                    "input_name": str(in_name),
                    "output_name": str(out_name),
                    "out_dtype": str(out_dtype),
                },
            )

            del model_ir.operators[int(cast_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_constant_input_cast_chains": int(rewritten)}


def _optimize_redundant_int64_to_int32_cast_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant CAST chains:
      CAST(* -> INT64)  -> CAST(INT64  -> INT32)
      CAST(* -> UINT64) -> CAST(UINT64 -> UINT32)

    This pattern appears in integer control/index subgraphs where upstream
    lowering temporarily widens to INT64 and the immediate consumer narrows
    back to 32-bit. Keep the first CAST and retarget it to the second cast
    output dtype, then remove the second CAST.
    """
    rewritten = 0
    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for cast2_idx, cast2_op in enumerate(model_ir.operators):
            if str(cast2_op.op_type) != "CAST" or len(cast2_op.inputs) != 1 or len(cast2_op.outputs) != 1:
                continue

            mid_name = str(cast2_op.inputs[0])
            cast2_out_name = str(cast2_op.outputs[0])
            if mid_name in model_outputs or cast2_out_name in model_outputs:
                continue

            cast2_in_dtype = str(cast2_op.options.get("inDataType", "")).upper()
            cast2_out_dtype = str(cast2_op.options.get("outDataType", "")).upper()
            target_dtype = ""
            if cast2_in_dtype == "INT64" and cast2_out_dtype == "INT32":
                target_dtype = "INT32"
            elif cast2_in_dtype == "UINT64" and cast2_out_dtype == "UINT32":
                target_dtype = "UINT32"
            if target_dtype == "":
                continue

            cast1_idx = producers.get(mid_name, None)
            if cast1_idx is None or int(cast1_idx) == int(cast2_idx):
                continue
            cast1_op = model_ir.operators[int(cast1_idx)]
            if str(cast1_op.op_type) != "CAST" or len(cast1_op.inputs) != 1 or len(cast1_op.outputs) != 1:
                continue
            if str(cast1_op.outputs[0]) != mid_name:
                continue

            mid_users = [int(v) for v in consumers.get(mid_name, [])]
            if len(mid_users) != 1 or int(mid_users[0]) != int(cast2_idx):
                continue

            cast1_out_dtype = str(cast1_op.options.get("outDataType", "")).upper()
            if cast1_out_dtype != cast2_in_dtype:
                continue

            cast1_op.options["outDataType"] = target_dtype
            mid_tensor = model_ir.tensors.get(mid_name, None)
            if mid_tensor is not None:
                mid_tensor.dtype = target_dtype
                src_tensor = model_ir.tensors.get(str(cast1_op.inputs[0]), None)
                if src_tensor is not None:
                    mid_tensor.quantization = _clone_quantization(src_tensor.quantization)

            _replace_tensor_inputs(
                model_ir=model_ir,
                src_name=cast2_out_name,
                dst_name=mid_name,
            )
            del model_ir.operators[int(cast2_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_redundant_int64_to_int32_cast_chains": int(rewritten)}


def _optimize_singleton_nms_maxpool_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Collapse SuperPoint-like singleton-channel NMS MAX_POOL reshape ladders into NHWC.

    Target (representative):
      ... -> base_nchw
      base_nchw --RESHAPE--> base_nhwc --MAX_POOL--> ...
      (EQUAL/CAST/GREATER/SELECT/NOT/AND/OR with repeated
       RESHAPE->MAX_POOL->RESHAPE wrappers in NCHW)

    Rewrite:
      Keep the whole logical NMS subgraph in NHWC and remove intermediate
      singleton layout wrappers; if downstream expects NCHW, add one terminal
      NHWC->NCHW adapter.
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _read_reshape_target_shape(op: OperatorIR) -> Optional[List[int]]:
        if str(op.op_type) != "RESHAPE" or len(op.inputs) < 2:
            return None
        vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        if vals is None or len(vals) != 4:
            return None
        return [int(v) for v in list(vals)]

    def _set_rank4_shape(tensor_name: str, shape: List[int]) -> None:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            return
        normalized = [int(v) for v in list(shape)]
        if len(normalized) != 4:
            return
        tensor.shape = [int(v) for v in list(normalized)]
        tensor.shape_signature = [int(v) for v in list(normalized)]

    def _find_single_consumer_op(
        tensor_name: str,
        consumers: Dict[str, List[int]],
        *,
        expected_op_type: str,
        operators: List[OperatorIR],
        required_input_index: Optional[int] = None,
    ) -> Optional[int]:
        user_indices = [int(v) for v in consumers.get(str(tensor_name), [])]
        matched_indices: List[int] = []
        for user_idx in user_indices:
            user_op = operators[int(user_idx)]
            if str(user_op.op_type) != str(expected_op_type):
                continue
            if required_input_index is not None:
                if (
                    int(required_input_index) < 0
                    or len(user_op.inputs) <= int(required_input_index)
                    or str(user_op.inputs[int(required_input_index)]) != str(tensor_name)
                ):
                    continue
            matched_indices.append(int(user_idx))
        if len(matched_indices) != 1:
            return None
        return int(matched_indices[0])

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post1_idx, post1_op in enumerate(model_ir.operators):
            # post1: pool1b_nhwc -> pool1b_nchw
            if str(post1_op.op_type) != "RESHAPE" or len(post1_op.inputs) < 2 or len(post1_op.outputs) != 1:
                continue
            post1_in_name = str(post1_op.inputs[0])
            post1_out_name = str(post1_op.outputs[0])
            if post1_out_name in model_outputs:
                continue
            post1_target_shape = _read_reshape_target_shape(post1_op)
            if post1_target_shape is None or int(post1_target_shape[1]) != 1:
                continue

            # pool1b and cast0 path
            pool1b_idx = producers.get(post1_in_name, None)
            if pool1b_idx is None:
                continue
            pool1b_op = model_ir.operators[int(pool1b_idx)]
            if str(pool1b_op.op_type) != "MAX_POOL_2D" or len(pool1b_op.inputs) != 1 or len(pool1b_op.outputs) != 1:
                continue
            if str(pool1b_op.outputs[0]) != post1_in_name:
                continue
            cast0_out_name = str(pool1b_op.inputs[0])
            cast0_idx = producers.get(cast0_out_name, None)
            if cast0_idx is None:
                continue
            cast0_op = model_ir.operators[int(cast0_idx)]
            if str(cast0_op.op_type) != "CAST" or len(cast0_op.inputs) != 1 or len(cast0_op.outputs) != 1:
                continue
            if str(cast0_op.outputs[0]) != cast0_out_name:
                continue
            eq0_nhwc_name = str(cast0_op.inputs[0])
            if eq0_nhwc_name in model_outputs:
                continue

            # eq0 path: eq0_nhwc = EQUAL(base_nhwc, pool1_nhwc)
            eq0_idx = producers.get(eq0_nhwc_name, None)
            if eq0_idx is None:
                continue
            eq0_op = model_ir.operators[int(eq0_idx)]
            if str(eq0_op.op_type) != "EQUAL" or len(eq0_op.inputs) != 2 or len(eq0_op.outputs) != 1:
                continue
            if str(eq0_op.outputs[0]) != eq0_nhwc_name:
                continue
            eq0_in0 = str(eq0_op.inputs[0])
            eq0_in1 = str(eq0_op.inputs[1])

            pool1_idx: Optional[int] = None
            pool1_out_name: Optional[str] = None
            base_nhwc_name: Optional[str] = None
            for candidate in [eq0_in0, eq0_in1]:
                cand_prod_idx = producers.get(str(candidate), None)
                if cand_prod_idx is None:
                    continue
                cand_op = model_ir.operators[int(cand_prod_idx)]
                if str(cand_op.op_type) == "MAX_POOL_2D" and len(cand_op.inputs) == 1 and len(cand_op.outputs) == 1:
                    if str(cand_op.outputs[0]) == str(candidate):
                        pool1_idx = int(cand_prod_idx)
                        pool1_out_name = str(candidate)
                        other = eq0_in1 if str(candidate) == eq0_in0 else eq0_in0
                        base_nhwc_name = str(other)
                        break
            if pool1_idx is None or pool1_out_name is None or base_nhwc_name is None:
                continue

            pool1_op = model_ir.operators[int(pool1_idx)]
            if str(pool1_op.inputs[0]) != str(base_nhwc_name):
                continue
            base_nhwc_shape = [int(v) for v in _read_reshape_target_shape(model_ir.operators[int(producers.get(base_nhwc_name, -1))]) or []]
            if len(base_nhwc_shape) != 4:
                base_nhwc_tensor = model_ir.tensors.get(base_nhwc_name, None)
                if base_nhwc_tensor is None or base_nhwc_tensor.shape is None or len(list(base_nhwc_tensor.shape)) != 4:
                    continue
                base_nhwc_shape = [int(v) for v in list(base_nhwc_tensor.shape)]
            if int(base_nhwc_shape[3]) != 1:
                continue

            pre_base_idx = producers.get(base_nhwc_name, None)
            if pre_base_idx is None:
                continue
            pre_base_op = model_ir.operators[int(pre_base_idx)]
            if str(pre_base_op.op_type) != "RESHAPE" or len(pre_base_op.inputs) < 2 or len(pre_base_op.outputs) != 1:
                continue
            if str(pre_base_op.outputs[0]) != str(base_nhwc_name):
                continue
            base_nchw_name = str(pre_base_op.inputs[0])
            if base_nchw_name in model_outputs:
                continue
            base_nchw_shape = _read_reshape_target_shape(post1_op)
            if base_nchw_shape is None:
                base_nchw_tensor = model_ir.tensors.get(base_nchw_name, None)
                if base_nchw_tensor is None or base_nchw_tensor.shape is None or len(list(base_nchw_tensor.shape)) != 4:
                    continue
                base_nchw_shape = [int(v) for v in list(base_nchw_tensor.shape)]
            if int(base_nchw_shape[1]) != 1:
                continue

            # greater0/select0/not0
            greater0_idx = _find_single_consumer_op(
                post1_out_name,
                consumers,
                expected_op_type="GREATER",
                operators=model_ir.operators,
            )
            if greater0_idx is None:
                continue
            greater0_op = model_ir.operators[int(greater0_idx)]
            if len(greater0_op.inputs) != 2 or len(greater0_op.outputs) != 1:
                continue
            cond0_name = str(greater0_op.outputs[0])
            cond0_users = [int(v) for v in consumers.get(cond0_name, [])]
            if len(cond0_users) != 2:
                continue
            select0_idx: Optional[int] = None
            not0_idx: Optional[int] = None
            for user_idx in cond0_users:
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) == "SELECT" and len(user_op.inputs) == 3 and str(user_op.inputs[0]) == cond0_name:
                    select0_idx = int(user_idx)
                elif str(user_op.op_type) == "LOGICAL_NOT" and len(user_op.inputs) == 1 and str(user_op.inputs[0]) == cond0_name:
                    not0_idx = int(user_idx)
            if select0_idx is None or not0_idx is None:
                continue
            select0_op = model_ir.operators[int(select0_idx)]
            not0_op = model_ir.operators[int(not0_idx)]
            select0_out_name = str(select0_op.outputs[0])
            not0_out_name = str(not0_op.outputs[0])
            zero_nchw_name = str(select0_op.inputs[1])
            base_nchw_from_select0 = str(select0_op.inputs[2])
            if base_nchw_from_select0 != base_nchw_name:
                continue

            # where0 -> reshape40 -> pool2 -> reshape42 -> equal1 -> and1
            reshape40_idx = _find_single_consumer_op(
                select0_out_name,
                consumers,
                expected_op_type="RESHAPE",
                operators=model_ir.operators,
                required_input_index=0,
            )
            if reshape40_idx is None:
                continue
            reshape40_op = model_ir.operators[int(reshape40_idx)]
            reshape40_out_name = str(reshape40_op.outputs[0])
            pool2_idx = _find_single_consumer_op(
                reshape40_out_name,
                consumers,
                expected_op_type="MAX_POOL_2D",
                operators=model_ir.operators,
            )
            if pool2_idx is None:
                continue
            pool2_op = model_ir.operators[int(pool2_idx)]
            if len(pool2_op.inputs) != 1 or len(pool2_op.outputs) != 1:
                continue
            pool2_out_name = str(pool2_op.outputs[0])
            reshape42_idx = _find_single_consumer_op(
                pool2_out_name,
                consumers,
                expected_op_type="RESHAPE",
                operators=model_ir.operators,
            )
            if reshape42_idx is None:
                continue
            reshape42_op = model_ir.operators[int(reshape42_idx)]
            reshape42_out_name = str(reshape42_op.outputs[0])
            equal1_idx = _find_single_consumer_op(
                reshape42_out_name,
                consumers,
                expected_op_type="EQUAL",
                operators=model_ir.operators,
            )
            if equal1_idx is None:
                continue
            equal1_op = model_ir.operators[int(equal1_idx)]
            if len(equal1_op.inputs) != 2 or len(equal1_op.outputs) != 1:
                continue
            if str(equal1_op.inputs[0]) != select0_out_name and str(equal1_op.inputs[1]) != select0_out_name:
                continue
            if str(equal1_op.inputs[0]) != reshape42_out_name and str(equal1_op.inputs[1]) != reshape42_out_name:
                continue
            equal1_out_name = str(equal1_op.outputs[0])
            and1_idx = _find_single_consumer_op(
                equal1_out_name,
                consumers,
                expected_op_type="LOGICAL_AND",
                operators=model_ir.operators,
            )
            if and1_idx is None:
                continue
            and1_op = model_ir.operators[int(and1_idx)]
            if len(and1_op.inputs) != 2 or len(and1_op.outputs) != 1:
                continue
            if not0_out_name not in set(str(v) for v in and1_op.inputs):
                continue
            and1_out_name = str(and1_op.outputs[0])

            # or0 with eq0_nchw adapter
            or0_idx = _find_single_consumer_op(
                and1_out_name,
                consumers,
                expected_op_type="LOGICAL_OR",
                operators=model_ir.operators,
            )
            if or0_idx is None:
                continue
            or0_op = model_ir.operators[int(or0_idx)]
            if len(or0_op.inputs) != 2 or len(or0_op.outputs) != 1:
                continue
            eq0_nchw_name = (
                str(or0_op.inputs[0])
                if str(or0_op.inputs[1]) == and1_out_name
                else str(or0_op.inputs[1])
            )
            eq0_nchw_adapter_idx = producers.get(eq0_nchw_name, None)
            if eq0_nchw_adapter_idx is None:
                continue
            eq0_nchw_adapter_op = model_ir.operators[int(eq0_nchw_adapter_idx)]
            if (
                str(eq0_nchw_adapter_op.op_type) != "RESHAPE"
                or len(eq0_nchw_adapter_op.inputs) < 1
                or len(eq0_nchw_adapter_op.outputs) != 1
                or str(eq0_nchw_adapter_op.outputs[0]) != eq0_nchw_name
                or str(eq0_nchw_adapter_op.inputs[0]) != eq0_nhwc_name
            ):
                continue
            or0_out_name = str(or0_op.outputs[0])

            # cast1 -> reshape47 -> pool3 -> reshape49 -> greater1 -> select1/not1
            cast1_idx = _find_single_consumer_op(
                or0_out_name,
                consumers,
                expected_op_type="CAST",
                operators=model_ir.operators,
                required_input_index=0,
            )
            if cast1_idx is None:
                continue
            cast1_op = model_ir.operators[int(cast1_idx)]
            if len(cast1_op.inputs) != 1 or len(cast1_op.outputs) != 1:
                continue
            cast1_out_name = str(cast1_op.outputs[0])
            reshape47_idx = _find_single_consumer_op(
                cast1_out_name,
                consumers,
                expected_op_type="RESHAPE",
                operators=model_ir.operators,
            )
            if reshape47_idx is None:
                continue
            reshape47_op = model_ir.operators[int(reshape47_idx)]
            reshape47_out_name = str(reshape47_op.outputs[0])
            pool3_idx = _find_single_consumer_op(
                reshape47_out_name,
                consumers,
                expected_op_type="MAX_POOL_2D",
                operators=model_ir.operators,
            )
            if pool3_idx is None:
                continue
            pool3_op = model_ir.operators[int(pool3_idx)]
            if len(pool3_op.inputs) != 1 or len(pool3_op.outputs) != 1:
                continue
            pool3_out_name = str(pool3_op.outputs[0])
            reshape49_idx = _find_single_consumer_op(
                pool3_out_name,
                consumers,
                expected_op_type="RESHAPE",
                operators=model_ir.operators,
            )
            if reshape49_idx is None:
                continue
            reshape49_op = model_ir.operators[int(reshape49_idx)]
            reshape49_out_name = str(reshape49_op.outputs[0])
            greater1_idx = _find_single_consumer_op(
                reshape49_out_name,
                consumers,
                expected_op_type="GREATER",
                operators=model_ir.operators,
            )
            if greater1_idx is None:
                continue
            greater1_op = model_ir.operators[int(greater1_idx)]
            if len(greater1_op.inputs) != 2 or len(greater1_op.outputs) != 1:
                continue
            cond1_name = str(greater1_op.outputs[0])
            cond1_users = [int(v) for v in consumers.get(cond1_name, [])]
            if len(cond1_users) != 2:
                continue
            select1_idx: Optional[int] = None
            not1_idx: Optional[int] = None
            for user_idx in cond1_users:
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) == "SELECT" and len(user_op.inputs) == 3 and str(user_op.inputs[0]) == cond1_name:
                    select1_idx = int(user_idx)
                elif str(user_op.op_type) == "LOGICAL_NOT" and len(user_op.inputs) == 1 and str(user_op.inputs[0]) == cond1_name:
                    not1_idx = int(user_idx)
            if select1_idx is None or not1_idx is None:
                continue
            select1_op = model_ir.operators[int(select1_idx)]
            not1_op = model_ir.operators[int(not1_idx)]
            select1_out_name = str(select1_op.outputs[0])
            not1_out_name = str(not1_op.outputs[0])
            if str(select1_op.inputs[1]) != zero_nchw_name or str(select1_op.inputs[2]) != base_nchw_name:
                continue

            # where1 -> reshape53 -> pool4 -> reshape55 -> equal2 -> and2
            reshape53_idx = _find_single_consumer_op(
                select1_out_name,
                consumers,
                expected_op_type="RESHAPE",
                operators=model_ir.operators,
                required_input_index=0,
            )
            if reshape53_idx is None:
                continue
            reshape53_op = model_ir.operators[int(reshape53_idx)]
            reshape53_out_name = str(reshape53_op.outputs[0])
            pool4_idx = _find_single_consumer_op(
                reshape53_out_name,
                consumers,
                expected_op_type="MAX_POOL_2D",
                operators=model_ir.operators,
            )
            if pool4_idx is None:
                continue
            pool4_op = model_ir.operators[int(pool4_idx)]
            if len(pool4_op.inputs) != 1 or len(pool4_op.outputs) != 1:
                continue
            pool4_out_name = str(pool4_op.outputs[0])
            reshape55_idx = _find_single_consumer_op(
                pool4_out_name,
                consumers,
                expected_op_type="RESHAPE",
                operators=model_ir.operators,
            )
            if reshape55_idx is None:
                continue
            reshape55_op = model_ir.operators[int(reshape55_idx)]
            reshape55_out_name = str(reshape55_op.outputs[0])
            equal2_idx = _find_single_consumer_op(
                reshape55_out_name,
                consumers,
                expected_op_type="EQUAL",
                operators=model_ir.operators,
            )
            if equal2_idx is None:
                continue
            equal2_op = model_ir.operators[int(equal2_idx)]
            if len(equal2_op.inputs) != 2 or len(equal2_op.outputs) != 1:
                continue
            if str(equal2_op.inputs[0]) != select1_out_name and str(equal2_op.inputs[1]) != select1_out_name:
                continue
            if str(equal2_op.inputs[0]) != reshape55_out_name and str(equal2_op.inputs[1]) != reshape55_out_name:
                continue
            equal2_out_name = str(equal2_op.outputs[0])
            and2_idx = _find_single_consumer_op(
                equal2_out_name,
                consumers,
                expected_op_type="LOGICAL_AND",
                operators=model_ir.operators,
            )
            if and2_idx is None:
                continue
            and2_op = model_ir.operators[int(and2_idx)]
            if len(and2_op.inputs) != 2 or len(and2_op.outputs) != 1:
                continue
            if not1_out_name not in set(str(v) for v in and2_op.inputs):
                continue
            and2_out_name = str(and2_op.outputs[0])

            # or1 -> select2
            or1_idx = _find_single_consumer_op(
                and2_out_name,
                consumers,
                expected_op_type="LOGICAL_OR",
                operators=model_ir.operators,
            )
            if or1_idx is None:
                continue
            or1_op = model_ir.operators[int(or1_idx)]
            if len(or1_op.inputs) != 2 or len(or1_op.outputs) != 1:
                continue
            if or0_out_name not in set(str(v) for v in or1_op.inputs):
                continue
            or1_out_name = str(or1_op.outputs[0])
            select2_idx = _find_single_consumer_op(
                or1_out_name,
                consumers,
                expected_op_type="SELECT",
                operators=model_ir.operators,
            )
            if select2_idx is None:
                continue
            select2_op = model_ir.operators[int(select2_idx)]
            if len(select2_op.inputs) != 3 or len(select2_op.outputs) != 1:
                continue
            if str(select2_op.inputs[0]) != or1_out_name:
                continue
            if set([str(select2_op.inputs[1]), str(select2_op.inputs[2])]) != set([base_nchw_name, zero_nchw_name]):
                continue
            select2_out_nchw_name = str(select2_op.outputs[0])
            if select2_out_nchw_name in model_outputs:
                continue

            # Create zero_nhwc adapter from zero_nchw.
            zero_nhwc_shape = [int(v) for v in list(base_nhwc_shape)]
            zero_nhwc_shape_name = _unique_tensor_name(f"{zero_nchw_name}_nhwc_shape")
            zero_nhwc_name = _unique_tensor_name(f"{zero_nchw_name}_nhwc")
            model_ir.tensors[zero_nhwc_shape_name] = TensorIR(
                name=zero_nhwc_shape_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray(zero_nhwc_shape, dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            zero_nchw_tensor = model_ir.tensors.get(zero_nchw_name, None)
            model_ir.tensors[zero_nhwc_name] = TensorIR(
                name=zero_nhwc_name,
                dtype=(
                    str(zero_nchw_tensor.dtype)
                    if zero_nchw_tensor is not None
                    else "FLOAT32"
                ),
                shape=[int(v) for v in list(zero_nhwc_shape)],
                shape_signature=[int(v) for v in list(zero_nhwc_shape)],
                data=None,
                is_variable=False,
                quantization=(
                    _clone_quantization(zero_nchw_tensor.quantization)
                    if zero_nchw_tensor is not None
                    else None
                ),
            )

            select0_current_idx = next((int(i) for i, op in enumerate(model_ir.operators) if op is select0_op), None)
            if select0_current_idx is None:
                continue
            model_ir.operators.insert(
                int(select0_current_idx),
                OperatorIR(
                    op_type="RESHAPE",
                    inputs=[str(zero_nchw_name), str(zero_nhwc_shape_name)],
                    outputs=[str(zero_nhwc_name)],
                    options={"newShape": [int(v) for v in list(zero_nhwc_shape)]},
                ),
            )

            # Rewire to NHWC.
            greater0_inputs = [str(v) for v in list(greater0_op.inputs)]
            greater0_inputs = [
                str(post1_in_name) if str(v) == str(post1_out_name) else str(v)
                for v in list(greater0_inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=greater0_op,
                new_inputs=greater0_inputs,
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=select0_op,
                new_inputs=[str(cond0_name), str(zero_nhwc_name), str(base_nhwc_name)],
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=pool2_op,
                new_inputs=[str(select0_out_name)],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=equal1_op,
                new_inputs=[str(select0_out_name), str(pool2_out_name)],
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=or0_op,
                new_inputs=[
                    str(eq0_nhwc_name) if str(v) == str(eq0_nchw_name) else str(v)
                    for v in list(or0_op.inputs)
                ],
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=pool3_op,
                new_inputs=[str(cast1_out_name)],
            )
            greater1_inputs = [str(v) for v in list(greater1_op.inputs)]
            greater1_inputs = [
                str(pool3_out_name) if str(v) == str(reshape49_out_name) else str(v)
                for v in list(greater1_inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=greater1_op,
                new_inputs=greater1_inputs,
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=select1_op,
                new_inputs=[str(cond1_name), str(zero_nhwc_name), str(base_nhwc_name)],
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=pool4_op,
                new_inputs=[str(select1_out_name)],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=equal2_op,
                new_inputs=[str(select1_out_name), str(pool4_out_name)],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=select2_op,
                new_inputs=[str(or1_out_name), str(base_nhwc_name), str(zero_nhwc_name)],
            )

            # Convert NHWC output of final SELECT back to NCHW for downstream.
            select2_out_nhwc_name = _unique_tensor_name(f"{select2_out_nchw_name}_nhwc")
            model_ir.tensors[select2_out_nhwc_name] = TensorIR(
                name=select2_out_nhwc_name,
                dtype=str(model_ir.tensors[select2_out_nchw_name].dtype)
                if select2_out_nchw_name in model_ir.tensors
                else "BOOL",
                shape=[int(v) for v in list(base_nhwc_shape)],
                shape_signature=[int(v) for v in list(base_nhwc_shape)],
                data=None,
                is_variable=False,
                quantization=(
                    _clone_quantization(model_ir.tensors[select2_out_nchw_name].quantization)
                    if select2_out_nchw_name in model_ir.tensors
                    else None
                ),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=select2_op,
                new_outputs=[str(select2_out_nhwc_name)],
            )

            final_nchw_shape_name = _unique_tensor_name(f"{select2_out_nchw_name}_shape_nchw")
            model_ir.tensors[final_nchw_shape_name] = TensorIR(
                name=final_nchw_shape_name,
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray(base_nchw_shape, dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            select2_current_idx = next((int(i) for i, op in enumerate(model_ir.operators) if op is select2_op), None)
            if select2_current_idx is None:
                continue
            model_ir.operators.insert(
                int(select2_current_idx) + 1,
                OperatorIR(
                    op_type="RESHAPE",
                    inputs=[str(select2_out_nhwc_name), str(final_nchw_shape_name)],
                    outputs=[str(select2_out_nchw_name)],
                    options={"newShape": [int(v) for v in list(base_nchw_shape)]},
                ),
            )

            # Update tensor metadata for NHWC-propagated logical tensors.
            nhwc_tensors = [
                str(cond0_name),
                str(select0_out_name),
                str(not0_out_name),
                str(equal1_out_name),
                str(and1_out_name),
                str(or0_out_name),
                str(cast1_out_name),
                str(cond1_name),
                str(select1_out_name),
                str(not1_out_name),
                str(equal2_out_name),
                str(and2_out_name),
                str(or1_out_name),
                str(select2_out_nhwc_name),
            ]
            for tensor_name in nhwc_tensors:
                _set_rank4_shape(tensor_name, [int(v) for v in list(base_nhwc_shape)])

            # Remove redundant middle NCHW adapters.
            remove_refs: List[OperatorIR] = [
                eq0_nchw_adapter_op,
                post1_op,
                reshape40_op,
                reshape42_op,
                reshape47_op,
                reshape49_op,
                reshape53_op,
                reshape55_op,
            ]
            remove_indices: List[int] = []
            for op_idx, op_ref in enumerate(model_ir.operators):
                if any(op_ref is target for target in remove_refs):
                    remove_indices.append(int(op_idx))
            for remove_idx in sorted(list(set(remove_indices)), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_singleton_nms_maxpool_nhwc_chains": int(rewritten)}


def _optimize_consecutive_reshape_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant RESHAPE chains.

    Target:
      x --RESHAPE(no-op)--> y
      x --RESHAPE--> y --RESHAPE--> z

    Rewrite:
      x ----------------> y
      x -----------RESHAPE--> z

    Safety:
    - No-op reshape removal requires matching input/output shape and shape_signature.
    - For graph-output tensors, preserve visible output names.
    - Middle tensor `y` can have fan-out. In that case the second RESHAPE is
      rewired to the original source while keeping the first RESHAPE for other users.
    - First RESHAPE removal still requires `y` to be consumed only by the second RESHAPE.
    - `y` must not be a graph output.
    - Source and destination element counts must match with fully known static shapes.
    """
    rewritten = 0
    rewritten_fanout_bypass = 0
    removed_noop = 0

    def _reshape_depends_on_input_dims(reshape_op: OperatorIR) -> bool:
        """
        Return True when RESHAPE target depends on source tensor dimensions.
        Such reshapes must not be bypassed across an intermediate reshape,
        otherwise ONNX semantics of 0/-1 placeholders can change.
        """
        options = (
            dict(reshape_op.options)
            if isinstance(reshape_op.options, dict)
            else {}
        )
        raw_shape = options.get("onnxRawNewShape", None)
        if raw_shape is not None:
            try:
                values = [int(v) for v in np.asarray(raw_shape).reshape(-1).tolist()]
            except Exception:
                return True
            if any(int(v) <= 0 for v in values):
                return True
            return False
        new_shape = options.get("newShape", None)
        if new_shape is None:
            return False
        try:
            values = [int(v) for v in np.asarray(new_shape).reshape(-1).tolist()]
        except Exception:
            return True
        return any(int(v) <= 0 for v in values)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        # 1) Remove single no-op reshape: input/output metadata is identical.
        for reshape_idx, reshape_op in enumerate(model_ir.operators):
            if str(reshape_op.op_type) != "RESHAPE" or len(reshape_op.inputs) < 1 or len(reshape_op.outputs) != 1:
                continue
            if bool(reshape_op.options.get("preserveDynamicShape", False)):
                continue

            src_name = str(reshape_op.inputs[0])
            dst_name = str(reshape_op.outputs[0])
            src_tensor = model_ir.tensors.get(src_name, None)
            dst_tensor = model_ir.tensors.get(dst_name, None)
            if src_tensor is None or dst_tensor is None:
                continue

            src_shape = [int(v) for v in list(src_tensor.shape)]
            dst_shape = [int(v) for v in list(dst_tensor.shape)]
            src_signature = (
                [int(v) for v in list(src_tensor.shape_signature)]
                if src_tensor.shape_signature is not None
                else [int(v) for v in list(src_shape)]
            )
            dst_signature = (
                [int(v) for v in list(dst_tensor.shape_signature)]
                if dst_tensor.shape_signature is not None
                else [int(v) for v in list(dst_shape)]
            )
            if src_shape != dst_shape or src_signature != dst_signature:
                continue

            # Keep graph-output tensor names stable when possible.
            if dst_name in model_outputs:
                if len(consumers.get(dst_name, [])) != 0:
                    continue
                producer_idx = producers.get(src_name, None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                producer_outputs = [str(v) for v in list(producer_op.outputs)]
                if src_name not in set(producer_outputs):
                    continue
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=producer_op,
                    new_outputs=[
                        str(dst_name) if str(v) == str(src_name) else str(v)
                        for v in producer_outputs
                    ],
                )
            else:
                _replace_tensor_inputs(
                    model_ir=model_ir,
                    src_name=str(dst_name),
                    dst_name=str(src_name),
                )

            del model_ir.operators[int(reshape_idx)]
            removed_noop += 1
            changed = True
            break

        if changed:
            continue

        # 2) Bypass intermediate reshape for RESHAPE->RESHAPE links, even with fan-out.
        for first_idx, first_op in enumerate(model_ir.operators):
            if str(first_op.op_type) != "RESHAPE" or len(first_op.inputs) < 1 or len(first_op.outputs) != 1:
                continue

            first_input_name = str(first_op.inputs[0])
            first_output_name = str(first_op.outputs[0])
            if first_output_name in model_outputs:
                continue

            first_input_tensor = model_ir.tensors.get(first_input_name, None)
            if first_input_tensor is None or not _is_fully_known_positive_shape(first_input_tensor.shape):
                continue
            first_input_shape = [int(v) for v in list(first_input_tensor.shape)]
            if len(first_input_shape) == 0:
                continue
            first_input_elements = int(np.prod(np.asarray(first_input_shape, dtype=np.int64)))

            first_users = [int(v) for v in consumers.get(first_output_name, []) if int(v) != int(first_idx)]
            if len(first_users) <= 1:
                continue

            for second_idx in first_users:
                second_op = model_ir.operators[int(second_idx)]
                if str(second_op.op_type) != "RESHAPE" or len(second_op.inputs) < 1 or len(second_op.outputs) != 1:
                    continue
                if str(second_op.inputs[0]) != first_output_name:
                    continue
                if _reshape_depends_on_input_dims(second_op):
                    continue

                second_output_name = str(second_op.outputs[0])
                second_output_tensor = model_ir.tensors.get(second_output_name, None)
                if second_output_tensor is None or not _is_fully_known_positive_shape(second_output_tensor.shape):
                    continue
                second_output_shape = [int(v) for v in list(second_output_tensor.shape)]
                if len(second_output_shape) == 0:
                    continue
                second_output_elements = int(np.prod(np.asarray(second_output_shape, dtype=np.int64)))
                if int(first_input_elements) != int(second_output_elements):
                    continue

                second_inputs = [str(v) for v in list(second_op.inputs)]
                if str(second_inputs[0]) == str(first_input_name):
                    continue
                second_inputs[0] = str(first_input_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=second_op,
                    new_inputs=second_inputs,
                )

                rewritten += 1
                rewritten_fanout_bypass += 1
                changed = True
                break

            if changed:
                break

        if changed:
            continue

        # 3) Remove consecutive reshape chains when the middle tensor has a single user.
        for first_idx, first_op in enumerate(model_ir.operators):
            if str(first_op.op_type) != "RESHAPE" or len(first_op.inputs) < 1 or len(first_op.outputs) != 1:
                continue

            first_input_name = str(first_op.inputs[0])
            first_output_name = str(first_op.outputs[0])
            if first_output_name in model_outputs:
                continue

            first_users = [int(v) for v in consumers.get(first_output_name, [])]
            if len(first_users) != 1:
                continue

            second_idx = int(first_users[0])
            if int(second_idx) == int(first_idx):
                continue
            second_op = model_ir.operators[int(second_idx)]
            if str(second_op.op_type) != "RESHAPE" or len(second_op.inputs) < 1 or len(second_op.outputs) != 1:
                continue
            if str(second_op.inputs[0]) != first_output_name:
                continue
            if _reshape_depends_on_input_dims(second_op):
                continue

            second_output_name = str(second_op.outputs[0])

            first_input_tensor = model_ir.tensors.get(first_input_name, None)
            second_output_tensor = model_ir.tensors.get(second_output_name, None)
            if (
                first_input_tensor is None
                or second_output_tensor is None
                or not _is_fully_known_positive_shape(first_input_tensor.shape)
                or not _is_fully_known_positive_shape(second_output_tensor.shape)
            ):
                continue

            first_input_shape = [int(v) for v in list(first_input_tensor.shape)]
            second_output_shape = [int(v) for v in list(second_output_tensor.shape)]
            if len(first_input_shape) == 0 or len(second_output_shape) == 0:
                continue
            if int(np.prod(np.asarray(first_input_shape, dtype=np.int64))) != int(
                np.prod(np.asarray(second_output_shape, dtype=np.int64))
            ):
                continue

            second_inputs = [str(v) for v in list(second_op.inputs)]
            second_inputs[0] = str(first_input_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=second_op,
                new_inputs=second_inputs,
            )

            del model_ir.operators[int(first_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0 or removed_noop > 0:
        _prune_unused_tensors(model_ir)
    return {
        "removed_noop_reshape_chains": int(removed_noop),
        "rewritten_consecutive_reshape_passthrough_chains": int(rewritten),
        "rewritten_fanout_bypass_reshape_passthrough_chains": int(rewritten_fanout_bypass),
    }


def _optimize_flatten_concat_expanddims_to_nhwc_concat(model_ir: ModelIR) -> Dict[str, int]:
    """
    Rewrite 4D->2D flatten concat and 2D->4D reshape back to direct NHWC concat.

    Target:
      a4d[N,1,1,Ca] --RESHAPE--> a2d[N,Ca]
      b2d[N,Cb]
      CONCAT(axis=1, [a2d,b2d]) -> c2d[N,Ca+Cb]
      RESHAPE -> c4d[N,1,1,Ca+Cb]

    Rewrite:
      b2d --RESHAPE--> b4d[N,1,1,Cb]
      CONCAT(axis=3, [a4d,b4d]) -> c4d
    """
    rewritten = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "RESHAPE" or len(pre_op.inputs) < 1 or len(pre_op.outputs) != 1:
                continue

            a4d_name = str(pre_op.inputs[0])
            a2d_name = str(pre_op.outputs[0])
            if a2d_name in model_outputs:
                continue

            concat_users = [int(v) for v in consumers.get(a2d_name, [])]
            if len(concat_users) != 1:
                continue
            concat_idx = int(concat_users[0])
            concat_op = model_ir.operators[int(concat_idx)]
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_axis = int(concat_op.options.get("axis", 1))
            if concat_axis < 0:
                concat_axis += 2
            if concat_axis != 1:
                continue
            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) != 2 or a2d_name not in concat_inputs:
                continue

            c2d_name = str(concat_op.outputs[0])
            if c2d_name in model_outputs:
                continue
            c2d_users = [int(v) for v in consumers.get(c2d_name, [])]
            if len(c2d_users) != 1:
                continue
            post_idx = int(c2d_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "RESHAPE" or len(post_op.inputs) < 1 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != c2d_name:
                continue
            c4d_name = str(post_op.outputs[0])

            a4d_tensor = model_ir.tensors.get(a4d_name, None)
            a2d_tensor = model_ir.tensors.get(a2d_name, None)
            c2d_tensor = model_ir.tensors.get(c2d_name, None)
            c4d_tensor = model_ir.tensors.get(c4d_name, None)
            if (
                a4d_tensor is None
                or a2d_tensor is None
                or c2d_tensor is None
                or c4d_tensor is None
                or not _is_fully_known_positive_shape(a4d_tensor.shape)
                or not _is_fully_known_positive_shape(a2d_tensor.shape)
                or not _is_fully_known_positive_shape(c2d_tensor.shape)
                or not _is_fully_known_positive_shape(c4d_tensor.shape)
            ):
                continue

            a4d_shape = [int(v) for v in list(a4d_tensor.shape)]
            a2d_shape = [int(v) for v in list(a2d_tensor.shape)]
            c2d_shape = [int(v) for v in list(c2d_tensor.shape)]
            c4d_shape = [int(v) for v in list(c4d_tensor.shape)]
            if len(a4d_shape) != 4 or len(a2d_shape) != 2 or len(c2d_shape) != 2 or len(c4d_shape) != 4:
                continue
            if int(a4d_shape[1]) != 1 or int(a4d_shape[2]) != 1:
                continue
            if int(c4d_shape[1]) != 1 or int(c4d_shape[2]) != 1:
                continue
            if int(a2d_shape[0]) != int(a4d_shape[0]) or int(a2d_shape[1]) != int(a4d_shape[3]):
                continue
            if int(c4d_shape[0]) != int(c2d_shape[0]) or int(c4d_shape[3]) != int(c2d_shape[1]):
                continue

            b2d_name = str(concat_inputs[0]) if str(concat_inputs[1]) == a2d_name else str(concat_inputs[1])
            b2d_tensor = model_ir.tensors.get(b2d_name, None)
            if (
                b2d_tensor is None
                or not _is_fully_known_positive_shape(b2d_tensor.shape)
                or len(list(b2d_tensor.shape)) != 2
            ):
                continue
            b2d_shape = [int(v) for v in list(b2d_tensor.shape)]
            if int(b2d_shape[0]) != int(a4d_shape[0]):
                continue

            expected_c = int(a4d_shape[3]) + int(b2d_shape[1])
            if int(c4d_shape[3]) != expected_c:
                continue

            b4d_shape = [int(b2d_shape[0]), 1, 1, int(b2d_shape[1])]
            b4d_name = _unique_tensor_name(f"{b2d_name}_nhwc")
            b4d_shape_name = _unique_tensor_name(f"{b4d_name}_reshape_shape")
            model_ir.tensors[b4d_shape_name] = TensorIR(
                name=str(b4d_shape_name),
                dtype="INT32",
                shape=[4],
                shape_signature=[4],
                data=np.asarray(b4d_shape, dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            model_ir.tensors[b4d_name] = TensorIR(
                name=str(b4d_name),
                dtype=str(b2d_tensor.dtype),
                shape=[int(v) for v in list(b4d_shape)],
                shape_signature=[int(v) for v in list(b4d_shape)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(b2d_tensor.quantization),
            )
            concat_pos = int(model_ir.operators.index(concat_op))
            model_ir.operators.insert(
                int(concat_pos),
                OperatorIR(
                    op_type="RESHAPE",
                    inputs=[str(b2d_name), str(b4d_shape_name)],
                    outputs=[str(b4d_name)],
                    options={"newShape": [int(v) for v in list(b4d_shape)]},
                ),
            )

            # Re-find after insertion.
            concat_op = next((op for op in model_ir.operators if op is concat_op), None)
            post_op = next((op for op in model_ir.operators if op is post_op), None)
            if concat_op is None or post_op is None:
                continue

            concat_in_4d = [str(a4d_name), str(b4d_name)]
            if str(concat_inputs[0]) != str(a2d_name):
                concat_in_4d = [str(b4d_name), str(a4d_name)]
            concat_op.options = {
                **(dict(concat_op.options) if isinstance(concat_op.options, dict) else {}),
                "axis": 3,
            }
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=concat_in_4d,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[str(c4d_name)],
            )
            _replace_tensor_inputs(model_ir, c2d_name, c4d_name)

            remove_indices: List[int] = []
            pre_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is pre_op), None)
            if pre_remove_idx is not None:
                remove_indices.append(int(pre_remove_idx))
            post_remove_idx = next((idx for idx, op in enumerate(model_ir.operators) if op is post_op), None)
            if post_remove_idx is not None:
                remove_indices.append(int(post_remove_idx))
            for remove_idx in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"rewritten_flatten_concat_expanddims_to_nhwc_concat": int(rewritten)}


def _optimize_singleton_spatial_nhwc_transpose_reshape_flatten(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove NHWC->NCHW transpose before 2D flatten reshape when spatial dims are 1x1.

    Pattern:
      x_nhwc [N,1,1,C] -> TRANSPOSE(0,3,1,2) -> x_nchw [N,C,1,1]
      x_nchw -> RESHAPE -> y [N,C]

    Extended pattern:
      x_nhwc [N,1,1,C] -> TRANSPOSE(0,3,1,2) -> x_nchw [N,C,1,1]
      x_nchw -> RESHAPE(identity) -> x_keep [N,C,1,1]
      x_keep -> RESHAPE -> y [N,C]

    For singleton spatial dims, flatten result is identical without transpose:
      RESHAPE(x_nhwc) -> y
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for transpose_idx, transpose_op in enumerate(model_ir.operators):
            if (
                str(transpose_op.op_type) != "TRANSPOSE"
                or len(transpose_op.inputs) < 2
                or len(transpose_op.outputs) != 1
                or _read_transpose_perm(model_ir, transpose_op) != perm_nhwc_to_nchw
            ):
                continue

            transpose_input_name = str(transpose_op.inputs[0])
            transpose_output_name = str(transpose_op.outputs[0])
            users = [int(v) for v in consumers.get(transpose_output_name, [])]
            if len(users) != 1:
                continue

            first_reshape_idx = int(users[0])
            first_reshape_op = model_ir.operators[int(first_reshape_idx)]
            if (
                str(first_reshape_op.op_type) != "RESHAPE"
                or len(first_reshape_op.inputs) < 1
                or str(first_reshape_op.inputs[0]) != transpose_output_name
                or len(first_reshape_op.outputs) != 1
            ):
                continue

            target_reshape_idx = int(first_reshape_idx)
            target_reshape_op = first_reshape_op
            target_reshape_output_name = str(target_reshape_op.outputs[0])
            intermediate_reshape_idx: Optional[int] = None
            intermediate_reshape_output_name = ""

            first_reshape_output_name = str(first_reshape_op.outputs[0])
            first_reshape_output_tensor = model_ir.tensors.get(first_reshape_output_name, None)
            if (
                first_reshape_output_tensor is not None
                and len(list(first_reshape_output_tensor.shape)) == 4
                and first_reshape_output_name not in model_outputs
            ):
                second_users = [int(v) for v in consumers.get(first_reshape_output_name, [])]
                if len(second_users) == 1:
                    second_reshape_idx = int(second_users[0])
                    second_reshape_op = model_ir.operators[int(second_reshape_idx)]
                    if (
                        str(second_reshape_op.op_type) == "RESHAPE"
                        and len(second_reshape_op.inputs) >= 1
                        and str(second_reshape_op.inputs[0]) == first_reshape_output_name
                        and len(second_reshape_op.outputs) == 1
                    ):
                        second_reshape_output_name = str(second_reshape_op.outputs[0])
                        second_reshape_output_tensor = model_ir.tensors.get(second_reshape_output_name, None)
                        if (
                            second_reshape_output_tensor is not None
                            and len(list(second_reshape_output_tensor.shape)) == 2
                        ):
                            target_reshape_idx = int(second_reshape_idx)
                            target_reshape_op = second_reshape_op
                            target_reshape_output_name = second_reshape_output_name
                            intermediate_reshape_idx = int(first_reshape_idx)
                            intermediate_reshape_output_name = first_reshape_output_name

            input_tensor = model_ir.tensors.get(transpose_input_name, None)
            output_tensor = model_ir.tensors.get(transpose_output_name, None)
            reshape_output_tensor = model_ir.tensors.get(target_reshape_output_name, None)
            if (
                input_tensor is None
                or output_tensor is None
                or reshape_output_tensor is None
                or not _is_fully_known_positive_shape(input_tensor.shape)
                or not _is_fully_known_positive_shape(output_tensor.shape)
                or not _is_fully_known_positive_shape(reshape_output_tensor.shape)
            ):
                continue

            input_shape = [int(v) for v in list(input_tensor.shape)]
            output_shape = [int(v) for v in list(output_tensor.shape)]
            reshape_output_shape = [int(v) for v in list(reshape_output_tensor.shape)]
            if len(input_shape) != 4 or len(output_shape) != 4 or len(reshape_output_shape) != 2:
                continue
            if int(input_shape[1]) != 1 or int(input_shape[2]) != 1:
                continue
            expected_output_shape = _permute_shape(input_shape, perm_nhwc_to_nchw)
            if expected_output_shape is None or [int(v) for v in list(expected_output_shape)] != output_shape:
                continue
            if intermediate_reshape_idx is not None:
                intermediate_tensor = model_ir.tensors.get(intermediate_reshape_output_name, None)
                if intermediate_tensor is None:
                    continue
                intermediate_shape = [int(v) for v in list(intermediate_tensor.shape)]
                if intermediate_shape != output_shape:
                    continue
            if int(reshape_output_shape[1]) != int(input_shape[3]):
                continue

            reshape_inputs = [str(v) for v in list(target_reshape_op.inputs)]
            reshape_inputs[0] = str(transpose_input_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=target_reshape_op,
                new_inputs=reshape_inputs,
            )
            remove_indices = [int(transpose_idx)]
            if (
                intermediate_reshape_idx is not None
                and intermediate_reshape_output_name not in model_outputs
                and set(int(v) for v in consumers.get(intermediate_reshape_output_name, [])) == {int(target_reshape_idx)}
            ):
                remove_indices.append(int(intermediate_reshape_idx))
            for remove_idx in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_singleton_spatial_nhwc_transpose_reshape_flatten": int(rewritten)}


def _optimize_singleton_reshape_concat_post_transpose_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Rewrite singleton-channel NCHW concat blocks back to NHWC and remove trailing transpose.

    Target:
      x_i_nhwc --RESHAPE([N,1,H,W])--> x_i_nchw (subset of inputs; some may be reused elsewhere)
      c_nchw --CONCAT(axis=1)--> y_nchw --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      - CONCAT inputs switch to NHWC aliases (reshape sources when available)
      - direct singleton NCHW inputs receive local NCHW->NHWC RESHAPE adapters
      - CONCAT axis becomes 3
      - trailing TRANSPOSE is removed and CONCAT directly produces y_nhwc
    """
    rewritten = 0
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue

            concat_output_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if concat_output_name in model_outputs:
                continue

            concat_idx = producers.get(concat_output_name, None)
            if concat_idx is None:
                continue
            concat_op = model_ir.operators[int(concat_idx)]
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            if str(concat_op.outputs[0]) != concat_output_name:
                continue
            concat_axis = int(concat_op.options.get("axis", 1))
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue
            if set(int(v) for v in consumers.get(concat_output_name, [])) != {int(post_idx)}:
                continue

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue

            mapped_concat_inputs: List[str] = []
            reshape_ops_to_remove: List[OperatorIR] = []
            adapters_to_insert: List[Tuple[str, str, str, List[int]]] = []
            valid = True

            for input_name in concat_inputs:
                input_tensor = model_ir.tensors.get(str(input_name), None)
                if (
                    input_tensor is None
                    or not _is_fully_known_positive_shape(input_tensor.shape)
                    or len(list(input_tensor.shape)) != 4
                ):
                    valid = False
                    break
                input_shape = [int(v) for v in list(input_tensor.shape)]

                input_prod_idx = producers.get(str(input_name), None)
                if input_prod_idx is not None:
                    input_prod_op = model_ir.operators[int(input_prod_idx)]
                    if (
                        str(input_prod_op.op_type) == "RESHAPE"
                        and len(input_prod_op.inputs) >= 1
                        and len(input_prod_op.outputs) == 1
                        and str(input_prod_op.outputs[0]) == str(input_name)
                    ):
                        source_name = str(input_prod_op.inputs[0])
                        source_tensor = model_ir.tensors.get(source_name, None)
                        if (
                            source_tensor is not None
                            and _is_fully_known_positive_shape(source_tensor.shape)
                            and len(list(source_tensor.shape)) == 4
                        ):
                            source_shape = [int(v) for v in list(source_tensor.shape)]
                            expected_shape = _permute_shape(source_shape, perm_nhwc_to_nchw)
                            if (
                                expected_shape is not None
                                and [int(v) for v in list(expected_shape)] == input_shape
                                and int(source_shape[3]) == 1
                                and int(input_shape[1]) == 1
                            ):
                                mapped_concat_inputs.append(str(source_name))
                                input_users = [int(v) for v in consumers.get(str(input_name), [])]
                                if set(input_users) == {int(concat_idx)}:
                                    reshape_ops_to_remove.append(input_prod_op)
                                continue

                # Direct NCHW singleton input: create local NCHW->NHWC reshape adapter.
                if int(input_shape[1]) != 1:
                    valid = False
                    break
                nhwc_shape = _permute_shape(input_shape, perm_nchw_to_nhwc)
                if nhwc_shape is None or len(nhwc_shape) != 4 or int(nhwc_shape[3]) != 1:
                    valid = False
                    break

                adapter_name = _unique_tensor_name(f"{input_name}_nhwc_adapter")
                shape_name = _unique_tensor_name(f"{adapter_name}_reshape_shape")
                adapters_to_insert.append(
                    (
                        str(input_name),
                        str(adapter_name),
                        str(shape_name),
                        [int(v) for v in list(nhwc_shape)],
                    )
                )
                mapped_concat_inputs.append(str(adapter_name))

            if not valid:
                continue

            # Insert local adapters right before concat.
            if len(adapters_to_insert) > 0:
                concat_pos = int(model_ir.operators.index(concat_op))
                insert_pos = int(concat_pos)
                for source_name, adapter_name, shape_name, nhwc_shape in adapters_to_insert:
                    source_tensor = model_ir.tensors.get(source_name, None)
                    source_dtype = str(source_tensor.dtype) if source_tensor is not None else "FLOAT32"
                    source_quant = (
                        _clone_quantization(source_tensor.quantization)
                        if source_tensor is not None
                        else None
                    )
                    model_ir.tensors[shape_name] = TensorIR(
                        name=str(shape_name),
                        dtype="INT32",
                        shape=[len(nhwc_shape)],
                        shape_signature=[len(nhwc_shape)],
                        data=np.asarray(nhwc_shape, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                    model_ir.tensors[adapter_name] = TensorIR(
                        name=str(adapter_name),
                        dtype=str(source_dtype),
                        shape=[int(v) for v in list(nhwc_shape)],
                        shape_signature=[int(v) for v in list(nhwc_shape)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(source_quant),
                    )
                    model_ir.operators.insert(
                        int(insert_pos),
                        OperatorIR(
                            op_type="RESHAPE",
                            inputs=[str(source_name), str(shape_name)],
                            outputs=[str(adapter_name)],
                            options={"newShape": [int(v) for v in list(nhwc_shape)]},
                        ),
                    )
                    insert_pos += 1

            # Re-find operators after insertion.
            concat_op = next((op for op in model_ir.operators if op is concat_op), None)
            post_op = next((op for op in model_ir.operators if op is post_op), None)
            if concat_op is None or post_op is None:
                continue

            concat_old_output_name = str(concat_op.outputs[0])
            concat_op.options = {
                **(dict(concat_op.options) if isinstance(concat_op.options, dict) else {}),
                "axis": 3,
            }
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in mapped_concat_inputs],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=concat_op,
                new_outputs=[str(post_output_name)],
            )
            _replace_tensor_inputs(model_ir, concat_old_output_name, post_output_name)

            post_out_tensor = model_ir.tensors.get(str(post_output_name), None)
            if post_out_tensor is not None and post_out_tensor.shape is not None and len(list(post_out_tensor.shape)) == 4:
                # Keep transpose output metadata authoritative.
                pass
            else:
                concat_old_tensor = model_ir.tensors.get(str(concat_old_output_name), None)
                if concat_old_tensor is not None and concat_old_tensor.shape is not None:
                    mapped_shape = _permute_shape(
                        [int(v) for v in list(concat_old_tensor.shape)],
                        perm_nchw_to_nhwc,
                    )
                    if mapped_shape is not None:
                        model_ir.tensors[str(post_output_name)] = TensorIR(
                            name=str(post_output_name),
                            dtype=str(concat_old_tensor.dtype),
                            shape=[int(v) for v in list(mapped_shape)],
                            shape_signature=[int(v) for v in list(mapped_shape)],
                            data=None,
                            is_variable=False,
                            quantization=_clone_quantization(concat_old_tensor.quantization),
                        )

            remove_indices: List[int] = []
            post_index = next((idx for idx, op in enumerate(model_ir.operators) if op is post_op), None)
            if post_index is not None:
                remove_indices.append(int(post_index))
            for reshape_op in reshape_ops_to_remove:
                reshape_index = next((idx for idx, op in enumerate(model_ir.operators) if op is reshape_op), None)
                if reshape_index is not None:
                    remove_indices.append(int(reshape_index))
            for remove_idx in sorted(set(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_singleton_reshape_concat_post_transpose_nhwc_chains": int(rewritten)}


def _optimize_transpose_unary_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around strict unary passthrough chains.

    Target:
      X --TRANSPOSE(P)--> A --(UNARY)*--> B --TRANSPOSE(inv(P))--> Y

    Rewrite:
      X --(UNARY)*--> Y

    Safety:
    - Unary chain must be strictly linear (single consumer per main-path output).
    - Only layout-agnostic unary ops are allowed.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "FLOOR",
        "CEIL",
    }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_ir.outputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            # Only fold the canonical NHWC->NCHW->NHWC bridge direction.
            # Reverse-direction folds can corrupt downstream layout contracts.
            if [int(v) for v in list(perm_pre)] != perm_nhwc_to_nchw:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            chain_indices: List[int] = []
            chain_ops: List[OperatorIR] = []
            chain_output_names: List[str] = []
            chain_main_input_index: Dict[int, int] = {}
            current_tensor = pre_output_name

            while True:
                current_users = [int(v) for v in consumers.get(current_tensor, [])]
                if len(current_users) != 1:
                    break
                op_idx = int(current_users[0])
                op = model_ir.operators[op_idx]
                op_type = str(op.op_type)
                if op_type not in unary_passthrough_ops:
                    break
                if len(op.inputs) != 1 or len(op.outputs) != 1:
                    break
                if str(op.inputs[0]) != current_tensor:
                    break
                out_name = str(op.outputs[0])
                if out_name in model_ir.outputs:
                    break
                chain_indices.append(int(op_idx))
                chain_ops.append(op)
                chain_output_names.append(out_name)
                current_tensor = out_name

            if len(chain_ops) == 0:
                continue

            tail_users = [int(v) for v in consumers.get(current_tensor, [])]
            if len(tail_users) != 1:
                continue
            post_idx = int(tail_users[0])
            if post_idx in set(chain_indices):
                continue
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "TRANSPOSE":
                continue
            if len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != current_tensor:
                continue
            perm_post = _read_transpose_perm(model_ir, post_op)
            if perm_post is None or perm_post != perm_post_expected:
                continue
            post_output_name = str(post_op.outputs[0])

            # Rewire chain head to the pre-transpose source.
            _set_operator_inputs(
                model_ir=model_ir,
                op=chain_ops[0],
                new_inputs=[pre_input_name],
            )
            # Preserve output name expected after post-transpose.
            _set_operator_outputs(
                model_ir=model_ir,
                op=chain_ops[-1],
                new_outputs=[post_output_name],
            )

            # Update intermediate metadata to post-transpose layout.
            for out_name in chain_output_names[:-1]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_post_expected,
                )

            old_last_name = str(chain_output_names[-1])
            old_last_tensor = model_ir.tensors.get(old_last_name, None)
            post_output_tensor = model_ir.tensors.get(post_output_name, None)
            if old_last_tensor is not None and post_output_tensor is not None:
                post_output_tensor.dtype = str(old_last_tensor.dtype)
                post_output_tensor.quantization = _clone_quantization(old_last_tensor.quantization)
                post_output_tensor.shape = [int(v) for v in list(old_last_tensor.shape)]
                post_output_tensor.shape_signature = (
                    [int(v) for v in list(old_last_tensor.shape_signature)]
                    if old_last_tensor.shape_signature is not None
                    else [int(v) for v in list(old_last_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    post_output_tensor,
                    perm_post_expected,
                )

            remove_indices = sorted(list({int(pre_idx), int(post_idx)}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_unary_passthrough_chains": int(rewritten)}


def _optimize_transpose_elementwise_roundtrip_nchw_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove NCHW->NHWC->NCHW bridge pairs around a layout-agnostic elementwise subgraph.

    Pattern:
      (one or more) pre_transpose_i:  T(0,2,3,1)(x_i_nchw) -> x_i_nhwc
      elementwise-only subgraph in NHWC
      root_nhwc --T(0,3,1,2)--> root_nchw

    Rewrite:
      - Bypass each pre-transpose into the NHWC elementwise subgraph.
      - Keep the subgraph in NCHW.
      - Remove pre-transposes and the trailing inverse post-transpose.
    """
    rewritten = 0
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    # Restrict to ops whose semantics are independent of channel axis placement.
    allowed_elementwise_ops = {
        "ABS",
        "ADD",
        "DIV",
        "EXP",
        "FLOOR",
        "LOG",
        "MAXIMUM",
        "MINIMUM",
        "MUL",
        "NEG",
        "POW",
        "RSQRT",
        "SIGN",
        "SQRT",
        "SUB",
    }

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nhwc_to_nchw:
                continue

            root_nhwc_name = str(post_op.inputs[0])
            post_output_name = str(post_op.outputs[0])
            if root_nhwc_name in model_outputs or post_output_name in model_outputs:
                continue

            root_idx = producers.get(root_nhwc_name, None)
            if root_idx is None:
                continue

            # Backward traverse from root producer to collect an elementwise-only subgraph.
            subgraph_indices: set[int] = set()
            pre_transpose_output_to_index: Dict[str, int] = {}
            stack: List[int] = [int(root_idx)]
            valid = True

            while len(stack) > 0 and valid:
                op_idx = int(stack.pop())
                if op_idx in subgraph_indices:
                    continue
                op = model_ir.operators[int(op_idx)]
                if str(op.op_type) not in allowed_elementwise_ops:
                    valid = False
                    break

                subgraph_indices.add(int(op_idx))
                for input_name_raw in list(op.inputs):
                    input_name = str(input_name_raw)
                    prod_idx = producers.get(input_name, None)
                    if prod_idx is None:
                        input_tensor = model_ir.tensors.get(input_name, None)
                        # Allow embedded constants only.
                        if input_tensor is not None and input_tensor.data is not None:
                            continue
                        valid = False
                        break

                    prod_op = model_ir.operators[int(prod_idx)]
                    if (
                        str(prod_op.op_type) == "TRANSPOSE"
                        and len(prod_op.inputs) >= 2
                        and len(prod_op.outputs) == 1
                        and str(prod_op.outputs[0]) == input_name
                        and _read_transpose_perm(model_ir, prod_op) == perm_nchw_to_nhwc
                        and input_name not in model_outputs
                        and str(prod_op.inputs[0]) not in model_outputs
                    ):
                        pre_transpose_output_to_index[str(input_name)] = int(prod_idx)
                        continue

                    if str(prod_op.op_type) in allowed_elementwise_ops:
                        stack.append(int(prod_idx))
                        continue

                    valid = False
                    break

            if not valid or len(pre_transpose_output_to_index) == 0:
                continue

            # Root output must be consumed only by the matched post-transpose.
            if set(int(v) for v in consumers.get(root_nhwc_name, [])) != {int(post_idx)}:
                continue

            # Subgraph outputs must not leak outside subgraph except root->post.
            valid_outputs = True
            for op_idx in list(subgraph_indices):
                op = model_ir.operators[int(op_idx)]
                for output_name_raw in list(op.outputs):
                    output_name = str(output_name_raw)
                    users = set(int(v) for v in consumers.get(output_name, []))
                    if output_name == root_nhwc_name:
                        if users != {int(post_idx)}:
                            valid_outputs = False
                            break
                    else:
                        if not users.issubset(set(int(v) for v in list(subgraph_indices))):
                            valid_outputs = False
                            break
                if not valid_outputs:
                    break
            if not valid_outputs:
                continue

            # Each pre-transpose output must be local to this subgraph.
            valid_pre_consumers = True
            for pre_output_name, _ in pre_transpose_output_to_index.items():
                users = set(int(v) for v in consumers.get(str(pre_output_name), []))
                if len(users) == 0:
                    valid_pre_consumers = False
                    break
                if not users.issubset(set(int(v) for v in list(subgraph_indices))):
                    valid_pre_consumers = False
                    break
            if not valid_pre_consumers:
                continue

            pre_input_map: Dict[str, str] = {}
            for pre_output_name, pre_op_idx in pre_transpose_output_to_index.items():
                pre_op = model_ir.operators[int(pre_op_idx)]
                pre_input_map[str(pre_output_name)] = str(pre_op.inputs[0])

            # Rewire all subgraph inputs from transpose outputs to original NCHW tensors.
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                new_inputs = [pre_input_map.get(str(v), str(v)) for v in list(op.inputs)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=op,
                    new_inputs=new_inputs,
                )

            # Subgraph tensors were NHWC; convert metadata to NCHW.
            for op_idx in list(subgraph_indices):
                op = model_ir.operators[int(op_idx)]
                for output_name_raw in list(op.outputs):
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(output_name_raw), None),
                        perm_nhwc_to_nchw,
                    )

            # Replace root NHWC output with canonical post-transpose output.
            root_op = model_ir.operators[int(root_idx)]
            _set_operator_outputs(
                model_ir=model_ir,
                op=root_op,
                new_outputs=[post_output_name],
            )
            _replace_tensor_inputs(model_ir, root_nhwc_name, post_output_name)

            old_root_tensor = model_ir.tensors.get(root_nhwc_name, None)
            canonical_tensor = model_ir.tensors.get(post_output_name, None)
            if old_root_tensor is not None and canonical_tensor is not None:
                canonical_tensor.dtype = str(old_root_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(old_root_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(old_root_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(old_root_tensor.shape_signature)]
                    if old_root_tensor.shape_signature is not None
                    else [int(v) for v in list(old_root_tensor.shape)]
                )
                _permute_tensor_metadata_if_rank_matches(
                    canonical_tensor,
                    perm_nhwc_to_nchw,
                )

            remove_indices = set([int(post_idx)])
            remove_indices.update(int(v) for v in list(pre_transpose_output_to_index.values()))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_elementwise_roundtrip_nchw_nhwc_chains": int(rewritten)}


def _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC->NCHW and inverse NCHW->NHWC transpose bridges around an
    elementwise-only fanout subgraph.

    Pattern:
      x_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw --(elementwise DAG; const side inputs allowed)--> y*_nchw
      y*_nchw --T(0,2,3,1)--> y*_nhwc

    Rewrite:
      - Keep the full elementwise DAG in NHWC by bypassing the leading transpose.
      - Remove all trailing inverse transposes at subgraph boundaries.

    Safety:
      - Subgraph ops are limited to layout-agnostic elementwise ops.
      - Per-channel rank-4 constants are remapped only when they are local to the subgraph.
      - Legacy NCHW users are preserved via local adapter TRANSPOSE ops.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    allowed_elementwise_ops = {
        "ABS",
        "ADD",
        "DIV",
        "ERF",
        "EXP",
        "FLOOR",
        "LOG",
        "MAXIMUM",
        "MINIMUM",
        "MUL",
        "NEG",
        "POW",
        "RSQRT",
        "SIGN",
        "SQRT",
        "SUB",
    }

    def _unique_tensor_name(base_name: str) -> str:
        candidate = str(base_name)
        serial = 0
        while candidate in model_ir.tensors:
            serial += 1
            candidate = f"{base_name}_{serial}"
        return candidate

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            # Forward-collect elementwise-only reachable ops.
            subgraph_indices: set[int] = set()
            frontier_tensors: List[str] = [pre_output_name]
            seen_tensors: set[str] = {pre_output_name}
            while len(frontier_tensors) > 0:
                current_tensor = str(frontier_tensors.pop())
                for user_idx_raw in consumers.get(current_tensor, []):
                    user_idx = int(user_idx_raw)
                    user_op = model_ir.operators[user_idx]
                    if str(user_op.op_type) not in allowed_elementwise_ops:
                        continue
                    if user_idx in subgraph_indices:
                        continue
                    subgraph_indices.add(user_idx)
                    for out_name_raw in list(user_op.outputs):
                        out_name = str(out_name_raw)
                        if out_name not in seen_tensors:
                            seen_tensors.add(out_name)
                            frontier_tensors.append(out_name)

            if len(subgraph_indices) == 0:
                continue

            # Validate local shape and collect runtime inputs coming from outside
            # the local elementwise region. Those external NCHW tensors are
            # adapted locally into NHWC.
            valid = True
            external_runtime_inputs: set[str] = set()
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[op_idx]
                if len(op.outputs) != 1:
                    valid = False
                    break
                for input_name_raw in list(op.inputs):
                    input_name = str(input_name_raw)
                    input_tensor = model_ir.tensors.get(input_name, None)
                    if input_tensor is not None and input_tensor.data is not None:
                        continue

                    if input_name == pre_output_name:
                        continue

                    prod_idx = producers.get(input_name, None)
                    if prod_idx is None or int(prod_idx) not in subgraph_indices:
                        external_runtime_inputs.add(str(input_name))
                if not valid:
                    break
            if not valid:
                continue

            # Ensure pre-transpose output is consumed only by this local subgraph.
            pre_users = set(int(v) for v in consumers.get(pre_output_name, []))
            if len(pre_users) == 0 or not pre_users.issubset(subgraph_indices):
                continue

            # Detect subgraph boundaries and classify external users:
            # - inverse post-transpose consumers (remove)
            # - legacy users (preserve via NHWC->NCHW adapter)
            boundary_posts: Dict[str, List[int]] = {}
            boundary_legacy_users: Dict[str, List[int]] = {}
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[op_idx]
                out_name = str(op.outputs[0])
                if out_name in model_outputs:
                    valid = False
                    break

                users = set(int(v) for v in consumers.get(out_name, []))
                external_users = sorted(list(users - subgraph_indices))
                if len(external_users) == 0:
                    continue

                for user_idx in external_users:
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) != "TRANSPOSE"
                        or len(user_op.inputs) < 2
                        or len(user_op.outputs) != 1
                        or str(user_op.inputs[0]) != out_name
                    ):
                        boundary_legacy_users.setdefault(out_name, []).append(int(user_idx))
                        continue
                    if (
                        _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        boundary_posts.setdefault(out_name, []).append(int(user_idx))
                    else:
                        boundary_legacy_users.setdefault(out_name, []).append(int(user_idx))
                if not valid:
                    break
            if not valid:
                continue
            if sum(len(v) for v in boundary_posts.values()) <= 0:
                continue
            # Conservative guard:
            # Rewriting mixed fanout islands that require external runtime
            # NCHW->NHWC adapter insertion has produced dangling `__to_nhwc`
            # tensors in some large graphs (e.g. network.7 branches).
            # Skip this candidate and leave it to other safer transpose passes.
            if len(external_runtime_inputs) > 0:
                continue

            # Safety guard:
            # This pass can be very aggressive on wide fanout graphs. Snapshot one
            # candidate rewrite and roll back if it introduces new unbound runtime
            # inputs (which later appears as "Input tensor N lacks data").
            candidate_snapshot = copy.deepcopy(model_ir)
            unbound_before_count = int(len(_find_unbound_nonconstant_operator_inputs(model_ir)))
            def _rollback_candidate() -> None:
                model_ir.tensors = candidate_snapshot.tensors
                model_ir.operators = candidate_snapshot.operators
                model_ir.inputs = candidate_snapshot.inputs
                model_ir.outputs = candidate_snapshot.outputs
                model_ir.subgraphs = candidate_snapshot.subgraphs
                model_ir.metadata = candidate_snapshot.metadata

            # Remap per-channel rank-4 constants from NCHW->NHWC where required.
            remapped_const_names: set[str] = set()
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[op_idx]
                runtime_input_names = []
                const_input_names = []
                for input_name_raw in list(op.inputs):
                    input_name = str(input_name_raw)
                    input_tensor = model_ir.tensors.get(input_name, None)
                    if input_tensor is not None and input_tensor.data is not None:
                        const_input_names.append(input_name)
                    else:
                        runtime_input_names.append(input_name)
                if len(runtime_input_names) == 0:
                    continue

                ref_name = next((name for name in runtime_input_names if name == pre_output_name), runtime_input_names[0])
                ref_tensor = model_ir.tensors.get(str(ref_name), None)
                if ref_tensor is None or ref_tensor.shape is None or len(list(ref_tensor.shape)) != 4:
                    continue
                ref_shape_nchw = [int(v) for v in list(ref_tensor.shape)]
                if any(int(v) <= 0 for v in ref_shape_nchw):
                    continue
                ref_shape_nhwc = _permute_shape(ref_shape_nchw, perm_nchw_to_nhwc)
                if ref_shape_nhwc is None:
                    continue

                for const_name in const_input_names:
                    if const_name in remapped_const_names:
                        continue
                    const_tensor = model_ir.tensors.get(str(const_name), None)
                    if const_tensor is None or const_tensor.data is None:
                        continue
                    const_data = np.asarray(const_tensor.data)
                    if int(const_data.ndim) != 4:
                        continue
                    const_shape = [int(v) for v in list(const_data.shape)]
                    # Strictly target NCHW per-channel broadcast constants [1,C,1,1].
                    if not (
                        const_shape[0] == 1
                        and const_shape[1] == ref_shape_nchw[1]
                        and const_shape[2] == 1
                        and const_shape[3] == 1
                    ):
                        continue
                    # Shared constants cannot be remapped safely because non-local
                    # users still expect NCHW semantics.
                    const_users = set(int(v) for v in consumers.get(str(const_name), []))
                    rotated_data = np.transpose(const_data, axes=perm_nchw_to_nhwc)
                    rotated_shape = [int(v) for v in list(rotated_data.shape)]
                    if _broadcast_static_shapes(ref_shape_nhwc, rotated_shape) is None:
                        valid = False
                        break
                    if const_users.issubset(subgraph_indices):
                        const_tensor.data = rotated_data.astype(const_data.dtype, copy=False)
                        const_tensor.shape = [int(v) for v in list(rotated_shape)]
                        const_tensor.shape_signature = (
                            [int(v) for v in list(rotated_shape)]
                            if const_tensor.shape_signature is None
                            else [int(v) for v in list(rotated_shape)]
                        )
                        remapped_const_names.add(str(const_name))
                    else:
                        # Shared constants are cloned for this rewritten branch.
                        clone_name = _unique_tensor_name(f"{const_name}__nhwc")
                        model_ir.tensors[str(clone_name)] = TensorIR(
                            name=str(clone_name),
                            dtype=str(const_tensor.dtype),
                            shape=[int(v) for v in list(rotated_shape)],
                            shape_signature=[int(v) for v in list(rotated_shape)],
                            data=rotated_data.astype(const_data.dtype, copy=False),
                            is_variable=False,
                            quantization=_clone_quantization(const_tensor.quantization),
                        )
                        op_inputs = [str(v) for v in list(op.inputs)]
                        op_inputs = [str(clone_name) if str(v) == str(const_name) else str(v) for v in op_inputs]
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=op,
                            new_inputs=op_inputs,
                        )
                        remapped_const_names.add(str(clone_name))
                if not valid:
                    break
            if not valid:
                _rollback_candidate()
                continue

            # Rewire all subgraph inputs from pre-transpose output to NHWC source.
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                new_inputs = [
                    pre_input_name if str(v) == pre_output_name else str(v)
                    for v in list(op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=op,
                    new_inputs=new_inputs,
                )

            # Bridge external NCHW runtime inputs into NHWC for rewritten ops.
            if len(external_runtime_inputs) > 0:
                for external_input_name in sorted(list(external_runtime_inputs)):
                    external_tensor = model_ir.tensors.get(str(external_input_name), None)
                    if (
                        external_tensor is None
                        or external_tensor.shape is None
                        or len(list(external_tensor.shape)) != 4
                    ):
                        valid = False
                        break
                    source_shape = [int(v) for v in list(external_tensor.shape)]
                    target_shape = _permute_shape(source_shape, perm_nchw_to_nhwc)
                    if target_shape is None:
                        valid = False
                        break
                    adapter_perm_name = _unique_tensor_name(
                        f"{external_input_name}__to_nhwc_perm",
                    )
                    adapter_output_name = _unique_tensor_name(
                        f"{external_input_name}__to_nhwc",
                    )
                    model_ir.tensors[str(adapter_perm_name)] = TensorIR(
                        name=str(adapter_perm_name),
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nchw_to_nhwc, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                    model_ir.tensors[str(adapter_output_name)] = TensorIR(
                        name=str(adapter_output_name),
                        dtype=str(external_tensor.dtype),
                        shape=[int(v) for v in list(target_shape)],
                        shape_signature=[int(v) for v in list(target_shape)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(external_tensor.quantization),
                    )
                    model_ir.operators.append(
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(external_input_name), str(adapter_perm_name)],
                            outputs=[str(adapter_output_name)],
                        )
                    )
                    for op_idx in sorted(list(subgraph_indices)):
                        op = model_ir.operators[int(op_idx)]
                        new_inputs = [
                            str(adapter_output_name) if str(v) == str(external_input_name) else str(v)
                            for v in list(op.inputs)
                        ]
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=op,
                            new_inputs=new_inputs,
                        )
            if not valid:
                _rollback_candidate()
                continue

            # Subgraph tensors were NCHW; convert metadata to NHWC.
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                out_name = str(op.outputs[0])
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_nchw_to_nhwc,
                )

            # Final broadcast-safety guard:
            # if an elementwise op now consumes NHWC data with a legacy
            # [1,C,1,1] constant, rotate (or clone+rotate) that constant.
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                if str(op.op_type) not in {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM", "POW"}:
                    continue
                input_names = [str(v) for v in list(op.inputs)]
                for const_input_name in list(input_names):
                    const_tensor = model_ir.tensors.get(str(const_input_name), None)
                    if const_tensor is None or const_tensor.data is None:
                        continue
                    const_data = np.asarray(const_tensor.data)
                    if int(const_data.ndim) != 4:
                        continue
                    data_input_name = next((name for name in input_names if str(name) != str(const_input_name)), None)
                    if data_input_name is None:
                        continue
                    data_tensor = model_ir.tensors.get(str(data_input_name), None)
                    if data_tensor is None or data_tensor.shape is None or len(list(data_tensor.shape)) != 4:
                        continue
                    data_shape = [int(v) for v in list(data_tensor.shape)]
                    if any(int(v) <= 0 for v in data_shape):
                        continue
                    const_shape = [int(v) for v in list(const_data.shape)]
                    if _broadcast_static_shapes(data_shape, const_shape) is not None:
                        continue
                    rotated_data = np.transpose(const_data, axes=perm_nchw_to_nhwc)
                    rotated_shape = [int(v) for v in list(rotated_data.shape)]
                    if _broadcast_static_shapes(data_shape, rotated_shape) is None:
                        continue
                    const_users = set(int(v) for v in consumers.get(str(const_input_name), []))
                    if const_users.issubset(subgraph_indices):
                        const_tensor.data = rotated_data.astype(const_data.dtype, copy=False)
                        const_tensor.shape = [int(v) for v in list(rotated_shape)]
                        const_tensor.shape_signature = [int(v) for v in list(rotated_shape)]
                    else:
                        clone_name = _unique_tensor_name(f"{const_input_name}__nhwc")
                        model_ir.tensors[str(clone_name)] = TensorIR(
                            name=str(clone_name),
                            dtype=str(const_tensor.dtype),
                            shape=[int(v) for v in list(rotated_shape)],
                            shape_signature=[int(v) for v in list(rotated_shape)],
                            data=rotated_data.astype(const_data.dtype, copy=False),
                            is_variable=False,
                            quantization=_clone_quantization(const_tensor.quantization),
                        )
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=op,
                            new_inputs=[
                                str(clone_name) if str(v) == str(const_input_name) else str(v)
                                for v in list(op.inputs)
                            ],
                        )

            # Replace each boundary with one canonical post-transpose output.
            post_indices_to_remove: set[int] = set()
            boundary_output_names = sorted(
                set(list(boundary_posts.keys()) + list(boundary_legacy_users.keys()))
            )
            for boundary_name in boundary_output_names:
                post_indices = [int(v) for v in list(boundary_posts.get(boundary_name, []))]
                legacy_users = [int(v) for v in list(boundary_legacy_users.get(boundary_name, []))]
                producer_idx = producers.get(str(boundary_name), None)
                if producer_idx is None or int(producer_idx) not in subgraph_indices:
                    valid = False
                    break
                producer_op = model_ir.operators[int(producer_idx)]
                if len(producer_op.outputs) != 1 or str(producer_op.outputs[0]) != str(boundary_name):
                    valid = False
                    break

                if len(post_indices) > 0:
                    canonical_post_idx = int(sorted(set(int(v) for v in post_indices))[0])
                    canonical_post_op = model_ir.operators[canonical_post_idx]
                    canonical_output_name = str(canonical_post_op.outputs[0])
                else:
                    canonical_post_idx = None
                    canonical_output_name = _unique_tensor_name(f"{boundary_name}__nhwc")
                    old_boundary_tensor = model_ir.tensors.get(str(boundary_name), None)
                    if old_boundary_tensor is not None:
                        model_ir.tensors[str(canonical_output_name)] = TensorIR(
                            name=str(canonical_output_name),
                            dtype=str(old_boundary_tensor.dtype),
                            shape=[int(v) for v in list(old_boundary_tensor.shape)],
                            shape_signature=(
                                [int(v) for v in list(old_boundary_tensor.shape_signature)]
                                if old_boundary_tensor.shape_signature is not None
                                else [int(v) for v in list(old_boundary_tensor.shape)]
                            ),
                            data=None,
                            is_variable=False,
                            quantization=_clone_quantization(old_boundary_tensor.quantization),
                        )

                _set_operator_outputs(
                    model_ir=model_ir,
                    op=producer_op,
                    new_outputs=[canonical_output_name],
                )
                _replace_tensor_inputs(model_ir, str(boundary_name), canonical_output_name)

                old_boundary_tensor = model_ir.tensors.get(str(boundary_name), None)
                canonical_tensor = model_ir.tensors.get(canonical_output_name, None)
                if old_boundary_tensor is not None and canonical_tensor is not None:
                    canonical_tensor.dtype = str(old_boundary_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(old_boundary_tensor.quantization)
                    canonical_tensor.shape = [int(v) for v in list(old_boundary_tensor.shape)]
                    canonical_tensor.shape_signature = (
                        [int(v) for v in list(old_boundary_tensor.shape_signature)]
                        if old_boundary_tensor.shape_signature is not None
                        else [int(v) for v in list(old_boundary_tensor.shape)]
                    )

                for post_idx in sorted(set(int(v) for v in list(post_indices))):
                    post_indices_to_remove.add(int(post_idx))
                    if canonical_post_idx is not None and int(post_idx) == int(canonical_post_idx):
                        continue
                    post_output_name = str(model_ir.operators[int(post_idx)].outputs[0])
                    _replace_tensor_inputs(model_ir, post_output_name, canonical_output_name)

                # Preserve external NCHW users by appending a local NHWC->NCHW adapter.
                if len(legacy_users) > 0:
                    adapter_perm_name = _unique_tensor_name(
                        f"{boundary_name}__to_nchw_perm",
                    )
                    adapter_output_name = str(boundary_name)
                    canonical_tensor = model_ir.tensors.get(str(canonical_output_name), None)
                    if canonical_tensor is None or canonical_tensor.shape is None:
                        valid = False
                        break
                    canonical_shape = [int(v) for v in list(canonical_tensor.shape)]
                    legacy_shape = _permute_shape(canonical_shape, perm_nhwc_to_nchw)
                    if legacy_shape is None:
                        valid = False
                        break
                    model_ir.tensors[str(adapter_perm_name)] = TensorIR(
                        name=str(adapter_perm_name),
                        dtype="INT32",
                        shape=[4],
                        shape_signature=[4],
                        data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                        is_variable=False,
                        quantization=None,
                    )
                    model_ir.tensors[str(adapter_output_name)] = TensorIR(
                        name=str(adapter_output_name),
                        dtype=str(canonical_tensor.dtype),
                        shape=[int(v) for v in list(legacy_shape)],
                        shape_signature=[int(v) for v in list(legacy_shape)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(canonical_tensor.quantization),
                    )
                    model_ir.operators.append(
                        OperatorIR(
                            op_type="TRANSPOSE",
                            inputs=[str(canonical_output_name), str(adapter_perm_name)],
                            outputs=[str(adapter_output_name)],
                        )
                    )
                    # Restore legacy users to NCHW tensor name.
                    for legacy_idx in list(legacy_users):
                        legacy_op = model_ir.operators[int(legacy_idx)]
                        new_inputs = [
                            str(adapter_output_name) if str(v) == str(canonical_output_name) else str(v)
                            for v in list(legacy_op.inputs)
                        ]
                        _set_operator_inputs(
                            model_ir=model_ir,
                            op=legacy_op,
                            new_inputs=new_inputs,
                        )

            if not valid:
                _rollback_candidate()
                continue

            remove_indices = set([int(pre_idx)])
            remove_indices.update(int(v) for v in list(post_indices_to_remove))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            unbound_after_count = int(len(_find_unbound_nonconstant_operator_inputs(model_ir)))
            if unbound_after_count > unbound_before_count:
                _rollback_candidate()
                continue

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains": int(rewritten)}


def _repair_rank4_channelwise_broadcast_constants_to_runtime_layout(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Repair channelwise constants that became layout-misaligned after
    transpose/layout rewrites.

    If a binary op consumes rank-4 runtime data tensor `X` and a rank-3/4
    channelwise constant `C` where:
      - as-is broadcast does not preserve `X` shape
      - NHWC-rotated broadcast preserves `X` shape
    then rotate `C` to NHWC layout. Shared constants are cloned.
    """
    repaired = 0
    binary_ops = {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM", "POW"}
    consumers = _build_tensor_consumer_map(model_ir)
    producers = _build_tensor_producer_map(model_ir)

    def _unique_tensor_name(base_name: str) -> str:
        candidate = str(base_name)
        serial = 0
        while candidate in model_ir.tensors:
            serial += 1
            candidate = f"{base_name}_{serial}"
        return candidate

    def _prefer_runtime_layout_for_rank4_data_tensor(data_tensor_name: str) -> Optional[str]:
        name = str(data_tensor_name)
        lower_name = name.lower()
        if lower_name.endswith("_nhwc"):
            return "NHWC"
        if lower_name.endswith("_nchw"):
            return "NCHW"
        producer_index = producers.get(name, None)
        if producer_index is None:
            return None
        if int(producer_index) < 0 or int(producer_index) >= len(model_ir.operators):
            return None
        producer_op = model_ir.operators[int(producer_index)]
        producer_op_type = str(producer_op.op_type)
        if producer_op_type in {
            "CONV_2D",
            "DEPTHWISE_CONV_2D",
            "TRANSPOSE_CONV",
            "AVERAGE_POOL_2D",
            "MAX_POOL_2D",
            "RESIZE_BILINEAR",
            "RESIZE_NEAREST_NEIGHBOR",
        }:
            return "NHWC"
        return None

    for op in list(model_ir.operators):
        if str(op.op_type) not in binary_ops or len(op.inputs) != 2:
            continue
        input_names = [str(v) for v in list(op.inputs)]

        for const_input_name in list(input_names):
            const_tensor = model_ir.tensors.get(str(const_input_name), None)
            if const_tensor is None or const_tensor.data is None:
                continue
            const_data = np.asarray(const_tensor.data)
            if int(const_data.ndim) not in {3, 4}:
                continue

            data_input_name = next((name for name in input_names if str(name) != str(const_input_name)), None)
            if data_input_name is None:
                continue
            data_tensor = model_ir.tensors.get(str(data_input_name), None)
            if data_tensor is None or data_tensor.shape is None or len(list(data_tensor.shape)) != 4:
                continue
            data_shape = [int(v) for v in list(data_tensor.shape)]
            if any(int(v) <= 0 for v in data_shape):
                continue

            const_shape = [int(v) for v in list(const_data.shape)]
            as_is_broadcast = _broadcast_static_shapes(data_shape, const_shape)

            rotated_data: Optional[np.ndarray] = None
            if int(const_data.ndim) == 4:
                rotated_data = np.transpose(const_data, axes=[0, 2, 3, 1])
            elif (
                int(const_data.ndim) == 3
                and int(const_shape[0]) > 0
                and int(const_shape[1]) == 1
                and int(const_shape[2]) == 1
            ):
                rotated_data = np.transpose(const_data, axes=[1, 2, 0])
            else:
                continue

            rotated_shape = [int(v) for v in list(rotated_data.shape)]
            rotated_broadcast = _broadcast_static_shapes(data_shape, rotated_shape)
            if rotated_broadcast is None:
                continue
            if rotated_broadcast != data_shape:
                continue
            force_rotate_even_if_ambiguous = False
            if as_is_broadcast == data_shape and int(const_data.ndim) == 4:
                # Ambiguous case:
                # [1,C,1,1] can broadcast to NHWC [N,H,W,C] when H==C, but this
                # semantically applies scale over height, not channel.
                # If runtime data layout is NHWC, prefer [1,1,1,C].
                preferred_layout = _prefer_runtime_layout_for_rank4_data_tensor(
                    data_input_name
                )
                if (
                    preferred_layout == "NHWC"
                    and int(const_shape[0]) == 1
                    and int(const_shape[1]) > 1
                    and int(const_shape[2]) == 1
                    and int(const_shape[3]) == 1
                    and int(rotated_shape[0]) == 1
                    and int(rotated_shape[1]) == 1
                    and int(rotated_shape[2]) == 1
                    and int(rotated_shape[3]) > 1
                ):
                    force_rotate_even_if_ambiguous = True
            if as_is_broadcast == data_shape and not force_rotate_even_if_ambiguous:
                continue

            const_users = set(int(v) for v in consumers.get(str(const_input_name), []))
            if len(const_users) <= 1:
                const_tensor.data = rotated_data.astype(const_data.dtype, copy=False)
                const_tensor.shape = [int(v) for v in list(rotated_shape)]
                const_tensor.shape_signature = [int(v) for v in list(rotated_shape)]
            else:
                clone_name = _unique_tensor_name(f"{const_input_name}__nhwc")
                model_ir.tensors[str(clone_name)] = TensorIR(
                    name=str(clone_name),
                    dtype=str(const_tensor.dtype),
                    shape=[int(v) for v in list(rotated_shape)],
                    shape_signature=[int(v) for v in list(rotated_shape)],
                    data=rotated_data.astype(const_data.dtype, copy=False),
                    is_variable=False,
                    quantization=_clone_quantization(const_tensor.quantization),
                )
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=op,
                    new_inputs=[
                        str(clone_name) if str(v) == str(const_input_name) else str(v)
                        for v in list(op.inputs)
                    ],
                )

            repaired += 1
            break

    return {"repaired_rank4_channelwise_broadcast_constants": int(repaired)}


def _optimize_convpool_output_transpose_nhwc_passthrough_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC->NCHW TRANSPOSE directly after Conv/Pool-family outputs and
    keep downstream layout contracts via local adapters only when required.

    Pattern:
      conv_or_pool_out_nhwc --T(0,3,1,2)--> x_nchw
      x_nchw -> elementwise-only subgraph

    Rewrite:
      - bypass leading transpose (subgraph runs in NHWC)
      - keep legacy external NCHW users through local NHWC->NCHW adapters
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    convpool_ops = {
        "CONV_2D",
        "DEPTHWISE_CONV_2D",
        "TRANSPOSE_CONV",
        "AVERAGE_POOL_2D",
        "MAX_POOL_2D",
    }
    elementwise_ops = {
        "ABS",
        "ADD",
        "DIV",
        "ERF",
        "EXP",
        "FLOOR",
        "LOG",
        "MAXIMUM",
        "MINIMUM",
        "MUL",
        "NEG",
        "POW",
        "RSQRT",
        "SIGN",
        "SQRT",
        "SUB",
    }

    def _unique_tensor_name(base_name: str) -> str:
        candidate = str(base_name)
        serial = 0
        while candidate in model_ir.tensors:
            serial += 1
            candidate = f"{base_name}_{serial}"
        return candidate

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue
            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_input_producer_idx = producers.get(pre_input_name, None)
            if pre_input_producer_idx is None:
                continue
            if str(model_ir.operators[int(pre_input_producer_idx)].op_type) not in convpool_ops:
                continue

            # Collect forward elementwise-only region from pre_output.
            subgraph_indices: set[int] = set()
            frontier: List[str] = [pre_output_name]
            seen_tensors: set[str] = {pre_output_name}
            while len(frontier) > 0:
                cur = str(frontier.pop())
                for user_idx_raw in consumers.get(cur, []):
                    user_idx = int(user_idx_raw)
                    user_op = model_ir.operators[int(user_idx)]
                    if str(user_op.op_type) not in elementwise_ops:
                        continue
                    if int(user_idx) in subgraph_indices:
                        continue
                    subgraph_indices.add(int(user_idx))
                    for out_name_raw in list(user_op.outputs):
                        out_name = str(out_name_raw)
                        if out_name not in seen_tensors:
                            seen_tensors.add(out_name)
                            frontier.append(out_name)
            if len(subgraph_indices) == 0:
                continue

            # Validate runtime inputs and gather external runtime dependencies.
            valid = True
            external_runtime_inputs: set[str] = set()
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                if len(op.outputs) != 1:
                    valid = False
                    break
                for input_name_raw in list(op.inputs):
                    input_name = str(input_name_raw)
                    input_tensor = model_ir.tensors.get(input_name, None)
                    if input_tensor is not None and input_tensor.data is not None:
                        continue
                    if input_name == pre_output_name:
                        continue
                    prod_idx = producers.get(input_name, None)
                    if prod_idx is None or int(prod_idx) not in subgraph_indices:
                        external_runtime_inputs.add(str(input_name))
            if not valid:
                continue

            # pre_output must be used only by local elementwise subgraph.
            pre_users = set(int(v) for v in consumers.get(pre_output_name, []))
            if len(pre_users) == 0 or not pre_users.issubset(subgraph_indices):
                continue

            # classify boundary users outside subgraph
            boundary_legacy_users: Dict[str, List[int]] = {}
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                out_name = str(op.outputs[0])
                if out_name in model_outputs:
                    valid = False
                    break
                users = set(int(v) for v in consumers.get(out_name, []))
                external_users = sorted(list(users - subgraph_indices))
                if len(external_users) == 0:
                    continue
                boundary_legacy_users[out_name] = [int(v) for v in list(external_users)]
            if not valid:
                continue

            # Rewire subgraph from pre_output -> pre_input.
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=op,
                    new_inputs=[
                        pre_input_name if str(v) == pre_output_name else str(v)
                        for v in list(op.inputs)
                    ],
                )

            # Adapt external runtime NCHW inputs used by subgraph.
            for external_input_name in sorted(list(external_runtime_inputs)):
                ext_tensor = model_ir.tensors.get(str(external_input_name), None)
                if ext_tensor is None or ext_tensor.shape is None or len(list(ext_tensor.shape)) != 4:
                    valid = False
                    break
                ext_shape = [int(v) for v in list(ext_tensor.shape)]
                ext_nhwc_shape = _permute_shape(ext_shape, perm_nchw_to_nhwc)
                if ext_nhwc_shape is None:
                    valid = False
                    break
                adapter_perm_name = _unique_tensor_name(f"{external_input_name}__to_nhwc_perm")
                adapter_output_name = _unique_tensor_name(f"{external_input_name}__to_nhwc")
                model_ir.tensors[str(adapter_perm_name)] = TensorIR(
                    name=str(adapter_perm_name),
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nchw_to_nhwc, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                model_ir.tensors[str(adapter_output_name)] = TensorIR(
                    name=str(adapter_output_name),
                    dtype=str(ext_tensor.dtype),
                    shape=[int(v) for v in list(ext_nhwc_shape)],
                    shape_signature=[int(v) for v in list(ext_nhwc_shape)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(ext_tensor.quantization),
                )
                model_ir.operators.append(
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(external_input_name), str(adapter_perm_name)],
                        outputs=[str(adapter_output_name)],
                    )
                )
                for op_idx in sorted(list(subgraph_indices)):
                    op = model_ir.operators[int(op_idx)]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=op,
                        new_inputs=[
                            str(adapter_output_name) if str(v) == str(external_input_name) else str(v)
                            for v in list(op.inputs)
                        ],
                    )
            if not valid:
                continue

            # Metadata update: subgraph tensors are now NHWC.
            for op_idx in sorted(list(subgraph_indices)):
                op = model_ir.operators[int(op_idx)]
                out_name = str(op.outputs[0])
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(out_name, None),
                    perm_nchw_to_nhwc,
                )

            # Keep external users with local NHWC->NCHW adapters at boundaries.
            for boundary_name, legacy_users in boundary_legacy_users.items():
                producer_idx = producers.get(str(boundary_name), None)
                if producer_idx is None or int(producer_idx) not in subgraph_indices:
                    valid = False
                    break
                boundary_nhwc_name = _unique_tensor_name(f"{boundary_name}__to_nhwc")
                producer_op = model_ir.operators[int(producer_idx)]
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=producer_op,
                    new_outputs=[str(boundary_nhwc_name)],
                )
                _replace_tensor_inputs(model_ir, str(boundary_name), str(boundary_nhwc_name))

                boundary_tensor = model_ir.tensors.get(str(boundary_name), None)
                boundary_nhwc_tensor = model_ir.tensors.get(str(boundary_nhwc_name), None)
                if boundary_tensor is not None and boundary_nhwc_tensor is None:
                    boundary_nhwc_tensor = TensorIR(
                        name=str(boundary_nhwc_name),
                        dtype=str(boundary_tensor.dtype),
                        shape=[int(v) for v in list(boundary_tensor.shape)],
                        shape_signature=(
                            [int(v) for v in list(boundary_tensor.shape_signature)]
                            if boundary_tensor.shape_signature is not None
                            else [int(v) for v in list(boundary_tensor.shape)]
                        ),
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(boundary_tensor.quantization),
                    )
                    model_ir.tensors[str(boundary_nhwc_name)] = boundary_nhwc_tensor

                # First, try to absorb NCHW-only MEAN consumers by remapping axes
                # and keeping them on NHWC input directly.
                remapped_legacy_users: set[int] = set()
                for legacy_idx in list(legacy_users):
                    legacy_op = model_ir.operators[int(legacy_idx)]
                    if (
                        str(legacy_op.op_type) != "MEAN"
                        or len(legacy_op.inputs) < 2
                        or len(legacy_op.outputs) != 1
                        or str(legacy_op.inputs[0]) != str(boundary_nhwc_name)
                        or not bool(legacy_op.options.get("keepDims", False))
                    ):
                        continue
                    if (
                        boundary_nhwc_tensor is None
                        or boundary_nhwc_tensor.shape is None
                        or len(list(boundary_nhwc_tensor.shape)) != 4
                    ):
                        continue

                    mean_axes_name = str(legacy_op.inputs[1])
                    mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
                    mean_axes_vals = _read_const_ints_from_tensor(mean_axes_tensor)
                    if mean_axes_vals is None or len(mean_axes_vals) == 0:
                        continue

                    normalized_axes: List[int] = []
                    valid_axes = True
                    for axis in mean_axes_vals:
                        a = int(axis)
                        if a < 0:
                            a += 4
                        if a < 0 or a >= 4:
                            valid_axes = False
                            break
                        normalized_axes.append(int(a))
                    if not valid_axes:
                        continue
                    mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]

                    _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in mapped_axes])
                    if isinstance(legacy_op.options, dict):
                        mean_options = dict(legacy_op.options)
                        for key in ["axis", "axes", "onnxRawAxes"]:
                            value = mean_options.get(key, None)
                            if isinstance(value, list) and len(value) == len(mapped_axes):
                                mean_options[key] = [int(v) for v in mapped_axes]
                        legacy_op.options = mean_options

                    mean_out_name = str(legacy_op.outputs[0])
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(mean_out_name, None),
                        perm_nchw_to_nhwc,
                    )

                    for mean_user_idx in consumers.get(mean_out_name, []):
                        mean_user_op = model_ir.operators[int(mean_user_idx)]
                        if (
                            str(mean_user_op.op_type) != "SQUEEZE"
                            or len(mean_user_op.inputs) != 1
                            or len(mean_user_op.outputs) != 1
                            or str(mean_user_op.inputs[0]) != mean_out_name
                            or not isinstance(mean_user_op.options, dict)
                            or "squeezeDims" not in mean_user_op.options
                        ):
                            continue
                        raw_axes = list(mean_user_op.options.get("squeezeDims", []))
                        remapped_squeeze_axes: List[int] = []
                        valid_squeeze_axes = True
                        for axis in raw_axes:
                            a = int(axis)
                            if a < 0:
                                a += 4
                            if a < 0 or a >= 4:
                                valid_squeeze_axes = False
                                break
                            remapped_squeeze_axes.append(int(perm_nhwc_to_nchw[int(a)]))
                        if not valid_squeeze_axes:
                            continue
                        squeeze_options = dict(mean_user_op.options)
                        squeeze_options["squeezeDims"] = sorted([int(v) for v in remapped_squeeze_axes])
                        mean_user_op.options = squeeze_options

                    remapped_legacy_users.add(int(legacy_idx))

                pending_legacy_users = [
                    int(v) for v in list(legacy_users)
                    if int(v) not in remapped_legacy_users
                ]
                if len(pending_legacy_users) <= 0:
                    continue

                perm_name = _unique_tensor_name(f"{boundary_name}__to_nchw_perm")
                model_ir.tensors[str(perm_name)] = TensorIR(
                    name=str(perm_name),
                    dtype="INT32",
                    shape=[4],
                    shape_signature=[4],
                    data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
                    is_variable=False,
                    quantization=None,
                )
                model_ir.operators.append(
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(boundary_nhwc_name), str(perm_name)],
                        outputs=[str(boundary_name)],
                    )
                )
                for legacy_idx in list(pending_legacy_users):
                    legacy_op = model_ir.operators[int(legacy_idx)]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=legacy_op,
                        new_inputs=[
                            str(boundary_name) if str(v) == str(boundary_nhwc_name) else str(v)
                            for v in list(legacy_op.inputs)
                        ],
                    )
            if not valid:
                continue

            # Remove leading transpose.
            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_convpool_output_transpose_nhwc_passthrough_chains": int(rewritten)}


def _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate transpose wrappers around unary fanout branches with inverse post-transpose consumers.

    Target:
      X --TRANSPOSE(P)--> A --(RELU|RELU6|LOGISTIC|TANH)--> B
      B --TRANSPOSE(inv(P))--> Y0
      B --TRANSPOSE(inv(P))--> Y1
      ...
      (optionally: B also has non-transpose legacy consumers)

    Rewrite:
      X --(UNARY)--> Y0
      (all uses of Y1... are rewired to Y0)
      If legacy consumers exist:
        keep one adapter TRANSPOSE(P): Y0 -> B for legacy consumers.

    Safety:
    - Pre-transpose output is consumed only by the unary op.
    - Inverse-post transpose consumers must use inv(P).
    - Fanout post-transpose outputs are not graph outputs.
    - If non-transpose legacy consumers exist, preserve B via one adapter.
    """
    rewritten = 0
    unary_ops = {"RELU", "RELU6", "RELU_0_TO_1", "HARD_SWISH", "LEAKY_RELU", "LOGISTIC", "TANH", "GELU"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            perm_post_expected = _invert_perm(perm_pre)
            if perm_post_expected is None:
                continue

            unary_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(unary_users) != 1:
                continue
            unary_idx = int(unary_users[0])
            unary_op = model_ir.operators[unary_idx]
            if str(unary_op.op_type) not in unary_ops:
                continue
            if len(unary_op.inputs) != 1 or len(unary_op.outputs) != 1:
                continue
            if str(unary_op.inputs[0]) != pre_output_name:
                continue

            unary_output_name = str(unary_op.outputs[0])
            if unary_output_name in model_outputs:
                continue
            output_users = [int(v) for v in consumers.get(unary_output_name, [])]
            if len(output_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in output_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == unary_output_name
                ):
                    perm_post = _read_transpose_perm(model_ir, user_op)
                    if perm_post is None or perm_post != perm_post_expected:
                        valid_users = False
                        break
                    post_output_name = str(user_op.outputs[0])
                    if post_output_name in model_outputs:
                        valid_users = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                else:
                    legacy_users.append(int(user_idx))
            if not valid_users or len(post_indices) == 0:
                continue

            representative_output_name = str(post_output_names[0])
            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[pre_input_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=unary_op,
                new_outputs=[representative_output_name],
            )

            for post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, post_output_name, representative_output_name)

            old_unary_tensor = model_ir.tensors.get(unary_output_name, None)
            representative_tensor = model_ir.tensors.get(representative_output_name, None)
            if old_unary_tensor is not None and representative_tensor is not None:
                representative_tensor.dtype = str(old_unary_tensor.dtype)
                representative_tensor.quantization = _clone_quantization(old_unary_tensor.quantization)

            if len(legacy_users) > 0:
                keep_post_idx = int(post_indices[0])
                keep_post_op = model_ir.operators[keep_post_idx]
                keep_perm_name = str(keep_post_op.inputs[1])
                keep_perm_tensor = model_ir.tensors.get(keep_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_pre, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[representative_output_name, keep_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[unary_output_name],
                )
                remove_post_indices = [int(v) for v in post_indices[1:]]
            else:
                remove_post_indices = [int(v) for v in post_indices]

            remove_indices = sorted(list({int(pre_idx), *remove_post_indices}), reverse=True)
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_unary_fanout_inverse_post_bridges": int(rewritten)}


def _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold transpose wrappers around unary->binary chains with inverse post-transpose fanout.

    Target:
      X --T(P)--> X_t --(RELU|RELU6|LOGISTIC|TANH)--> U_t
      Y --T(P)--> Y_t
      U_t,Y_t --(ADD|SUB|MUL|DIV)--> Z_t
      Z_t --T(inv(P))--> Z0
      Z_t --T(inv(P))--> Z1
      ...

    Rewrite:
      X --(UNARY)--> U
      U,Y --(BINARY)--> Z0
      (all uses of Z1... are rewired to Z0)

    Safety:
    - The unary-side transpose output is consumed only by the unary op.
    - The unary output and the other transpose output are consumed only by the binary op.
    - All binary consumers are inverse post-transpose ops.
    - Intermediate tensors and post outputs are not graph outputs.
    """
    rewritten = 0
    unary_ops = {"RELU", "RELU6", "RELU_0_TO_1", "HARD_SWISH", "LOGISTIC", "TANH", "GELU"}
    binary_ops = {"ADD", "SUB", "MUL", "DIV"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for mid_idx, mid_op in enumerate(model_ir.operators):
            if str(mid_op.op_type) not in binary_ops:
                continue
            if len(mid_op.inputs) != 2 or len(mid_op.outputs) != 1:
                continue

            in0_name = str(mid_op.inputs[0])
            in1_name = str(mid_op.inputs[1])
            out_name = str(mid_op.outputs[0])

            if in0_name in model_outputs or in1_name in model_outputs or out_name in model_outputs:
                continue

            # One binary input must be produced by unary-from-transpose,
            # and the other by a matching transpose.
            candidate = None
            for unary_on_lhs in [True, False]:
                unary_input_name = in0_name if unary_on_lhs else in1_name
                other_input_name = in1_name if unary_on_lhs else in0_name

                unary_idx = producers.get(unary_input_name, None)
                if unary_idx is None:
                    continue
                unary_op = model_ir.operators[int(unary_idx)]
                if (
                    str(unary_op.op_type) not in unary_ops
                    or len(unary_op.inputs) != 1
                    or len(unary_op.outputs) != 1
                    or str(unary_op.outputs[0]) != unary_input_name
                ):
                    continue

                pre_unary_out_name = str(unary_op.inputs[0])
                pre_unary_idx = producers.get(pre_unary_out_name, None)
                if pre_unary_idx is None:
                    continue
                pre_unary_op = model_ir.operators[int(pre_unary_idx)]
                if (
                    str(pre_unary_op.op_type) != "TRANSPOSE"
                    or len(pre_unary_op.inputs) < 2
                    or len(pre_unary_op.outputs) != 1
                    or str(pre_unary_op.outputs[0]) != pre_unary_out_name
                ):
                    continue

                pre_other_idx = producers.get(other_input_name, None)
                if pre_other_idx is None:
                    continue
                pre_other_op = model_ir.operators[int(pre_other_idx)]
                if (
                    str(pre_other_op.op_type) != "TRANSPOSE"
                    or len(pre_other_op.inputs) < 2
                    or len(pre_other_op.outputs) != 1
                    or str(pre_other_op.outputs[0]) != other_input_name
                ):
                    continue

                if set(consumers.get(pre_unary_out_name, [])) != {int(unary_idx)}:
                    continue
                if set(consumers.get(unary_input_name, [])) != {int(mid_idx)}:
                    continue
                if set(consumers.get(other_input_name, [])) != {int(mid_idx)}:
                    continue

                perm_pre_unary = _read_transpose_perm(model_ir, pre_unary_op)
                perm_pre_other = _read_transpose_perm(model_ir, pre_other_op)
                if perm_pre_unary is None or perm_pre_other is None:
                    continue
                if perm_pre_unary != perm_pre_other:
                    continue

                candidate = {
                    "unary_on_lhs": bool(unary_on_lhs),
                    "unary_idx": int(unary_idx),
                    "unary_op": unary_op,
                    "unary_input_name": unary_input_name,
                    "pre_unary_idx": int(pre_unary_idx),
                    "pre_unary_op": pre_unary_op,
                    "pre_unary_out_name": pre_unary_out_name,
                    "other_input_name": other_input_name,
                    "pre_other_idx": int(pre_other_idx),
                    "pre_other_op": pre_other_op,
                    "perm_pre": perm_pre_unary,
                }
                break

            if candidate is None:
                continue

            unary_idx = int(candidate["unary_idx"])
            unary_op = candidate["unary_op"]
            unary_input_name = str(candidate["unary_input_name"])
            pre_unary_idx = int(candidate["pre_unary_idx"])
            pre_unary_op = candidate["pre_unary_op"]
            pre_unary_out_name = str(candidate["pre_unary_out_name"])
            other_input_name = str(candidate["other_input_name"])
            pre_other_idx = int(candidate["pre_other_idx"])
            pre_other_op = candidate["pre_other_op"]
            perm_pre = candidate["perm_pre"]
            unary_on_lhs = bool(candidate["unary_on_lhs"])

            out_users = [int(v) for v in consumers.get(out_name, []) if int(v) != int(mid_idx)]
            if len(out_users) == 0:
                continue

            post_indices: List[int] = []
            post_output_names: List[str] = []
            legacy_users: List[int] = []
            valid_users = True
            for user_idx in out_users:
                post_op = model_ir.operators[int(user_idx)]
                if (
                    str(post_op.op_type) == "TRANSPOSE"
                    and len(post_op.inputs) >= 2
                    and len(post_op.outputs) == 1
                    and str(post_op.inputs[0]) == out_name
                ):
                    perm_post = _read_transpose_perm(model_ir, post_op)
                    if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                        valid_users = False
                        break
                    post_output_name = str(post_op.outputs[0])
                    if post_output_name in model_outputs:
                        valid_users = False
                        break
                    post_indices.append(int(user_idx))
                    post_output_names.append(post_output_name)
                else:
                    legacy_users.append(int(user_idx))

            if not valid_users or len(post_indices) == 0:
                continue

            if pre_unary_out_name in model_outputs:
                continue

            raw_unary_input_name = str(pre_unary_op.inputs[0])
            raw_other_input_name = str(pre_other_op.inputs[0])
            raw_unary_tensor = model_ir.tensors.get(raw_unary_input_name, None)
            raw_other_tensor = model_ir.tensors.get(raw_other_input_name, None)
            unary_input_tensor = model_ir.tensors.get(unary_input_name, None)
            other_input_tensor = model_ir.tensors.get(other_input_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            if not _all_per_tensor_quantized(
                [raw_unary_tensor, raw_other_tensor, unary_input_tensor, other_input_tensor, out_tensor]
            ):
                continue

            raw_shape_unary = list(raw_unary_tensor.shape) if raw_unary_tensor is not None else None
            raw_shape_other = list(raw_other_tensor.shape) if raw_other_tensor is not None else None
            raw_broadcast_shape = _broadcast_static_shapes(raw_shape_unary, raw_shape_other)
            if raw_broadcast_shape is not None:
                mid_expected_shape = _permute_shape(list(raw_broadcast_shape), perm_pre)
                if mid_expected_shape is None:
                    continue

                in0_tensor = model_ir.tensors.get(in0_name, None)
                in1_tensor = model_ir.tensors.get(in1_name, None)
                in_shape0 = list(in0_tensor.shape) if in0_tensor is not None else None
                in_shape1 = list(in1_tensor.shape) if in1_tensor is not None else None
                mid_broadcast_shape = _broadcast_static_shapes(in_shape0, in_shape1)
                if mid_broadcast_shape is not None and not _shapes_match_if_known(mid_broadcast_shape, mid_expected_shape):
                    continue

                out_shape = list(out_tensor.shape) if out_tensor is not None else None
                if not _shapes_match_if_known(out_shape, mid_expected_shape):
                    continue

                valid_post_shapes = True
                for post_output_name in post_output_names:
                    post_tensor = model_ir.tensors.get(post_output_name, None)
                    post_shape = list(post_tensor.shape) if post_tensor is not None else None
                    if not _shapes_match_if_known(post_shape, raw_broadcast_shape):
                        valid_post_shapes = False
                        break
                if not valid_post_shapes:
                    continue

            canonical_post_output = str(post_output_names[0])
            canonical_tensor = model_ir.tensors.get(canonical_post_output, None)
            keep_post_idx = int(post_indices[0])
            keep_post_op = model_ir.operators[keep_post_idx]
            keep_post_perm_name = str(keep_post_op.inputs[1])

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[raw_unary_input_name],
            )

            if unary_on_lhs:
                new_binary_inputs = [unary_input_name, raw_other_input_name]
            else:
                new_binary_inputs = [raw_other_input_name, unary_input_name]
            _set_operator_inputs(
                model_ir=model_ir,
                op=mid_op,
                new_inputs=new_binary_inputs,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=mid_op,
                new_outputs=[canonical_post_output],
            )

            unary_out_tensor = model_ir.tensors.get(unary_input_name, None)
            if unary_out_tensor is not None and raw_unary_tensor is not None:
                unary_out_tensor.shape = [int(v) for v in list(raw_unary_tensor.shape)]
                unary_out_tensor.shape_signature = (
                    [int(v) for v in list(raw_unary_tensor.shape_signature)]
                    if raw_unary_tensor.shape_signature is not None
                    else [int(v) for v in list(raw_unary_tensor.shape)]
                )

            if canonical_tensor is not None and raw_broadcast_shape is not None:
                canonical_tensor.shape = [int(v) for v in list(raw_broadcast_shape)]
                raw_sig_unary = (
                    list(raw_unary_tensor.shape_signature)
                    if raw_unary_tensor is not None and raw_unary_tensor.shape_signature is not None
                    else list(raw_shape_unary) if raw_shape_unary is not None else None
                )
                raw_sig_other = (
                    list(raw_other_tensor.shape_signature)
                    if raw_other_tensor is not None and raw_other_tensor.shape_signature is not None
                    else list(raw_shape_other) if raw_shape_other is not None else None
                )
                raw_broadcast_signature = _broadcast_shape_signatures(raw_sig_unary, raw_sig_other)
                if raw_broadcast_signature is None:
                    raw_broadcast_signature = [int(v) for v in list(raw_broadcast_shape)]
                canonical_tensor.shape_signature = [int(v) for v in list(raw_broadcast_signature)]
                if out_tensor is not None:
                    canonical_tensor.dtype = str(out_tensor.dtype)
                    canonical_tensor.quantization = _clone_quantization(out_tensor.quantization)

            for post_output_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, post_output_name, canonical_post_output)

            if len(legacy_users) > 0:
                # Keep one adapter transpose for legacy NCHW consumers.
                keep_perm_tensor = model_ir.tensors.get(keep_post_perm_name, None)
                if keep_perm_tensor is not None:
                    keep_perm_tensor.data = np.asarray(perm_pre, dtype=np.int32)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_inputs=[canonical_post_output, keep_post_perm_name],
                )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=keep_post_op,
                    new_outputs=[out_name],
                )
                post_remove_indices = [int(v) for v in post_indices[1:]]
            else:
                post_remove_indices = [int(v) for v in post_indices]

            remove_indices = sorted(
                list(set([int(pre_unary_idx), int(pre_other_idx)] + post_remove_indices)),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_transpose_unary_binary_full_post_fanout_bridges": int(rewritten)}


def _optimize_trailing_output_transpose_passthrough_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate trailing layout transposes on output-side passthrough chains.

    Targets:
      X --TRANSPOSE(P)--> Y_out(graph output)
      X --TRANSPOSE(P)--> A --(UNARY|BINARY_WITH_CONST)*--> Y_out(graph output)

    Rewrite:
      X --(same chain without TRANSPOSE)--> Y_out

    Safety:
    - Main path must be strictly linear.
    - Binary ops must use singleton constant as side input.
    - Layout metadata is updated by applying inverse(P).
    """
    rewritten = 0
    preserve_layout_boundary_marker = "__preserve_layout_boundary__"
    layout_perms = {
        (0, 2, 1),       # NCW -> NWC
        (0, 3, 1, 2),    # NHWC -> NCHW
        (0, 2, 3, 1),    # NCHW -> NHWC
        (0, 4, 1, 2, 3), # NDHWC -> NCDHW
        (0, 2, 3, 4, 1), # NCDHW -> NDHWC
    }
    unary_passthrough_ops = {
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LEAKY_RELU",
        "LOGISTIC",
        "TANH",
        "GELU",
        "ABS",
        "NEG",
        "SQRT",
        "EXP",
        "FLOOR",
        "CEIL",
    }
    binary_passthrough_ops = {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM"}

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE":
                continue
            if len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            pre_opts = dict(pre_op.options) if isinstance(pre_op.options, dict) else {}
            preserve_layout_boundary = bool(pre_opts.get(preserve_layout_boundary_marker, False))

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])

            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue
            if tuple(int(v) for v in list(perm_pre)) not in layout_perms:
                # Keep non-layout terminal transposes (e.g. [1,0]) to preserve
                # ONNX output contracts such as [N,C] descriptor tensors.
                continue
            perm_inv = _invert_perm(perm_pre)
            if perm_inv is None:
                continue

            chain_indices: List[int] = []
            chain_ops: List[OperatorIR] = []
            chain_output_names: List[str] = []
            chain_main_input_index: Dict[int, int] = {}
            current_tensor = pre_output_name

            while True:
                current_users = [int(v) for v in consumers.get(current_tensor, [])]
                if len(current_users) != 1:
                    break
                op_idx = int(current_users[0])
                op = model_ir.operators[int(op_idx)]
                op_type = str(op.op_type)

                if op_type in unary_passthrough_ops:
                    if len(op.inputs) != 1 or len(op.outputs) != 1:
                        break
                    if str(op.inputs[0]) != current_tensor:
                        break
                    out_name = str(op.outputs[0])
                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    current_tensor = out_name
                    continue

                if op_type in binary_passthrough_ops:
                    if len(op.inputs) != 2 or len(op.outputs) != 1:
                        break
                    input_0 = str(op.inputs[0])
                    input_1 = str(op.inputs[1])
                    if input_0 == current_tensor:
                        side_input_name = input_1
                        main_input_index = 0
                    elif input_1 == current_tensor:
                        side_input_name = input_0
                        main_input_index = 1
                    else:
                        break
                    if not _is_singleton_constant_tensor(model_ir, side_input_name):
                        break
                    out_name = str(op.outputs[0])
                    chain_indices.append(int(op_idx))
                    chain_ops.append(op)
                    chain_output_names.append(out_name)
                    chain_main_input_index[int(op_idx)] = int(main_input_index)
                    current_tensor = out_name
                    continue

                break

            # Case 1: direct terminal transpose output.
            if len(chain_ops) == 0:
                if preserve_layout_boundary:
                    continue
                # Leave terminal transpose for symmetric transpose-binary patterns so
                # dedicated bridge passes can remove both pre-transposes safely.
                terminal_binary_bridge_candidate = False
                producer_idx = producers.get(pre_input_name, None)
                if producer_idx is not None:
                    producer_op = model_ir.operators[int(producer_idx)]
                    producer_type = str(producer_op.op_type)
                    if (
                        producer_type in {"ADD", "SUB", "MUL", "DIV"}
                        and len(producer_op.inputs) == 2
                        and len(producer_op.outputs) == 1
                        and str(producer_op.outputs[0]) == pre_input_name
                    ):
                        mid_in0 = str(producer_op.inputs[0])
                        mid_in1 = str(producer_op.inputs[1])
                        pre0_idx = producers.get(mid_in0, None)
                        pre1_idx = producers.get(mid_in1, None)
                        if pre0_idx is not None and pre1_idx is not None:
                            pre0_op = model_ir.operators[int(pre0_idx)]
                            pre1_op = model_ir.operators[int(pre1_idx)]
                            if (
                                str(pre0_op.op_type) == "TRANSPOSE"
                                and str(pre1_op.op_type) == "TRANSPOSE"
                                and len(pre0_op.inputs) >= 2
                                and len(pre1_op.inputs) >= 2
                                and len(pre0_op.outputs) == 1
                                and len(pre1_op.outputs) == 1
                                and str(pre0_op.outputs[0]) == mid_in0
                                and str(pre1_op.outputs[0]) == mid_in1
                                and set(int(v) for v in consumers.get(mid_in0, [])) == {int(producer_idx)}
                                and set(int(v) for v in consumers.get(mid_in1, [])) == {int(producer_idx)}
                            ):
                                perm_mid0 = _read_transpose_perm(model_ir, pre0_op)
                                perm_mid1 = _read_transpose_perm(model_ir, pre1_op)
                                if (
                                    perm_mid0 is not None
                                    and perm_mid1 is not None
                                    and perm_mid0 == perm_mid1
                                    and _is_inverse_perm(perm_mid0, perm_pre)
                                ):
                                    terminal_binary_bridge_candidate = True
                if terminal_binary_bridge_candidate:
                    continue

                if pre_output_name not in model_outputs:
                    continue
                if len(consumers.get(pre_output_name, [])) != 0:
                    continue
                producer_idx = producers.get(pre_input_name, None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                # Keep explicit output transpose after SOFTMAX to preserve graph
                # output layout contract (axis semantics are layout-sensitive).
                if str(producer_op.op_type) == "SOFTMAX":
                    continue
                producer_outputs = [str(v) for v in list(producer_op.outputs)]
                if pre_input_name not in set(producer_outputs):
                    continue
                new_outputs = [
                    str(pre_output_name) if str(v) == str(pre_input_name) else str(v)
                    for v in producer_outputs
                ]
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=producer_op,
                    new_outputs=new_outputs,
                )
                src_tensor = model_ir.tensors.get(pre_input_name, None)
                dst_tensor = model_ir.tensors.get(pre_output_name, None)
                if src_tensor is not None and dst_tensor is not None:
                    dst_tensor.dtype = str(src_tensor.dtype)
                    dst_tensor.quantization = _clone_quantization(src_tensor.quantization)
                    dst_tensor.shape = [int(v) for v in list(src_tensor.shape)]
                    dst_tensor.shape_signature = (
                        [int(v) for v in list(src_tensor.shape_signature)]
                        if src_tensor.shape_signature is not None
                        else [int(v) for v in list(src_tensor.shape)]
                    )
                del model_ir.operators[int(pre_idx)]
                rewritten += 1
                changed = True
                break

            # Case 2: transpose + passthrough chain ending at graph output.
            if current_tensor not in model_outputs:
                continue
            if preserve_layout_boundary:
                # Keep protected boundary transposes by default.
                # Exception: single TANH tail where dropping NHWC->NCHW boundary
                # transpose is explicitly safe and requested by optimization users.
                if len(chain_ops) != 1 or str(chain_ops[0].op_type) != "TANH":
                    continue
            # If the graph-output tensor is still consumed internally, removing only
            # the pre-transpose would shift layout on the internal path and can
            # leave a mismatched NCHW->NHWC adapter behind.
            if len(consumers.get(current_tensor, [])) != 0:
                continue

            # Rewire chain head to bypass pre-transpose.
            first_op = chain_ops[0]
            first_inputs = [str(v) for v in list(first_op.inputs)]
            if str(first_op.op_type) in unary_passthrough_ops:
                if len(first_inputs) != 1 or first_inputs[0] != pre_output_name:
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=[pre_input_name],
                )
            else:
                if len(first_inputs) != 2:
                    continue
                main_input_index = int(chain_main_input_index.get(int(chain_indices[0]), -1))
                if main_input_index not in [0, 1]:
                    if first_inputs[0] == pre_output_name:
                        main_input_index = 0
                    elif first_inputs[1] == pre_output_name:
                        main_input_index = 1
                    else:
                        continue
                first_inputs[main_input_index] = str(pre_input_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=first_op,
                    new_inputs=first_inputs,
                )

            # Adjust metadata from transposed layout back to source layout.
            for out_name in chain_output_names:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(out_name), None),
                    perm_inv,
                )

            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"rewritten_trailing_output_transpose_passthrough_chains": int(rewritten)}


def _optimize_dequant_prelu_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->PRELU->QUANTIZE into quantized PRELU.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --PRELU(alpha_f)--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --PRELU(alpha_q)--> Yq
    """
    folded = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            prelu_users = consumers.get(f_in_name, [])
            if len(prelu_users) != 1:
                continue
            prelu_idx = int(prelu_users[0])
            prelu_op = model_ir.operators[prelu_idx]
            if str(prelu_op.op_type) != "PRELU" or len(prelu_op.inputs) != 2 or len(prelu_op.outputs) != 1:
                continue
            if str(prelu_op.inputs[0]) != f_in_name:
                continue

            f_out_name = str(prelu_op.outputs[0])
            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or f_out_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            alpha_name = str(prelu_op.inputs[1])
            alpha_tensor = model_ir.tensors.get(alpha_name, None)
            if q_in_tensor is None or q_out_tensor is None or alpha_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue

            target_dtype = str(q_in_tensor.dtype)
            if target_dtype not in {"INT8", "UINT8"}:
                continue
            if str(q_out_tensor.dtype) != target_dtype:
                continue
            if not isinstance(alpha_tensor.data, np.ndarray):
                continue

            try:
                alpha_q, alpha_qparams = _quantize_prelu_alpha(alpha_tensor.data, target_dtype)
            except Exception:
                continue

            alpha_users = consumers.get(alpha_name, [])
            if len(alpha_users) == 1 and int(alpha_users[0]) == int(prelu_idx):
                alpha_q_name = alpha_name
                alpha_tensor.data = np.asarray(alpha_q)
                alpha_tensor.dtype = target_dtype
                alpha_tensor.shape = [int(v) for v in alpha_q.shape]
                alpha_tensor.shape_signature = [int(v) for v in alpha_q.shape]
                alpha_tensor.quantization = alpha_qparams
            else:
                alpha_q_name = _unique_tensor_name(f"{alpha_name}_q")
                model_ir.tensors[alpha_q_name] = TensorIR(
                    name=alpha_q_name,
                    dtype=target_dtype,
                    shape=[int(v) for v in alpha_q.shape],
                    shape_signature=[int(v) for v in alpha_q.shape],
                    data=np.asarray(alpha_q),
                    is_variable=False,
                    quantization=alpha_qparams,
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=prelu_op,
                new_inputs=[q_in_name, alpha_q_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=prelu_op,
                new_outputs=[q_out_name],
            )

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_prelu_quantize_chains": int(folded)}


def _optimize_dequant_prelu_depthwise_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->PRELU->DEPTHWISE_CONV_2D->QUANTIZE into quantized PRELU+DEPTHWISE_CONV_2D.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --PRELU(alpha_f)--> Pf --DEPTHWISE_CONV_2D(w_f,b_f)--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --PRELU(alpha_q)--> Pq --DEPTHWISE_CONV_2D(w_q,b_q)--> Yq
    """
    folded = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            prelu_users = consumers.get(f_in_name, [])
            if len(prelu_users) != 1:
                continue
            prelu_idx = int(prelu_users[0])
            prelu_op = model_ir.operators[prelu_idx]
            if str(prelu_op.op_type) != "PRELU" or len(prelu_op.inputs) != 2 or len(prelu_op.outputs) != 1:
                continue
            if str(prelu_op.inputs[0]) != f_in_name:
                continue
            p_f_name = str(prelu_op.outputs[0])

            dw_users = consumers.get(p_f_name, [])
            if len(dw_users) != 1:
                continue
            dw_idx = int(dw_users[0])
            dw_op = model_ir.operators[dw_idx]
            if str(dw_op.op_type) != "DEPTHWISE_CONV_2D" or len(dw_op.inputs) != 3 or len(dw_op.outputs) != 1:
                continue
            if str(dw_op.inputs[0]) != p_f_name:
                continue
            y_f_name = str(dw_op.outputs[0])

            q_users = consumers.get(y_f_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != y_f_name:
                continue
            y_q_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or p_f_name in model_ir.outputs or y_f_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            y_q_tensor = model_ir.tensors.get(y_q_name, None)
            alpha_tensor = model_ir.tensors.get(str(prelu_op.inputs[1]), None)
            w_f_tensor = model_ir.tensors.get(str(dw_op.inputs[1]), None)
            b_f_tensor = model_ir.tensors.get(str(dw_op.inputs[2]), None)
            if (
                q_in_tensor is None
                or y_q_tensor is None
                or alpha_tensor is None
                or w_f_tensor is None
                or b_f_tensor is None
            ):
                continue
            if not _all_per_tensor_quantized([q_in_tensor, y_q_tensor]):
                continue

            target_dtype = str(q_in_tensor.dtype)
            if target_dtype not in {"INT8", "UINT8"}:
                continue
            if str(y_q_tensor.dtype) != target_dtype:
                continue
            if not isinstance(alpha_tensor.data, np.ndarray):
                continue
            if not isinstance(w_f_tensor.data, np.ndarray):
                continue
            if not isinstance(b_f_tensor.data, np.ndarray):
                continue

            x_qparams = _get_per_tensor_scale_zero_point(q_in_tensor.quantization)
            if x_qparams is None:
                continue
            x_scale, _x_zero_point = x_qparams

            weights_f = np.asarray(w_f_tensor.data, dtype=np.float32)
            if weights_f.ndim != 4:
                continue
            bias_f = np.asarray(b_f_tensor.data, dtype=np.float32).reshape(-1)
            if bias_f.size != int(weights_f.shape[-1]):
                continue

            try:
                alpha_q, alpha_qparams = _quantize_prelu_alpha(alpha_tensor.data, target_dtype)
                w_q, w_qparams = _quantize_tensor_per_tensor(weights_f, target_dtype)
            except Exception:
                continue

            w_scale = float(w_qparams.scale[0])
            bias_scale = max(float(x_scale * w_scale), 1e-12)
            bias_q = np.clip(
                np.round(bias_f / bias_scale),
                np.iinfo(np.int32).min,
                np.iinfo(np.int32).max,
            ).astype(np.int32)
            b_qparams = QuantParamIR(
                scale=[float(bias_scale)],
                zero_point=[0],
                quantized_dimension=0,
            )

            alpha_q_name = _unique_tensor_name(f"{prelu_op.inputs[1]}_q")
            model_ir.tensors[alpha_q_name] = TensorIR(
                name=alpha_q_name,
                dtype=target_dtype,
                shape=[int(v) for v in alpha_q.shape],
                shape_signature=[int(v) for v in alpha_q.shape],
                data=np.asarray(alpha_q),
                is_variable=False,
                quantization=alpha_qparams,
            )

            w_q_name = _unique_tensor_name(f"{dw_op.inputs[1]}_q")
            model_ir.tensors[w_q_name] = TensorIR(
                name=w_q_name,
                dtype=target_dtype,
                shape=[int(v) for v in w_q.shape],
                shape_signature=[int(v) for v in w_q.shape],
                data=np.asarray(w_q),
                is_variable=False,
                quantization=w_qparams,
            )

            b_q_name = _unique_tensor_name(f"{dw_op.inputs[2]}_q")
            model_ir.tensors[b_q_name] = TensorIR(
                name=b_q_name,
                dtype="INT32",
                shape=[int(v) for v in bias_q.shape],
                shape_signature=[int(v) for v in bias_q.shape],
                data=np.asarray(bias_q),
                is_variable=False,
                quantization=b_qparams,
            )

            p_q_tensor = model_ir.tensors.get(p_f_name, None)
            if p_q_tensor is not None:
                p_q_tensor.dtype = target_dtype
                p_q_tensor.quantization = _clone_quantization(q_in_tensor.quantization)

            _set_operator_inputs(
                model_ir=model_ir,
                op=prelu_op,
                new_inputs=[q_in_name, alpha_q_name],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=dw_op,
                new_inputs=[p_f_name, w_q_name, b_q_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=dw_op,
                new_outputs=[y_q_name],
            )

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_prelu_depthwise_quantize_chains": int(folded)}


def _optimize_dequant_transposeconv_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->TRANSPOSE_CONV->QUANTIZE into quantized TRANSPOSE_CONV.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --TRANSPOSE_CONV(Wf)--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --TRANSPOSE_CONV(Wq)--> Yq

    Safety conditions:
    - Linear chain topology.
    - TRANSPOSE_CONV data input is the DEQUANTIZE output.
    - INT8 per-tensor activation quantization on Xq/Yq.
    - Weight tensor is a constant rank-4 tensor.
    """
    folded = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            tc_users = consumers.get(f_in_name, [])
            if len(tc_users) != 1:
                continue
            tc_idx = int(tc_users[0])
            tc_op = model_ir.operators[tc_idx]
            if str(tc_op.op_type) != "TRANSPOSE_CONV" or len(tc_op.inputs) != 3 or len(tc_op.outputs) != 1:
                continue
            # TRANSPOSE_CONV inputs: [output_shape, filter, input]
            if str(tc_op.inputs[2]) != f_in_name:
                continue
            f_out_name = str(tc_op.outputs[0])

            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or f_out_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            f_out_tensor = model_ir.tensors.get(f_out_name, None)
            if q_in_tensor is None or q_out_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue

            q_in_dtype = str(q_in_tensor.dtype)
            q_out_dtype = str(q_out_tensor.dtype)
            if q_in_dtype != "INT8" or q_out_dtype != "INT8":
                continue
            if _get_per_tensor_scale_zero_point(q_in_tensor.quantization) is None:
                continue
            if _get_per_tensor_scale_zero_point(q_out_tensor.quantization) is None:
                continue

            w_name = str(tc_op.inputs[1])
            w_tensor = model_ir.tensors.get(w_name, None)
            if w_tensor is None or not isinstance(w_tensor.data, np.ndarray):
                continue
            w_data = np.asarray(w_tensor.data)
            if w_data.ndim != 4:
                continue

            if str(w_tensor.dtype) == "INT8" and _is_per_tensor_quantization(w_tensor.quantization):
                w_q_name = w_name
            elif str(w_tensor.dtype) in {"FLOAT32", "FLOAT16", "FLOAT64"}:
                try:
                    w_q, w_qparams = _quantize_tensor_per_tensor(
                        np.asarray(w_tensor.data, dtype=np.float32),
                        "INT8",
                    )
                except Exception:
                    continue

                w_users = [int(v) for v in consumers.get(w_name, [])]
                if len(w_users) == 1 and int(w_users[0]) == int(tc_idx):
                    w_q_name = w_name
                    w_tensor.data = np.asarray(w_q)
                    w_tensor.dtype = "INT8"
                    w_tensor.shape = [int(v) for v in list(w_q.shape)]
                    w_tensor.shape_signature = [int(v) for v in list(w_q.shape)]
                    w_tensor.quantization = w_qparams
                else:
                    w_q_name = _unique_tensor_name(f"{w_name}_q")
                    model_ir.tensors[w_q_name] = TensorIR(
                        name=w_q_name,
                        dtype="INT8",
                        shape=[int(v) for v in list(w_q.shape)],
                        shape_signature=[int(v) for v in list(w_q.shape)],
                        data=np.asarray(w_q),
                        is_variable=False,
                        quantization=w_qparams,
                    )
            else:
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=tc_op,
                new_inputs=[str(tc_op.inputs[0]), w_q_name, q_in_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=tc_op,
                new_outputs=[q_out_name],
            )
            # Quantized TRANSPOSE_CONV kernels are exposed with newer opcode versions.
            tc_op.version = max(int(tc_op.version), 3)

            q_out_tensor.dtype = "INT8"
            q_out_tensor.quantization = _clone_quantization(q_out_tensor.quantization)
            if f_out_tensor is not None:
                q_out_tensor.shape = [int(v) for v in list(f_out_tensor.shape)]
                q_out_tensor.shape_signature = (
                    [int(v) for v in list(f_out_tensor.shape_signature)]
                    if f_out_tensor.shape_signature is not None
                    else [int(v) for v in list(f_out_tensor.shape)]
                )

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_transposeconv_quantize_chains": int(folded)}


def _optimize_dequant_reshape_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->RESHAPE->QUANTIZE into quantized RESHAPE.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --RESHAPE(shape)--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --RESHAPE(shape)--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - input/output quantized tensors use equivalent per-tensor quantization
    - input/output quantized dtypes are identical
    """
    folded = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            reshape_users = consumers.get(f_in_name, [])
            if len(reshape_users) != 1:
                continue
            reshape_idx = int(reshape_users[0])
            reshape_op = model_ir.operators[reshape_idx]
            if str(reshape_op.op_type) != "RESHAPE" or len(reshape_op.inputs) < 1 or len(reshape_op.outputs) != 1:
                continue
            if str(reshape_op.inputs[0]) != f_in_name:
                continue
            f_out_name = str(reshape_op.outputs[0])

            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or f_out_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            f_out_tensor = model_ir.tensors.get(f_out_name, None)
            if q_in_tensor is None or q_out_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue
            if not _is_same_per_tensor_quantization(
                q_in_tensor.quantization,
                q_out_tensor.quantization,
            ):
                continue

            q_in_dtype = str(q_in_tensor.dtype)
            q_out_dtype = str(q_out_tensor.dtype)
            if q_in_dtype != q_out_dtype:
                continue
            if q_in_dtype in {"FLOAT16", "FLOAT32", "FLOAT64", "BOOL", "STRING"}:
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=reshape_op,
                input_index=0,
                new_input_name=q_in_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=reshape_op,
                new_outputs=[q_out_name],
            )

            q_out_tensor.dtype = q_in_dtype
            q_out_tensor.quantization = _clone_quantization(q_in_tensor.quantization)
            if f_out_tensor is not None:
                q_out_tensor.shape = [int(v) for v in list(f_out_tensor.shape)]
                if f_out_tensor.shape_signature is not None:
                    q_out_tensor.shape_signature = [
                        int(v) for v in list(f_out_tensor.shape_signature)
                    ]
                else:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape)]

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_reshape_quantize_chains": int(folded)}


def _optimize_dequant_hardsigmoid_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->HardSigmoid(expanded)->QUANTIZE into quantized ops.

    Target pattern:
      Xq --DEQUANTIZE--> Xf
         --MUL(alpha)--> A
         --ADD(beta)--> B
         --MAXIMUM(lo)--> C
         --MINIMUM(hi)--> Yf
      Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --MUL(alpha_q)--> A_q
         --ADD(beta_q)--> B_q
         --MAXIMUM(lo_q)--> C_q
         --MINIMUM(hi_q)--> Yq

    Safety conditions:
    - Strict linear chain topology (single consumer at each bridge tensor).
    - Side inputs are singleton constants.
    - input/output quantized tensors are per-tensor INT8/UINT8.
    """
    folded = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _scalar_from_const_tensor(tensor_name: str) -> Optional[float]:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return None
        try:
            arr = np.asarray(tensor.data, dtype=np.float32).reshape(-1)
        except Exception:
            return None
        if int(arr.size) != 1:
            return None
        return float(arr[0])

    def _quantize_values_with_qparams(
        *,
        values: np.ndarray,
        dtype: str,
        quantization: Any,
    ) -> Optional[np.ndarray]:
        qparams = _get_per_tensor_scale_zero_point(quantization)
        if qparams is None:
            return None
        scale, zero_point = qparams
        if float(scale) <= 0.0:
            return None
        values_f = np.asarray(values, dtype=np.float32)
        if str(dtype) == "INT8":
            quantized = np.round(values_f / float(scale)) + int(zero_point)
            quantized = np.clip(quantized, -128, 127)
            return np.asarray(quantized, dtype=np.int8)
        if str(dtype) == "UINT8":
            quantized = np.round(values_f / float(scale)) + int(zero_point)
            quantized = np.clip(quantized, 0, 255)
            return np.asarray(quantized, dtype=np.uint8)
        return None

    def _retarget_const_input_to_quantized(
        *,
        op_idx: int,
        op: OperatorIR,
        input_index: int,
        target_dtype: str,
        target_qparams: Any,
        consumers: Dict[str, List[int]],
    ) -> bool:
        if input_index < 0 or input_index >= len(op.inputs):
            return False
        old_name = str(op.inputs[input_index])
        old_tensor = model_ir.tensors.get(old_name, None)
        if old_tensor is None or old_tensor.data is None:
            return False
        old_values = np.asarray(old_tensor.data, dtype=np.float32)
        quantized_values = _quantize_values_with_qparams(
            values=old_values,
            dtype=target_dtype,
            quantization=target_qparams,
        )
        if quantized_values is None:
            return False

        # If the const tensor is exclusively used here, update in place.
        const_users = [int(v) for v in consumers.get(old_name, [])]
        if len(const_users) == 1 and int(const_users[0]) == int(op_idx):
            old_tensor.data = quantized_values
            old_tensor.dtype = str(target_dtype)
            old_tensor.quantization = _clone_quantization(target_qparams)
            return True

        # Otherwise clone and bind the cloned const only for this op input.
        new_name = _unique_tensor_name(f"{old_name}_q")
        model_ir.tensors[new_name] = TensorIR(
            name=new_name,
            dtype=str(target_dtype),
            shape=list(old_tensor.shape),
            shape_signature=(
                list(old_tensor.shape_signature)
                if old_tensor.shape_signature is not None
                else list(old_tensor.shape)
            ),
            data=quantized_values,
            is_variable=False,
            quantization=_clone_quantization(target_qparams),
        )
        _replace_operator_input_at(
            model_ir=model_ir,
            op=op,
            input_index=int(input_index),
            new_input_name=str(new_name),
        )
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            mul_users = consumers.get(f_in_name, [])
            if len(mul_users) != 1:
                continue
            mul_idx = int(mul_users[0])
            mul_op = model_ir.operators[mul_idx]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue

            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            if mul_inputs[0] == f_in_name:
                mul_data_input_index = 0
                alpha_name = mul_inputs[1]
                alpha_input_index = 1
            elif mul_inputs[1] == f_in_name:
                mul_data_input_index = 1
                alpha_name = mul_inputs[0]
                alpha_input_index = 0
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, alpha_name):
                continue
            mul_out_name = str(mul_op.outputs[0])
            if mul_out_name in model_ir.outputs:
                continue

            add_users = consumers.get(mul_out_name, [])
            if len(add_users) != 1:
                continue
            add_idx = int(add_users[0])
            add_op = model_ir.operators[add_idx]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_inputs = [str(v) for v in list(add_op.inputs)]
            if add_inputs[0] == mul_out_name:
                beta_name = add_inputs[1]
                beta_input_index = 1
            elif add_inputs[1] == mul_out_name:
                beta_name = add_inputs[0]
                beta_input_index = 0
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, beta_name):
                continue
            add_out_name = str(add_op.outputs[0])
            if add_out_name in model_ir.outputs:
                continue

            max_users = consumers.get(add_out_name, [])
            if len(max_users) != 1:
                continue
            max_idx = int(max_users[0])
            max_op = model_ir.operators[max_idx]
            if str(max_op.op_type) != "MAXIMUM" or len(max_op.inputs) != 2 or len(max_op.outputs) != 1:
                continue
            max_inputs = [str(v) for v in list(max_op.inputs)]
            if max_inputs[0] == add_out_name:
                lo_name = max_inputs[1]
                lo_input_index = 1
            elif max_inputs[1] == add_out_name:
                lo_name = max_inputs[0]
                lo_input_index = 0
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, lo_name):
                continue
            max_out_name = str(max_op.outputs[0])
            if max_out_name in model_ir.outputs:
                continue

            min_users = consumers.get(max_out_name, [])
            if len(min_users) != 1:
                continue
            min_idx = int(min_users[0])
            min_op = model_ir.operators[min_idx]
            if str(min_op.op_type) != "MINIMUM" or len(min_op.inputs) != 2 or len(min_op.outputs) != 1:
                continue
            min_inputs = [str(v) for v in list(min_op.inputs)]
            if min_inputs[0] == max_out_name:
                hi_name = min_inputs[1]
                hi_input_index = 1
            elif min_inputs[1] == max_out_name:
                hi_name = min_inputs[0]
                hi_input_index = 0
            else:
                continue
            if not _is_singleton_constant_tensor(model_ir, hi_name):
                continue
            f_out_name = str(min_op.outputs[0])
            if f_out_name in model_ir.outputs:
                continue

            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            f_in_tensor = model_ir.tensors.get(f_in_name, None)
            f_out_tensor = model_ir.tensors.get(f_out_name, None)
            mul_out_tensor = model_ir.tensors.get(mul_out_name, None)
            add_out_tensor = model_ir.tensors.get(add_out_name, None)
            max_out_tensor = model_ir.tensors.get(max_out_name, None)
            if (
                q_in_tensor is None
                or q_out_tensor is None
                or mul_out_tensor is None
                or add_out_tensor is None
                or max_out_tensor is None
            ):
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue

            target_dtype = str(q_in_tensor.dtype).upper()
            if target_dtype not in {"INT8", "UINT8"}:
                continue
            if str(q_out_tensor.dtype).upper() != target_dtype:
                continue

            q_in_qparams = _get_per_tensor_scale_zero_point(q_in_tensor.quantization)
            q_out_qparams = _get_per_tensor_scale_zero_point(q_out_tensor.quantization)
            if q_in_qparams is None or q_out_qparams is None:
                continue
            if not _is_same_per_tensor_quantization(
                q_in_tensor.quantization,
                q_out_tensor.quantization,
            ):
                continue
            x_scale, x_zero_point = q_in_qparams
            if float(x_scale) <= 0.0:
                continue

            alpha = _scalar_from_const_tensor(alpha_name)
            beta = _scalar_from_const_tensor(beta_name)
            lo = _scalar_from_const_tensor(lo_name)
            hi = _scalar_from_const_tensor(hi_name)
            if any(v is None for v in [alpha, beta, lo, hi]):
                continue
            alpha = float(alpha)
            beta = float(beta)
            lo = float(lo)
            hi = float(hi)
            base_qparams = _clone_quantization(q_in_tensor.quantization)
            if base_qparams is None:
                continue

            # Conservative gating:
            # apply this int8ization only when scalar constants are representable
            # by the shared activation quantization with small error.
            def _is_scalar_representable(value: float) -> bool:
                q = _quantize_values_with_qparams(
                    values=np.asarray([value], dtype=np.float32),
                    dtype=target_dtype,
                    quantization=base_qparams,
                )
                if q is None:
                    return False
                recon = (float(np.asarray(q).reshape(-1)[0]) - float(x_zero_point)) * float(x_scale)
                tol = max(float(x_scale) * 0.25, 1e-3)
                return bool(abs(float(recon) - float(value)) <= tol)

            if not (
                _is_scalar_representable(alpha)
                and _is_scalar_representable(beta)
                and _is_scalar_representable(lo)
                and _is_scalar_representable(hi)
            ):
                continue

            if f_in_tensor is not None and not _shapes_match_if_known(q_in_tensor.shape, f_in_tensor.shape):
                continue
            if f_out_tensor is not None and not _shapes_match_if_known(q_out_tensor.shape, f_out_tensor.shape):
                continue

            if not _retarget_const_input_to_quantized(
                op_idx=int(mul_idx),
                op=mul_op,
                input_index=int(alpha_input_index),
                target_dtype=target_dtype,
                target_qparams=base_qparams,
                consumers=consumers,
            ):
                continue
            if not _retarget_const_input_to_quantized(
                op_idx=int(add_idx),
                op=add_op,
                input_index=int(beta_input_index),
                target_dtype=target_dtype,
                target_qparams=base_qparams,
                consumers=consumers,
            ):
                continue
            if not _retarget_const_input_to_quantized(
                op_idx=int(max_idx),
                op=max_op,
                input_index=int(lo_input_index),
                target_dtype=target_dtype,
                target_qparams=base_qparams,
                consumers=consumers,
            ):
                continue
            if not _retarget_const_input_to_quantized(
                op_idx=int(min_idx),
                op=min_op,
                input_index=int(hi_input_index),
                target_dtype=target_dtype,
                target_qparams=base_qparams,
                consumers=consumers,
            ):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul_op,
                input_index=int(mul_data_input_index),
                new_input_name=q_in_name,
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=min_op,
                new_outputs=[q_out_name],
            )

            mul_out_tensor.dtype = str(target_dtype)
            mul_out_tensor.quantization = _clone_quantization(base_qparams)
            add_out_tensor.dtype = str(target_dtype)
            add_out_tensor.quantization = _clone_quantization(base_qparams)
            max_out_tensor.dtype = str(target_dtype)
            max_out_tensor.quantization = _clone_quantization(base_qparams)
            q_out_tensor.dtype = str(target_dtype)
            q_out_tensor.quantization = _clone_quantization(base_qparams)

            if f_out_tensor is not None:
                q_out_tensor.shape = [int(v) for v in list(f_out_tensor.shape)]
                if f_out_tensor.shape_signature is not None:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape_signature)]
                else:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape)]

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_hardsigmoid_quantize_chains": int(folded)}


def _optimize_dequant_maxpool_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->MAX_POOL_2D->QUANTIZE into quantized MAX_POOL_2D.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --MAX_POOL_2D--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --MAX_POOL_2D--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - input/output quantized tensors use equivalent per-tensor quantization
    - input/output quantized dtypes are identical INT8/UINT8
    """
    folded = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            pool_users = consumers.get(f_in_name, [])
            if len(pool_users) != 1:
                continue
            pool_idx = int(pool_users[0])
            pool_op = model_ir.operators[pool_idx]
            if str(pool_op.op_type) != "MAX_POOL_2D" or len(pool_op.inputs) != 1 or len(pool_op.outputs) != 1:
                continue
            if str(pool_op.inputs[0]) != f_in_name:
                continue
            f_out_name = str(pool_op.outputs[0])

            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or f_out_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            f_in_tensor = model_ir.tensors.get(f_in_name, None)
            f_out_tensor = model_ir.tensors.get(f_out_name, None)
            if q_in_tensor is None or q_out_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue
            if not _is_same_per_tensor_quantization(
                q_in_tensor.quantization,
                q_out_tensor.quantization,
            ):
                continue

            q_in_dtype = str(q_in_tensor.dtype).upper()
            q_out_dtype = str(q_out_tensor.dtype).upper()
            if q_in_dtype not in {"INT8", "UINT8"} or q_out_dtype != q_in_dtype:
                continue

            if f_in_tensor is not None and not _shapes_match_if_known(q_in_tensor.shape, f_in_tensor.shape):
                continue
            if f_out_tensor is not None and not _shapes_match_if_known(q_out_tensor.shape, f_out_tensor.shape):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=pool_op,
                new_inputs=[q_in_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=pool_op,
                new_outputs=[q_out_name],
            )

            q_out_tensor.dtype = q_in_dtype
            q_out_tensor.quantization = _clone_quantization(q_in_tensor.quantization)
            if f_out_tensor is not None:
                q_out_tensor.shape = [int(v) for v in list(f_out_tensor.shape)]
                if f_out_tensor.shape_signature is not None:
                    q_out_tensor.shape_signature = [
                        int(v) for v in list(f_out_tensor.shape_signature)
                    ]
                else:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape)]

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_maxpool_quantize_chains": int(folded)}


def _optimize_dequant_softmax_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->SOFTMAX->QUANTIZE into quantized SOFTMAX when TFLite-compatible.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --SOFTMAX(beta=1.0)--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --SOFTMAX(int8/uint8)--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - input/output quantized tensors are per-tensor INT8/UINT8
    - output quantization matches TFLite quantized softmax canonical params
      (INT8: scale=1/256, zp=-128; UINT8: scale=1/256, zp=0)
    - Softmax beta is approximately 1.0
    """
    folded = 0

    def _is_supported_softmax_qparams(dtype: str, quantization: Any) -> bool:
        qparams = _get_per_tensor_scale_zero_point(quantization)
        if qparams is None:
            return False
        scale, zero_point = qparams
        if not np.isclose(float(scale), 1.0 / 256.0, rtol=0.0, atol=1e-7):
            return False
        dtype_u = str(dtype).upper()
        if dtype_u == "INT8":
            return int(zero_point) == -128
        if dtype_u == "UINT8":
            return int(zero_point) == 0
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            softmax_users = consumers.get(f_in_name, [])
            if len(softmax_users) != 1:
                continue
            softmax_idx = int(softmax_users[0])
            softmax_op = model_ir.operators[softmax_idx]
            if str(softmax_op.op_type) != "SOFTMAX" or len(softmax_op.inputs) != 1 or len(softmax_op.outputs) != 1:
                continue
            if str(softmax_op.inputs[0]) != f_in_name:
                continue
            f_out_name = str(softmax_op.outputs[0])

            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or f_out_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            f_in_tensor = model_ir.tensors.get(f_in_name, None)
            f_out_tensor = model_ir.tensors.get(f_out_name, None)
            if q_in_tensor is None or q_out_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue

            q_in_dtype = str(q_in_tensor.dtype).upper()
            q_out_dtype = str(q_out_tensor.dtype).upper()
            if q_in_dtype not in {"INT8", "UINT8"} or q_out_dtype != q_in_dtype:
                continue
            if not _is_supported_softmax_qparams(q_out_dtype, q_out_tensor.quantization):
                continue

            beta = float(softmax_op.options.get("beta", 1.0))
            if not np.isclose(beta, 1.0, rtol=0.0, atol=1e-6):
                continue

            if f_in_tensor is not None and not _shapes_match_if_known(q_in_tensor.shape, f_in_tensor.shape):
                continue
            if f_out_tensor is not None and not _shapes_match_if_known(q_out_tensor.shape, f_out_tensor.shape):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=softmax_op,
                new_inputs=[q_in_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=softmax_op,
                new_outputs=[q_out_name],
            )
            # TFLite quantized SOFTMAX opcode version differs by dtype.
            softmax_op.version = 2 if q_out_dtype == "INT8" else 1

            # Keep output tensor metadata aligned to the pre-quantized softmax output shape.
            if f_out_tensor is not None:
                q_out_tensor.shape = [int(v) for v in list(f_out_tensor.shape)]
                if f_out_tensor.shape_signature is not None:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape_signature)]
                else:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape)]

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_softmax_quantize_chains": int(folded)}


def _optimize_dequant_logistic_quantize_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold DEQUANTIZE->LOGISTIC->QUANTIZE into quantized LOGISTIC when TFLite-compatible.

    Target pattern:
      Xq --DEQUANTIZE--> Xf --LOGISTIC--> Yf --QUANTIZE--> Yq

    Rewritten:
      Xq --LOGISTIC(int8/uint8)--> Yq

    Safety conditions:
    - Chain is linear (single consumer at each bridge tensor)
    - input/output quantized tensors are per-tensor INT8/UINT8
    - output quantization matches TFLite quantized logistic canonical params
      (INT8: scale=1/256, zp=-128; UINT8: scale=1/256, zp=0)
    """
    folded = 0

    def _is_supported_logistic_qparams(dtype: str, quantization: Any) -> bool:
        qparams = _get_per_tensor_scale_zero_point(quantization)
        if qparams is None:
            return False
        scale, zero_point = qparams
        if not np.isclose(float(scale), 1.0 / 256.0, rtol=0.0, atol=1e-7):
            return False
        dtype_u = str(dtype).upper()
        if dtype_u == "INT8":
            return int(zero_point) == -128
        if dtype_u == "UINT8":
            return int(zero_point) == 0
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            q_in_name = str(dq_op.inputs[0])
            f_in_name = str(dq_op.outputs[0])

            logistic_users = consumers.get(f_in_name, [])
            if len(logistic_users) != 1:
                continue
            logistic_idx = int(logistic_users[0])
            logistic_op = model_ir.operators[logistic_idx]
            if str(logistic_op.op_type) != "LOGISTIC" or len(logistic_op.inputs) != 1 or len(logistic_op.outputs) != 1:
                continue
            if str(logistic_op.inputs[0]) != f_in_name:
                continue
            f_out_name = str(logistic_op.outputs[0])

            q_users = consumers.get(f_out_name, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != f_out_name:
                continue
            q_out_name = str(q_op.outputs[0])

            if f_in_name in model_ir.outputs or f_out_name in model_ir.outputs:
                continue

            q_in_tensor = model_ir.tensors.get(q_in_name, None)
            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            f_in_tensor = model_ir.tensors.get(f_in_name, None)
            f_out_tensor = model_ir.tensors.get(f_out_name, None)
            if q_in_tensor is None or q_out_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_in_tensor, q_out_tensor]):
                continue

            q_in_dtype = str(q_in_tensor.dtype).upper()
            q_out_dtype = str(q_out_tensor.dtype).upper()
            if q_in_dtype not in {"INT8", "UINT8"} or q_out_dtype != q_in_dtype:
                continue
            if not _is_supported_logistic_qparams(q_out_dtype, q_out_tensor.quantization):
                continue

            if f_in_tensor is not None and not _shapes_match_if_known(q_in_tensor.shape, f_in_tensor.shape):
                continue
            if f_out_tensor is not None and not _shapes_match_if_known(q_out_tensor.shape, f_out_tensor.shape):
                continue

            _set_operator_inputs(
                model_ir=model_ir,
                op=logistic_op,
                new_inputs=[q_in_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=logistic_op,
                new_outputs=[q_out_name],
            )
            # TFLite quantized LOGISTIC uses version 2 for INT8.
            logistic_op.version = 2 if q_out_dtype == "INT8" else 1

            if f_out_tensor is not None:
                q_out_tensor.shape = [int(v) for v in list(f_out_tensor.shape)]
                if f_out_tensor.shape_signature is not None:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape_signature)]
                else:
                    q_out_tensor.shape_signature = [int(v) for v in list(f_out_tensor.shape)]

            for remove_idx in sorted([dq_idx, q_idx], reverse=True):
                del model_ir.operators[remove_idx]
            folded += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"folded_dequant_logistic_quantize_chains": int(folded)}


def _optimize_terminal_quantize_dequantize(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove terminal QUANTIZE->DEQUANTIZE pairs for float outputs when safe.

    Target pattern:
      x_float -> QUANTIZE -> q -> DEQUANTIZE -> y_float (graph output)

    Safety conditions:
    - q is consumed only by that DEQUANTIZE
    - y_float is a graph output and has no internal consumers
    - x_float is produced by another op and consumed only by QUANTIZE
      (allows preserving output name by renaming x_float -> y_float)
    """
    removed_pairs = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for q_idx, q_op in enumerate(model_ir.operators):
            if str(q_op.op_type) != "QUANTIZE":
                continue
            if len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue

            float_input_name = q_op.inputs[0]
            quantized_name = q_op.outputs[0]
            if float_input_name in model_ir.inputs:
                continue

            quantized_users = consumers.get(quantized_name, [])
            if len(quantized_users) != 1:
                continue

            dq_idx = int(quantized_users[0])
            if dq_idx == q_idx:
                continue
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE":
                continue
            if len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if dq_op.inputs[0] != quantized_name:
                continue

            float_output_name = dq_op.outputs[0]
            if float_output_name not in model_ir.outputs:
                continue
            if len(consumers.get(float_output_name, [])) > 0:
                continue

            float_input_users = consumers.get(float_input_name, [])
            if len(float_input_users) != 1 or int(float_input_users[0]) != q_idx:
                continue
            if float_input_name not in producers:
                continue
            if float_input_name in model_ir.outputs:
                continue

            _rename_tensor_globally(
                model_ir=model_ir,
                old_name=float_input_name,
                new_name=float_output_name,
            )

            for remove_idx in sorted([q_idx, dq_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_pairs += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_terminal_quantize_dequantize_pairs": int(removed_pairs),
    }


def _optimize_fuse_conv_activation_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fuse producer -> Activation chains into producer fusedActivationFunction in flatbuffer_direct IR.

    Target:
      (CONV_2D|DEPTHWISE_CONV_2D, fusedActivationFunction=NONE) -> (RELU|RELU6)
      (ADD|SUB|MUL|DIV, fusedActivationFunction=NONE) -> (RELU|RELU6)

    Rewrite:
      (producer, fusedActivationFunction=<activation>)

    Safety:
    - Producer output must have exactly one consumer.
    - Activation must be a unary op with one input and one output.
    - Producer output tensor and activation output tensor must have matching dtype when both exist.
    """
    fused = 0
    fused_conv = 0
    fused_add = 0
    fused_sub = 0
    fused_mul = 0
    fused_div = 0
    skip_add_activation_fuse_marker = "__skip_add_activation_fuse__"
    binary_activation_map = {
        "RELU": "RELU",
        "RELU6": "RELU6",
    }
    activation_map_by_producer = {
        "CONV_2D": {
            "RELU": "RELU",
            "RELU6": "RELU6",
            "RELU_N1_TO_1": "RELU_N1_TO_1",
        },
        "DEPTHWISE_CONV_2D": {
            "RELU": "RELU",
            "RELU6": "RELU6",
            "RELU_N1_TO_1": "RELU_N1_TO_1",
        },
        "ADD": dict(binary_activation_map),
        "SUB": dict(binary_activation_map),
        "MUL": dict(binary_activation_map),
        "DIV": dict(binary_activation_map),
    }
    protected_boundary_tensor_names = _get_protected_boundary_tensor_names(model_ir)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for producer_idx, producer_op in enumerate(model_ir.operators):
            producer_type = str(producer_op.op_type).upper()
            activation_map = activation_map_by_producer.get(producer_type, None)
            if activation_map is None:
                continue
            if len(producer_op.outputs) != 1:
                continue

            producer_opts = dict(producer_op.options) if isinstance(producer_op.options, dict) else {}
            fused_act = str(producer_opts.get("fusedActivationFunction", "NONE")).upper()
            if fused_act != "NONE":
                continue
            producer_out_name = str(producer_op.outputs[0])
            producer_users = [int(v) for v in consumers.get(producer_out_name, [])]
            if len(producer_users) != 1:
                continue

            act_idx = int(producer_users[0])
            if act_idx < 0 or act_idx >= len(model_ir.operators):
                continue
            if act_idx == int(producer_idx):
                continue
            act_op = model_ir.operators[act_idx]
            act_type = str(act_op.op_type).upper()
            fused_target = activation_map.get(act_type, None)
            if fused_target is None:
                continue
            if len(act_op.inputs) != 1 or len(act_op.outputs) != 1:
                continue
            if str(act_op.inputs[0]) != producer_out_name:
                continue

            act_out_name = str(act_op.outputs[0])
            producer_out_tensor = model_ir.tensors.get(producer_out_name, None)
            act_out_tensor = model_ir.tensors.get(act_out_name, None)
            if producer_out_tensor is not None and act_out_tensor is not None:
                if str(producer_out_tensor.dtype).upper() != str(act_out_tensor.dtype).upper():
                    continue
            if producer_out_name in protected_boundary_tensor_names:
                continue
            if act_out_name in protected_boundary_tensor_names:
                continue
            # Keep explicit activation node when its output is both a graph output and
            # an internal bridge tensor. Fusing in this case can relabel NHWC bridge
            # tensors to ONNX/NCHW names and later trigger wrong layout adapters.
            if act_out_name in model_ir.outputs and len(consumers.get(act_out_name, [])) > 0:
                continue

            if producer_type == "ADD":
                # Transpose bridge rewrites may leave this marker behind. For strict
                # single-consumer Add->Activation chains, fusing is still safe.
                producer_opts.pop(skip_add_activation_fuse_marker, None)
            producer_opts["fusedActivationFunction"] = str(fused_target)
            producer_op.options = producer_opts
            _set_operator_outputs(
                model_ir=model_ir,
                op=producer_op,
                new_outputs=[act_out_name],
            )
            if producer_type in {"CONV_2D", "DEPTHWISE_CONV_2D"}:
                _append_tensor_lineage_event(
                    model_ir=model_ir,
                    event={
                        "kind": "fuse_conv_activation",
                        "conv_op_type": str(producer_op.op_type),
                        "activation_op_type": str(act_op.op_type),
                        "fused_activation": str(fused_target),
                        "conv_output": str(producer_out_name),
                        "fused_output": str(act_out_name),
                    },
                )
            elif producer_type == "ADD":
                _append_tensor_lineage_event(
                    model_ir=model_ir,
                    event={
                        "kind": "fuse_add_activation",
                        "add_op_type": str(producer_op.op_type),
                        "activation_op_type": str(act_op.op_type),
                        "fused_activation": str(fused_target),
                        "add_output": str(producer_out_name),
                        "fused_output": str(act_out_name),
                    },
                )
            elif producer_type in {"SUB", "MUL", "DIV"}:
                _append_tensor_lineage_event(
                    model_ir=model_ir,
                    event={
                        "kind": "fuse_binary_activation",
                        "binary_op_type": str(producer_op.op_type),
                        "activation_op_type": str(act_op.op_type),
                        "fused_activation": str(fused_target),
                        "binary_output": str(producer_out_name),
                        "fused_output": str(act_out_name),
                    },
                )

            del model_ir.operators[int(act_idx)]
            fused += 1
            if producer_type in {"CONV_2D", "DEPTHWISE_CONV_2D"}:
                fused_conv += 1
            elif producer_type == "ADD":
                fused_add += 1
            elif producer_type == "SUB":
                fused_sub += 1
            elif producer_type == "MUL":
                fused_mul += 1
            elif producer_type == "DIV":
                fused_div += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "fused_conv_activation_chains": int(fused_conv),
        "fused_add_activation_chains": int(fused_add),
        "fused_sub_activation_chains": int(fused_sub),
        "fused_mul_activation_chains": int(fused_mul),
        "fused_div_activation_chains": int(fused_div),
        "fused_binary_activation_chains": int(fused_add + fused_sub + fused_mul + fused_div),
        "fused_activation_chains_total": int(fused),
    }


def _optimize_fold_conv_mul_add_affine_chains(
    model_ir: ModelIR,
    *,
    enable_conv_add_only_fold: bool = True,
) -> Dict[str, int]:
    """
    Fold single-path affine chains after CONV_2D into CONV_2D.

    Targets:
      x --CONV_2D(w,b)--> y
      y --ADD(a)--> o
        or
      y --MUL(m)--> z
      (optional) z --ADD(a)--> o

    Rewrite 1:
      x --CONV_2D(w,b')--> o
      where:
        b'[oc] = b[oc] + a[oc]

    Rewrite 2:
      x --CONV_2D(w',b')--> z

    Rewrite 3:
      z --ADD(a)--> o
      x --CONV_2D(w',b')--> o
      where:
        w'[oc,kh,kw,ic] = w[oc,kh,kw,ic] * m[oc]
        b'[oc]          = b[oc] * m[oc] + a[oc]

    Rewrite 4 (RELU-preserving partial fold):
      x --CONV_2D(w,b, fused=RELU)--> y
      y --MUL(m)--> z
      (optional) z --ADD(a)--> o
      x --CONV_2D(w',b', fused=RELU)--> z
      (optional) z --ADD(a)--> o
      where m[oc] >= 0 and:
        w'[oc,kh,kw,ic] = w[oc,kh,kw,ic] * m[oc]
        b'[oc]          = b[oc] * m[oc]

      NOTE: ADD is intentionally kept because a post-RELU affine offset
      cannot, in general, be represented by CONV_2D fused activations.

    Safety:
    - Conv output must have exactly one consumer (single path).
    - ADD/MUL side input must be constant representable as scalar or NHWC-channelwise.
    - When MUL is followed by ADD, ADD side input must also be constant representable as scalar or NHWC-channelwise.
    - Conv filter/bias must be constant and not shared by other operators.
    - Chains with Conv output fan-out (>=2) are excluded.
    """

    def _extract_nhwc_channelwise_coeff(
        *,
        tensor: Optional[TensorIR],
        out_channels: int,
        allow_leading_channel_axis: bool = True,
    ) -> Optional[np.ndarray]:
        if tensor is None or tensor.data is None:
            return None
        if int(out_channels) <= 0:
            return None
        raw = np.asarray(tensor.data)
        if raw.size == 0:
            return None
        coeff = np.asarray(raw, dtype=np.float32)
        if coeff.size == 1:
            return np.full((int(out_channels),), float(coeff.reshape(-1)[0]), dtype=np.float32)
        shape = [int(v) for v in list(coeff.shape)]
        if len(shape) == 1 and int(shape[0]) == int(out_channels):
            return coeff.reshape(int(out_channels)).astype(np.float32, copy=False)
        if bool(allow_leading_channel_axis) and len(shape) >= 2 and int(shape[0]) == int(out_channels):
            if int(np.prod(shape[1:])) == 1:
                return coeff.reshape(int(out_channels)).astype(np.float32, copy=False)
        if len(shape) >= 2 and int(shape[-1]) == int(out_channels):
            if int(np.prod(shape[:-1])) == 1:
                return coeff.reshape(int(out_channels)).astype(np.float32, copy=False)
        return None

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    folded = 0
    folded_conv_add_only = 0
    folded_conv_mul_only = 0
    folded_conv_mul_add = 0

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for conv_idx, conv_op in enumerate(model_ir.operators):
            if str(conv_op.op_type) != "CONV_2D" or len(conv_op.inputs) < 2 or len(conv_op.outputs) != 1:
                continue
            conv_opts = dict(conv_op.options) if isinstance(conv_op.options, dict) else {}
            conv_fused_activation = str(
                conv_opts.get("fusedActivationFunction", "NONE")
            ).upper()
            if conv_fused_activation not in {"NONE", "RELU"}:
                continue
            conv_out_name = str(conv_op.outputs[0])
            conv_users = [int(v) for v in consumers.get(conv_out_name, [])]
            if len(conv_users) != 1:
                continue

            first_idx = int(conv_users[0])
            first_op = model_ir.operators[int(first_idx)]
            first_type = str(first_op.op_type)

            add_only_mode = False
            mul_side_tensor: Optional[TensorIR] = None
            mul_out_name: Optional[str] = None
            mul_idx: Optional[int] = None

            add_only_idx: Optional[int] = None
            add_only_op: Optional[OperatorIR] = None
            add_only_side_tensor: Optional[TensorIR] = None
            add_only_out_name: Optional[str] = None
            add_only_fused_activation: str = "NONE"

            if first_type == "MUL":
                if len(first_op.inputs) != 2 or len(first_op.outputs) != 1:
                    continue
                mul_idx = int(first_idx)
                mul_op = first_op
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                if mul_inputs[0] == conv_out_name:
                    mul_side_name = str(mul_inputs[1])
                elif mul_inputs[1] == conv_out_name:
                    mul_side_name = str(mul_inputs[0])
                else:
                    continue
                mul_side_tensor = model_ir.tensors.get(mul_side_name, None)
                if mul_side_tensor is None or mul_side_tensor.data is None:
                    continue
                mul_out_name = str(mul_op.outputs[0])
                mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            elif first_type == "ADD":
                if conv_fused_activation != "NONE":
                    continue
                if not bool(enable_conv_add_only_fold):
                    continue
                if len(first_op.inputs) != 2 or len(first_op.outputs) != 1:
                    continue
                add_only_opts = dict(first_op.options) if isinstance(first_op.options, dict) else {}
                add_only_fused_activation = str(
                    add_only_opts.get("fusedActivationFunction", "NONE")
                ).upper()
                if add_only_fused_activation not in {"NONE", "RELU", "RELU6"}:
                    continue
                add_only_idx = int(first_idx)
                add_only_op = first_op
                add_only_inputs = [str(v) for v in list(add_only_op.inputs)]
                if add_only_inputs[0] == conv_out_name:
                    add_only_side_name = str(add_only_inputs[1])
                elif add_only_inputs[1] == conv_out_name:
                    add_only_side_name = str(add_only_inputs[0])
                else:
                    continue
                add_only_side_tensor = model_ir.tensors.get(add_only_side_name, None)
                if add_only_side_tensor is None or add_only_side_tensor.data is None:
                    continue
                add_only_out_name = str(add_only_op.outputs[0])
                add_only_users = [int(v) for v in consumers.get(add_only_out_name, [])]
                # Keep this rewrite strict to single-path chains.
                if len(add_only_users) != 1:
                    continue
                add_only_mode = True
                mul_users = []
            else:
                continue

            has_add = False
            add_idx: Optional[int] = None
            add_op: Optional[OperatorIR] = None
            add_side_name: Optional[str] = None
            add_side_tensor: Optional[TensorIR] = None
            keep_add_after_mul_fold = False
            if not add_only_mode and len(mul_users) == 1:
                candidate_add_idx = int(mul_users[0])
                candidate_add_op = model_ir.operators[int(candidate_add_idx)]
                if str(candidate_add_op.op_type) == "ADD" and len(candidate_add_op.inputs) == 2 and len(candidate_add_op.outputs) == 1:
                    candidate_add_opts = dict(candidate_add_op.options) if isinstance(candidate_add_op.options, dict) else {}
                    if str(candidate_add_opts.get("fusedActivationFunction", "NONE")).upper() != "NONE":
                        continue
                    add_inputs = [str(v) for v in list(candidate_add_op.inputs)]
                    candidate_add_side_name: Optional[str] = None
                    if add_inputs[0] == str(mul_out_name):
                        candidate_add_side_name = str(add_inputs[1])
                    elif add_inputs[1] == str(mul_out_name):
                        candidate_add_side_name = str(add_inputs[0])
                    if candidate_add_side_name is not None:
                        candidate_add_side_tensor = model_ir.tensors.get(candidate_add_side_name, None)
                        if candidate_add_side_tensor is not None and candidate_add_side_tensor.data is not None:
                            has_add = True
                            add_idx = int(candidate_add_idx)
                            add_op = candidate_add_op
                            add_side_name = str(candidate_add_side_name)
                            add_side_tensor = candidate_add_side_tensor

            filter_name = str(conv_op.inputs[1])
            filter_tensor = model_ir.tensors.get(filter_name, None)
            if (
                filter_tensor is None
                or filter_tensor.data is None
                or bool(filter_tensor.is_variable)
                or filter_tensor.quantization is not None
            ):
                continue
            filter_users = [int(v) for v in consumers.get(filter_name, [])]
            if set(int(v) for v in filter_users) != {int(conv_idx)}:
                continue
            filter_data = np.asarray(filter_tensor.data)
            if filter_data.ndim != 4:
                continue
            out_channels = int(filter_data.shape[0])
            if out_channels <= 0:
                continue

            bias_name: Optional[str] = None
            bias_tensor: Optional[TensorIR] = None
            if len(conv_op.inputs) >= 3:
                bias_name = str(conv_op.inputs[2])
                bias_tensor = model_ir.tensors.get(bias_name, None)
                if (
                    bias_tensor is None
                    or bias_tensor.data is None
                    or bool(bias_tensor.is_variable)
                    or bias_tensor.quantization is not None
                ):
                    continue
                bias_users = [int(v) for v in consumers.get(str(bias_name), [])]
                if set(int(v) for v in bias_users) != {int(conv_idx)}:
                    continue
                bias_values = np.asarray(bias_tensor.data).reshape(-1)
                if int(bias_values.size) != int(out_channels):
                    continue
            else:
                bias_name = _unique_tensor_name(f"{conv_out_name}_affine_fold_bias")
                bias_tensor = TensorIR(
                    name=str(bias_name),
                    dtype=str(filter_tensor.dtype),
                    shape=[int(out_channels)],
                    shape_signature=[int(out_channels)],
                    data=np.zeros((int(out_channels),), dtype=np.float32),
                    is_variable=False,
                    quantization=None,
                )
                model_ir.tensors[str(bias_name)] = bias_tensor
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=conv_op,
                    new_inputs=[str(conv_op.inputs[0]), str(filter_name), str(bias_name)],
                )

            if add_only_mode:
                mul_coeff = np.ones((int(out_channels),), dtype=np.float32)
                add_coeff = _extract_nhwc_channelwise_coeff(
                    tensor=add_only_side_tensor,
                    out_channels=int(out_channels),
                    # Keep add-only folding strict: accept only scalar,
                    # [C], or trailing-channel broadcast forms (..., C)
                    # with all non-channel dims == 1.
                    allow_leading_channel_axis=False,
                )
                if add_coeff is None:
                    continue
                folded_output_name = str(add_only_out_name)
                remove_indices = [int(add_only_idx)]
            else:
                mul_coeff = _extract_nhwc_channelwise_coeff(
                    tensor=mul_side_tensor,
                    out_channels=int(out_channels),
                )
                if mul_coeff is None:
                    continue
                if conv_fused_activation == "RELU" and bool(np.any(mul_coeff < 0.0)):
                    # For RELU-preserving rewrite, scale must be non-negative:
                    # relu(z) * m == relu(z * m) only when m >= 0.
                    continue
                add_coeff = np.zeros((int(out_channels),), dtype=np.float32)
                if has_add:
                    if conv_fused_activation == "RELU":
                        # Keep ADD in place; fold only MUL into CONV_2D.
                        keep_add_after_mul_fold = True
                        folded_output_name = str(mul_out_name)
                        remove_indices = [int(mul_idx)]
                    else:
                        add_coeff = _extract_nhwc_channelwise_coeff(
                            tensor=add_side_tensor,
                            out_channels=int(out_channels),
                        )
                        if add_coeff is None:
                            continue
                        assert add_op is not None
                        assert add_idx is not None
                        folded_output_name = str(add_op.outputs[0])
                        remove_indices = [int(mul_idx), int(add_idx)]
                else:
                    folded_output_name = str(mul_out_name)
                    remove_indices = [int(mul_idx)]
            if not np.isfinite(mul_coeff).all() or not np.isfinite(add_coeff).all():
                continue

            filter_dtype = np.asarray(filter_tensor.data).dtype
            folded_filter = (
                np.asarray(filter_tensor.data, dtype=np.float32)
                * np.asarray(mul_coeff, dtype=np.float32).reshape(int(out_channels), 1, 1, 1)
            )
            filter_tensor.data = np.asarray(folded_filter, dtype=filter_dtype)

            assert bias_tensor is not None
            bias_dtype = np.asarray(bias_tensor.data).dtype
            bias_values = np.asarray(bias_tensor.data, dtype=np.float32).reshape(int(out_channels))
            folded_bias = bias_values * np.asarray(mul_coeff, dtype=np.float32) + np.asarray(add_coeff, dtype=np.float32)
            bias_tensor.data = np.asarray(folded_bias, dtype=bias_dtype)
            bias_tensor.shape = [int(out_channels)]
            bias_tensor.shape_signature = [int(out_channels)]

            _set_operator_outputs(
                model_ir=model_ir,
                op=conv_op,
                new_outputs=[folded_output_name],
            )
            if add_only_mode and str(add_only_fused_activation) != "NONE":
                if not isinstance(conv_op.options, dict):
                    conv_op.options = {}
                conv_op.options["fusedActivationFunction"] = str(add_only_fused_activation)

            for remove_idx in sorted(remove_indices, reverse=True):
                del model_ir.operators[int(remove_idx)]

            folded += 1
            if add_only_mode:
                folded_conv_add_only += 1
            else:
                if has_add and not bool(keep_add_after_mul_fold):
                    folded_conv_mul_add += 1
                else:
                    folded_conv_mul_only += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "folded_conv_mul_add_affine_chains": int(folded),
        "folded_conv_add_only_affine_chains": int(folded_conv_add_only),
        "folded_conv_mul_only_affine_chains": int(folded_conv_mul_only),
        "folded_conv_mul_add_only_affine_chains": int(folded_conv_mul_add),
    }


def _sanitize_terminal_transpose_before_dequantize(model_ir: ModelIR) -> Dict[str, int]:
    """
    Sanitize terminal TRANSPOSE->DEQUANTIZE by swapping to DEQUANTIZE->TRANSPOSE.

    Target pattern:
      q --TRANSPOSE(P)--> q_t --DEQUANTIZE--> y(graph output)

    Rewritten:
      q --DEQUANTIZE--> y_pre --TRANSPOSE(P)--> y(graph output)

    Safety conditions:
    - y is a graph output without internal consumers
    - q_t is consumed only by that DEQUANTIZE
    - q/q_t quantization is per-tensor (to avoid axis remap ambiguity)
    """
    sanitized = 0
    removed_terminal_dequantize_transpose = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for dq_idx, dq_op in enumerate(model_ir.operators):
            if str(dq_op.op_type) != "DEQUANTIZE":
                continue
            if len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue

            dq_input_name = str(dq_op.inputs[0])
            dq_output_name = str(dq_op.outputs[0])
            if dq_output_name not in model_ir.outputs:
                continue
            if len(consumers.get(dq_output_name, [])) > 0:
                continue
            if dq_input_name in model_ir.outputs:
                continue

            transpose_idx = producers.get(dq_input_name, None)
            if transpose_idx is None:
                continue
            transpose_op = model_ir.operators[int(transpose_idx)]
            if str(transpose_op.op_type) != "TRANSPOSE":
                continue
            if len(transpose_op.inputs) < 2 or len(transpose_op.outputs) != 1:
                continue
            if str(transpose_op.outputs[0]) != dq_input_name:
                continue

            q_input_name = str(transpose_op.inputs[0])
            if q_input_name in model_ir.outputs:
                continue

            dq_input_users = [int(v) for v in consumers.get(dq_input_name, [])]
            if len(dq_input_users) != 1 or int(dq_input_users[0]) != int(dq_idx):
                continue

            q_tensor = model_ir.tensors.get(q_input_name, None)
            qt_tensor = model_ir.tensors.get(dq_input_name, None)
            y_tensor = model_ir.tensors.get(dq_output_name, None)
            if q_tensor is None or qt_tensor is None or y_tensor is None:
                continue
            if not _all_per_tensor_quantized([q_tensor, qt_tensor]):
                continue

            q_shape = list(q_tensor.shape)
            q_signature = (
                list(q_tensor.shape_signature)
                if q_tensor.shape_signature is not None
                else list(q_shape)
            )
            if len(q_shape) == 0:
                continue
            perm = _read_transpose_perm(model_ir, transpose_op)
            if perm is None:
                continue
            expected_out_shape = _permute_shape(q_shape, perm)
            expected_out_signature = _permute_shape(q_signature, perm)
            if expected_out_shape is None or expected_out_signature is None:
                continue

            pre_transpose_name = _unique_tensor_name(f"{dq_output_name}_before_transpose")
            model_ir.tensors[pre_transpose_name] = TensorIR(
                name=pre_transpose_name,
                dtype=str(y_tensor.dtype),
                shape=[int(v) for v in q_shape],
                shape_signature=[int(v) for v in q_signature],
                data=None,
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_input_name],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=dq_op,
                new_outputs=[pre_transpose_name],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=transpose_op,
                new_inputs=[pre_transpose_name, str(transpose_op.inputs[1])],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=transpose_op,
                new_outputs=[dq_output_name],
            )

            y_tensor.shape = [int(v) for v in expected_out_shape]
            y_tensor.shape_signature = [int(v) for v in expected_out_signature]

            transpose_idx_int = int(transpose_idx)
            if transpose_idx_int < int(dq_idx):
                moved = model_ir.operators.pop(transpose_idx_int)
                insert_idx = int(dq_idx) - 1
                model_ir.operators.insert(insert_idx + 1, moved)

            sanitized += 1
            changed = True
            break

        if changed:
            continue

        # Pattern B:
        #   q --DEQUANTIZE--> y_pre --TRANSPOSE(P)--> y(graph output)
        #   => q --DEQUANTIZE--> y(graph output)
        # Keep user-visible output tensor name stable by renaming y_pre -> y.
        for transpose_idx, transpose_op in enumerate(model_ir.operators):
            if str(transpose_op.op_type) != "TRANSPOSE":
                continue
            if len(transpose_op.inputs) < 2 or len(transpose_op.outputs) != 1:
                continue

            pre_output_name = str(transpose_op.inputs[0])
            final_output_name = str(transpose_op.outputs[0])
            if final_output_name not in model_ir.outputs:
                continue
            if len(consumers.get(final_output_name, [])) > 0:
                continue
            if pre_output_name in model_ir.outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1 or int(pre_users[0]) != int(transpose_idx):
                continue

            dq_idx = producers.get(pre_output_name, None)
            if dq_idx is None:
                continue
            dq_op = model_ir.operators[int(dq_idx)]
            if str(dq_op.op_type) != "DEQUANTIZE":
                continue
            if len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.outputs[0]) != pre_output_name:
                continue

            _rename_tensor_globally(
                model_ir=model_ir,
                old_name=pre_output_name,
                new_name=final_output_name,
            )
            del model_ir.operators[int(transpose_idx)]
            removed_terminal_dequantize_transpose += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "sanitized_terminal_transpose_before_dequantize": int(sanitized),
        "removed_terminal_dequantize_transpose": int(removed_terminal_dequantize_transpose),
    }


def _optimize_concat_pre_quantize_dequantize(model_ir: ModelIR) -> Dict[str, int]:
    """
    Bypass redundant QUANTIZE->DEQUANTIZE immediately before CONCATENATION inputs.

    Target pattern:
      ... -> x_float --QUANTIZE--> q --DEQUANTIZE--> x_dq --CONCATENATION--> ...

    Rewritten:
      ... -> x_float ---------------------------------> --CONCATENATION--> ...

    Safety conditions:
    - q is consumed only by that DEQUANTIZE
    - neither q nor x_dq is a graph output
    - x_float is already in a dequantized flow:
      - producer(x_float) is DEQUANTIZE, or
      - producer(x_float) has at least one DEQUANTIZE input
    """
    bypassed = 0

    def _has_dequantized_origin(
        tensor_name: str,
        producers: Dict[str, int],
        max_depth: int = 6,
    ) -> bool:
        traceable_ops = {
            "TRANSPOSE",
            "RESHAPE",
            "MAX_POOL_2D",
            "AVERAGE_POOL_2D",
            "PAD",
            "ADD",
            "SUB",
            "MUL",
            "DIV",
            "CONCATENATION",
            "MEAN",
        }
        visited = set()
        stack: List[Tuple[str, int]] = [(str(tensor_name), 0)]
        while len(stack) > 0:
            current_name, depth = stack.pop()
            key = (str(current_name), int(depth))
            if key in visited:
                continue
            visited.add(key)
            if int(depth) > int(max_depth):
                continue

            producer_idx = producers.get(str(current_name), None)
            if producer_idx is None:
                continue
            producer = model_ir.operators[int(producer_idx)]
            producer_type = str(producer.op_type)
            if producer_type == "DEQUANTIZE":
                return True
            if producer_type not in traceable_ops:
                continue
            for input_name in list(producer.inputs):
                stack.append((str(input_name), int(depth) + 1))
        return False

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION":
                continue
            if len(concat_op.inputs) == 0:
                continue

            concat_inputs = list(concat_op.inputs)
            for input_pos, concat_input_name in enumerate(concat_inputs):
                dq_idx = producers.get(str(concat_input_name), None)
                if dq_idx is None:
                    continue
                dq_op = model_ir.operators[int(dq_idx)]
                if str(dq_op.op_type) != "DEQUANTIZE":
                    continue
                if len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                    continue
                if str(dq_op.outputs[0]) != str(concat_input_name):
                    continue

                quantized_name = str(dq_op.inputs[0])
                if quantized_name in model_ir.outputs or str(concat_input_name) in model_ir.outputs:
                    continue

                q_idx = producers.get(quantized_name, None)
                if q_idx is None:
                    continue
                q_op = model_ir.operators[int(q_idx)]
                if str(q_op.op_type) != "QUANTIZE":
                    continue
                if len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                    continue
                if str(q_op.outputs[0]) != quantized_name:
                    continue

                q_users = [int(v) for v in consumers.get(quantized_name, [])]
                if len(q_users) != 1 or int(q_users[0]) != int(dq_idx):
                    continue

                float_name = str(q_op.inputs[0])
                float_tensor = model_ir.tensors.get(float_name, None)
                dq_tensor = model_ir.tensors.get(str(concat_input_name), None)
                if float_tensor is None or dq_tensor is None:
                    continue
                if str(float_tensor.dtype).upper().startswith("INT"):
                    continue
                if not _shapes_equal(list(float_tensor.shape), list(dq_tensor.shape)):
                    continue
                if not _has_dequantized_origin(float_name, producers):
                    continue

                updated_concat_inputs = [str(v) for v in list(concat_op.inputs)]
                if int(input_pos) >= 0 and int(input_pos) < len(updated_concat_inputs):
                    updated_concat_inputs[int(input_pos)] = str(float_name)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=concat_op,
                        new_inputs=updated_concat_inputs,
                    )
                bypassed += 1
                changed = True
                break

            if changed:
                break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "bypassed_concat_pre_quantize_dequantize": int(bypassed),
    }


def _optimize_transpose_dequantize_mean_quantize_bridges(model_ir: ModelIR) -> Dict[str, int]:
    """
    Rewrite TRANSPOSE->DEQUANTIZE->MEAN->QUANTIZE bridges by moving TRANSPOSE after MEAN.

    Target pattern:
      X --TRANSPOSE(P)--> A --DEQUANTIZE--> B --MEAN(axes=K, keepDims=True)--> C --QUANTIZE--> Y

    Rewritten:
      X --DEQUANTIZE--> B --MEAN(axes=map(P,K), keepDims=True)--> C'
        --TRANSPOSE(P)--> C --QUANTIZE--> Y

    This preserves Y layout while avoiding transposing a large activation tensor.
    """
    moved_bridges = 0

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre is None:
                continue

            pre_input = str(pre_op.inputs[0])
            pre_output = str(pre_op.outputs[0])
            if pre_output in model_ir.outputs:
                continue

            dq_users = consumers.get(pre_output, [])
            if len(dq_users) != 1:
                continue
            dq_idx = int(dq_users[0])
            dq_op = model_ir.operators[dq_idx]
            if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != pre_output:
                continue
            dq_output = str(dq_op.outputs[0])
            if dq_output in model_ir.outputs:
                continue

            mean_users = consumers.get(dq_output, [])
            if len(mean_users) != 1:
                continue
            mean_idx = int(mean_users[0])
            mean_op = model_ir.operators[mean_idx]
            if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                continue
            if str(mean_op.inputs[0]) != dq_output:
                continue
            if not bool(mean_op.options.get("keepDims", False)):
                continue

            axes_name = str(mean_op.inputs[1])
            axes_users = consumers.get(axes_name, [])
            if len(axes_users) != 1:
                continue
            axes_tensor = model_ir.tensors.get(axes_name, None)
            if axes_tensor is None or axes_tensor.data is None:
                continue
            old_axes_raw = [int(v) for v in np.asarray(axes_tensor.data).reshape(-1).tolist()]
            if len(old_axes_raw) == 0:
                continue

            mean_output = str(mean_op.outputs[0])
            if mean_output in model_ir.outputs:
                continue

            q_users = consumers.get(mean_output, [])
            if len(q_users) != 1:
                continue
            q_idx = int(q_users[0])
            q_op = model_ir.operators[q_idx]
            if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                continue
            if str(q_op.inputs[0]) != mean_output:
                continue

            input_tensor = model_ir.tensors.get(pre_input, None)
            dq_tensor = model_ir.tensors.get(dq_output, None)
            mean_tensor = model_ir.tensors.get(mean_output, None)
            if input_tensor is None or dq_tensor is None or mean_tensor is None:
                continue
            rank = len(list(input_tensor.shape))
            if rank <= 0 or len(perm_pre) != rank:
                continue

            old_axes: List[int] = []
            valid_axes = True
            for axis in old_axes_raw:
                norm_axis = int(axis)
                if norm_axis < 0:
                    norm_axis += rank
                if norm_axis < 0 or norm_axis >= rank:
                    valid_axes = False
                    break
                old_axes.append(norm_axis)
            if not valid_axes:
                continue

            new_axes = [int(perm_pre[axis]) for axis in old_axes]

            # 1) Bypass pre-transpose for DEQUANTIZE.
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[pre_input],
            )
            dq_tensor.shape = list(input_tensor.shape)
            dq_tensor.shape_signature = (
                list(input_tensor.shape_signature)
                if input_tensor.shape_signature is not None
                else list(input_tensor.shape)
            )

            # 2) Update MEAN axes to match input layout.
            axes_arr = np.asarray(new_axes, dtype=np.int32)
            axes_tensor.data = axes_arr
            axes_tensor.dtype = "INT32"
            axes_tensor.shape = [int(len(new_axes))]
            axes_tensor.shape_signature = [int(len(new_axes))]

            # 3) Recompute MEAN output tensor metadata in new layout.
            reduced_shape = list(dq_tensor.shape)
            reduced_signature = (
                list(dq_tensor.shape_signature)
                if dq_tensor.shape_signature is not None
                else list(dq_tensor.shape)
            )
            for axis in new_axes:
                if 0 <= int(axis) < len(reduced_shape):
                    reduced_shape[int(axis)] = 1
                if 0 <= int(axis) < len(reduced_signature):
                    reduced_signature[int(axis)] = 1
            mean_tensor.shape = list(reduced_shape)
            mean_tensor.shape_signature = list(reduced_signature)

            # 4) Insert a transpose before QUANTIZE to preserve QUANTIZE input layout.
            q_bridge_name = _unique_tensor_name(f"{mean_output}_for_quant")
            q_bridge_shape = _permute_shape(list(mean_tensor.shape), perm_pre)
            q_bridge_signature = _permute_shape(list(mean_tensor.shape_signature), perm_pre)
            if q_bridge_shape is None or q_bridge_signature is None:
                continue
            model_ir.tensors[q_bridge_name] = TensorIR(
                name=q_bridge_name,
                dtype=str(mean_tensor.dtype),
                shape=[int(v) for v in q_bridge_shape],
                shape_signature=[int(v) for v in q_bridge_signature],
                data=None,
            )

            perm_name = _unique_tensor_name(f"{q_bridge_name}_perm")
            perm_data = np.asarray(perm_pre, dtype=np.int32)
            model_ir.tensors[perm_name] = TensorIR(
                name=perm_name,
                dtype="INT32",
                shape=[int(perm_data.size)],
                shape_signature=[int(perm_data.size)],
                data=perm_data,
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=q_op,
                new_inputs=[q_bridge_name],
            )
            model_ir.operators.insert(
                q_idx,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[mean_output, perm_name],
                    outputs=[q_bridge_name],
                ),
            )

            # 5) Remove the original pre-transpose.
            remove_idx = int(pre_idx if pre_idx < q_idx else pre_idx + 1)
            del model_ir.operators[remove_idx]

            moved_bridges += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "moved_transpose_dequantize_mean_quantize_bridges": int(moved_bridges),
    }


def _optimize_transpose_mean_hardsigmoid_muladd_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate redundant transpose bridges around mixed MEAN and HardSigmoid branches.

    Target pattern (rank-4, quantized):
      q0_raw --TRANSPOSE(0,3,1,2)--> q0_nchw
             --DEQUANTIZE--> ...MEAN(axes=[2,3])... --QUANTIZE--> qm_nchw --TRANSPOSE(0,2,3,1)--> qm_nhwc -> CONV
      q1_raw --TRANSPOSE(0,3,1,2)--> q1_nchw
             --DEQUANTIZE--> MUL(c)->ADD(c)->MAXIMUM(c)->MINIMUM(c) --QUANTIZE--> qsig_nchw
      MUL(q0_nchw, qsig_nchw) -> m
      ADD(q0_nchw, m) -> y

    Rewrite:
    - Bypass both pre-transposes on the two dequantize branches.
    - Convert MEAN axes to NHWC-side and remove post-quantize transpose.
    - Rewire MUL/ADD main inputs to q0_raw (NHWC).
    - Remove the three bridge transposes.
    - Insert NHWC->NCHW adapters only for legacy consumers of `y`.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw_const_name = "__nhwc_to_nchw_perm_rank4__"

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _tensor_signature(tensor: TensorIR) -> List[int]:
        if tensor.shape_signature is not None:
            return [int(v) for v in list(tensor.shape_signature)]
        return [int(v) for v in list(tensor.shape)]

    if perm_nhwc_to_nchw_const_name not in model_ir.tensors:
        model_ir.tensors[perm_nhwc_to_nchw_const_name] = TensorIR(
            name=perm_nhwc_to_nchw_const_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for pre0_idx, pre0_op in enumerate(model_ir.operators):
            if str(pre0_op.op_type) != "TRANSPOSE" or len(pre0_op.inputs) < 2 or len(pre0_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre0_op) != perm_nhwc_to_nchw:
                continue

            q0_raw_name = str(pre0_op.inputs[0])
            q0_nchw_name = str(pre0_op.outputs[0])
            if q0_raw_name in model_ir.outputs or q0_nchw_name in model_ir.outputs:
                continue

            q0_users = [int(v) for v in consumers.get(q0_nchw_name, [])]
            if len(q0_users) != 3:
                continue

            dq0_idx: Optional[int] = None
            mul0_idx: Optional[int] = None
            add0_idx: Optional[int] = None
            for u_idx in q0_users:
                u_op = model_ir.operators[int(u_idx)]
                u_type = str(u_op.op_type)
                if u_type == "DEQUANTIZE":
                    if dq0_idx is not None:
                        dq0_idx = None
                        break
                    dq0_idx = int(u_idx)
                elif u_type == "MUL":
                    if mul0_idx is not None:
                        mul0_idx = None
                        break
                    mul0_idx = int(u_idx)
                elif u_type == "ADD":
                    if add0_idx is not None:
                        add0_idx = None
                        break
                    add0_idx = int(u_idx)
            if dq0_idx is None or mul0_idx is None or add0_idx is None:
                continue

            dq0_op = model_ir.operators[int(dq0_idx)]
            if len(dq0_op.inputs) != 1 or len(dq0_op.outputs) != 1 or str(dq0_op.inputs[0]) != q0_nchw_name:
                continue
            dq0_out_name = str(dq0_op.outputs[0])
            if dq0_out_name in model_ir.outputs:
                continue

            mean_users = [int(v) for v in consumers.get(dq0_out_name, [])]
            if len(mean_users) != 1:
                continue
            mean_idx = int(mean_users[0])
            mean_op = model_ir.operators[int(mean_idx)]
            if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                continue
            if str(mean_op.inputs[0]) != dq0_out_name or not bool(mean_op.options.get("keepDims", False)):
                continue
            mean_out_name = str(mean_op.outputs[0])
            if mean_out_name in model_ir.outputs:
                continue

            mean_axes_name = str(mean_op.inputs[1])
            mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
            mean_axes = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes is None or len(mean_axes) == 0:
                continue

            qmean_users = [int(v) for v in consumers.get(mean_out_name, [])]
            if len(qmean_users) != 1:
                continue
            qmean_idx = int(qmean_users[0])
            qmean_op = model_ir.operators[int(qmean_idx)]
            if str(qmean_op.op_type) != "QUANTIZE" or len(qmean_op.inputs) != 1 or len(qmean_op.outputs) != 1:
                continue
            if str(qmean_op.inputs[0]) != mean_out_name:
                continue
            qmean_out_name = str(qmean_op.outputs[0])
            if qmean_out_name in model_ir.outputs:
                continue

            post_users = [int(v) for v in consumers.get(qmean_out_name, [])]
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != qmean_out_name:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue
            post_out_name = str(post_op.outputs[0])
            if post_out_name in model_ir.outputs:
                continue

            # Ensure post-transpose is a pure adapter to downstream ops.
            post_out_users = [int(v) for v in consumers.get(post_out_name, [])]
            if len(post_out_users) == 0:
                continue

            mul0_op = model_ir.operators[int(mul0_idx)]
            add0_op = model_ir.operators[int(add0_idx)]
            if len(mul0_op.inputs) != 2 or len(mul0_op.outputs) != 1:
                continue
            if len(add0_op.inputs) != 2 or len(add0_op.outputs) != 1:
                continue

            mul0_inputs = [str(v) for v in list(mul0_op.inputs)]
            if q0_nchw_name == mul0_inputs[0]:
                mul0_q0_input_idx = 0
                sig_q_name = mul0_inputs[1]
            elif q0_nchw_name == mul0_inputs[1]:
                mul0_q0_input_idx = 1
                sig_q_name = mul0_inputs[0]
            else:
                continue
            if sig_q_name in model_ir.outputs:
                continue

            add0_inputs = [str(v) for v in list(add0_op.inputs)]
            mul0_out_name = str(mul0_op.outputs[0])
            if q0_nchw_name == add0_inputs[0] and mul0_out_name == add0_inputs[1]:
                add0_q0_input_idx = 0
            elif q0_nchw_name == add0_inputs[1] and mul0_out_name == add0_inputs[0]:
                add0_q0_input_idx = 1
            else:
                continue

            sig_q_users = [int(v) for v in consumers.get(sig_q_name, [])]
            if len(sig_q_users) != 1 or int(sig_q_users[0]) != int(mul0_idx):
                continue
            mul0_users = [int(v) for v in consumers.get(mul0_out_name, [])]
            if len(mul0_users) != 1 or int(mul0_users[0]) != int(add0_idx):
                continue

            qsig_prod_idx = producers.get(sig_q_name, None)
            if qsig_prod_idx is None:
                continue
            qsig_op = model_ir.operators[int(qsig_prod_idx)]
            if str(qsig_op.op_type) != "QUANTIZE" or len(qsig_op.inputs) != 1 or len(qsig_op.outputs) != 1:
                continue
            hsig_min_out_name = str(qsig_op.inputs[0])

            hsig_min_idx = producers.get(hsig_min_out_name, None)
            if hsig_min_idx is None:
                continue
            hsig_min_op = model_ir.operators[int(hsig_min_idx)]
            if str(hsig_min_op.op_type) != "MINIMUM" or len(hsig_min_op.inputs) != 2 or len(hsig_min_op.outputs) != 1:
                continue
            hsig_max_out_name = (
                str(hsig_min_op.inputs[0])
                if str(hsig_min_op.inputs[0]) != str(hsig_min_op.inputs[1])
                else str(hsig_min_op.inputs[0])
            )
            # Identify main path input of MINIMUM.
            if producers.get(str(hsig_min_op.inputs[0]), None) is not None:
                hsig_max_out_name = str(hsig_min_op.inputs[0])
            elif producers.get(str(hsig_min_op.inputs[1]), None) is not None:
                hsig_max_out_name = str(hsig_min_op.inputs[1])
            else:
                continue

            hsig_max_idx = producers.get(hsig_max_out_name, None)
            if hsig_max_idx is None:
                continue
            hsig_max_op = model_ir.operators[int(hsig_max_idx)]
            if str(hsig_max_op.op_type) != "MAXIMUM" or len(hsig_max_op.inputs) != 2 or len(hsig_max_op.outputs) != 1:
                continue

            hsig_add_out_name = (
                str(hsig_max_op.inputs[0])
                if producers.get(str(hsig_max_op.inputs[0]), None) is not None
                else str(hsig_max_op.inputs[1])
            )
            hsig_add_idx = producers.get(hsig_add_out_name, None)
            if hsig_add_idx is None:
                continue
            hsig_add_op = model_ir.operators[int(hsig_add_idx)]
            if str(hsig_add_op.op_type) != "ADD" or len(hsig_add_op.inputs) != 2 or len(hsig_add_op.outputs) != 1:
                continue

            hsig_mul_out_name = (
                str(hsig_add_op.inputs[0])
                if producers.get(str(hsig_add_op.inputs[0]), None) is not None
                else str(hsig_add_op.inputs[1])
            )
            hsig_mul_idx = producers.get(hsig_mul_out_name, None)
            if hsig_mul_idx is None:
                continue
            hsig_mul_op = model_ir.operators[int(hsig_mul_idx)]
            if str(hsig_mul_op.op_type) != "MUL" or len(hsig_mul_op.inputs) != 2 or len(hsig_mul_op.outputs) != 1:
                continue

            dq1_out_name = (
                str(hsig_mul_op.inputs[0])
                if producers.get(str(hsig_mul_op.inputs[0]), None) is not None
                else str(hsig_mul_op.inputs[1])
            )
            dq1_idx = producers.get(dq1_out_name, None)
            if dq1_idx is None:
                continue
            dq1_op = model_ir.operators[int(dq1_idx)]
            if str(dq1_op.op_type) != "DEQUANTIZE" or len(dq1_op.inputs) != 1 or len(dq1_op.outputs) != 1:
                continue
            if str(dq1_op.outputs[0]) != dq1_out_name:
                continue
            q1_nchw_name = str(dq1_op.inputs[0])
            if q1_nchw_name in model_ir.outputs:
                continue

            pre1_idx = producers.get(q1_nchw_name, None)
            if pre1_idx is None:
                continue
            pre1_op = model_ir.operators[int(pre1_idx)]
            if str(pre1_op.op_type) != "TRANSPOSE" or len(pre1_op.inputs) < 2 or len(pre1_op.outputs) != 1:
                continue
            if str(pre1_op.outputs[0]) != q1_nchw_name:
                continue
            if _read_transpose_perm(model_ir, pre1_op) != perm_nhwc_to_nchw:
                continue
            q1_raw_name = str(pre1_op.inputs[0])
            if q1_raw_name in model_ir.outputs:
                continue
            q1_users = [int(v) for v in consumers.get(q1_nchw_name, [])]
            if len(q1_users) != 1 or int(q1_users[0]) != int(dq1_idx):
                continue

            q0_raw_tensor = model_ir.tensors.get(q0_raw_name, None)
            q0_nchw_tensor = model_ir.tensors.get(q0_nchw_name, None)
            q1_raw_tensor = model_ir.tensors.get(q1_raw_name, None)
            q1_nchw_tensor = model_ir.tensors.get(q1_nchw_name, None)
            dq0_out_tensor = model_ir.tensors.get(dq0_out_name, None)
            mean_out_tensor = model_ir.tensors.get(mean_out_name, None)
            qmean_out_tensor = model_ir.tensors.get(qmean_out_name, None)
            dq1_out_tensor = model_ir.tensors.get(dq1_out_name, None)
            hsig_mul_out_tensor = model_ir.tensors.get(hsig_mul_out_name, None)
            hsig_add_out_tensor = model_ir.tensors.get(hsig_add_out_name, None)
            hsig_max_out_tensor = model_ir.tensors.get(hsig_max_out_name, None)
            hsig_min_out_tensor = model_ir.tensors.get(hsig_min_out_name, None)
            sig_q_tensor = model_ir.tensors.get(sig_q_name, None)
            mul0_out_tensor = model_ir.tensors.get(mul0_out_name, None)
            add0_out_name = str(add0_op.outputs[0])
            add0_out_tensor = model_ir.tensors.get(add0_out_name, None)
            if any(
                t is None
                for t in [
                    q0_raw_tensor,
                    q0_nchw_tensor,
                    q1_raw_tensor,
                    q1_nchw_tensor,
                    dq0_out_tensor,
                    mean_out_tensor,
                    qmean_out_tensor,
                    dq1_out_tensor,
                    hsig_mul_out_tensor,
                    hsig_add_out_tensor,
                    hsig_max_out_tensor,
                    hsig_min_out_tensor,
                    sig_q_tensor,
                    mul0_out_tensor,
                    add0_out_tensor,
                ]
            ):
                continue

            # Avoid per-axis qdim remap complexity in this rewrite.
            if not _all_per_tensor_quantized(
                [
                    q0_raw_tensor,
                    q0_nchw_tensor,
                    q1_raw_tensor,
                    q1_nchw_tensor,
                    qmean_out_tensor,
                    sig_q_tensor,
                    mul0_out_tensor,
                    add0_out_tensor,
                ]
            ):
                continue

            # 1) Bypass pre0 transpose on mean branch.
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq0_op,
                new_inputs=[q0_raw_name],
            )
            dq0_out_tensor.shape = [int(v) for v in list(q0_raw_tensor.shape)]
            dq0_out_tensor.shape_signature = _tensor_signature(q0_raw_tensor)

            # 2) Remap mean axes from NCHW to NHWC side.
            rank = len(list(q0_raw_tensor.shape))
            normalized_axes: List[int] = []
            valid_axes = True
            for axis in mean_axes:
                a = int(axis)
                if a < 0:
                    a += rank
                if a < 0 or a >= rank:
                    valid_axes = False
                    break
                normalized_axes.append(int(a))
            if not valid_axes:
                continue
            new_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]
            if not _write_const_ints_to_tensor(mean_axes_tensor, [int(v) for v in new_axes]):
                continue

            mean_shape = [int(v) for v in list(q0_raw_tensor.shape)]
            mean_sig = _tensor_signature(q0_raw_tensor)
            for axis in new_axes:
                if 0 <= int(axis) < len(mean_shape):
                    mean_shape[int(axis)] = 1
                if 0 <= int(axis) < len(mean_sig):
                    mean_sig[int(axis)] = 1
            mean_out_tensor.shape = [int(v) for v in list(mean_shape)]
            mean_out_tensor.shape_signature = [int(v) for v in list(mean_sig)]

            # qmean now follows MEAN (NHWC-side), so remove post transpose.
            qmean_out_tensor.shape = [int(v) for v in list(mean_shape)]
            qmean_out_tensor.shape_signature = [int(v) for v in list(mean_sig)]
            for post_user_idx in post_out_users:
                post_user_op = model_ir.operators[int(post_user_idx)]
                updated_inputs = [
                    str(qmean_out_name) if str(inp) == str(post_out_name) else str(inp)
                    for inp in list(post_user_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=post_user_op,
                    new_inputs=updated_inputs,
                )

            # 3) Bypass pre1 transpose on HardSigmoid branch.
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq1_op,
                new_inputs=[q1_raw_name],
            )
            dq1_out_tensor.shape = [int(v) for v in list(q1_raw_tensor.shape)]
            dq1_out_tensor.shape_signature = _tensor_signature(q1_raw_tensor)

            # 4) Rewire MUL/ADD main path to q0_raw.
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mul0_op,
                input_index=int(mul0_q0_input_idx),
                new_input_name=q0_raw_name,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=add0_op,
                input_index=int(add0_q0_input_idx),
                new_input_name=q0_raw_name,
            )

            # 5) Metadata now follows NHWC on rewritten branches.
            for tensor in [
                dq1_out_tensor,
                hsig_mul_out_tensor,
                hsig_add_out_tensor,
                hsig_max_out_tensor,
                hsig_min_out_tensor,
                sig_q_tensor,
                mul0_out_tensor,
                add0_out_tensor,
            ]:
                _permute_tensor_metadata_if_rank_matches(tensor, perm_nchw_to_nhwc)

            # Preserve legacy NCHW semantics for add0 downstream via adapters.
            legacy_add0_shape = [int(v) for v in list(add0_out_tensor.shape)]
            legacy_add0_sig = _tensor_signature(add0_out_tensor)
            adapter_insertions: List[Tuple[int, OperatorIR]] = []
            consumers_after = _build_tensor_consumer_map(model_ir)
            add0_consumers = [int(v) for v in consumers_after.get(add0_out_name, [])]
            for legacy_idx in add0_consumers:
                if int(legacy_idx) == int(add0_idx):
                    continue
                legacy_op = model_ir.operators[int(legacy_idx)]
                slots = [
                    int(i)
                    for i, inp in enumerate(list(legacy_op.inputs))
                    if str(inp) == str(add0_out_name)
                ]
                if len(slots) == 0:
                    continue
                adapter_name = _unique_tensor_name(f"{add0_out_name}_nchw_adapter")
                model_ir.tensors[adapter_name] = TensorIR(
                    name=adapter_name,
                    dtype=str(add0_out_tensor.dtype),
                    shape=[int(v) for v in legacy_add0_shape],
                    shape_signature=[int(v) for v in legacy_add0_sig],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(add0_out_tensor.quantization),
                )
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[add0_out_name, perm_nhwc_to_nchw_const_name],
                    outputs=[adapter_name],
                )
                adapter_insertions.append((int(legacy_idx), adapter_op))
                updated_inputs = [str(v) for v in list(legacy_op.inputs)]
                for slot in slots:
                    if 0 <= int(slot) < len(updated_inputs):
                        updated_inputs[int(slot)] = str(adapter_name)
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=legacy_op,
                    new_inputs=updated_inputs,
                )

            if len(adapter_insertions) > 0:
                inserted = 0
                for insert_idx, adapter_op in sorted(adapter_insertions, key=lambda v: int(v[0])):
                    model_ir.operators.insert(int(insert_idx + inserted), adapter_op)
                    inserted += 1

            # 6) Remove bridge transposes.
            remove_indices = sorted(
                set([int(pre0_idx), int(post_idx), int(pre1_idx)]),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_mean_hardsigmoid_muladd_chains": int(optimized),
    }


def _optimize_nhwc_propagation_qlinear_concat_conv(model_ir: ModelIR) -> Dict[str, int]:
    """
    Propagate NHWC layout through a QLinearConcat -> QLinearConv bridge.

    Target IR shape:
      (NCHW quantized inputs)
        -> DEQUANTIZE[*]
        -> CONCATENATION(axis=1)
        -> QUANTIZE(q_out_nchw)
        -> TRANSPOSE([0,2,3,1])  # q_out_nhwc
        -> (CONV_2D|DEPTHWISE_CONV_2D)

    Rewrites:
    - Fold per-input NCHW adapters (TRANSPOSE [0,3,1,2]) into DQ inputs when safe.
    - Convert CONCAT axis from NCHW to NHWC.
    - Remove post-QUANTIZE TRANSPOSE adapters to CONV inputs.
    - Allow additional CONCAT consumers that are NCHW->NHWC TRANSPOSE adapters
      and remove them after CONCAT becomes NHWC.
    """
    propagated = 0
    rank4_perm_nchw_to_nhwc = [0, 2, 3, 1]
    rank4_perm_nhwc_to_nchw = [0, 3, 1, 2]
    inv_nchw_to_nhwc = _invert_perm(rank4_perm_nchw_to_nhwc)
    if inv_nchw_to_nhwc is None:
        return {"propagated_qlinear_concat_conv_nhwc_chains": 0}

    def _remap_qdim_for_permute(tensor: Optional[TensorIR], perm: List[int]) -> None:
        if tensor is None or tensor.quantization is None:
            return
        if _quant_scale_count(tensor.quantization) <= 1:
            return
        inv = _invert_perm(perm)
        if inv is None:
            return
        old_qdim = int(tensor.quantization.quantized_dimension)
        if 0 <= old_qdim < len(inv):
            tensor.quantization.quantized_dimension = int(inv[old_qdim])

    def _is_nchw_nhwc_reinterpret_safe(tensor: Optional[TensorIR]) -> bool:
        if tensor is None:
            return False
        shape = list(tensor.shape)
        if len(shape) != 4:
            return False
        signature = (
            list(tensor.shape_signature)
            if tensor.shape_signature is not None
            else list(shape)
        )
        if len(signature) != 4:
            return False
        # Reinterpretation is safe when spatial dimensions are statically singleton.
        return int(shape[2]) == 1 and int(shape[3]) == 1 and int(signature[2]) == 1 and int(signature[3]) == 1

    def _permute_tensor_shape_signature(tensor: Optional[TensorIR], perm: List[int]) -> bool:
        if tensor is None:
            return False
        new_shape = _permute_shape(list(tensor.shape), perm)
        signature_source = (
            list(tensor.shape_signature)
            if tensor.shape_signature is not None
            else list(tensor.shape)
        )
        new_signature = _permute_shape(signature_source, perm)
        if new_shape is None or new_signature is None:
            return False
        tensor.shape = [int(v) for v in new_shape]
        tensor.shape_signature = [int(v) for v in new_signature]
        _remap_qdim_for_permute(tensor, perm)
        return True

    def _permuted_shape_signature(
        tensor: Optional[TensorIR],
        perm: List[int],
    ) -> Optional[Tuple[List[int], List[int]]]:
        if tensor is None:
            return None
        new_shape = _permute_shape(list(tensor.shape), perm)
        signature_source = (
            list(tensor.shape_signature)
            if tensor.shape_signature is not None
            else list(tensor.shape)
        )
        new_signature = _permute_shape(signature_source, perm)
        if new_shape is None or new_signature is None:
            return None
        return [int(v) for v in new_shape], [int(v) for v in new_signature]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.inputs) == 0 or len(concat_op.outputs) != 1:
                continue
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                continue

            q_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(q_users) == 0:
                continue
            q_idx_candidates: List[int] = []
            removable_concat_post_indices: List[int] = []
            valid_concat_users = True
            for user_idx in q_users:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "QUANTIZE"
                    and len(user_op.inputs) == 1
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                ):
                    q_idx_candidates.append(int(user_idx))
                    continue
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == concat_out_name
                    and _read_transpose_perm(model_ir, user_op) == rank4_perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    removable_concat_post_indices.append(int(user_idx))
                    continue
                valid_concat_users = False
                break
            if not valid_concat_users:
                continue
            if len(q_idx_candidates) != 1:
                continue
            q_idx = int(q_idx_candidates[0])
            q_op = model_ir.operators[q_idx]
            q_out_name = str(q_op.outputs[0])
            if q_out_name in model_outputs:
                continue

            post_users = [int(v) for v in consumers.get(q_out_name, [])]
            if len(post_users) == 0:
                continue

            removable_post_indices: List[int] = []
            post_output_names: List[str] = []
            valid_posts = True
            for post_idx in post_users:
                post_op = model_ir.operators[post_idx]
                if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                    valid_posts = False
                    break
                if str(post_op.inputs[0]) != q_out_name:
                    valid_posts = False
                    break
                perm_post = _read_transpose_perm(model_ir, post_op)
                if perm_post != rank4_perm_nchw_to_nhwc:
                    valid_posts = False
                    break
                post_output_name = str(post_op.outputs[0])
                if post_output_name in model_outputs:
                    valid_posts = False
                    break
                removable_post_indices.append(int(post_idx))
                post_output_names.append(post_output_name)
            if not valid_posts:
                continue

            # Collect and rewrite concat input DQ adapters.
            removable_pre_indices: List[int] = []
            convertible = True
            pending_dq_input_rewrites: Dict[int, str] = {}
            pending_quant_input_rewrites: Dict[int, str] = {}
            pending_tensor_shape_updates: Dict[str, Tuple[List[int], List[int]]] = {}
            pending_qdim_remaps: set[str] = set()
            for dq_input_name in list(concat_op.inputs):
                dq_idx = producers.get(str(dq_input_name), None)
                if dq_idx is None:
                    convertible = False
                    break
                dq_op = model_ir.operators[int(dq_idx)]
                if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                    convertible = False
                    break
                if str(dq_op.outputs[0]) != str(dq_input_name):
                    convertible = False
                    break

                q_in_name = str(dq_op.inputs[0])
                if q_in_name in model_ir.outputs:
                    convertible = False
                    break

                q_in_tensor = model_ir.tensors.get(q_in_name, None)
                if q_in_tensor is None or len(list(q_in_tensor.shape)) != 4:
                    convertible = False
                    break

                q_in_producer_idx = producers.get(q_in_name, None)
                if q_in_producer_idx is None:
                    convertible = False
                    break
                q_in_producer = model_ir.operators[int(q_in_producer_idx)]

                # Pattern 1:
                #   q_raw --TRANSPOSE(0,3,1,2)--> q_nchw --DEQUANTIZE--> dq
                if (
                    str(q_in_producer.op_type) == "TRANSPOSE"
                    and len(q_in_producer.inputs) >= 2
                    and len(q_in_producer.outputs) == 1
                    and _read_transpose_perm(model_ir, q_in_producer) == rank4_perm_nhwc_to_nchw
                ):
                    q_raw_name = str(q_in_producer.inputs[0])
                    q_in_users = [int(v) for v in consumers.get(q_in_name, [])]
                    if len(q_in_users) != 1 or int(q_in_users[0]) != int(dq_idx):
                        convertible = False
                        break
                    if q_in_name in model_ir.outputs:
                        convertible = False
                        break

                    q_raw_tensor = model_ir.tensors.get(q_raw_name, None)
                    dq_out_tensor = model_ir.tensors.get(str(dq_op.outputs[0]), None)
                    if q_raw_tensor is None or dq_out_tensor is None:
                        convertible = False
                        break
                    pending_dq_input_rewrites[int(dq_idx)] = str(q_raw_name)
                    pending_tensor_shape_updates[str(dq_op.outputs[0])] = (
                        [int(v) for v in list(q_raw_tensor.shape)],
                        [int(v) for v in list(q_raw_tensor.shape_signature)]
                        if q_raw_tensor.shape_signature is not None
                        else [int(v) for v in list(q_raw_tensor.shape)]
                    )
                    removable_pre_indices.append(int(q_in_producer_idx))
                    continue

                # Pattern 2:
                #   x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw --QUANTIZE--> q_nchw --DEQUANTIZE--> dq
                if (
                    str(q_in_producer.op_type) == "QUANTIZE"
                    and len(q_in_producer.inputs) == 1
                    and len(q_in_producer.outputs) == 1
                    and str(q_in_producer.outputs[0]) == q_in_name
                ):
                    q_float_name = str(q_in_producer.inputs[0])
                    q_float_producer_idx = producers.get(q_float_name, None)
                    if q_float_producer_idx is not None:
                        q_float_producer = model_ir.operators[int(q_float_producer_idx)]
                        if (
                            str(q_float_producer.op_type) == "TRANSPOSE"
                            and len(q_float_producer.inputs) >= 2
                            and len(q_float_producer.outputs) == 1
                            and _read_transpose_perm(model_ir, q_float_producer) == rank4_perm_nhwc_to_nchw
                        ):
                            q_float_users = [int(v) for v in consumers.get(q_float_name, [])]
                            if len(q_float_users) != 1 or int(q_float_users[0]) != int(q_in_producer_idx):
                                convertible = False
                                break
                            if q_float_name in model_ir.outputs:
                                convertible = False
                                break

                            permuted_q_in = _permuted_shape_signature(q_in_tensor, rank4_perm_nchw_to_nhwc)
                            if permuted_q_in is None:
                                convertible = False
                                break
                            pending_quant_input_rewrites[int(q_in_producer_idx)] = str(q_float_producer.inputs[0])
                            pending_tensor_shape_updates[q_in_name] = permuted_q_in
                            pending_qdim_remaps.add(q_in_name)

                            dq_out_tensor = model_ir.tensors.get(str(dq_op.outputs[0]), None)
                            if dq_out_tensor is None:
                                convertible = False
                                break
                            pending_tensor_shape_updates[str(dq_op.outputs[0])] = (
                                list(permuted_q_in[0]),
                                list(permuted_q_in[1]),
                            )
                            removable_pre_indices.append(int(q_float_producer_idx))
                            continue

                # Pattern 3:
                #   q_nchw --DEQUANTIZE--> dq, where q_nchw is effectively layout-invariant
                #   (e.g., NCHW N,C,1,1), so we can reinterpret metadata without data movement.
                q_in_users = [int(v) for v in consumers.get(q_in_name, [])]
                if len(q_in_users) == 1 and int(q_in_users[0]) == int(dq_idx):
                    # q_in produced directly by QUANTIZE is shape-reconciled from its float input
                    # in later passes. Metadata-only reinterpretation here is unstable and can
                    # recreate NCHW/NHWC mismatches at CONCATENATION prepare time.
                    if str(q_in_producer.op_type) == "QUANTIZE":
                        convertible = False
                        break
                    if _is_nchw_nhwc_reinterpret_safe(q_in_tensor):
                        permuted_q_in = _permuted_shape_signature(q_in_tensor, rank4_perm_nchw_to_nhwc)
                        if permuted_q_in is None:
                            convertible = False
                            break
                        pending_tensor_shape_updates[q_in_name] = permuted_q_in
                        pending_qdim_remaps.add(q_in_name)
                        dq_out_tensor = model_ir.tensors.get(str(dq_op.outputs[0]), None)
                        if dq_out_tensor is None:
                            convertible = False
                            break
                        pending_tensor_shape_updates[str(dq_op.outputs[0])] = (
                            list(permuted_q_in[0]),
                            list(permuted_q_in[1]),
                        )
                        continue

                convertible = False
                break

            if not convertible:
                continue

            concat_axis_old = int(concat_op.options.get("axis", 1))
            if concat_axis_old < 0 or concat_axis_old >= len(inv_nchw_to_nhwc):
                continue
            concat_axis_new = int(inv_nchw_to_nhwc[concat_axis_old])

            # Validate prospective concat input layout compatibility before committing rewrites.
            prospective_concat_shapes: List[List[int]] = []
            prospective_concat_signatures: List[List[int]] = []
            prospective_valid = True
            for concat_input_name in list(concat_op.inputs):
                tensor = model_ir.tensors.get(str(concat_input_name), None)
                if tensor is None:
                    prospective_valid = False
                    break
                pending_shape_sig = pending_tensor_shape_updates.get(str(concat_input_name), None)
                if pending_shape_sig is not None:
                    shape_i, sig_i = pending_shape_sig
                else:
                    shape_i = [int(v) for v in list(tensor.shape)]
                    sig_i = (
                        [int(v) for v in list(tensor.shape_signature)]
                        if tensor.shape_signature is not None
                        else [int(v) for v in list(tensor.shape)]
                    )
                if len(list(shape_i)) != 4 or len(list(sig_i)) != 4:
                    prospective_valid = False
                    break
                prospective_concat_shapes.append([int(v) for v in list(shape_i)])
                prospective_concat_signatures.append([int(v) for v in list(sig_i)])
            if not prospective_valid or len(prospective_concat_shapes) == 0:
                continue
            ref_shape = list(prospective_concat_shapes[0])
            for shape_i in prospective_concat_shapes[1:]:
                for dim_idx in range(4):
                    if int(dim_idx) == int(concat_axis_new):
                        continue
                    if (
                        int(ref_shape[dim_idx]) >= 0
                        and int(shape_i[dim_idx]) >= 0
                        and int(ref_shape[dim_idx]) != int(shape_i[dim_idx])
                    ):
                        prospective_valid = False
                        break
                if not prospective_valid:
                    break
            if not prospective_valid:
                continue

            for op_idx, new_input_name in pending_dq_input_rewrites.items():
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=model_ir.operators[int(op_idx)],
                    new_inputs=[str(new_input_name)],
                )
            for op_idx, new_input_name in pending_quant_input_rewrites.items():
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=model_ir.operators[int(op_idx)],
                    new_inputs=[str(new_input_name)],
                )
            for tensor_name, (new_shape, new_signature) in pending_tensor_shape_updates.items():
                tensor = model_ir.tensors.get(str(tensor_name), None)
                if tensor is None:
                    continue
                tensor.shape = [int(v) for v in list(new_shape)]
                tensor.shape_signature = [int(v) for v in list(new_signature)]
            for tensor_name in pending_qdim_remaps:
                tensor = model_ir.tensors.get(str(tensor_name), None)
                _remap_qdim_for_permute(tensor, rank4_perm_nchw_to_nhwc)

            concat_op.options["axis"] = int(concat_axis_new)

            concat_input_tensors: List[Optional[TensorIR]] = [
                model_ir.tensors.get(str(input_name), None)
                for input_name in concat_op.inputs
            ]
            if any(t is None for t in concat_input_tensors):
                continue
            first_tensor = concat_input_tensors[0]
            if first_tensor is None:
                continue
            rank = len(list(first_tensor.shape))
            if rank != 4:
                continue

            concat_shape = [int(v) for v in list(first_tensor.shape)]
            concat_signature = (
                [int(v) for v in list(first_tensor.shape_signature)]
                if first_tensor.shape_signature is not None
                else [int(v) for v in list(first_tensor.shape)]
            )
            dynamic_concat_axis = False
            for tensor in concat_input_tensors[1:]:
                if tensor is None:
                    continue
                if len(list(tensor.shape)) != rank:
                    convertible = False
                    break
                for dim_idx in range(rank):
                    if int(dim_idx) == int(concat_axis_new):
                        continue
                    if (
                        int(concat_shape[dim_idx]) >= 0
                        and int(tensor.shape[dim_idx]) >= 0
                        and int(concat_shape[dim_idx]) != int(tensor.shape[dim_idx])
                    ):
                        convertible = False
                        break
                if not convertible:
                    break
                concat_shape[concat_axis_new] += int(tensor.shape[concat_axis_new])
            if not convertible:
                continue
            for tensor in concat_input_tensors:
                if tensor is None:
                    continue
                sig = (
                    [int(v) for v in list(tensor.shape_signature)]
                    if tensor.shape_signature is not None
                    else [int(v) for v in list(tensor.shape)]
                )
                if int(sig[concat_axis_new]) < 0:
                    dynamic_concat_axis = True
                    break
            if dynamic_concat_axis:
                concat_signature[concat_axis_new] = -1
            else:
                concat_signature[concat_axis_new] = int(
                    sum(
                        int(
                            (
                                t.shape_signature[concat_axis_new]
                                if t is not None and t.shape_signature is not None
                                else t.shape[concat_axis_new]
                            )
                        )
                        for t in concat_input_tensors
                        if t is not None
                    )
                )

            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            if concat_out_tensor is None:
                continue
            concat_out_tensor.shape = [int(v) for v in concat_shape]
            concat_out_tensor.shape_signature = [int(v) for v in concat_signature]

            q_out_tensor = model_ir.tensors.get(q_out_name, None)
            if q_out_tensor is None:
                continue
            q_out_tensor.shape = [int(v) for v in concat_shape]
            q_out_tensor.shape_signature = [int(v) for v in concat_signature]
            _remap_qdim_for_permute(q_out_tensor, rank4_perm_nchw_to_nhwc)

            # Remove output transpose adapters and reconnect their consumers to q_out directly.
            for post_idx in removable_post_indices:
                post_op = model_ir.operators[int(post_idx)]
                post_out_name = str(post_op.outputs[0])
                _replace_tensor_inputs(model_ir, post_out_name, q_out_name)
            for post_idx in removable_concat_post_indices:
                post_op = model_ir.operators[int(post_idx)]
                post_out_name = str(post_op.outputs[0])
                _replace_tensor_inputs(model_ir, post_out_name, concat_out_name)

            remove_indices = sorted(
                set(int(v) for v in (removable_pre_indices + removable_post_indices + removable_concat_post_indices)),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            propagated += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "propagated_qlinear_concat_conv_nhwc_chains": int(propagated),
    }


def _optimize_nhwc_prefix_qlinear_silu_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Propagate NHWC through early QLinear SiLU chains and remove unnecessary
    NCHW round-trips around Conv outputs.

    Target pattern (rank-4):
      q_raw_nhwc --TRANSPOSE(0,3,1,2)--> q_nchw
      q_nchw --DEQUANTIZE--> dq --LOGISTIC--> sig_f --QUANTIZE--> sig_q
      MUL(q_nchw, sig_q) --> y_q
    or
      q_raw_nhwc --TRANSPOSE(0,3,1,2)--> q_nchw
      q_nchw --DEQUANTIZE--> dq --MUL(c1)--> --ADD(c2)--> --MAXIMUM(c3)--> --MINIMUM(c4)--> --QUANTIZE--> sig_q
      MUL(q_nchw, sig_q) --> y_q

    Rewrite:
      q_raw_nhwc --DEQUANTIZE--> dq --LOGISTIC--> sig_f --QUANTIZE--> sig_q
      MUL(q_raw_nhwc, sig_q) --> y_q

    Plus optional removal of:
      y_q --TRANSPOSE(0,2,3,1)--> y_nhwc

    Safety:
    - Strict SiLU topology
      : single DEQUANTIZE->LOGISTIC->QUANTIZE branch
      : or DEQUANTIZE->(MUL/ADD/MAXIMUM/MINIMUM)->QUANTIZE with singleton constants.
    - Per-tensor quantization only on quantized tensors touched by the rewrite.
    - For non-transpose consumers, preserve legacy NCHW semantics by inserting
      NHWC->NCHW adapters at the consumer input slots when the op is considered
      layout-safe (axis count preserving / non-layout-sensitive).
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw_const_name = "__nhwc_to_nchw_perm_rank4__"
    # Ops excluded from generic legacy-adapter rewiring because they are
    # layout-sensitive or modify axis semantics directly.
    legacy_adapter_blocklist = {
        "CONCATENATION",
        "CONV_2D",
        "DEPTHWISE_CONV_2D",
        "MAX_POOL_2D",
        "AVERAGE_POOL_2D",
        "RESIZE_BILINEAR",
        "RESIZE_NEAREST_NEIGHBOR",
        "RESHAPE",
        "TRANSPOSE",
        "STRIDED_SLICE",
        "SLICE",
        "SOFTMAX",
        "FULLY_CONNECTED",
        "MEAN",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    if perm_nhwc_to_nchw_const_name not in model_ir.tensors:
        model_ir.tensors[perm_nhwc_to_nchw_const_name] = TensorIR(
            name=perm_nhwc_to_nchw_const_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            q_raw_name = str(pre_op.inputs[0])
            q_nchw_name = str(pre_op.outputs[0])
            if q_raw_name in model_outputs or q_nchw_name in model_outputs:
                continue

            q_raw_tensor = model_ir.tensors.get(q_raw_name, None)
            q_nchw_tensor = model_ir.tensors.get(q_nchw_name, None)
            if q_raw_tensor is None or q_nchw_tensor is None:
                continue
            if len(list(q_raw_tensor.shape)) != 4 or len(list(q_nchw_tensor.shape)) != 4:
                continue
            if not _all_per_tensor_quantized([q_raw_tensor, q_nchw_tensor]):
                continue

            q_nchw_users = [int(v) for v in consumers.get(q_nchw_name, [])]
            if len(q_nchw_users) != 2:
                continue

            dq_idx: Optional[int] = None
            mul_idx: Optional[int] = None
            for user_idx in q_nchw_users:
                user_op = model_ir.operators[int(user_idx)]
                if str(user_op.op_type) == "DEQUANTIZE":
                    if dq_idx is not None:
                        dq_idx = None
                        break
                    dq_idx = int(user_idx)
                elif str(user_op.op_type) == "MUL":
                    if mul_idx is not None:
                        mul_idx = None
                        break
                    mul_idx = int(user_idx)
                else:
                    dq_idx = None
                    mul_idx = None
                    break
            if dq_idx is None or mul_idx is None:
                continue

            dq_op = model_ir.operators[int(dq_idx)]
            mul_op = model_ir.operators[int(mul_idx)]
            if len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                continue
            if str(dq_op.inputs[0]) != q_nchw_name:
                continue
            dq_out_name = str(dq_op.outputs[0])
            if dq_out_name in model_outputs:
                continue

            dq_users = [int(v) for v in consumers.get(dq_out_name, [])]
            if len(dq_users) != 1:
                continue

            sig_head_idx = int(dq_users[0])
            sig_head_op = model_ir.operators[int(sig_head_idx)]
            sig_head_type = str(sig_head_op.op_type)
            sig_intermediate_names: List[str] = []
            sig_q_name: Optional[str] = None

            # Pattern A:
            #   DEQUANTIZE -> LOGISTIC -> QUANTIZE
            if sig_head_type == "LOGISTIC":
                if len(sig_head_op.inputs) != 1 or len(sig_head_op.outputs) != 1:
                    continue
                if str(sig_head_op.inputs[0]) != dq_out_name:
                    continue
                logistic_out_name = str(sig_head_op.outputs[0])
                if logistic_out_name in model_outputs:
                    continue
                logistic_users = [int(v) for v in consumers.get(logistic_out_name, [])]
                if len(logistic_users) != 1:
                    continue
                qsig_idx = int(logistic_users[0])
                qsig_op = model_ir.operators[int(qsig_idx)]
                if str(qsig_op.op_type) != "QUANTIZE" or len(qsig_op.inputs) != 1 or len(qsig_op.outputs) != 1:
                    continue
                if str(qsig_op.inputs[0]) != logistic_out_name:
                    continue
                sig_intermediate_names = [logistic_out_name]
                sig_q_name = str(qsig_op.outputs[0])

            # Pattern B:
            #   DEQUANTIZE -> MUL(const) -> ADD(const) -> MAXIMUM(const) -> MINIMUM(const) -> QUANTIZE
            else:
                if sig_head_type != "MUL" or len(sig_head_op.inputs) != 2 or len(sig_head_op.outputs) != 1:
                    continue
                sig_head_inputs = [str(v) for v in list(sig_head_op.inputs)]
                if dq_out_name == sig_head_inputs[0]:
                    hs_mul_side = sig_head_inputs[1]
                elif dq_out_name == sig_head_inputs[1]:
                    hs_mul_side = sig_head_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_mul_side):
                    continue
                hs_mul_out_name = str(sig_head_op.outputs[0])
                if hs_mul_out_name in model_outputs:
                    continue

                hs_mul_users = [int(v) for v in consumers.get(hs_mul_out_name, [])]
                if len(hs_mul_users) != 1:
                    continue
                hs_add_idx = int(hs_mul_users[0])
                hs_add_op = model_ir.operators[int(hs_add_idx)]
                if str(hs_add_op.op_type) != "ADD" or len(hs_add_op.inputs) != 2 or len(hs_add_op.outputs) != 1:
                    continue
                hs_add_inputs = [str(v) for v in list(hs_add_op.inputs)]
                if hs_mul_out_name == hs_add_inputs[0]:
                    hs_add_side = hs_add_inputs[1]
                elif hs_mul_out_name == hs_add_inputs[1]:
                    hs_add_side = hs_add_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_add_side):
                    continue
                hs_add_out_name = str(hs_add_op.outputs[0])
                if hs_add_out_name in model_outputs:
                    continue

                hs_add_users = [int(v) for v in consumers.get(hs_add_out_name, [])]
                if len(hs_add_users) != 1:
                    continue
                hs_max_idx = int(hs_add_users[0])
                hs_max_op = model_ir.operators[int(hs_max_idx)]
                if str(hs_max_op.op_type) != "MAXIMUM" or len(hs_max_op.inputs) != 2 or len(hs_max_op.outputs) != 1:
                    continue
                hs_max_inputs = [str(v) for v in list(hs_max_op.inputs)]
                if hs_add_out_name == hs_max_inputs[0]:
                    hs_max_side = hs_max_inputs[1]
                elif hs_add_out_name == hs_max_inputs[1]:
                    hs_max_side = hs_max_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_max_side):
                    continue
                hs_max_out_name = str(hs_max_op.outputs[0])
                if hs_max_out_name in model_outputs:
                    continue

                hs_max_users = [int(v) for v in consumers.get(hs_max_out_name, [])]
                if len(hs_max_users) != 1:
                    continue
                hs_min_idx = int(hs_max_users[0])
                hs_min_op = model_ir.operators[int(hs_min_idx)]
                if str(hs_min_op.op_type) != "MINIMUM" or len(hs_min_op.inputs) != 2 or len(hs_min_op.outputs) != 1:
                    continue
                hs_min_inputs = [str(v) for v in list(hs_min_op.inputs)]
                if hs_max_out_name == hs_min_inputs[0]:
                    hs_min_side = hs_min_inputs[1]
                elif hs_max_out_name == hs_min_inputs[1]:
                    hs_min_side = hs_min_inputs[0]
                else:
                    continue
                if not _is_singleton_constant_tensor(model_ir, hs_min_side):
                    continue
                hs_min_out_name = str(hs_min_op.outputs[0])
                if hs_min_out_name in model_outputs:
                    continue

                hs_min_users = [int(v) for v in consumers.get(hs_min_out_name, [])]
                if len(hs_min_users) != 1:
                    continue
                qsig_idx = int(hs_min_users[0])
                qsig_op = model_ir.operators[int(qsig_idx)]
                if str(qsig_op.op_type) != "QUANTIZE" or len(qsig_op.inputs) != 1 or len(qsig_op.outputs) != 1:
                    continue
                if str(qsig_op.inputs[0]) != hs_min_out_name:
                    continue
                sig_intermediate_names = [
                    hs_mul_out_name,
                    hs_add_out_name,
                    hs_max_out_name,
                    hs_min_out_name,
                ]
                sig_q_name = str(qsig_op.outputs[0])

            if sig_q_name is None or sig_q_name in model_outputs:
                continue

            if len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            if q_nchw_name not in set(mul_inputs):
                continue
            if sig_q_name not in set(mul_inputs):
                continue
            sig_q_users = [int(v) for v in consumers.get(sig_q_name, [])]
            if sig_q_users != [int(mul_idx)] and set(sig_q_users) != {int(mul_idx)}:
                continue

            sig_q_tensor = model_ir.tensors.get(sig_q_name, None)
            mul_out_name = str(mul_op.outputs[0])
            mul_out_tensor = model_ir.tensors.get(mul_out_name, None)
            if mul_out_name in model_outputs or sig_q_tensor is None or mul_out_tensor is None:
                continue
            if not _all_per_tensor_quantized([sig_q_tensor, mul_out_tensor]):
                continue

            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            removable_post_ops: List[OperatorIR] = []
            legacy_user_input_slots: List[Tuple[int, int]] = []
            users_supported = True
            for user_idx in mul_users:
                user_op = model_ir.operators[int(user_idx)]
                user_type = str(user_op.op_type)
                if user_type == "TRANSPOSE":
                    if len(user_op.inputs) < 2 or len(user_op.outputs) != 1:
                        users_supported = False
                        break
                    if str(user_op.inputs[0]) != mul_out_name:
                        users_supported = False
                        break
                    if _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc:
                        users_supported = False
                        break
                    if str(user_op.outputs[0]) in model_outputs:
                        users_supported = False
                        break
                    removable_post_ops.append(user_op)
                else:
                    if user_type in legacy_adapter_blocklist:
                        users_supported = False
                        break
                    input_indices = [
                        int(input_idx)
                        for input_idx, input_name in enumerate(user_op.inputs)
                        if str(input_name) == str(mul_out_name)
                    ]
                    if len(input_indices) == 0:
                        users_supported = False
                        break
                    if len(user_op.outputs) == 0:
                        users_supported = False
                        break
                    if any(str(output_name) in model_outputs for output_name in list(user_op.outputs)):
                        users_supported = False
                        break
                    for input_idx in input_indices:
                        legacy_user_input_slots.append((int(user_idx), int(input_idx)))
            if not users_supported:
                continue

            # Rewire SiLU branch to consume NHWC source directly.
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_op,
                new_inputs=[q_raw_name],
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=[
                    q_raw_name if str(input_name) == q_nchw_name else str(input_name)
                    for input_name in mul_inputs
                ],
            )

            # For DEQUANTIZE users that should keep NCHW semantics downstream,
            # insert NHWC->NCHW adapters and redirect those DEQUANTIZE inputs.
            legacy_mul_shape = [int(v) for v in list(mul_out_tensor.shape)]
            legacy_mul_signature = (
                [int(v) for v in list(mul_out_tensor.shape_signature)]
                if mul_out_tensor.shape_signature is not None
                else [int(v) for v in list(mul_out_tensor.shape)]
            )
            adapter_insertions: List[Tuple[int, OperatorIR]] = []
            for legacy_user_idx, legacy_input_index in legacy_user_input_slots:
                legacy_user_op = model_ir.operators[int(legacy_user_idx)]
                adapter_name = _unique_tensor_name(f"{mul_out_name}_nchw_adapter")
                model_ir.tensors[adapter_name] = TensorIR(
                    name=adapter_name,
                    dtype=str(mul_out_tensor.dtype),
                    shape=list(legacy_mul_shape),
                    shape_signature=list(legacy_mul_signature),
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(mul_out_tensor.quantization),
                )
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[mul_out_name, perm_nhwc_to_nchw_const_name],
                    outputs=[adapter_name],
                )
                adapter_insertions.append((int(legacy_user_idx), adapter_op))
                updated_inputs = [str(v) for v in list(legacy_user_op.inputs)]
                if 0 <= int(legacy_input_index) < len(updated_inputs):
                    updated_inputs[int(legacy_input_index)] = str(adapter_name)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=legacy_user_op,
                        new_inputs=updated_inputs,
                    )

            if len(adapter_insertions) > 0:
                inserted = 0
                for insert_idx, adapter_op in sorted(adapter_insertions, key=lambda v: int(v[0])):
                    model_ir.operators.insert(int(insert_idx + inserted), adapter_op)
                    inserted += 1

            # Metadata now follows NHWC for the rewritten chain.
            metadata_targets = [dq_out_name] + list(sig_intermediate_names) + [sig_q_name, mul_out_name]
            for tensor_name in metadata_targets:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(tensor_name, None),
                    perm_nchw_to_nhwc,
                )

            # Remove redundant output transpose adapters.
            for post_op in removable_post_ops:
                post_out_name = str(post_op.outputs[0])
                _replace_tensor_inputs(model_ir, post_out_name, mul_out_name)

            removable_post_op_ids = {id(op) for op in removable_post_ops}
            removable_post_indices = [
                int(op_idx)
                for op_idx, op in enumerate(model_ir.operators)
                if id(op) in removable_post_op_ids
            ]
            remove_indices = sorted(
                set([int(pre_idx)] + [int(v) for v in removable_post_indices]),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_nhwc_prefix_qlinear_silu_chains": int(optimized),
    }


def _optimize_transpose_mean_maxpool_concat_conv_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate redundant NCHW/NHWC transpose chains around
    DEQUANTIZE->MEAN and MAX_POOL_2D branches before CONCATENATION->CONV.

    Target pattern (rank-4):
      q_raw_nhwc --TRANSPOSE(0,3,1,2)--> q_nchw --DEQUANTIZE--> dq_mean --MEAN(axes=[2,3])--> mean_nchw
      q_raw_nhwc --DEQUANTIZE--> dq_pool --MAX_POOL_2D--> pool_nhwc --TRANSPOSE(0,3,1,2)--> pool_nchw
      CONCAT(axis=1, [mean_nchw, pool_nchw]) -> QUANTIZE -> q_cat_nchw --TRANSPOSE(0,2,3,1)--> q_cat_nhwc -> CONV

    Rewritten:
      q_raw_nhwc --DEQUANTIZE--> dq_mean --MEAN(axes=[1,2])--> mean_nhwc
      q_raw_nhwc --DEQUANTIZE--> dq_pool --MAX_POOL_2D--> pool_nhwc
      CONCAT(axis=3, [mean_nhwc, pool_nhwc]) -> QUANTIZE -> q_cat_nhwc -> CONV
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    old_to_new_axis = perm_nhwc_to_nchw

    def _normalize_axis(axis: int, rank: int) -> Optional[int]:
        a = int(axis)
        if a < 0:
            a += int(rank)
        if a < 0 or a >= int(rank):
            return None
        return int(a)

    def _tensor_shape_signature(tensor: TensorIR) -> List[int]:
        if tensor.shape_signature is not None:
            return [int(v) for v in list(tensor.shape_signature)]
        return [int(v) for v in list(tensor.shape)]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            q_raw_name = str(pre_op.inputs[0])
            q_nchw_name = str(pre_op.outputs[0])
            if q_raw_name in model_ir.outputs or q_nchw_name in model_ir.outputs:
                continue

            q_raw_tensor = model_ir.tensors.get(q_raw_name, None)
            q_nchw_tensor = model_ir.tensors.get(q_nchw_name, None)
            if q_raw_tensor is None or q_nchw_tensor is None:
                continue
            q_raw_shape = [int(v) for v in list(q_raw_tensor.shape)]
            if len(q_raw_shape) != 4:
                continue

            q_nchw_users = [int(v) for v in consumers.get(q_nchw_name, [])]
            if len(q_nchw_users) != 1:
                continue
            dq_mean_idx = int(q_nchw_users[0])
            dq_mean_op = model_ir.operators[dq_mean_idx]
            if str(dq_mean_op.op_type) != "DEQUANTIZE" or len(dq_mean_op.inputs) != 1 or len(dq_mean_op.outputs) != 1:
                continue
            dq_mean_out_name = str(dq_mean_op.outputs[0])
            if dq_mean_out_name in model_ir.outputs:
                continue

            dq_mean_out_users = [int(v) for v in consumers.get(dq_mean_out_name, [])]
            if len(dq_mean_out_users) != 1:
                continue
            mean_idx = int(dq_mean_out_users[0])
            mean_op = model_ir.operators[mean_idx]
            if str(mean_op.op_type) != "MEAN" or len(mean_op.inputs) < 2 or len(mean_op.outputs) != 1:
                continue
            if not bool(mean_op.options.get("keepDims", True)):
                continue
            mean_out_name = str(mean_op.outputs[0])
            if mean_out_name in model_ir.outputs:
                continue

            mean_axes_name = str(mean_op.inputs[1])
            mean_axes_tensor = model_ir.tensors.get(mean_axes_name, None)
            mean_axes_raw = _read_const_ints_from_tensor(mean_axes_tensor)
            if mean_axes_raw is None:
                continue
            mean_axes_old: List[int] = []
            axes_ok = True
            for axis in mean_axes_raw:
                norm_axis = _normalize_axis(int(axis), 4)
                if norm_axis is None:
                    axes_ok = False
                    break
                mean_axes_old.append(int(norm_axis))
            if not axes_ok:
                continue
            if sorted(mean_axes_old) != [2, 3]:
                continue

            q_raw_users = [int(v) for v in consumers.get(q_raw_name, [])]
            q_raw_users_wo_pre = [int(v) for v in q_raw_users if int(v) != int(pre_idx)]
            if len(q_raw_users_wo_pre) != 1:
                continue
            dq_pool_idx = int(q_raw_users_wo_pre[0])
            dq_pool_op = model_ir.operators[dq_pool_idx]
            if str(dq_pool_op.op_type) != "DEQUANTIZE" or len(dq_pool_op.inputs) != 1 or len(dq_pool_op.outputs) != 1:
                continue
            if str(dq_pool_op.inputs[0]) != q_raw_name:
                continue
            dq_pool_out_name = str(dq_pool_op.outputs[0])
            if dq_pool_out_name in model_ir.outputs:
                continue

            dq_pool_out_users = [int(v) for v in consumers.get(dq_pool_out_name, [])]
            if len(dq_pool_out_users) != 1:
                continue
            pool_idx = int(dq_pool_out_users[0])
            pool_op = model_ir.operators[pool_idx]
            if str(pool_op.op_type) != "MAX_POOL_2D" or len(pool_op.inputs) != 1 or len(pool_op.outputs) != 1:
                continue
            if str(pool_op.inputs[0]) != dq_pool_out_name:
                continue
            pool_out_name = str(pool_op.outputs[0])
            if pool_out_name in model_ir.outputs:
                continue

            pool_out_users = [int(v) for v in consumers.get(pool_out_name, [])]
            if len(pool_out_users) != 1:
                continue
            pool_post_idx = int(pool_out_users[0])
            pool_post_op = model_ir.operators[pool_post_idx]
            if str(pool_post_op.op_type) != "TRANSPOSE" or len(pool_post_op.inputs) < 2 or len(pool_post_op.outputs) != 1:
                continue
            if str(pool_post_op.inputs[0]) != pool_out_name:
                continue
            if _read_transpose_perm(model_ir, pool_post_op) != perm_nhwc_to_nchw:
                continue
            pool_nchw_name = str(pool_post_op.outputs[0])
            if pool_nchw_name in model_ir.outputs:
                continue

            mean_users = [int(v) for v in consumers.get(mean_out_name, [])]
            pool_nchw_users = [int(v) for v in consumers.get(pool_nchw_name, [])]
            if len(mean_users) != 1 or len(pool_nchw_users) != 1:
                continue
            concat_idx = int(mean_users[0])
            if int(pool_nchw_users[0]) != int(concat_idx):
                continue
            concat_op = model_ir.operators[concat_idx]
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                continue
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_ir.outputs:
                continue
            concat_axis_old = _normalize_axis(int(concat_op.options.get("axis", 1)), 4)
            if concat_axis_old is None or int(concat_axis_old) != 1:
                continue
            if mean_out_name not in set(str(v) for v in concat_op.inputs):
                continue
            if pool_nchw_name not in set(str(v) for v in concat_op.inputs):
                continue

            concat_out_users = [int(v) for v in consumers.get(concat_out_name, [])]
            if len(concat_out_users) != 1:
                continue
            q_cat_idx = int(concat_out_users[0])
            q_cat_op = model_ir.operators[q_cat_idx]
            if str(q_cat_op.op_type) != "QUANTIZE" or len(q_cat_op.inputs) != 1 or len(q_cat_op.outputs) != 1:
                continue
            if str(q_cat_op.inputs[0]) != concat_out_name:
                continue
            q_cat_name = str(q_cat_op.outputs[0])
            if q_cat_name in model_ir.outputs:
                continue

            post_users = [int(v) for v in consumers.get(q_cat_name, [])]
            if len(post_users) == 0:
                continue
            removable_post_indices: List[int] = []
            valid_posts = True
            for post_idx in post_users:
                post_op = model_ir.operators[post_idx]
                if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                    valid_posts = False
                    break
                if str(post_op.inputs[0]) != q_cat_name:
                    valid_posts = False
                    break
                if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                    valid_posts = False
                    break
                post_out_name = str(post_op.outputs[0])
                if post_out_name in model_ir.outputs:
                    valid_posts = False
                    break
                removable_post_indices.append(int(post_idx))
            if not valid_posts:
                continue

            dq_mean_out_tensor = model_ir.tensors.get(dq_mean_out_name, None)
            mean_out_tensor = model_ir.tensors.get(mean_out_name, None)
            pool_out_tensor = model_ir.tensors.get(pool_out_name, None)
            concat_out_tensor = model_ir.tensors.get(concat_out_name, None)
            q_cat_tensor = model_ir.tensors.get(q_cat_name, None)
            if any(t is None for t in [dq_mean_out_tensor, mean_out_tensor, pool_out_tensor, concat_out_tensor, q_cat_tensor]):
                continue

            # Keep the QDIM semantics deterministic for this metadata rewrite.
            if q_cat_tensor is not None and q_cat_tensor.quantization is not None and _quant_scale_count(q_cat_tensor.quantization) > 1:
                if int(q_cat_tensor.quantization.quantized_dimension) != 1:
                    continue

            # Rewire mean branch to consume NHWC quantized source directly.
            _set_operator_inputs(
                model_ir=model_ir,
                op=dq_mean_op,
                new_inputs=[q_raw_name],
            )
            dq_mean_out_tensor.shape = [int(v) for v in list(q_raw_tensor.shape)]
            dq_mean_out_tensor.shape_signature = _tensor_shape_signature(q_raw_tensor)

            mean_axes_new = [int(old_to_new_axis[int(axis)]) for axis in mean_axes_old]
            _write_const_ints_to_tensor(mean_axes_tensor, mean_axes_new)

            q_raw_signature = _tensor_shape_signature(q_raw_tensor)
            mean_shape = [int(v) for v in list(q_raw_tensor.shape)]
            mean_signature = [int(v) for v in list(q_raw_signature)]
            for axis in mean_axes_new:
                mean_shape[int(axis)] = 1
                mean_signature[int(axis)] = 1
            mean_out_tensor.shape = [int(v) for v in list(mean_shape)]
            mean_out_tensor.shape_signature = [int(v) for v in list(mean_signature)]

            # Replace pool NCHW adapter output with NHWC pool output.
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[
                    str(pool_out_name) if str(input_name) == str(pool_nchw_name) else str(input_name)
                    for input_name in list(concat_op.inputs)
                ],
            )
            concat_op.options["axis"] = 3

            concat_inputs = [model_ir.tensors.get(str(input_name), None) for input_name in list(concat_op.inputs)]
            if any(t is None for t in concat_inputs):
                continue
            first_tensor = concat_inputs[0]
            if first_tensor is None:
                continue
            if len(list(first_tensor.shape)) != 4:
                continue
            concat_shape = [int(v) for v in list(first_tensor.shape)]
            concat_signature = _tensor_shape_signature(first_tensor)
            dynamic_concat_axis = False
            for tensor in concat_inputs[1:]:
                if tensor is None:
                    continue
                shape_i = [int(v) for v in list(tensor.shape)]
                if len(shape_i) != 4:
                    valid_posts = False
                    break
                concat_shape[3] += int(shape_i[3])
            if not valid_posts:
                continue
            for tensor in concat_inputs:
                if tensor is None:
                    continue
                sig_i = _tensor_shape_signature(tensor)
                if int(sig_i[3]) < 0:
                    dynamic_concat_axis = True
                    break
            if dynamic_concat_axis:
                concat_signature[3] = -1
            else:
                concat_signature[3] = int(
                    sum(int(_tensor_shape_signature(tensor)[3]) for tensor in concat_inputs if tensor is not None)
                )

            concat_out_tensor.shape = [int(v) for v in list(concat_shape)]
            concat_out_tensor.shape_signature = [int(v) for v in list(concat_signature)]
            q_cat_tensor.shape = [int(v) for v in list(concat_shape)]
            q_cat_tensor.shape_signature = [int(v) for v in list(concat_signature)]
            if q_cat_tensor.quantization is not None and _quant_scale_count(q_cat_tensor.quantization) > 1:
                q_cat_tensor.quantization.quantized_dimension = 3

            # Remove post-quantize transpose adapters and reconnect their consumers.
            for post_idx in removable_post_indices:
                post_op = model_ir.operators[int(post_idx)]
                post_out_name = str(post_op.outputs[0])
                _replace_tensor_inputs(model_ir, post_out_name, q_cat_name)

            remove_indices = sorted(
                set([int(pre_idx), int(pool_post_idx)] + [int(v) for v in removable_post_indices]),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "optimized_transpose_mean_maxpool_concat_conv_chains": int(optimized),
    }


def _canonicalize_softmax_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Canonicalize ONNX-driven transpose chains around SOFTMAX to expose removable pairs.

    Target template:
      T0(0,3,1,2) -> T1(0,3,2,1) -> SOFTMAX -> T2(0,3,2,1)

    Rewritten to:
      T0(0,3,1,2) -> T1(0,2,3,1) -> SOFTMAX -> T2(0,3,1,2)

    Then T0/T1 become inverse pairs and are removed by `_optimize_layout_transpose_chains`.
    """
    rewritten = 0
    softmax_nhwc_propagated_marker = "__softmax_nhwc_propagated__"
    rank4_perm_nhwc_to_nchw = [0, 3, 1, 2]
    rank4_perm_nchw_to_nhwc = [0, 2, 3, 1]
    rank4_perm_nchw_to_nwhc = [0, 3, 2, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _write_transpose_perm_for_op(
        op_idx: int,
        new_perm: List[int],
        consumers: Dict[str, List[int]],
    ) -> bool:
        op = model_ir.operators[int(op_idx)]
        if str(op.op_type) != "TRANSPOSE" or len(op.inputs) < 2:
            return False
        perm_name = str(op.inputs[1])
        perm_users = [int(v) for v in consumers.get(perm_name, [])]
        perm_data = np.asarray(new_perm, dtype=np.int32)

        # If perm tensor is shared, clone it to avoid side effects on other TRANSPOSE ops.
        if len(perm_users) != 1 or int(perm_users[0]) != int(op_idx):
            new_perm_name = _unique_tensor_name(f"{perm_name}_canon")
            model_ir.tensors[new_perm_name] = TensorIR(
                name=new_perm_name,
                dtype="INT32",
                shape=[int(len(new_perm))],
                shape_signature=[int(len(new_perm))],
                data=perm_data,
                is_variable=False,
                quantization=None,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=op,
                input_index=1,
                new_input_name=new_perm_name,
            )
            return True

        perm_tensor = model_ir.tensors.get(perm_name, None)
        if perm_tensor is None:
            return False
        perm_tensor.data = perm_data
        perm_tensor.dtype = "INT32"
        perm_tensor.shape = [int(len(new_perm))]
        perm_tensor.shape_signature = [int(len(new_perm))]
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for softmax_idx, softmax_op in enumerate(model_ir.operators):
            if str(softmax_op.op_type) != "SOFTMAX" or len(softmax_op.inputs) != 1 or len(softmax_op.outputs) != 1:
                continue
            softmax_in = str(softmax_op.inputs[0])
            softmax_out = str(softmax_op.outputs[0])
            if softmax_in in model_ir.outputs or softmax_out in model_ir.outputs:
                continue

            pre_idx = producers.get(softmax_in, None)
            if pre_idx is None:
                continue
            pre_op = model_ir.operators[int(pre_idx)]
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if str(pre_op.outputs[0]) != softmax_in:
                continue
            if _read_transpose_perm(model_ir, pre_op) != rank4_perm_nchw_to_nwhc:
                continue
            pre_users = [int(v) for v in consumers.get(softmax_in, [])]
            if pre_users != [int(softmax_idx)]:
                continue

            pre_input = str(pre_op.inputs[0])
            pre_prev_idx = producers.get(pre_input, None)
            if pre_prev_idx is None:
                continue
            pre_prev_op = model_ir.operators[int(pre_prev_idx)]
            if (
                str(pre_prev_op.op_type) != "TRANSPOSE"
                or len(pre_prev_op.inputs) < 2
                or len(pre_prev_op.outputs) != 1
                or str(pre_prev_op.outputs[0]) != pre_input
            ):
                continue
            if _read_transpose_perm(model_ir, pre_prev_op) != rank4_perm_nhwc_to_nchw:
                continue
            pre_prev_users = [int(v) for v in consumers.get(pre_input, [])]
            if pre_prev_users != [int(pre_idx)]:
                continue

            post_users = [int(v) for v in consumers.get(softmax_out, [])]
            if len(post_users) != 1:
                continue
            post_idx = int(post_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != softmax_out:
                continue
            if _read_transpose_perm(model_ir, post_op) != rank4_perm_nchw_to_nwhc:
                continue
            post_out_name = str(post_op.outputs[0])
            # Terminal graph outputs are safe as long as they are not internally consumed.
            # This enables removing redundant pre-softmax transpose pairs while preserving
            # the final output layout adapter.
            if post_out_name in model_ir.outputs and len(consumers.get(post_out_name, [])) != 0:
                continue

            # Avoid qdim remapping complexity for per-axis activation quantization.
            pre_in_tensor = model_ir.tensors.get(pre_input, None)
            softmax_out_tensor = model_ir.tensors.get(softmax_out, None)
            post_out_tensor = model_ir.tensors.get(str(post_op.outputs[0]), None)
            if not _all_per_tensor_quantized([pre_in_tensor, softmax_out_tensor, post_out_tensor]):
                continue

            if not _write_transpose_perm_for_op(int(pre_idx), rank4_perm_nchw_to_nhwc, consumers):
                continue
            if not _write_transpose_perm_for_op(int(post_idx), rank4_perm_nhwc_to_nchw, consumers):
                continue

            # Refresh static metadata conservatively; final reconcile pass will refine.
            if pre_in_tensor is not None:
                pre_out_tensor = model_ir.tensors.get(softmax_in, None)
                if pre_out_tensor is not None:
                    new_shape = _permute_shape(list(pre_in_tensor.shape), rank4_perm_nchw_to_nhwc)
                    sig_src = (
                        list(pre_in_tensor.shape_signature)
                        if pre_in_tensor.shape_signature is not None
                        else list(pre_in_tensor.shape)
                    )
                    new_sig = _permute_shape(sig_src, rank4_perm_nchw_to_nhwc)
                    if new_shape is not None:
                        pre_out_tensor.shape = [int(v) for v in new_shape]
                    if new_sig is not None:
                        pre_out_tensor.shape_signature = [int(v) for v in new_sig]

            if softmax_out_tensor is not None:
                post_out_tensor = model_ir.tensors.get(str(post_op.outputs[0]), None)
                if post_out_tensor is not None:
                    new_shape = _permute_shape(list(softmax_out_tensor.shape), rank4_perm_nhwc_to_nchw)
                    sig_src = (
                        list(softmax_out_tensor.shape_signature)
                        if softmax_out_tensor.shape_signature is not None
                        else list(softmax_out_tensor.shape)
                    )
                    new_sig = _permute_shape(sig_src, rank4_perm_nhwc_to_nchw)
                    if new_shape is not None:
                        post_out_tensor.shape = [int(v) for v in new_shape]
                    if new_sig is not None:
                        post_out_tensor.shape_signature = [int(v) for v in new_sig]

            softmax_opts = (
                dict(softmax_op.options)
                if isinstance(softmax_op.options, dict)
                else {}
            )
            softmax_opts[softmax_nhwc_propagated_marker] = True
            softmax_op.options = softmax_opts

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"canonicalized_softmax_transpose_chains": int(rewritten)}


def _optimize_terminal_softmax_transpose_after_nhwc_propagation(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove terminal SOFTMAX->TRANSPOSE(NHWC->NCHW) when SOFTMAX already carries
    NHWC-propagated semantics from `_canonicalize_softmax_transpose_chains`.

    Target:
      ... -> SOFTMAX(s) -> TRANSPOSE(0,3,1,2) -> y(graph output)

    Rewrite:
      ... -> SOFTMAX(y)

    Safety:
    - y must be a graph output with no internal consumers.
    - SOFTMAX output is consumed only by that terminal TRANSPOSE.
    - SOFTMAX is marked as NHWC-propagated by canonicalize pass.
    """
    rewritten = 0
    rank4_perm_nhwc_to_nchw = [0, 3, 1, 2]
    softmax_nhwc_propagated_marker = "__softmax_nhwc_propagated__"

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)

        for output_name in list(model_ir.outputs):
            output_name = str(output_name)
            if len(consumers.get(output_name, [])) != 0:
                continue

            post_idx = producers.get(output_name, None)
            if post_idx is None:
                continue
            post_op = model_ir.operators[int(post_idx)]
            if (
                str(post_op.op_type) != "TRANSPOSE"
                or len(post_op.inputs) < 2
                or len(post_op.outputs) != 1
                or str(post_op.outputs[0]) != output_name
            ):
                continue
            if _read_transpose_perm(model_ir, post_op) != rank4_perm_nhwc_to_nchw:
                continue

            softmax_out_name = str(post_op.inputs[0])
            softmax_idx = producers.get(softmax_out_name, None)
            if softmax_idx is None:
                continue
            softmax_op = model_ir.operators[int(softmax_idx)]
            if (
                str(softmax_op.op_type) != "SOFTMAX"
                or len(softmax_op.inputs) != 1
                or len(softmax_op.outputs) != 1
                or str(softmax_op.outputs[0]) != softmax_out_name
            ):
                continue
            softmax_users = [int(v) for v in consumers.get(softmax_out_name, [])]
            if softmax_users != [int(post_idx)]:
                continue

            softmax_opts = (
                dict(softmax_op.options)
                if isinstance(softmax_op.options, dict)
                else {}
            )
            if not bool(softmax_opts.get(softmax_nhwc_propagated_marker, False)):
                continue
            softmax_opts.pop(softmax_nhwc_propagated_marker, None)
            softmax_op.options = softmax_opts

            _set_operator_outputs(
                model_ir=model_ir,
                op=softmax_op,
                new_outputs=[output_name],
            )
            src_tensor = model_ir.tensors.get(softmax_out_name, None)
            dst_tensor = model_ir.tensors.get(output_name, None)
            if src_tensor is not None and dst_tensor is not None:
                dst_tensor.dtype = str(src_tensor.dtype)
                dst_tensor.quantization = _clone_quantization(src_tensor.quantization)
                dst_tensor.shape = [int(v) for v in list(src_tensor.shape)]
                dst_tensor.shape_signature = (
                    [int(v) for v in list(src_tensor.shape_signature)]
                    if src_tensor.shape_signature is not None
                    else [int(v) for v in list(src_tensor.shape)]
                )

            del model_ir.operators[int(post_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"removed_terminal_softmax_transpose_after_nhwc_propagation": int(rewritten)}


def _optimize_transpose_pre_argmax_nhwc_terminal_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove redundant NHWC->NCHW transpose immediately before ARG_MAX by remapping axis.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw --ARG_MAX(axis=1)--> y

    Rewrite:
      x_nhwc --ARG_MAX(axis=3)--> y

    Safety:
    - Transpose output has single consumer (ARG_MAX).
    - ARG_MAX axis is constant and resolves to channel axis=1 in NCHW.
    - Rank-4 transpose permutation is exactly [0,3,1,2].
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_input_name in model_ir.outputs or pre_output_name in model_ir.outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            argmax_idx = int(pre_users[0])
            argmax_op = model_ir.operators[int(argmax_idx)]
            if str(argmax_op.op_type) != "ARG_MAX" or len(argmax_op.inputs) < 2 or len(argmax_op.outputs) != 1:
                continue
            if str(argmax_op.inputs[0]) != pre_output_name:
                continue

            axis_name = str(argmax_op.inputs[1])
            axis_tensor = model_ir.tensors.get(axis_name, None)
            axis_vals = _read_const_ints_from_tensor(axis_tensor)
            if axis_vals is None or len(axis_vals) != 1:
                continue

            axis_old = int(axis_vals[0])
            if axis_old < 0:
                axis_old += 4
            # Restrict to channel-axis argmax in NCHW to preserve output ordering.
            if int(axis_old) != 1:
                continue
            axis_new = int(perm_nhwc_to_nchw[int(axis_old)])

            axis_users = [int(v) for v in consumers.get(axis_name, [])]
            axis_input_name = str(axis_name)
            if axis_users == [int(argmax_idx)]:
                if not _write_const_ints_to_tensor(axis_tensor, [int(axis_new)]):
                    # value may already be rewritten in prior sweep
                    pass
            else:
                axis_input_name = _unique_tensor_name(f"{axis_name}_nhwc")
                np_dtype = np.int32
                if axis_tensor is not None and axis_tensor.data is not None:
                    try:
                        np_dtype = np.asarray(axis_tensor.data).dtype
                    except Exception:
                        np_dtype = np.int32
                axis_data = np.asarray([int(axis_new)], dtype=np_dtype)
                model_ir.tensors[axis_input_name] = TensorIR(
                    name=axis_input_name,
                    dtype=str(axis_tensor.dtype) if axis_tensor is not None else "INT32",
                    shape=[1],
                    shape_signature=[1],
                    data=axis_data,
                    is_variable=False,
                    quantization=_clone_quantization(axis_tensor.quantization if axis_tensor is not None else None),
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=argmax_op,
                new_inputs=[str(pre_input_name), str(axis_input_name)],
            )

            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_pre_argmax_nhwc_terminal_chains": int(rewritten)}


def _optimize_transpose_gather_transpose_nhwc_channel_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Fold NHWC<->NCHW adapters around channel-axis Gather into a single NHWC Gather.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
             --GATHER(axis=1,batchDims=0)--> y_nchw
             --TRANSPOSE(0,2,3,1)--> y_nhwc

    Rewrite:
      x_nhwc --GATHER(axis=3,batchDims=0)--> y_nhwc
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_input_name = str(pre_op.inputs[0])
            pre_output_name = str(pre_op.outputs[0])
            if pre_output_name in model_outputs:
                continue

            pre_users = [int(v) for v in consumers.get(pre_output_name, [])]
            if len(pre_users) != 1:
                continue
            gather_idx = int(pre_users[0])
            gather_op = model_ir.operators[int(gather_idx)]
            if str(gather_op.op_type) != "GATHER" or len(gather_op.inputs) < 2 or len(gather_op.outputs) != 1:
                continue
            if str(gather_op.inputs[0]) != pre_output_name:
                continue

            gather_output_name = str(gather_op.outputs[0])
            if gather_output_name in model_outputs:
                continue

            gather_users = [int(v) for v in consumers.get(gather_output_name, [])]
            if len(gather_users) != 1:
                continue
            post_idx = int(gather_users[0])
            post_op = model_ir.operators[int(post_idx)]
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if str(post_op.inputs[0]) != gather_output_name:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            gather_in_tensor = model_ir.tensors.get(pre_output_name, None)
            if gather_in_tensor is None:
                continue
            gather_in_rank = len(list(gather_in_tensor.shape))
            if gather_in_rank != 4:
                continue

            gather_options = dict(gather_op.options) if isinstance(gather_op.options, dict) else {}
            gather_axis = int(gather_options.get("axis", 0))
            if gather_axis < 0:
                gather_axis += int(gather_in_rank)
            if int(gather_axis) != 1:
                continue
            gather_batch_dims = int(
                gather_options.get(
                    "batchDims",
                    gather_options.get("batch_dims", 0),
                )
            )
            if int(gather_batch_dims) != 0:
                continue

            gather_options["axis"] = 3
            gather_options["batchDims"] = int(gather_batch_dims)
            gather_op.options = gather_options

            _set_operator_inputs(
                model_ir=model_ir,
                op=gather_op,
                new_inputs=[str(pre_input_name), str(gather_op.inputs[1])],
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=gather_op,
                new_outputs=[str(post_op.outputs[0])],
            )

            for remove_idx in sorted([int(post_idx), int(pre_idx)], reverse=True):
                del model_ir.operators[int(remove_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_gather_transpose_nhwc_channel_chains": int(rewritten)}


def _optimize_layout_transpose_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate redundant TRANSPOSE chains introduced by channel-first/channel-last bridging.

    This pass removes:
    - identity transpose: X --Transpose(identity)--> Y
    - inverse transpose pairs with single-edge bridge:
        A --Transpose(P)--> B --Transpose(inv(P))--> C
      by directly reconnecting C consumers to A.

    The optimization is output-safe: graph outputs are never renamed.
    """
    removed_identity = 0
    removed_inverse_pairs = 0
    removed_inverse_fanout_branches = 0
    composed_consecutive_pairs = 0
    iterations = 0
    preserve_layout_boundary_marker = "__preserve_layout_boundary__"
    rank4_perm_nhwc_to_nchw = [0, 3, 1, 2]
    rank4_perm_nchw_to_nwhc = [0, 3, 2, 1]

    def _compose_permutations(perm_pre: List[int], perm_post: List[int]) -> Optional[List[int]]:
        if len(perm_pre) != len(perm_post) or len(perm_pre) == 0:
            return None
        rank = len(perm_pre)
        if sorted(int(v) for v in perm_pre) != list(range(rank)):
            return None
        if sorted(int(v) for v in perm_post) != list(range(rank)):
            return None
        try:
            return [int(perm_pre[int(v)]) for v in perm_post]
        except Exception:
            return None

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    while True:
        iterations += 1
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)

        # 1) Remove identity transpose.
        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE":
                continue
            if len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            perm = _read_transpose_perm(model_ir, op)
            if perm is None or not _is_identity_perm(perm):
                continue
            transposed_output = op.outputs[0]
            transposed_input = op.inputs[0]
            if transposed_output in model_ir.outputs:
                continue
            _replace_tensor_inputs(model_ir, transposed_output, transposed_input)
            del model_ir.operators[op_idx]
            removed_identity += 1
            changed = True
            break
        if changed:
            continue

        # 2) Remove inverse transpose pair.
        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE":
                continue
            if len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            bridge_tensor = op.outputs[0]
            bridge_users = consumers.get(bridge_tensor, [])
            if len(bridge_users) != 1:
                continue
            next_op_idx = int(bridge_users[0])
            if next_op_idx == op_idx:
                continue
            next_op = model_ir.operators[next_op_idx]
            if str(next_op.op_type) != "TRANSPOSE":
                continue
            if len(next_op.inputs) < 2 or len(next_op.outputs) != 1:
                continue
            if next_op.inputs[0] != bridge_tensor:
                continue

            perm_1 = _read_transpose_perm(model_ir, op)
            perm_2 = _read_transpose_perm(model_ir, next_op)
            if perm_1 is None or perm_2 is None:
                continue
            if not _is_inverse_perm(perm_1, perm_2):
                continue

            transpose_1_input = op.inputs[0]
            transpose_1_output = op.outputs[0]
            transpose_2_output = next_op.outputs[0]

            # Keep user-visible output names stable.
            if transpose_1_output in model_ir.outputs or transpose_2_output in model_ir.outputs:
                continue

            _replace_tensor_inputs(model_ir, transpose_2_output, transpose_1_input)
            for remove_idx in sorted([op_idx, next_op_idx], reverse=True):
                del model_ir.operators[remove_idx]
            removed_inverse_pairs += 1
            changed = True
            break

        if not changed:
            # 2b) Remove inverse transpose branches even when the first transpose
            # fans out to multiple consumers.
            for op_idx, op in enumerate(model_ir.operators):
                if str(op.op_type) != "TRANSPOSE":
                    continue
                if len(op.inputs) < 2 or len(op.outputs) != 1:
                    continue

                bridge_tensor = str(op.outputs[0])
                bridge_users = [int(v) for v in consumers.get(bridge_tensor, [])]
                if len(bridge_users) <= 1:
                    continue

                perm_pre = _read_transpose_perm(model_ir, op)
                if perm_pre is None:
                    continue

                for post_idx in bridge_users:
                    if int(post_idx) == int(op_idx):
                        continue
                    post_op = model_ir.operators[int(post_idx)]
                    if str(post_op.op_type) != "TRANSPOSE":
                        continue
                    if len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                        continue
                    if str(post_op.inputs[0]) != bridge_tensor:
                        continue

                    post_opts = dict(post_op.options) if isinstance(post_op.options, dict) else {}
                    if bool(post_opts.get(preserve_layout_boundary_marker, False)):
                        continue

                    perm_post = _read_transpose_perm(model_ir, post_op)
                    if perm_post is None or not _is_inverse_perm(perm_pre, perm_post):
                        continue

                    post_output = str(post_op.outputs[0])
                    if bridge_tensor in model_ir.outputs or post_output in model_ir.outputs:
                        continue

                    _replace_tensor_inputs(model_ir, post_output, str(op.inputs[0]))
                    del model_ir.operators[int(post_idx)]
                    removed_inverse_fanout_branches += 1
                    changed = True
                    break
                if changed:
                    break

        if not changed:
            # 3) Compose consecutive transpose pairs into one transpose.
            for op_idx, op in enumerate(model_ir.operators):
                if str(op.op_type) != "TRANSPOSE":
                    continue
                if len(op.inputs) < 2 or len(op.outputs) != 1:
                    continue
                op_opts = dict(op.options) if isinstance(op.options, dict) else {}
                if bool(op_opts.get(preserve_layout_boundary_marker, False)):
                    continue

                bridge_tensor = str(op.outputs[0])
                bridge_users = [int(v) for v in consumers.get(bridge_tensor, [])]
                if len(bridge_users) != 1:
                    continue

                next_op_idx = int(bridge_users[0])
                if next_op_idx == int(op_idx):
                    continue
                next_op = model_ir.operators[int(next_op_idx)]
                if str(next_op.op_type) != "TRANSPOSE":
                    continue
                if len(next_op.inputs) < 2 or len(next_op.outputs) != 1:
                    continue
                if str(next_op.inputs[0]) != bridge_tensor:
                    continue
                next_opts = dict(next_op.options) if isinstance(next_op.options, dict) else {}
                if bool(next_opts.get(preserve_layout_boundary_marker, False)):
                    continue

                perm_1 = _read_transpose_perm(model_ir, op)
                perm_2 = _read_transpose_perm(model_ir, next_op)
                if perm_1 is None or perm_2 is None:
                    continue
                # Keep this pattern intact for `_canonicalize_softmax_transpose_chains`.
                if (
                    [int(v) for v in perm_1] == rank4_perm_nhwc_to_nchw
                    and [int(v) for v in perm_2] == rank4_perm_nchw_to_nwhc
                ):
                    next_out_name = str(next_op.outputs[0])
                    next_out_users = [int(v) for v in consumers.get(next_out_name, [])]
                    if len(next_out_users) == 1:
                        softmax_idx = int(next_out_users[0])
                        if softmax_idx != int(next_op_idx):
                            softmax_op = model_ir.operators[softmax_idx]
                            if (
                                str(softmax_op.op_type) == "SOFTMAX"
                                and len(softmax_op.inputs) == 1
                                and len(softmax_op.outputs) == 1
                                and str(softmax_op.inputs[0]) == next_out_name
                            ):
                                softmax_out_name = str(softmax_op.outputs[0])
                                softmax_out_users = [int(v) for v in consumers.get(softmax_out_name, [])]
                                if len(softmax_out_users) == 1:
                                    post_idx = int(softmax_out_users[0])
                                    if post_idx != int(softmax_idx):
                                        post_op = model_ir.operators[post_idx]
                                        if (
                                            str(post_op.op_type) == "TRANSPOSE"
                                            and len(post_op.inputs) >= 2
                                            and len(post_op.outputs) == 1
                                            and str(post_op.inputs[0]) == softmax_out_name
                                            and _read_transpose_perm(model_ir, post_op) == rank4_perm_nchw_to_nwhc
                                        ):
                                            continue
                composed_perm = _compose_permutations(
                    [int(v) for v in perm_1],
                    [int(v) for v in perm_2],
                )
                if composed_perm is None:
                    continue

                next_perm_name = str(next_op.inputs[1])
                next_perm_users = [int(v) for v in consumers.get(next_perm_name, [])]
                if len(next_perm_users) == 1 and int(next_perm_users[0]) == int(next_op_idx):
                    perm_tensor = model_ir.tensors.get(next_perm_name, None)
                    if perm_tensor is None:
                        continue
                    perm_arr = np.asarray(composed_perm, dtype=np.int32)
                    perm_tensor.data = perm_arr
                    perm_tensor.dtype = "INT32"
                    perm_tensor.shape = [int(perm_arr.size)]
                    perm_tensor.shape_signature = [int(perm_arr.size)]
                    new_perm_name = next_perm_name
                else:
                    new_perm_name = _unique_tensor_name(f"{next_perm_name}_composed")
                    perm_arr = np.asarray(composed_perm, dtype=np.int32)
                    model_ir.tensors[new_perm_name] = TensorIR(
                        name=new_perm_name,
                        dtype="INT32",
                        shape=[int(perm_arr.size)],
                        shape_signature=[int(perm_arr.size)],
                        data=perm_arr,
                        is_variable=False,
                        quantization=None,
                    )

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=next_op,
                    new_inputs=[str(op.inputs[0]), str(new_perm_name)],
                )
                del model_ir.operators[int(op_idx)]
                composed_consecutive_pairs += 1
                changed = True
                break

            if not changed:
                break

    _prune_unused_tensors(model_ir)
    return {
        "iterations": int(iterations),
        "removed_identity_transpose": int(removed_identity),
        "removed_inverse_transpose_pairs": int(removed_inverse_pairs),
        "removed_inverse_transpose_fanout_branches": int(removed_inverse_fanout_branches),
        "composed_consecutive_transpose_pairs": int(composed_consecutive_pairs),
    }


def _optimize_boundary_input_transpose_channel_slice_blocks(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove shared boundary input TRANSPOSE while preserving downstream NCHW semantics.

    Target:
      input(NHWC) --TRANSPOSE(0,3,1,2)--> input_onnx_ncx_internal
        |- SLICE(channel axis=1)
        |- ...
        `- other consumers

    Rewrite:
    - For eligible channel SLICE consumers, rewrite slice source to NHWC input and
      move NHWC->NCHW transpose after each slice output.
    - Localize remaining NCHW requirements by inserting per-consumer transposes.
    - Remove shared boundary TRANSPOSE so identifier=0 is no longer TRANSPOSE.
    """
    removed_boundary = 0
    rewritten_channel_slices = 0
    inserted_local_transposes = 0
    rewritten_axis_ops = 0

    layout_passthrough_unary = {
        "CAST",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LOGISTIC",
        "TANH",
        "NEG",
        "EXP",
        "ABS",
        "SQRT",
    }
    layout_passthrough_binary = {
        "ADD",
        "SUB",
        "MUL",
        "DIV",
        "MAXIMUM",
        "MINIMUM",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _is_channel_slice_axis1(
        *,
        op: OperatorIR,
        source_shape_nchw: Optional[List[int]] = None,
    ) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3 or len(op.outputs) != 1:
            return False
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        if int(size_vals[1]) != 1:
            return False
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return False
        if source_shape_nchw is not None and len(source_shape_nchw) == 4:
            if int(begin_vals[1]) < 0 or int(begin_vals[1]) >= int(source_shape_nchw[1]):
                return False
            h = int(source_shape_nchw[2])
            w = int(source_shape_nchw[3])
            if int(size_vals[2]) not in {h, -1}:
                return False
            if int(size_vals[3]) not in {w, -1}:
                return False
        else:
            if int(begin_vals[1]) < 0:
                return False
        return True

    def _rewrite_channel_slice_axis1_to_axis3(
        *,
        op: OperatorIR,
    ) -> bool:
        begin_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
        new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(str(v) for v in model_ir.inputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            input_name = str(pre_op.inputs[0])
            internal_name = str(pre_op.outputs[0])
            if input_name not in model_inputs:
                continue
            if not str(internal_name).endswith("_onnx_ncx_internal"):
                continue
            perm_pre = _read_transpose_perm(model_ir, pre_op)
            if perm_pre != [0, 3, 1, 2]:
                continue

            input_tensor = model_ir.tensors.get(input_name, None)
            internal_tensor = model_ir.tensors.get(internal_name, None)
            if input_tensor is None or internal_tensor is None:
                continue
            if len(list(internal_tensor.shape)) != 4:
                continue

            internal_users = [int(v) for v in consumers.get(internal_name, [])]
            if len(internal_users) == 0:
                continue

            # 1) Rewrite direct channel-slice consumers to NHWC source.
            converted_ops: set = set()
            nhwc_tensors: set = set()
            rewritten_any_slice = False
            for user_idx in sorted(internal_users):
                user_op = model_ir.operators[int(user_idx)]
                if not _is_channel_slice_axis1(
                    op=user_op,
                    source_shape_nchw=[int(v) for v in list(internal_tensor.shape)],
                ):
                    continue
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=user_op,
                    new_inputs=[input_name, str(user_op.inputs[1]), str(user_op.inputs[2])],
                )
                if _rewrite_channel_slice_axis1_to_axis3(op=user_op):
                    rewritten_channel_slices += 1
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(user_op.outputs[0]), None),
                    [0, 2, 3, 1],
                )
                converted_ops.add(int(user_idx))
                nhwc_tensors.add(str(user_op.outputs[0]))
                rewritten_any_slice = True

            if not rewritten_any_slice:
                continue

            # 2) Propagate NHWC through local axis-sensitive block.
            while True:
                propagated = False
                for op_idx, op in enumerate(model_ir.operators):
                    if int(op_idx) in converted_ops:
                        continue
                    if len(op.outputs) != 1:
                        continue
                    out_name = str(op.outputs[0])
                    op_type = str(op.op_type)

                    if op_type == "CONCATENATION":
                        axis = int(op.options.get("axis", -1))
                        if axis == 1 and len(op.inputs) > 0 and all(str(inp) in nhwc_tensors for inp in op.inputs):
                            op.options["axis"] = 3
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                [0, 2, 3, 1],
                            )
                            rewritten_axis_ops += 1
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type == "SLICE" and len(op.inputs) >= 3:
                        in_name = str(op.inputs[0])
                        if in_name in nhwc_tensors and _is_channel_slice_axis1(op=op, source_shape_nchw=None):
                            if _rewrite_channel_slice_axis1_to_axis3(op=op):
                                rewritten_axis_ops += 1
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                [0, 2, 3, 1],
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type in layout_passthrough_unary and len(op.inputs) >= 1:
                        if str(op.inputs[0]) in nhwc_tensors:
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                [0, 2, 3, 1],
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type in layout_passthrough_binary and len(op.inputs) >= 2:
                        dynamic_inputs: List[str] = []
                        for inp in list(op.inputs[:2]):
                            inp_name = str(inp)
                            inp_tensor = model_ir.tensors.get(inp_name, None)
                            if inp_tensor is not None and inp_tensor.data is not None:
                                continue
                            dynamic_inputs.append(inp_name)
                        if len(dynamic_inputs) > 0 and all(inp_name in nhwc_tensors for inp_name in dynamic_inputs):
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                [0, 2, 3, 1],
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type in {"CONV_2D", "DEPTHWISE_CONV_2D"} and len(op.inputs) >= 1:
                        # TFLite CONV_2D/DEPTHWISE_CONV_2D always produce NHWC outputs.
                        if str(op.inputs[0]) in nhwc_tensors:
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                [0, 2, 3, 1],
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                if not propagated:
                    break

            # 3) Bridge NHWC tensors back to NCHW for consumers outside converted block.
            consumers_nhwc = _build_tensor_consumer_map(model_ir)
            bridge_plans: List[Tuple[int, str, str, List[int]]] = []
            for tensor_name in sorted(list(nhwc_tensors)):
                users = [
                    int(v)
                    for v in consumers_nhwc.get(str(tensor_name), [])
                    if int(v) not in converted_ops
                ]
                if len(users) == 0:
                    continue
                tensor = model_ir.tensors.get(str(tensor_name), None)
                if tensor is None or len(list(tensor.shape)) != 4:
                    continue
                bridge_name = _unique_tensor_name(f"{tensor_name}_nchw_bridge")
                bridge_shape = _permute_shape(list(tensor.shape), [0, 3, 1, 2])
                signature_src = (
                    list(tensor.shape_signature)
                    if tensor.shape_signature is not None
                    else list(tensor.shape)
                )
                bridge_signature = _permute_shape(signature_src, [0, 3, 1, 2])
                if bridge_shape is None or bridge_signature is None:
                    continue
                model_ir.tensors[bridge_name] = TensorIR(
                    name=bridge_name,
                    dtype=str(tensor.dtype),
                    shape=[int(v) for v in list(bridge_shape)],
                    shape_signature=[int(v) for v in list(bridge_signature)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(tensor.quantization),
                )
                for user_idx in users:
                    user_op = model_ir.operators[int(user_idx)]
                    updated_inputs = [
                        bridge_name if str(inp) == str(tensor_name) else str(inp)
                        for inp in list(user_op.inputs)
                    ]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=user_op,
                        new_inputs=updated_inputs,
                    )
                bridge_plans.append((int(min(users)), str(tensor_name), str(bridge_name), [int(v) for v in users]))

            inserted = 0
            for insert_idx, source_name, bridge_name, _users in sorted(bridge_plans, key=lambda v: int(v[0])):
                model_ir.operators.insert(
                    int(insert_idx + inserted),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(source_name), str(pre_op.inputs[1])],
                        outputs=[str(bridge_name)],
                    ),
                )
                inserted += 1
                inserted_local_transposes += 1

            # 4) Localize any remaining NCHW uses for this boundary tensor.
            consumers_after = _build_tensor_consumer_map(model_ir)
            remaining_users = [int(v) for v in consumers_after.get(internal_name, [])]
            inserted = 0
            for user_idx in sorted(remaining_users):
                user_op = model_ir.operators[int(user_idx + inserted)]
                local_name = _unique_tensor_name(f"{internal_name}_local")
                model_ir.tensors[local_name] = TensorIR(
                    name=local_name,
                    dtype=str(internal_tensor.dtype),
                    shape=[int(v) for v in list(internal_tensor.shape)],
                    shape_signature=(
                        [int(v) for v in list(internal_tensor.shape_signature)]
                        if internal_tensor.shape_signature is not None
                        else [int(v) for v in list(internal_tensor.shape)]
                    ),
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(internal_tensor.quantization),
                )
                model_ir.operators.insert(
                    int(user_idx + inserted),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[input_name, str(pre_op.inputs[1])],
                        outputs=[local_name],
                    ),
                )
                inserted += 1
                updated_inputs = [
                    local_name if str(inp) == internal_name else str(inp)
                    for inp in list(user_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=user_op,
                    new_inputs=updated_inputs,
                )
                inserted_local_transposes += 1

            # 5) Remove shared boundary transpose.
            del model_ir.operators[int(pre_idx)]
            removed_boundary += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_boundary_input_transpose": int(removed_boundary),
        "rewritten_boundary_channel_slices": int(rewritten_channel_slices),
        "rewritten_boundary_axis_ops": int(rewritten_axis_ops),
        "inserted_local_boundary_transposes": int(inserted_local_transposes),
    }


def _optimize_internal_transpose_channel_slice_nhwc_propagation_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove internal NHWC->NCHW transpose stems that only feed channel-slice branches.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
        |- SLICE(axis=1) ...
        `- SLICE(axis=1) ...

    Rewrite:
    - Rewrite slice sources/params to NHWC (axis=3 equivalent).
    - Propagate NHWC through local unary/binary/conv/pool block.
    - Bypass local TRANSPOSE wrappers fed by already-NHWC tensors.
    - Reinsert NHWC->NCHW bridge only for remaining legacy consumers.
    """
    removed_internal = 0
    rewritten_channel_slices = 0
    rewritten_axis_ops = 0
    inserted_local_transposes = 0

    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    layout_passthrough_unary = {
        "CAST",
        "RELU",
        "RELU6",
        "RELU_0_TO_1",
        "HARD_SWISH",
        "LOGISTIC",
        "TANH",
        "LEAKY_RELU",
        "NEG",
        "EXP",
        "ABS",
        "SQRT",
    }
    layout_passthrough_binary = {
        "ADD",
        "SUB",
        "MUL",
        "DIV",
        "MAXIMUM",
        "MINIMUM",
    }

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _is_channel_slice_axis1(
        *,
        op: OperatorIR,
        source_shape_nchw: Optional[List[int]] = None,
    ) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3 or len(op.outputs) != 1:
            return False
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        if int(size_vals[1]) <= 0:
            return False
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return False
        if source_shape_nchw is not None and len(source_shape_nchw) == 4:
            if int(begin_vals[1]) < 0:
                return False
            c = int(source_shape_nchw[1])
            if int(begin_vals[1]) + int(size_vals[1]) > c:
                return False
            h = int(source_shape_nchw[2])
            w = int(source_shape_nchw[3])
            if int(size_vals[2]) not in {h, -1}:
                return False
            if int(size_vals[3]) not in {w, -1}:
                return False
        return True

    def _rewrite_channel_slice_axis1_to_axis3(
        *,
        op: OperatorIR,
    ) -> bool:
        begin_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
        new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            input_name = str(pre_op.inputs[0])
            internal_name = str(pre_op.outputs[0])
            if input_name in model_outputs or internal_name in model_outputs:
                continue

            input_tensor = model_ir.tensors.get(input_name, None)
            internal_tensor = model_ir.tensors.get(internal_name, None)
            if input_tensor is None or internal_tensor is None:
                continue
            if len(list(input_tensor.shape)) != 4 or len(list(internal_tensor.shape)) != 4:
                continue

            internal_users = [int(v) for v in consumers.get(internal_name, [])]
            if len(internal_users) == 0:
                continue
            if not all(
                _is_channel_slice_axis1(
                    op=model_ir.operators[int(user_idx)],
                    source_shape_nchw=[int(v) for v in list(internal_tensor.shape)],
                )
                for user_idx in internal_users
            ):
                continue

            converted_ops: set = set()
            nhwc_tensors: set = set()
            removable_transposes: set = set()

            for user_idx in sorted(internal_users):
                user_op = model_ir.operators[int(user_idx)]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=user_op,
                    new_inputs=[input_name, str(user_op.inputs[1]), str(user_op.inputs[2])],
                )
                if _rewrite_channel_slice_axis1_to_axis3(op=user_op):
                    rewritten_channel_slices += 1
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(user_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
                converted_ops.add(int(user_idx))
                nhwc_tensors.add(str(user_op.outputs[0]))

            while True:
                propagated = False
                for op_idx, op in enumerate(model_ir.operators):
                    if int(op_idx) in converted_ops:
                        continue
                    if len(op.outputs) != 1:
                        continue
                    op_type = str(op.op_type)
                    out_name = str(op.outputs[0])
                    if out_name in model_outputs:
                        continue

                    if op_type == "CONCATENATION":
                        axis = int(op.options.get("axis", -1))
                        if axis == 1 and len(op.inputs) > 0 and all(str(inp) in nhwc_tensors for inp in op.inputs):
                            op.options["axis"] = 3
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                perm_nchw_to_nhwc,
                            )
                            rewritten_axis_ops += 1
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type == "SLICE" and len(op.inputs) >= 3:
                        in_name = str(op.inputs[0])
                        if in_name in nhwc_tensors and _is_channel_slice_axis1(op=op, source_shape_nchw=None):
                            if _rewrite_channel_slice_axis1_to_axis3(op=op):
                                rewritten_axis_ops += 1
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                perm_nchw_to_nhwc,
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type in layout_passthrough_unary and len(op.inputs) >= 1:
                        if str(op.inputs[0]) in nhwc_tensors:
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                perm_nchw_to_nhwc,
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                    if op_type in layout_passthrough_binary and len(op.inputs) >= 2:
                        input_names = [str(v) for v in list(op.inputs)]
                        dynamic_inputs: List[str] = []
                        for inp_name in input_names[:2]:
                            inp_tensor = model_ir.tensors.get(inp_name, None)
                            if inp_tensor is not None and inp_tensor.data is not None:
                                continue
                            dynamic_inputs.append(inp_name)

                        if len(dynamic_inputs) == 0 or not all(inp_name in nhwc_tensors for inp_name in dynamic_inputs):
                            continue

                        target_nhwc_shape: Optional[List[int]] = None
                        for dyn_name in dynamic_inputs:
                            dyn_tensor = model_ir.tensors.get(str(dyn_name), None)
                            if (
                                dyn_tensor is not None
                                and _is_fully_known_positive_shape(list(dyn_tensor.shape))
                                and len(list(dyn_tensor.shape)) == 4
                            ):
                                target_nhwc_shape = [int(v) for v in list(dyn_tensor.shape)]
                                break

                        # Rewrite NCHW channelwise constants to NHWC when needed.
                        safe_binary = True
                        new_inputs = [str(v) for v in input_names]
                        for input_index, const_name in enumerate(input_names[:2]):
                            const_tensor = model_ir.tensors.get(str(const_name), None)
                            if const_tensor is None or const_tensor.data is None:
                                continue
                            const_data = np.asarray(const_tensor.data)
                            if int(const_data.size) == 1:
                                continue

                            if target_nhwc_shape is None:
                                safe_binary = False
                                break

                            const_shape = [int(v) for v in list(const_data.shape)]
                            if _broadcast_static_shapes(target_nhwc_shape, const_shape) is not None:
                                continue

                            rotated: Optional[np.ndarray] = None
                            if int(const_data.ndim) == 4:
                                candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
                                if _broadcast_static_shapes(
                                    target_nhwc_shape,
                                    [int(v) for v in list(candidate.shape)],
                                ) is not None:
                                    rotated = np.asarray(candidate)
                            elif int(const_data.ndim) == 3:
                                candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
                                if _broadcast_static_shapes(
                                    target_nhwc_shape,
                                    [int(v) for v in list(candidate.shape)],
                                ) is not None:
                                    rotated = np.asarray(candidate)

                            if rotated is None:
                                safe_binary = False
                                break

                            const_users = [int(v) for v in consumers.get(str(const_name), [])]
                            if any(int(v) != int(op_idx) for v in const_users):
                                rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
                                model_ir.tensors[rotated_name] = TensorIR(
                                    name=rotated_name,
                                    dtype=str(const_tensor.dtype),
                                    shape=[int(v) for v in list(rotated.shape)],
                                    shape_signature=[int(v) for v in list(rotated.shape)],
                                    data=np.asarray(rotated),
                                    is_variable=False,
                                    quantization=_clone_quantization(const_tensor.quantization),
                                )
                                new_inputs[int(input_index)] = str(rotated_name)
                            else:
                                const_tensor.data = np.asarray(rotated)
                                const_tensor.shape = [int(v) for v in list(rotated.shape)]
                                const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]

                        if not safe_binary:
                            continue

                        if new_inputs != input_names:
                            _set_operator_inputs(
                                model_ir=model_ir,
                                op=op,
                                new_inputs=new_inputs,
                            )

                        _permute_tensor_metadata_if_rank_matches(
                            model_ir.tensors.get(out_name, None),
                            perm_nchw_to_nhwc,
                        )
                        nhwc_tensors.add(out_name)
                        converted_ops.add(int(op_idx))
                        propagated = True
                        continue

                    if op_type in {"CONV_2D", "DEPTHWISE_CONV_2D", "AVERAGE_POOL_2D", "MAX_POOL_2D"} and len(op.inputs) >= 1:
                        if str(op.inputs[0]) in nhwc_tensors:
                            _permute_tensor_metadata_if_rank_matches(
                                model_ir.tensors.get(out_name, None),
                                perm_nchw_to_nhwc,
                            )
                            nhwc_tensors.add(out_name)
                            converted_ops.add(int(op_idx))
                            propagated = True
                        continue

                if not propagated:
                    break

            consumers_nhwc = _build_tensor_consumer_map(model_ir)
            bridge_plans: List[Tuple[int, str, str, List[int]]] = []
            for tensor_name in sorted(list(nhwc_tensors)):
                users = [
                    int(v)
                    for v in consumers_nhwc.get(str(tensor_name), [])
                    if int(v) not in converted_ops
                ]
                if len(users) == 0:
                    continue
                tensor = model_ir.tensors.get(str(tensor_name), None)
                if tensor is None or len(list(tensor.shape)) != 4:
                    continue
                bridge_name = _unique_tensor_name(f"{tensor_name}_nchw_bridge")
                bridge_shape = _permute_shape(list(tensor.shape), perm_nhwc_to_nchw)
                signature_src = (
                    list(tensor.shape_signature)
                    if tensor.shape_signature is not None
                    else list(tensor.shape)
                )
                bridge_signature = _permute_shape(signature_src, perm_nhwc_to_nchw)
                if bridge_shape is None or bridge_signature is None:
                    continue
                model_ir.tensors[bridge_name] = TensorIR(
                    name=bridge_name,
                    dtype=str(tensor.dtype),
                    shape=[int(v) for v in list(bridge_shape)],
                    shape_signature=[int(v) for v in list(bridge_signature)],
                    data=None,
                    is_variable=False,
                    quantization=_clone_quantization(tensor.quantization),
                )
                for user_idx in users:
                    user_op = model_ir.operators[int(user_idx)]
                    updated_inputs = [
                        bridge_name if str(inp) == str(tensor_name) else str(inp)
                        for inp in list(user_op.inputs)
                    ]
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=user_op,
                        new_inputs=updated_inputs,
                    )
                bridge_plans.append((int(min(users)), str(tensor_name), str(bridge_name), [int(v) for v in users]))

            inserted = 0
            for insert_idx, source_name, bridge_name, _users in sorted(bridge_plans, key=lambda v: int(v[0])):
                model_ir.operators.insert(
                    int(insert_idx + inserted),
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(source_name), str(pre_op.inputs[1])],
                        outputs=[str(bridge_name)],
                    ),
                )
                inserted += 1
                inserted_local_transposes += 1

            remove_indices = sorted(
                list({int(pre_idx)} | {int(v) for v in removable_transposes}),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]
            removed_internal += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_internal_transpose_channel_slice_stems": int(removed_internal),
        "rewritten_internal_channel_slices": int(rewritten_channel_slices),
        "rewritten_internal_axis_ops": int(rewritten_axis_ops),
        "inserted_internal_local_transposes": int(inserted_local_transposes),
    }


def _optimize_transpose_channel_slice_muladd_nhwc_bridge_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC->NCHW transpose stems whose channel slices immediately bridge back to NHWC.

    Target:
      x_nhwc --TRANSPOSE(0,3,1,2)--> x_nchw
      x_nchw --SLICE(axis=1)--> s_i_nchw
      s_i_nchw --TRANSPOSE(0,2,3,1)--> s_i_nhwc
      or
      s_i_nchw --MUL(const)--> m_i_nchw --TRANSPOSE(0,2,3,1)--> m_i_nhwc

    Rewrite:
      x_nhwc --SLICE(axis=3)--> s_i_nhwc
      keep downstream NHWC consumers and remove bridge transposes.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        name = str(base)
        suffix = 1
        while name in model_ir.tensors:
            name = f"{base}_{suffix}"
            suffix += 1
        return name

    def _is_channel_slice_axis1(
        *,
        op: OperatorIR,
        source_shape_nchw: Optional[List[int]] = None,
    ) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3 or len(op.outputs) != 1:
            return False
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        if int(size_vals[1]) <= 0:
            return False
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return False
        if source_shape_nchw is not None and len(source_shape_nchw) == 4:
            c = int(source_shape_nchw[1])
            if int(begin_vals[1]) < 0 or int(begin_vals[1]) + int(size_vals[1]) > c:
                return False
        return True

    def _rewrite_slice_axis1_to_axis3(slice_op: OperatorIR) -> bool:
        begin_tensor = model_ir.tensors.get(str(slice_op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(slice_op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
        new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_idx: int,
        mul_op: OperatorIR,
        data_input_name: str,
        consumers: Dict[str, List[int]],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_index: Optional[int] = None
        for i, name in enumerate(mul_inputs):
            if str(name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(name), None)
            if tensor is not None and tensor.data is not None:
                const_index = int(i)
                break
        if const_index is None:
            return False

        const_name = str(mul_inputs[int(const_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        data_tensor = model_ir.tensors.get(str(data_input_name), None)
        if data_tensor is None or not _is_fully_known_positive_shape(list(data_tensor.shape)):
            return False
        target_shape = [int(v) for v in list(data_tensor.shape)]

        const_shape = [int(v) for v in list(const_data.shape)]
        if _broadcast_static_shapes(target_shape, const_shape) is not None:
            return True

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        elif int(const_data.ndim) == 3:
            candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        if rotated is None:
            return False

        const_users = [int(v) for v in consumers.get(const_name, [])]
        if any(int(v) != int(mul_idx) for v in const_users):
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return True

    def _is_safe_direct_post_rewire(
        *,
        post_output_name: str,
        replacement_name: str,
        consumers: Dict[str, List[int]],
    ) -> bool:
        replacement_tensor = model_ir.tensors.get(str(replacement_name), None)
        replacement_shape = (
            [int(v) for v in list(replacement_tensor.shape)]
            if replacement_tensor is not None
            else None
        )
        binary_ops = {
            "ADD",
            "SUB",
            "MUL",
            "DIV",
            "MAXIMUM",
            "MINIMUM",
            "SQUARED_DIFFERENCE",
            "POW",
        }
        for consumer_idx in [int(v) for v in consumers.get(str(post_output_name), [])]:
            consumer_op = model_ir.operators[int(consumer_idx)]
            consumer_inputs = [str(v) for v in list(consumer_op.inputs)]
            replace_positions = [i for i, n in enumerate(consumer_inputs) if str(n) == str(post_output_name)]
            if len(replace_positions) == 0:
                continue
            if len(consumer_inputs) <= 1:
                continue
            if str(consumer_op.op_type) not in binary_ops:
                return False
            if replacement_shape is None or not _is_fully_known_positive_shape(replacement_shape):
                return False
            for pos in replace_positions:
                for other_pos, other_name in enumerate(consumer_inputs):
                    if int(other_pos) == int(pos):
                        continue
                    other_tensor = model_ir.tensors.get(str(other_name), None)
                    if other_tensor is None or not _is_fully_known_positive_shape(list(other_tensor.shape)):
                        return False
                    other_shape = [int(v) for v in list(other_tensor.shape)]
                    if (
                        _broadcast_static_shapes(replacement_shape, other_shape) is None
                        and _broadcast_static_shapes(other_shape, replacement_shape) is None
                    ):
                        return False
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            pre_out_tensor = model_ir.tensors.get(pre_out_name, None)
            if pre_out_tensor is None or len(list(pre_out_tensor.shape)) != 4:
                continue

            slice_indices = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(slice_indices) == 0:
                continue
            if not all(
                _is_channel_slice_axis1(
                    op=model_ir.operators[int(slice_idx)],
                    source_shape_nchw=[int(v) for v in list(pre_out_tensor.shape)],
                )
                for slice_idx in slice_indices
            ):
                continue

            remove_indices: set[int] = {int(pre_idx)}
            rewritable = True
            for slice_idx in slice_indices:
                slice_op = model_ir.operators[int(slice_idx)]
                slice_out_name = str(slice_op.outputs[0])
                if slice_out_name in model_outputs:
                    rewritable = False
                    break

                if not _rewrite_slice_axis1_to_axis3(slice_op):
                    rewritable = False
                    break
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=slice_op,
                    input_index=0,
                    new_input_name=pre_in_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(slice_out_name, None),
                    perm_nchw_to_nhwc,
                )

                slice_users = [int(v) for v in consumers.get(slice_out_name, []) if int(v) != int(slice_idx)]
                if len(slice_users) == 0:
                    continue

                if len(slice_users) == 1:
                    user_idx = int(slice_users[0])
                    user_op = model_ir.operators[int(user_idx)]
                    if (
                        str(user_op.op_type) == "TRANSPOSE"
                        and len(user_op.inputs) >= 2
                        and len(user_op.outputs) == 1
                        and str(user_op.inputs[0]) == slice_out_name
                        and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        if not _is_safe_direct_post_rewire(
                            post_output_name=str(user_op.outputs[0]),
                            replacement_name=slice_out_name,
                            consumers=consumers,
                        ):
                            rewritable = False
                            break
                        _replace_tensor_inputs(model_ir, str(user_op.outputs[0]), slice_out_name)
                        remove_indices.add(int(user_idx))
                        continue

                    if (
                        str(user_op.op_type) == "MUL"
                        and len(user_op.inputs) == 2
                        and len(user_op.outputs) == 1
                        and slice_out_name in [str(v) for v in list(user_op.inputs)]
                        and str(user_op.outputs[0]) not in model_outputs
                    ):
                        mul_out_name = str(user_op.outputs[0])
                        mul_users = [int(v) for v in consumers.get(mul_out_name, []) if int(v) != int(user_idx)]
                        if len(mul_users) != 1:
                            rewritable = False
                            break
                        post_idx = int(mul_users[0])
                        post_op = model_ir.operators[int(post_idx)]
                        if (
                            str(post_op.op_type) != "TRANSPOSE"
                            or len(post_op.inputs) < 2
                            or len(post_op.outputs) != 1
                            or str(post_op.inputs[0]) != mul_out_name
                            or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                            or str(post_op.outputs[0]) in model_outputs
                        ):
                            rewritable = False
                            break
                        if not _rewrite_mul_const_to_nhwc_if_needed(
                            mul_idx=int(user_idx),
                            mul_op=user_op,
                            data_input_name=slice_out_name,
                            consumers=consumers,
                        ):
                            rewritable = False
                            break
                        _permute_tensor_metadata_if_rank_matches(
                            model_ir.tensors.get(mul_out_name, None),
                            perm_nchw_to_nhwc,
                        )
                        _replace_tensor_inputs(model_ir, str(post_op.outputs[0]), mul_out_name)
                        remove_indices.add(int(post_idx))
                        continue

                rewritable = False
                break

            if not rewritable:
                continue

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_channel_slice_muladd_nhwc_bridge_chains": int(optimized)}


def _optimize_transpose_channel_slice_dual_add_bridges_strict(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Strictly fold:
      NHWC->NCHW TRANSPOSE -> 2x channel SLICE
        branch-A: SLICE -> MUL(const) -> TRANSPOSE -> ADD
        branch-B: SLICE -> TRANSPOSE -> ADD
    back to NHWC.

    This is intentionally strict to avoid unsafe late-stage rewrites.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _is_axis1_channel_slice(op: OperatorIR) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3 or len(op.outputs) != 1:
            return False
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        if int(size_vals[1]) <= 0:
            return False
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return False
        return True

    def _rewrite_slice_axis1_to_axis3(slice_op: OperatorIR) -> bool:
        begin_tensor = model_ir.tensors.get(str(slice_op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(slice_op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
        new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_idx: int,
        mul_op: OperatorIR,
        data_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_index: Optional[int] = None
        for i, name in enumerate(mul_inputs):
            if str(name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(name), None)
            if tensor is not None and tensor.data is not None:
                const_index = int(i)
                break
        if const_index is None:
            return False

        const_name = str(mul_inputs[int(const_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        if not _is_fully_known_positive_shape(target_shape_nhwc):
            return False
        target_shape = [int(v) for v in list(target_shape_nhwc)]

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        elif int(const_data.ndim) == 3:
            candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        else:
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(const_data.shape)]) is not None:
                rotated = np.asarray(const_data)
        if rotated is None:
            return False

        const_users = [int(v) for v in consumers.get(const_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in const_users)
        if shared_outside_chain:
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]

        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_in_name in model_outputs or pre_out_name in model_outputs:
                continue

            slice_indices = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(slice_indices) != 2:
                continue
            if not all(_is_axis1_channel_slice(model_ir.operators[int(slice_idx)]) for slice_idx in slice_indices):
                continue

            branch_plans: List[Dict[str, Any]] = []
            branch_kind_set: set[str] = set()
            valid = True

            for slice_idx in sorted(slice_indices):
                slice_op = model_ir.operators[int(slice_idx)]
                slice_out_name = str(slice_op.outputs[0])
                if slice_out_name in model_outputs:
                    valid = False
                    break

                slice_users = [int(v) for v in consumers.get(slice_out_name, []) if int(v) != int(slice_idx)]
                if len(slice_users) != 1:
                    valid = False
                    break
                first_idx = int(slice_users[0])
                first_op = model_ir.operators[int(first_idx)]

                if (
                    str(first_op.op_type) == "TRANSPOSE"
                    and len(first_op.inputs) >= 2
                    and len(first_op.outputs) == 1
                    and str(first_op.inputs[0]) == slice_out_name
                    and _read_transpose_perm(model_ir, first_op) == perm_nchw_to_nhwc
                    and str(first_op.outputs[0]) not in model_outputs
                ):
                    post_out_name = str(first_op.outputs[0])
                    post_users = [int(v) for v in consumers.get(post_out_name, [])]
                    if len(post_users) != 1:
                        valid = False
                        break
                    add_idx = int(post_users[0])
                    add_op = model_ir.operators[int(add_idx)]
                    if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                        valid = False
                        break
                    add_inputs = [str(v) for v in list(add_op.inputs)]
                    if str(post_out_name) not in add_inputs:
                        valid = False
                        break
                    other_name = add_inputs[1] if add_inputs[0] == str(post_out_name) else add_inputs[0]
                    other_tensor = model_ir.tensors.get(str(other_name), None)
                    slice_tensor = model_ir.tensors.get(slice_out_name, None)
                    target_shape_nhwc = (
                        _permute_shape(list(slice_tensor.shape), perm_nchw_to_nhwc)
                        if slice_tensor is not None and len(list(slice_tensor.shape)) == 4
                        else None
                    )
                    if (
                        other_tensor is None
                        or target_shape_nhwc is None
                        or not _is_fully_known_positive_shape(list(other_tensor.shape))
                        or _broadcast_static_shapes(
                            [int(v) for v in list(target_shape_nhwc)],
                            [int(v) for v in list(other_tensor.shape)],
                        ) is None
                    ):
                        valid = False
                        break
                    branch_kind_set.add("direct")
                    branch_plans.append(
                        {
                            "kind": "direct",
                            "slice_idx": int(slice_idx),
                            "slice_out_name": str(slice_out_name),
                            "post_idx": int(first_idx),
                            "post_out_name": str(post_out_name),
                        }
                    )
                    continue

                if (
                    str(first_op.op_type) == "MUL"
                    and len(first_op.inputs) == 2
                    and len(first_op.outputs) == 1
                    and slice_out_name in [str(v) for v in list(first_op.inputs)]
                    and str(first_op.outputs[0]) not in model_outputs
                ):
                    mul_idx = int(first_idx)
                    mul_op = first_op
                    mul_out_name = str(mul_op.outputs[0])
                    mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
                    if len(mul_users) != 1:
                        valid = False
                        break
                    post_idx = int(mul_users[0])
                    post_op = model_ir.operators[int(post_idx)]
                    if (
                        str(post_op.op_type) != "TRANSPOSE"
                        or len(post_op.inputs) < 2
                        or len(post_op.outputs) != 1
                        or str(post_op.inputs[0]) != mul_out_name
                        or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
                        or str(post_op.outputs[0]) in model_outputs
                    ):
                        valid = False
                        break
                    post_out_name = str(post_op.outputs[0])
                    post_users = [int(v) for v in consumers.get(post_out_name, [])]
                    if len(post_users) != 1:
                        valid = False
                        break
                    add_idx = int(post_users[0])
                    add_op = model_ir.operators[int(add_idx)]
                    if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                        valid = False
                        break
                    add_inputs = [str(v) for v in list(add_op.inputs)]
                    if str(post_out_name) not in add_inputs:
                        valid = False
                        break
                    other_name = add_inputs[1] if add_inputs[0] == str(post_out_name) else add_inputs[0]
                    other_tensor = model_ir.tensors.get(str(other_name), None)
                    slice_tensor = model_ir.tensors.get(slice_out_name, None)
                    target_shape_nhwc = (
                        _permute_shape(list(slice_tensor.shape), perm_nchw_to_nhwc)
                        if slice_tensor is not None and len(list(slice_tensor.shape)) == 4
                        else None
                    )
                    if (
                        other_tensor is None
                        or target_shape_nhwc is None
                        or not _is_fully_known_positive_shape(list(other_tensor.shape))
                        or _broadcast_static_shapes(
                            [int(v) for v in list(target_shape_nhwc)],
                            [int(v) for v in list(other_tensor.shape)],
                        ) is None
                    ):
                        valid = False
                        break
                    branch_kind_set.add("mul")
                    branch_plans.append(
                        {
                            "kind": "mul",
                            "slice_idx": int(slice_idx),
                            "slice_out_name": str(slice_out_name),
                            "mul_idx": int(mul_idx),
                            "mul_out_name": str(mul_out_name),
                            "post_idx": int(post_idx),
                            "post_out_name": str(post_out_name),
                            "target_shape_nhwc": [int(v) for v in list(target_shape_nhwc)],
                        }
                    )
                    continue

                valid = False
                break

            if not valid:
                continue
            if branch_kind_set != {"direct", "mul"}:
                continue

            remove_indices: set[int] = {int(pre_idx)}
            chain_indices: set[int] = {int(pre_idx)}
            chain_indices.update(int(v) for v in slice_indices)
            for plan in branch_plans:
                if "mul_idx" in plan:
                    chain_indices.add(int(plan["mul_idx"]))
                chain_indices.add(int(plan["post_idx"]))

            apply_ok = True
            for slice_idx in sorted(slice_indices):
                slice_op = model_ir.operators[int(slice_idx)]
                if not _rewrite_slice_axis1_to_axis3(slice_op):
                    apply_ok = False
                    break
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=slice_op,
                    input_index=0,
                    new_input_name=pre_in_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(slice_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
            if not apply_ok:
                continue

            for plan in branch_plans:
                if str(plan["kind"]) == "mul":
                    mul_idx = int(plan["mul_idx"])
                    mul_op = model_ir.operators[int(mul_idx)]
                    if not _rewrite_mul_const_to_nhwc_if_needed(
                        mul_idx=int(mul_idx),
                        mul_op=mul_op,
                        data_input_name=str(plan["slice_out_name"]),
                        target_shape_nhwc=[int(v) for v in list(plan["target_shape_nhwc"])],
                        consumers=consumers,
                        chain_indices=chain_indices,
                    ):
                        apply_ok = False
                        break
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(plan["mul_out_name"]), None),
                        perm_nchw_to_nhwc,
                    )
                    _replace_tensor_inputs(
                        model_ir,
                        str(plan["post_out_name"]),
                        str(plan["mul_out_name"]),
                    )
                    remove_indices.add(int(plan["post_idx"]))
                else:
                    _replace_tensor_inputs(
                        model_ir,
                        str(plan["post_out_name"]),
                        str(plan["slice_out_name"]),
                    )
                    remove_indices.add(int(plan["post_idx"]))

            if not apply_ok:
                continue

            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_channel_slice_dual_add_bridges_strict": int(optimized)}


def _optimize_transpose_slice_muladd_conv_mergeadd_strict(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Strictly fold the following into a single post-merge adapter:
      T0: NHWC->NCHW
      S0/S1: channel slices from T0 output
      S0 -> MUL -> T1(NCHW->NHWC) -> ADD(const) -> CONV_2D -> T2(NHWC->NCHW)
      merge ADD(T2_out, S1_out) -> y_nchw

    Rewrite:
      S0/S1 read from T0 input in NHWC (axis=3 channel slices)
      S0 branch stays NHWC through MUL/ADD/CONV
      merge ADD runs in NHWC
      one new NHWC->NCHW transpose is inserted after merge ADD.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _find_or_create_nhwc_to_nchw_perm_tensor() -> str:
        target = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
        for tensor_name, tensor in model_ir.tensors.items():
            if (
                tensor is None
                or tensor.data is None
                or len(list(tensor.shape)) != 1
                or int(tensor.shape[0]) != 4
            ):
                continue
            data = np.asarray(tensor.data)
            if data.shape == (4,) and data.dtype == np.int32 and np.array_equal(data, target):
                return str(tensor_name)
        perm_name = _unique_tensor_name("__nhwc_to_nchw_perm_rank4__")
        model_ir.tensors[perm_name] = TensorIR(
            name=perm_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(target, dtype=np.int32),
            is_variable=False,
        )
        return perm_name

    def _is_axis1_channel_slice(op: OperatorIR) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3 or len(op.outputs) != 1:
            return False
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        if int(size_vals[1]) <= 0:
            return False
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return False
        return True

    def _rewrite_slice_axis1_to_axis3(slice_op: OperatorIR) -> bool:
        begin_tensor = model_ir.tensors.get(str(slice_op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(slice_op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
        new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_idx: int,
        mul_op: OperatorIR,
        data_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_index: Optional[int] = None
        for i, name in enumerate(mul_inputs):
            if str(name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(name), None)
            if tensor is not None and tensor.data is not None:
                const_index = int(i)
                break
        if const_index is None:
            return False

        const_name = str(mul_inputs[int(const_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        if not _is_fully_known_positive_shape(target_shape_nhwc):
            return False
        target_shape = [int(v) for v in list(target_shape_nhwc)]

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        elif int(const_data.ndim) == 3:
            candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        else:
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(const_data.shape)]) is not None:
                rotated = np.asarray(const_data)
        if rotated is None:
            return False

        const_users = [int(v) for v in consumers.get(const_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in const_users)
        if shared_outside_chain:
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue

            slice_indices = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(slice_indices) != 2:
                continue
            if not all(_is_axis1_channel_slice(model_ir.operators[int(v)]) for v in slice_indices):
                continue

            branch_a_plan: Optional[Dict[str, Any]] = None
            branch_b_plan: Optional[Dict[str, Any]] = None
            merge_add_idx: Optional[int] = None

            valid = True
            for slice_idx in sorted(slice_indices):
                slice_op = model_ir.operators[int(slice_idx)]
                slice_out_name = str(slice_op.outputs[0])
                if slice_out_name in model_outputs:
                    valid = False
                    break

                users1 = [int(v) for v in consumers.get(slice_out_name, []) if int(v) != int(slice_idx)]
                if len(users1) != 1:
                    valid = False
                    break
                first_idx = int(users1[0])
                first_op = model_ir.operators[int(first_idx)]

                # Branch-B: slice -> merge ADD directly
                if str(first_op.op_type) == "ADD" and len(first_op.inputs) == 2 and len(first_op.outputs) == 1:
                    if merge_add_idx is None:
                        merge_add_idx = int(first_idx)
                    elif int(merge_add_idx) != int(first_idx):
                        valid = False
                        break
                    if str(first_op.outputs[0]) in model_outputs:
                        valid = False
                        break
                    if branch_b_plan is not None:
                        valid = False
                        break
                    branch_b_plan = {
                        "slice_idx": int(slice_idx),
                        "slice_out_name": str(slice_out_name),
                        "merge_add_idx": int(first_idx),
                    }
                    continue

                # Branch-A: slice -> MUL -> T -> ADD -> CONV_2D -> T -> merge ADD
                if str(first_op.op_type) != "MUL" or len(first_op.inputs) != 2 or len(first_op.outputs) != 1:
                    valid = False
                    break
                mul_idx = int(first_idx)
                mul_out_name = str(first_op.outputs[0])
                if mul_out_name in model_outputs:
                    valid = False
                    break
                mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
                if len(mul_users) != 1:
                    valid = False
                    break
                t1_idx = int(mul_users[0])
                t1_op = model_ir.operators[int(t1_idx)]
                if (
                    str(t1_op.op_type) != "TRANSPOSE"
                    or len(t1_op.inputs) < 2
                    or len(t1_op.outputs) != 1
                    or str(t1_op.inputs[0]) != mul_out_name
                    or _read_transpose_perm(model_ir, t1_op) != perm_nchw_to_nhwc
                ):
                    valid = False
                    break
                t1_out_name = str(t1_op.outputs[0])
                if t1_out_name in model_outputs:
                    valid = False
                    break
                t1_users = [int(v) for v in consumers.get(t1_out_name, [])]
                if len(t1_users) != 1:
                    valid = False
                    break
                add0_idx = int(t1_users[0])
                add0_op = model_ir.operators[int(add0_idx)]
                if str(add0_op.op_type) != "ADD" or len(add0_op.inputs) != 2 or len(add0_op.outputs) != 1:
                    valid = False
                    break
                add0_out_name = str(add0_op.outputs[0])
                if add0_out_name in model_outputs:
                    valid = False
                    break
                add0_inputs = [str(v) for v in list(add0_op.inputs)]
                if str(t1_out_name) not in add0_inputs:
                    valid = False
                    break
                add0_side_name = add0_inputs[1] if add0_inputs[0] == str(t1_out_name) else add0_inputs[0]
                add0_side_tensor = model_ir.tensors.get(str(add0_side_name), None)
                if add0_side_tensor is None or add0_side_tensor.data is None:
                    valid = False
                    break

                conv_users = [int(v) for v in consumers.get(add0_out_name, [])]
                if len(conv_users) != 1:
                    valid = False
                    break
                conv_idx = int(conv_users[0])
                conv_op = model_ir.operators[int(conv_idx)]
                if (
                    str(conv_op.op_type) != "CONV_2D"
                    or len(conv_op.inputs) < 2
                    or len(conv_op.outputs) != 1
                    or str(conv_op.inputs[0]) != add0_out_name
                ):
                    valid = False
                    break
                conv_out_name = str(conv_op.outputs[0])
                if conv_out_name in model_outputs:
                    valid = False
                    break
                conv_out_users = [int(v) for v in consumers.get(conv_out_name, [])]
                if len(conv_out_users) != 1:
                    valid = False
                    break
                t2_idx = int(conv_out_users[0])
                t2_op = model_ir.operators[int(t2_idx)]
                if (
                    str(t2_op.op_type) != "TRANSPOSE"
                    or len(t2_op.inputs) < 2
                    or len(t2_op.outputs) != 1
                    or str(t2_op.inputs[0]) != conv_out_name
                    or _read_transpose_perm(model_ir, t2_op) != perm_nhwc_to_nchw
                ):
                    valid = False
                    break
                t2_out_name = str(t2_op.outputs[0])
                if t2_out_name in model_outputs:
                    valid = False
                    break
                merge_users = [int(v) for v in consumers.get(t2_out_name, [])]
                if len(merge_users) != 1:
                    valid = False
                    break
                cand_merge_add_idx = int(merge_users[0])
                merge_add_op = model_ir.operators[int(cand_merge_add_idx)]
                if str(merge_add_op.op_type) != "ADD" or len(merge_add_op.inputs) != 2 or len(merge_add_op.outputs) != 1:
                    valid = False
                    break
                if merge_add_idx is None:
                    merge_add_idx = int(cand_merge_add_idx)
                elif int(merge_add_idx) != int(cand_merge_add_idx):
                    valid = False
                    break
                if branch_a_plan is not None:
                    valid = False
                    break
                branch_a_plan = {
                    "slice_idx": int(slice_idx),
                    "slice_out_name": str(slice_out_name),
                    "mul_idx": int(mul_idx),
                    "mul_op": first_op,
                    "mul_out_name": str(mul_out_name),
                    "t1_idx": int(t1_idx),
                    "t1_out_name": str(t1_out_name),
                    "add0_idx": int(add0_idx),
                    "add0_side_name": str(add0_side_name),
                    "conv_idx": int(conv_idx),
                    "conv_out_name": str(conv_out_name),
                    "t2_idx": int(t2_idx),
                    "t2_out_name": str(t2_out_name),
                    "merge_add_idx": int(cand_merge_add_idx),
                }

            if not valid or branch_a_plan is None or branch_b_plan is None or merge_add_idx is None:
                continue

            merge_add_op = model_ir.operators[int(merge_add_idx)]
            merge_out_name = str(merge_add_op.outputs[0])
            if merge_out_name in model_outputs:
                continue
            if set(int(v) for v in consumers.get(pre_out_name, [])) != set(int(v) for v in slice_indices):
                continue
            if set(int(v) for v in consumers.get(str(branch_b_plan["slice_out_name"]), [])) != {int(merge_add_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(branch_a_plan["t2_out_name"]), [])) != {int(merge_add_idx)}:
                continue

            a_slice_tensor = model_ir.tensors.get(str(branch_a_plan["slice_out_name"]), None)
            conv_out_tensor = model_ir.tensors.get(str(branch_a_plan["conv_out_name"]), None)
            if (
                a_slice_tensor is None
                or conv_out_tensor is None
                or len(list(a_slice_tensor.shape)) != 4
                or len(list(conv_out_tensor.shape)) != 4
            ):
                continue
            a_slice_shape_nhwc = _permute_shape(list(a_slice_tensor.shape), perm_nchw_to_nhwc)
            b_slice_tensor = model_ir.tensors.get(str(branch_b_plan["slice_out_name"]), None)
            if b_slice_tensor is None or len(list(b_slice_tensor.shape)) != 4:
                continue
            b_slice_shape_nhwc = _permute_shape(list(b_slice_tensor.shape), perm_nchw_to_nhwc)
            conv_shape_nhwc = [int(v) for v in list(conv_out_tensor.shape)]
            if _broadcast_static_shapes(conv_shape_nhwc, b_slice_shape_nhwc) is None:
                continue
            add0_side_tensor = model_ir.tensors.get(str(branch_a_plan["add0_side_name"]), None)
            if (
                add0_side_tensor is None
                or not _is_fully_known_positive_shape(list(add0_side_tensor.shape))
                or _broadcast_static_shapes(
                    a_slice_shape_nhwc,
                    [int(v) for v in list(add0_side_tensor.shape)],
                ) is None
            ):
                continue

            chain_indices = {
                int(pre_idx),
                int(branch_a_plan["slice_idx"]),
                int(branch_b_plan["slice_idx"]),
                int(branch_a_plan["mul_idx"]),
                int(branch_a_plan["t1_idx"]),
                int(branch_a_plan["add0_idx"]),
                int(branch_a_plan["conv_idx"]),
                int(branch_a_plan["t2_idx"]),
                int(merge_add_idx),
            }
            if not _rewrite_mul_const_to_nhwc_if_needed(
                mul_idx=int(branch_a_plan["mul_idx"]),
                mul_op=branch_a_plan["mul_op"],
                data_input_name=str(branch_a_plan["slice_out_name"]),
                target_shape_nhwc=[int(v) for v in list(a_slice_shape_nhwc)],
                consumers=consumers,
                chain_indices=chain_indices,
            ):
                continue

            apply_ok = True
            for slice_idx in sorted(slice_indices):
                slice_op = model_ir.operators[int(slice_idx)]
                if not _rewrite_slice_axis1_to_axis3(slice_op):
                    apply_ok = False
                    break
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=slice_op,
                    input_index=0,
                    new_input_name=pre_in_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(slice_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
            if not apply_ok:
                continue

            # Branch-A rewires: remove T1/T2 and keep path in NHWC.
            _replace_tensor_inputs(
                model_ir,
                str(branch_a_plan["t1_out_name"]),
                str(branch_a_plan["mul_out_name"]),
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(branch_a_plan["mul_out_name"]), None),
                perm_nchw_to_nhwc,
            )
            _replace_tensor_inputs(
                model_ir,
                str(branch_a_plan["t2_out_name"]),
                str(branch_a_plan["conv_out_name"]),
            )

            # Merge ADD output is now NHWC, then insert one post adapter back to NCHW.
            merge_out_name = str(merge_add_op.outputs[0])
            merge_out_tensor = model_ir.tensors.get(str(merge_out_name), None)
            if merge_out_tensor is None or len(list(merge_out_tensor.shape)) != 4:
                continue
            nhwc_merge_out_name = _unique_tensor_name(f"{merge_out_name}_nhwc")
            nhwc_shape = _permute_shape(list(merge_out_tensor.shape), perm_nchw_to_nhwc)
            model_ir.tensors[nhwc_merge_out_name] = TensorIR(
                name=nhwc_merge_out_name,
                dtype=str(merge_out_tensor.dtype),
                shape=[int(v) for v in list(nhwc_shape)],
                shape_signature=[int(v) for v in list(nhwc_shape)],
                data=None,
                is_variable=False,
                quantization=_clone_quantization(merge_out_tensor.quantization),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=merge_add_op,
                new_outputs=[nhwc_merge_out_name],
            )

            perm_tensor_name = _find_or_create_nhwc_to_nchw_perm_tensor()
            post_adapter = OperatorIR(
                op_type="TRANSPOSE",
                inputs=[nhwc_merge_out_name, perm_tensor_name],
                outputs=[merge_out_name],
                options={},
            )

            remove_indices = {
                int(pre_idx),
                int(branch_a_plan["t1_idx"]),
                int(branch_a_plan["t2_idx"]),
            }
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]
            model_ir.operators.append(post_adapter)

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_slice_muladd_conv_mergeadd_strict": int(optimized)}


def _optimize_transpose_slice_muladd_mergeadd_posttranspose_strict(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Strictly fold:
      T0(NHWC->NCHW) -> two channel SLICE branches
        branch-A: SLICE -> MUL -> T1(NCHW->NHWC) -> ...
        branch-B: SLICE ------------------------------+
      ... -> T2(NHWC->NCHW) -------------------------+-> ADD(y_nchw)
      y_nchw -> T3(NCHW->NHWC) (plus optional legacy NCHW consumers)

    into NHWC and keep a single localized legacy adapter when required.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _find_or_create_nhwc_to_nchw_perm_tensor() -> str:
        target = np.asarray(perm_nhwc_to_nchw, dtype=np.int32)
        for tensor_name, tensor in model_ir.tensors.items():
            if (
                tensor is None
                or tensor.data is None
                or len(list(tensor.shape)) != 1
                or int(tensor.shape[0]) != 4
            ):
                continue
            data = np.asarray(tensor.data)
            if data.shape == (4,) and data.dtype == np.int32 and np.array_equal(data, target):
                return str(tensor_name)
        perm_name = _unique_tensor_name("__nhwc_to_nchw_perm_rank4__")
        model_ir.tensors[perm_name] = TensorIR(
            name=perm_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(target, dtype=np.int32),
            is_variable=False,
        )
        return perm_name

    def _is_axis1_channel_slice(op: OperatorIR) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3 or len(op.outputs) != 1:
            return False
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        size_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        if int(size_vals[1]) <= 0:
            return False
        if int(begin_vals[2]) != 0 or int(begin_vals[3]) != 0:
            return False
        return True

    def _rewrite_slice_axis1_to_axis3(slice_op: OperatorIR) -> bool:
        begin_tensor = model_ir.tensors.get(str(slice_op.inputs[1]), None)
        size_tensor = model_ir.tensors.get(str(slice_op.inputs[2]), None)
        begin_vals = _read_const_ints_from_tensor(begin_tensor)
        size_vals = _read_const_ints_from_tensor(size_tensor)
        if begin_vals is None or size_vals is None or len(begin_vals) != 4 or len(size_vals) != 4:
            return False
        new_begin = [int(begin_vals[0]), int(begin_vals[2]), int(begin_vals[3]), int(begin_vals[1])]
        new_size = [int(size_vals[0]), int(size_vals[2]), int(size_vals[3]), int(size_vals[1])]
        _write_const_ints_to_tensor(begin_tensor, new_begin)
        _write_const_ints_to_tensor(size_tensor, new_size)
        return True

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_idx: int,
        mul_op: OperatorIR,
        data_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_index: Optional[int] = None
        for i, name in enumerate(mul_inputs):
            if str(name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(name), None)
            if tensor is not None and tensor.data is not None:
                const_index = int(i)
                break
        if const_index is None:
            return False

        const_name = str(mul_inputs[int(const_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        if not _is_fully_known_positive_shape(target_shape_nhwc):
            return False
        target_shape = [int(v) for v in list(target_shape_nhwc)]

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        elif int(const_data.ndim) == 3:
            candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(candidate.shape)]) is not None:
                rotated = np.asarray(candidate)
        else:
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(const_data.shape)]) is not None:
                rotated = np.asarray(const_data)
        if rotated is None:
            return False

        const_users = [int(v) for v in consumers.get(const_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in const_users)
        if shared_outside_chain:
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return True

    def _find_reachable_to_nchw_merge_candidate(
        *,
        start_tensor_name: str,
        expected_merge_add_idx: Optional[int],
        model_outputs: set[str],
    ) -> Optional[Dict[str, Any]]:
        visited_tensors: set[str] = set()
        stack: List[str] = [str(start_tensor_name)]
        candidate: Optional[Dict[str, Any]] = None
        steps = 0

        while len(stack) > 0 and steps < 4096:
            steps += 1
            tensor_name = str(stack.pop())
            if tensor_name in visited_tensors:
                continue
            visited_tensors.add(tensor_name)

            user_indices = [int(v) for v in consumers.get(tensor_name, [])]
            for user_idx in user_indices:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == str(tensor_name)
                    and _read_transpose_perm(model_ir, user_op) == perm_nhwc_to_nchw
                ):
                    to_nchw_out_name = str(user_op.outputs[0])
                    to_nchw_out_users = [int(v) for v in consumers.get(to_nchw_out_name, [])]
                    if len(to_nchw_out_users) != 1:
                        continue
                    cand_merge_add_idx = int(to_nchw_out_users[0])
                    cand_merge_add_op = model_ir.operators[int(cand_merge_add_idx)]
                    if (
                        str(cand_merge_add_op.op_type) != "ADD"
                        or len(cand_merge_add_op.inputs) != 2
                        or len(cand_merge_add_op.outputs) != 1
                    ):
                        continue
                    if expected_merge_add_idx is not None and int(cand_merge_add_idx) != int(expected_merge_add_idx):
                        continue
                    this_candidate = {
                        "to_nchw_idx": int(user_idx),
                        "to_nchw_out_name": str(to_nchw_out_name),
                        "path_nhwc_tensor_name": str(tensor_name),
                        "merge_add_idx": int(cand_merge_add_idx),
                    }
                    if candidate is None:
                        candidate = this_candidate
                    else:
                        if int(candidate["to_nchw_idx"]) != int(this_candidate["to_nchw_idx"]):
                            return None

                for produced_name in [str(v) for v in list(user_op.outputs)]:
                    if produced_name in visited_tensors or produced_name in model_outputs:
                        continue
                    stack.append(str(produced_name))

        return candidate

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw:
                continue

            pre_in_name = str(pre_op.inputs[0])
            pre_out_name = str(pre_op.outputs[0])
            if pre_out_name in model_outputs:
                continue

            slice_indices = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(slice_indices) != 2:
                continue
            if not all(_is_axis1_channel_slice(model_ir.operators[int(v)]) for v in slice_indices):
                continue

            branch_a: Optional[Dict[str, Any]] = None
            branch_b: Optional[Dict[str, Any]] = None
            merge_add_idx: Optional[int] = None
            valid = True

            for slice_idx in sorted(slice_indices):
                slice_op = model_ir.operators[int(slice_idx)]
                slice_out_name = str(slice_op.outputs[0])
                if slice_out_name in model_outputs:
                    valid = False
                    break
                users1 = [int(v) for v in consumers.get(slice_out_name, []) if int(v) != int(slice_idx)]
                if len(users1) != 1:
                    valid = False
                    break
                first_idx = int(users1[0])
                first_op = model_ir.operators[int(first_idx)]

                if str(first_op.op_type) == "ADD" and len(first_op.inputs) == 2 and len(first_op.outputs) == 1:
                    if merge_add_idx is None:
                        merge_add_idx = int(first_idx)
                    elif int(merge_add_idx) != int(first_idx):
                        valid = False
                        break
                    if branch_b is not None:
                        valid = False
                        break
                    branch_b = {
                        "slice_idx": int(slice_idx),
                        "slice_out_name": str(slice_out_name),
                        "merge_add_idx": int(first_idx),
                    }
                    continue

                if str(first_op.op_type) != "MUL" or len(first_op.inputs) != 2 or len(first_op.outputs) != 1:
                    valid = False
                    break
                mul_idx = int(first_idx)
                mul_out_name = str(first_op.outputs[0])
                mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
                if len(mul_users) != 1:
                    valid = False
                    break
                t1_idx = int(mul_users[0])
                t1_op = model_ir.operators[int(t1_idx)]
                if (
                    str(t1_op.op_type) != "TRANSPOSE"
                    or len(t1_op.inputs) < 2
                    or len(t1_op.outputs) != 1
                    or str(t1_op.inputs[0]) != mul_out_name
                    or _read_transpose_perm(model_ir, t1_op) != perm_nchw_to_nhwc
                ):
                    valid = False
                    break
                t1_out_name = str(t1_op.outputs[0])
                t1_users = [int(v) for v in consumers.get(t1_out_name, [])]
                if len(t1_users) != 1:
                    valid = False
                    break
                reachable = _find_reachable_to_nchw_merge_candidate(
                    start_tensor_name=str(t1_out_name),
                    expected_merge_add_idx=merge_add_idx,
                    model_outputs=model_outputs,
                )
                if reachable is None:
                    valid = False
                    break
                nhwc_to_nchw_idx = int(reachable["to_nchw_idx"])
                nhwc_to_nchw_out_name = str(reachable["to_nchw_out_name"])
                cand_merge_add_idx = int(reachable["merge_add_idx"])
                if merge_add_idx is None:
                    merge_add_idx = int(cand_merge_add_idx)
                elif int(merge_add_idx) != int(cand_merge_add_idx):
                    valid = False
                    break

                if branch_a is not None:
                    valid = False
                    break
                branch_a = {
                    "slice_idx": int(slice_idx),
                    "slice_out_name": str(slice_out_name),
                    "mul_idx": int(mul_idx),
                    "mul_op": first_op,
                    "mul_out_name": str(mul_out_name),
                    "t1_idx": int(t1_idx),
                    "t1_out_name": str(t1_out_name),
                    "path_nhwc_tensor_name": str(reachable["path_nhwc_tensor_name"]),
                    "to_nchw_idx": int(nhwc_to_nchw_idx),
                    "to_nchw_out_name": str(nhwc_to_nchw_out_name),
                    "merge_add_idx": int(cand_merge_add_idx),
                }

            if not valid or branch_a is None or branch_b is None or merge_add_idx is None:
                continue
            merge_add_op = model_ir.operators[int(merge_add_idx)]
            merge_out_name = str(merge_add_op.outputs[0])
            if merge_out_name in model_outputs:
                continue

            if set(int(v) for v in consumers.get(pre_out_name, [])) != set(int(v) for v in slice_indices):
                continue
            if set(int(v) for v in consumers.get(str(branch_b["slice_out_name"]), [])) != {int(merge_add_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(branch_a["to_nchw_out_name"]), [])) != {int(merge_add_idx)}:
                continue

            # Merge output must have at least one NCHW->NHWC tail to fold.
            post_transpose_indices: List[int] = []
            post_transpose_out_names: List[str] = []
            post_mul_transpose_plans: List[Dict[str, Any]] = []
            legacy_consumer_indices: List[int] = []
            for user_idx in [int(v) for v in consumers.get(merge_out_name, [])]:
                user_op = model_ir.operators[int(user_idx)]
                if (
                    str(user_op.op_type) == "TRANSPOSE"
                    and len(user_op.inputs) >= 2
                    and len(user_op.outputs) == 1
                    and str(user_op.inputs[0]) == merge_out_name
                    and _read_transpose_perm(model_ir, user_op) == perm_nchw_to_nhwc
                    and str(user_op.outputs[0]) not in model_outputs
                ):
                    post_transpose_indices.append(int(user_idx))
                    post_transpose_out_names.append(str(user_op.outputs[0]))
                elif (
                    str(user_op.op_type) == "MUL"
                    and len(user_op.inputs) == 2
                    and len(user_op.outputs) == 1
                ):
                    mul_inputs = [str(v) for v in list(user_op.inputs)]
                    if str(merge_out_name) == mul_inputs[0]:
                        mul_data_input_index = 0
                        mul_const_input_index = 1
                    elif str(merge_out_name) == mul_inputs[1]:
                        mul_data_input_index = 1
                        mul_const_input_index = 0
                    else:
                        legacy_consumer_indices.append(int(user_idx))
                        continue
                    mul_const_name = str(mul_inputs[int(mul_const_input_index)])
                    mul_const_tensor = model_ir.tensors.get(str(mul_const_name), None)
                    if mul_const_tensor is None or mul_const_tensor.data is None:
                        legacy_consumer_indices.append(int(user_idx))
                        continue
                    mul_out_name = str(user_op.outputs[0])
                    if mul_out_name in model_outputs:
                        legacy_consumer_indices.append(int(user_idx))
                        continue
                    mul_out_users = [int(v) for v in consumers.get(mul_out_name, [])]
                    if len(mul_out_users) != 1:
                        legacy_consumer_indices.append(int(user_idx))
                        continue
                    mul_post_idx = int(mul_out_users[0])
                    mul_post_op = model_ir.operators[int(mul_post_idx)]
                    if (
                        str(mul_post_op.op_type) != "TRANSPOSE"
                        or len(mul_post_op.inputs) < 2
                        or len(mul_post_op.outputs) != 1
                        or str(mul_post_op.inputs[0]) != str(mul_out_name)
                        or _read_transpose_perm(model_ir, mul_post_op) != perm_nchw_to_nhwc
                        or str(mul_post_op.outputs[0]) in model_outputs
                    ):
                        legacy_consumer_indices.append(int(user_idx))
                        continue
                    post_mul_transpose_plans.append(
                        {
                            "mul_idx": int(user_idx),
                            "mul_data_input_index": int(mul_data_input_index),
                            "mul_const_name": str(mul_const_name),
                            "mul_out_name": str(mul_out_name),
                            "post_idx": int(mul_post_idx),
                            "post_out_name": str(mul_post_op.outputs[0]),
                        }
                    )
                else:
                    legacy_consumer_indices.append(int(user_idx))
            # Accept strict legacy-only tails as well. Even when merge output has
            # no NCHW->NHWC post adapters, we can still fold pre/T1/T2 transposes
            # and preserve existing NCHW consumers through one localized adapter.
            if (
                len(post_transpose_indices) == 0
                and len(post_mul_transpose_plans) == 0
                and len(legacy_consumer_indices) == 0
            ):
                continue

            b_slice_tensor = model_ir.tensors.get(str(branch_b["slice_out_name"]), None)
            a_nhwc_tensor = model_ir.tensors.get(str(branch_a["path_nhwc_tensor_name"]), None)
            if (
                b_slice_tensor is None
                or a_nhwc_tensor is None
                or len(list(b_slice_tensor.shape)) != 4
                or len(list(a_nhwc_tensor.shape)) != 4
            ):
                continue
            b_shape_nhwc = _permute_shape(list(b_slice_tensor.shape), perm_nchw_to_nhwc)
            a_shape_nhwc = [int(v) for v in list(a_nhwc_tensor.shape)]
            if _broadcast_static_shapes(a_shape_nhwc, b_shape_nhwc) is None:
                continue

            chain_indices: set[int] = {
                int(pre_idx),
                int(branch_a["slice_idx"]),
                int(branch_b["slice_idx"]),
                int(branch_a["mul_idx"]),
                int(branch_a["t1_idx"]),
                int(branch_a["to_nchw_idx"]),
                int(merge_add_idx),
            }
            chain_indices.update(int(v) for v in post_transpose_indices)
            chain_indices.update(int(v["mul_idx"]) for v in post_mul_transpose_plans)
            chain_indices.update(int(v["post_idx"]) for v in post_mul_transpose_plans)

            a_slice_tensor = model_ir.tensors.get(str(branch_a["slice_out_name"]), None)
            a_slice_shape_nhwc = (
                _permute_shape(list(a_slice_tensor.shape), perm_nchw_to_nhwc)
                if a_slice_tensor is not None and len(list(a_slice_tensor.shape)) == 4
                else None
            )
            if not _rewrite_mul_const_to_nhwc_if_needed(
                mul_idx=int(branch_a["mul_idx"]),
                mul_op=branch_a["mul_op"],
                data_input_name=str(branch_a["slice_out_name"]),
                target_shape_nhwc=a_slice_shape_nhwc,
                consumers=consumers,
                chain_indices=chain_indices,
            ):
                continue
            post_mul_rewrite_ok = True
            for post_mul_plan in post_mul_transpose_plans:
                post_mul_op = model_ir.operators[int(post_mul_plan["mul_idx"])]
                if not _rewrite_mul_const_to_nhwc_if_needed(
                    mul_idx=int(post_mul_plan["mul_idx"]),
                    mul_op=post_mul_op,
                    data_input_name=str(merge_out_name),
                    target_shape_nhwc=[int(v) for v in list(a_shape_nhwc)],
                    consumers=consumers,
                    chain_indices=chain_indices,
                ):
                    post_mul_rewrite_ok = False
                    break
            if not post_mul_rewrite_ok:
                continue

            apply_ok = True
            for slice_idx in sorted(slice_indices):
                slice_op = model_ir.operators[int(slice_idx)]
                if not _rewrite_slice_axis1_to_axis3(slice_op):
                    apply_ok = False
                    break
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=slice_op,
                    input_index=0,
                    new_input_name=pre_in_name,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(slice_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
            if not apply_ok:
                continue

            _replace_tensor_inputs(
                model_ir,
                str(branch_a["t1_out_name"]),
                str(branch_a["mul_out_name"]),
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(branch_a["mul_out_name"]), None),
                perm_nchw_to_nhwc,
            )

            # Merge ADD in NHWC.
            merge_out_tensor_nchw = model_ir.tensors.get(merge_out_name, None)
            merge_out_shape_nchw = (
                [int(v) for v in list(merge_out_tensor_nchw.shape)]
                if merge_out_tensor_nchw is not None and merge_out_tensor_nchw.shape is not None
                else None
            )
            merge_out_sig_nchw = (
                [int(v) for v in list(merge_out_tensor_nchw.shape_signature)]
                if merge_out_tensor_nchw is not None and merge_out_tensor_nchw.shape_signature is not None
                else None
            )
            merge_out_dtype = str(merge_out_tensor_nchw.dtype) if merge_out_tensor_nchw is not None else None
            merge_out_quant = (
                _clone_quantization(merge_out_tensor_nchw.quantization)
                if merge_out_tensor_nchw is not None
                else None
            )
            _replace_tensor_inputs(
                model_ir,
                str(branch_a["to_nchw_out_name"]),
                str(branch_a["path_nhwc_tensor_name"]),
            )

            if len(post_transpose_out_names) > 0:
                canonical_nhwc_out_name = str(post_transpose_out_names[0])
            else:
                canonical_nhwc_out_name = _unique_tensor_name(f"{merge_out_name}_nhwc")
                if canonical_nhwc_out_name not in model_ir.tensors:
                    model_ir.tensors[canonical_nhwc_out_name] = TensorIR(
                        name=canonical_nhwc_out_name,
                        dtype=str(merge_out_dtype) if merge_out_dtype is not None else "FLOAT32",
                        shape=[],
                        shape_signature=[],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(merge_out_quant),
                    )
            _set_operator_outputs(
                model_ir=model_ir,
                op=merge_add_op,
                new_outputs=[canonical_nhwc_out_name],
            )
            for alias_name in post_transpose_out_names[1:]:
                _replace_tensor_inputs(model_ir, alias_name, canonical_nhwc_out_name)
            for post_mul_plan in post_mul_transpose_plans:
                post_mul_idx = int(post_mul_plan["mul_idx"])
                post_mul_op = model_ir.operators[int(post_mul_idx)]
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=post_mul_op,
                    input_index=int(post_mul_plan["mul_data_input_index"]),
                    new_input_name=str(canonical_nhwc_out_name),
                )
                old_mul_out_name = str(post_mul_plan["mul_out_name"])
                post_out_name = str(post_mul_plan["post_out_name"])
                old_mul_tensor = model_ir.tensors.get(old_mul_out_name, None)
                post_out_tensor = model_ir.tensors.get(post_out_name, None)
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=post_mul_op,
                    new_outputs=[post_out_name],
                )
                if old_mul_tensor is not None and post_out_tensor is not None:
                    post_out_tensor.dtype = str(old_mul_tensor.dtype)
                    post_out_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                    post_out_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                    post_out_tensor.shape_signature = (
                        [int(v) for v in list(old_mul_tensor.shape_signature)]
                        if old_mul_tensor.shape_signature is not None
                        else [int(v) for v in list(old_mul_tensor.shape)]
                    )
                    _permute_tensor_metadata_if_rank_matches(
                        post_out_tensor,
                        perm_nchw_to_nhwc,
                    )

            canonical_nhwc_tensor = model_ir.tensors.get(canonical_nhwc_out_name, None)
            if canonical_nhwc_tensor is not None:
                if merge_out_dtype is not None:
                    canonical_nhwc_tensor.dtype = str(merge_out_dtype)
                if merge_out_quant is not None:
                    canonical_nhwc_tensor.quantization = _clone_quantization(merge_out_quant)
                if merge_out_shape_nchw is not None and len(merge_out_shape_nchw) == 4:
                    canonical_nhwc_tensor.shape = _permute_shape(
                        list(merge_out_shape_nchw),
                        perm_nchw_to_nhwc,
                    )
                    base_sig = (
                        list(merge_out_sig_nchw)
                        if merge_out_sig_nchw is not None and len(merge_out_sig_nchw) == 4
                        else list(merge_out_shape_nchw)
                    )
                    canonical_nhwc_tensor.shape_signature = _permute_shape(
                        list(base_sig),
                        perm_nchw_to_nhwc,
                    )

            # Keep legacy NCHW consumers via one localized adapter.
            if len(legacy_consumer_indices) > 0 or merge_out_name in model_outputs:
                perm_tensor_name = _find_or_create_nhwc_to_nchw_perm_tensor()
                adapter_op = OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[canonical_nhwc_out_name, perm_tensor_name],
                    outputs=[merge_out_name],
                    options={},
                )
                model_ir.operators.append(adapter_op)

            remove_indices = {
                int(pre_idx),
                int(branch_a["t1_idx"]),
                int(branch_a["to_nchw_idx"]),
            }
            remove_indices.update(int(v) for v in post_transpose_indices)
            remove_indices.update(int(v["post_idx"]) for v in post_mul_transpose_plans)
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_slice_muladd_mergeadd_posttranspose_strict": int(optimized)}


def _optimize_concat_mul_add_transpose_nhwc_bridge_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate concat->mul->transpose->add NHWC bridge tails when concat inputs
    are NHWC->NCHW adapters.

    Target:
      x0_nhwc --TRANSPOSE(0,3,1,2)--> x0_nchw \
      x1_nhwc --TRANSPOSE(0,3,1,2)--> x1_nchw  +-> CONCAT(axis=1) -> c_nchw
      ...                                       /
      MUL(c_nchw, const_nchw) -> m_nchw
      TRANSPOSE(m_nchw, 0,2,3,1) -> m_nhwc
      ADD(m_nhwc, const_nhwc) -> y_nhwc

    Rewrite:
      CONCAT(axis=3, [x0_nhwc, x1_nhwc, ...]) -> c_nhwc
      MUL(c_nhwc, const_nhwc) -> m_nhwc
      ADD(m_nhwc, const_nhwc) -> y_nhwc
      remove adapter/pre/post TRANSPOSE operators.

    Notes:
    - If CONCAT output has legacy NCHW consumers, preserve them via one local
      adapter TRANSPOSE(cat_nhwc->cat_nchw).
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _find_or_create_nhwc_to_nchw_perm_tensor() -> str:
        base_name = "__concat_mul_tail_nhwc_to_nchw_perm_rank4__"
        tensor = model_ir.tensors.get(base_name, None)
        if tensor is not None and tensor.data is not None:
            arr = np.asarray(tensor.data).reshape(-1)
            if arr.size == 4 and [int(v) for v in arr.tolist()] == perm_nhwc_to_nchw:
                return str(base_name)
        model_ir.tensors[base_name] = TensorIR(
            name=base_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )
        return str(base_name)

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_idx: int,
        mul_op: OperatorIR,
        data_input_name: str,
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
        target_shape_nhwc: Optional[List[int]],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_input_index: Optional[int] = None
        for i, input_name in enumerate(mul_inputs):
            if str(input_name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(input_name), None)
            if tensor is not None and tensor.data is not None:
                const_input_index = int(i)
                break
        if const_input_index is None:
            return False

        const_name = str(mul_inputs[int(const_input_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False

        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        if _is_fully_known_positive_shape(target_shape_nhwc):
            target_shape = [int(v) for v in list(target_shape_nhwc)]
        else:
            data_tensor = model_ir.tensors.get(str(data_input_name), None)
            target_shape = (
                [int(v) for v in list(data_tensor.shape)]
                if data_tensor is not None and _is_fully_known_positive_shape(list(data_tensor.shape))
                else None
            )

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            if target_shape is None or _broadcast_static_shapes(
                target_shape,
                [int(v) for v in list(rotated_candidate.shape)],
            ) is not None:
                rotated = np.asarray(rotated_candidate)
        elif int(const_data.ndim) == 3:
            rotated_candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            if target_shape is None or _broadcast_static_shapes(
                target_shape,
                [int(v) for v in list(rotated_candidate.shape)],
            ) is not None:
                rotated = np.asarray(rotated_candidate)
        elif target_shape is not None:
            if _broadcast_static_shapes(
                target_shape,
                [int(v) for v in list(const_data.shape)],
            ) is not None:
                rotated = np.asarray(const_data)

        if rotated is None:
            return False

        const_users = [int(v) for v in consumers.get(const_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in const_users)
        if shared_outside_chain:
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_input_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]

        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            mul_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if mul_out_name in model_outputs or post_out_name in model_outputs:
                continue

            mul_idx = producers.get(mul_out_name, None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(mul_op.outputs[0]) != mul_out_name:
                continue
            if set(int(v) for v in consumers.get(mul_out_name, [])) != {int(post_idx)}:
                continue

            add_users = [int(v) for v in consumers.get(post_out_name, [])]
            if len(add_users) != 1:
                continue
            add_idx = int(add_users[0])
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_inputs = [str(v) for v in list(add_op.inputs)]
            if str(post_out_name) not in add_inputs:
                continue
            add_out_name = str(add_op.outputs[0])

            add_side_name = add_inputs[1] if add_inputs[0] == str(post_out_name) else add_inputs[0]
            add_side_tensor = model_ir.tensors.get(str(add_side_name), None)
            if add_side_tensor is None or add_side_tensor.data is None:
                continue

            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            concat_idx: Optional[int] = None
            concat_out_name: Optional[str] = None
            concat_data_input_name: Optional[str] = None
            concat_data_input_index: Optional[int] = None
            for input_index, mul_input_name in enumerate(mul_inputs):
                producer_idx = producers.get(str(mul_input_name), None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if str(producer_op.op_type) != "CONCATENATION" or len(producer_op.outputs) != 1:
                    continue
                axis = int(producer_op.options.get("axis", 1))
                if axis < 0:
                    axis += 4
                if axis != 1:
                    continue
                if str(producer_op.outputs[0]) != str(mul_input_name):
                    continue
                concat_idx = int(producer_idx)
                concat_out_name = str(mul_input_name)
                concat_data_input_name = str(mul_input_name)
                concat_data_input_index = int(input_index)
                break

            if (
                concat_idx is None
                or concat_out_name is None
                or concat_data_input_name is None
                or concat_data_input_index is None
            ):
                continue

            concat_op = model_ir.operators[int(concat_idx)]
            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue

            pre_indices: List[int] = []
            new_concat_inputs: List[str] = []
            valid = True
            for concat_input_name in concat_inputs:
                pre_idx = producers.get(str(concat_input_name), None)
                if pre_idx is None:
                    valid = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(concat_input_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                ):
                    valid = False
                    break
                if str(concat_input_name) in model_outputs:
                    valid = False
                    break
                if set(int(v) for v in consumers.get(str(concat_input_name), [])) != {int(concat_idx)}:
                    valid = False
                    break
                source_name = str(pre_op.inputs[0])
                source_tensor = model_ir.tensors.get(source_name, None)
                if source_tensor is None:
                    valid = False
                    break
                pre_indices.append(int(pre_idx))
                new_concat_inputs.append(str(source_name))
            if not valid:
                continue

            chain_indices = {int(post_idx), int(mul_idx), int(add_idx), int(concat_idx)} | set(int(v) for v in pre_indices)
            target_shape_nhwc: Optional[List[int]] = None
            concat_out_tensor = model_ir.tensors.get(str(concat_out_name), None)
            if concat_out_tensor is not None and len(list(concat_out_tensor.shape)) == 4:
                target_shape_nhwc = _permute_shape(list(concat_out_tensor.shape), perm_nchw_to_nhwc)
            add_side_data = np.asarray(add_side_tensor.data)
            add_side_shape = [int(v) for v in list(add_side_data.shape)]
            if int(add_side_data.size) != 1:
                if int(add_side_data.ndim) != 4:
                    continue
                if target_shape_nhwc is not None and _broadcast_static_shapes(target_shape_nhwc, add_side_shape) is None:
                    continue
                # Keep rewrite strict: accept only NHWC channelwise bias constants.
                if not (
                    int(add_side_shape[0]) == 1
                    and int(add_side_shape[1]) == 1
                    and int(add_side_shape[2]) == 1
                    and int(add_side_shape[3]) > 0
                ):
                    continue
            if not _rewrite_mul_const_to_nhwc_if_needed(
                mul_idx=int(mul_idx),
                mul_op=mul_op,
                data_input_name=str(concat_data_input_name),
                consumers=consumers,
                chain_indices=chain_indices,
                target_shape_nhwc=target_shape_nhwc,
            ):
                continue

            concat_users = [int(v) for v in consumers.get(str(concat_out_name), [])]
            legacy_concat_users = [int(v) for v in concat_users if int(v) != int(mul_idx)]
            canonical_concat_out_name = str(concat_out_name)
            if len(legacy_concat_users) > 0 or str(concat_out_name) in model_outputs:
                canonical_concat_out_name = _unique_tensor_name(f"{concat_out_name}_nhwc")
                old_concat_tensor = model_ir.tensors.get(str(concat_out_name), None)
                if old_concat_tensor is not None:
                    old_shape = [int(v) for v in list(old_concat_tensor.shape)]
                    old_sig = (
                        [int(v) for v in list(old_concat_tensor.shape_signature)]
                        if old_concat_tensor.shape_signature is not None
                        else [int(v) for v in list(old_concat_tensor.shape)]
                    )
                    new_shape = (
                        _permute_shape(old_shape, perm_nchw_to_nhwc)
                        if len(old_shape) == 4
                        else [int(v) for v in list(old_shape)]
                    )
                    new_sig = (
                        _permute_shape(old_sig, perm_nchw_to_nhwc)
                        if len(old_sig) == 4
                        else [int(v) for v in list(old_sig)]
                    )
                    model_ir.tensors[canonical_concat_out_name] = TensorIR(
                        name=canonical_concat_out_name,
                        dtype=str(old_concat_tensor.dtype),
                        shape=[int(v) for v in list(new_shape)],
                        shape_signature=[int(v) for v in list(new_sig)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(old_concat_tensor.quantization),
                    )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=concat_op,
                    new_outputs=[str(canonical_concat_out_name)],
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=mul_op,
                    input_index=int(concat_data_input_index),
                    new_input_name=str(canonical_concat_out_name),
                )
            else:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(concat_out_name), None),
                    perm_nchw_to_nhwc,
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=new_concat_inputs,
            )
            concat_op.options["axis"] = 3

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(mul_out_name, None),
                perm_nchw_to_nhwc,
            )

            if str(add_op.inputs[0]) == str(post_out_name):
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=add_op,
                    input_index=0,
                    new_input_name=mul_out_name,
                )
            else:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=add_op,
                    input_index=1,
                    new_input_name=mul_out_name,
                )

            remove_indices = set(int(v) for v in pre_indices)
            remove_indices.add(int(post_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            if len(legacy_concat_users) > 0 or str(concat_out_name) in model_outputs:
                adapter_perm_name = _find_or_create_nhwc_to_nchw_perm_tensor()
                model_ir.operators.append(
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(canonical_concat_out_name), str(adapter_perm_name)],
                        outputs=[str(concat_out_name)],
                    )
                )

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_concat_mul_add_transpose_nhwc_bridge_chains": int(optimized)}


def _optimize_concat_mul_add_transpose_add_nhwc_bridge_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate DenseNet-like NHWC<->NCHW bridge tails:

      leaf_i_nhwc --TRANSPOSE(0,3,1,2)--> leaf_i_nchw
      CONCAT(axis=1, leaf_i_nchw) -> cat_nchw
      MUL(cat_nchw, const_nchw) -> m_nchw
      ADD(m_nchw, const_nchw) -> a_nchw
      TRANSPOSE(a_nchw, 0,2,3,1) -> a_nhwc
      ADD(a_nhwc, const_nhwc) -> y_nhwc

    Rewrite:
      CONCAT(axis=3, leaf_i_nhwc) -> cat_nhwc
      MUL(cat_nhwc, const_nhwc) -> m_nhwc
      ADD(m_nhwc, const_nhwc) -> a_nhwc
      ADD(a_nhwc, const_nhwc) -> y_nhwc

    Notes:
    - If CONCAT output has legacy NCHW consumers, preserve them via one local
      adapter TRANSPOSE(cat_nhwc->cat_nchw).
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _find_or_create_nhwc_to_nchw_perm_tensor() -> str:
        base_name = "__concat_affine_tail_nhwc_to_nchw_perm_rank4__"
        tensor = model_ir.tensors.get(base_name, None)
        if tensor is not None and tensor.data is not None:
            arr = np.asarray(tensor.data).reshape(-1)
            if arr.size == 4 and [int(v) for v in arr.tolist()] == perm_nhwc_to_nchw:
                return str(base_name)
        model_ir.tensors[base_name] = TensorIR(
            name=base_name,
            dtype="INT32",
            shape=[4],
            shape_signature=[4],
            data=np.asarray(perm_nhwc_to_nchw, dtype=np.int32),
            is_variable=False,
            quantization=None,
        )
        return str(base_name)

    def _rewrite_binary_const_to_nhwc_if_needed(
        *,
        op: OperatorIR,
        const_input_index: int,
        const_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> Optional[str]:
        const_tensor = model_ir.tensors.get(str(const_input_name), None)
        if const_tensor is None or const_tensor.data is None:
            return None
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return str(const_input_name)

        target_shape = (
            [int(v) for v in list(target_shape_nhwc)]
            if _is_fully_known_positive_shape(target_shape_nhwc)
            else None
        )

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            as_is_shape = [int(v) for v in list(const_data.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, as_is_shape) is not None:
                return str(const_input_name)
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, rotated_shape) is None:
                return None
            rotated = np.asarray(rotated_candidate)
        elif int(const_data.ndim) == 3:
            as_is_shape = [int(v) for v in list(const_data.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, as_is_shape) is not None:
                return str(const_input_name)
            rotated_candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, rotated_shape) is None:
                return None
            rotated = np.asarray(rotated_candidate)
        else:
            if target_shape is not None and _broadcast_static_shapes(
                target_shape,
                [int(v) for v in list(const_data.shape)],
            ) is not None:
                return str(const_input_name)
            return None

        if rotated is None:
            return None

        const_users = [int(v) for v in consumers.get(str(const_input_name), [])]
        shared_outside = any(int(v) not in chain_indices for v in const_users)
        if shared_outside:
            rotated_name = _unique_tensor_name(f"{const_input_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            new_inputs = [str(v) for v in list(op.inputs)]
            new_inputs[int(const_input_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=op,
                new_inputs=new_inputs,
            )
            return str(rotated_name)

        const_tensor.data = np.asarray(rotated)
        const_tensor.shape = [int(v) for v in list(rotated.shape)]
        const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return str(const_input_name)

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            pre_add_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if pre_add_out_name in model_outputs or post_out_name in model_outputs:
                continue

            pre_add_idx = producers.get(str(pre_add_out_name), None)
            if pre_add_idx is None:
                continue
            pre_add_op = model_ir.operators[int(pre_add_idx)]
            if (
                str(pre_add_op.op_type) != "ADD"
                or len(pre_add_op.inputs) != 2
                or len(pre_add_op.outputs) != 1
                or str(pre_add_op.outputs[0]) != str(pre_add_out_name)
            ):
                continue
            if set(int(v) for v in consumers.get(str(pre_add_out_name), [])) != {int(post_idx)}:
                continue

            pre_add_inputs = [str(v) for v in list(pre_add_op.inputs)]
            pre_add_const_input_index: Optional[int] = None
            pre_add_const_name: Optional[str] = None
            pre_mul_data_name: Optional[str] = None
            pre_mul_idx: Optional[int] = None
            pre_mul_op: Optional[OperatorIR] = None
            pre_mul_data_input_index: Optional[int] = None
            pre_mul_const_input_index: Optional[int] = None
            pre_mul_const_name: Optional[str] = None
            concat_idx: Optional[int] = None
            concat_op: Optional[OperatorIR] = None
            concat_out_name: Optional[str] = None

            for pre_add_input_index, pre_add_input_name in enumerate(pre_add_inputs):
                candidate_mul_idx = producers.get(str(pre_add_input_name), None)
                if candidate_mul_idx is None:
                    continue
                candidate_mul_op = model_ir.operators[int(candidate_mul_idx)]
                if (
                    str(candidate_mul_op.op_type) != "MUL"
                    or len(candidate_mul_op.inputs) != 2
                    or len(candidate_mul_op.outputs) != 1
                    or str(candidate_mul_op.outputs[0]) != str(pre_add_input_name)
                ):
                    continue
                if set(int(v) for v in consumers.get(str(pre_add_input_name), [])) != {int(pre_add_idx)}:
                    continue
                side_name = pre_add_inputs[1 - int(pre_add_input_index)]
                side_tensor = model_ir.tensors.get(str(side_name), None)
                if side_tensor is None or side_tensor.data is None:
                    continue

                mul_inputs = [str(v) for v in list(candidate_mul_op.inputs)]
                for mul_data_input_index, mul_data_input_name in enumerate(mul_inputs):
                    candidate_concat_idx = producers.get(str(mul_data_input_name), None)
                    if candidate_concat_idx is None:
                        continue
                    candidate_concat_op = model_ir.operators[int(candidate_concat_idx)]
                    if (
                        str(candidate_concat_op.op_type) != "CONCATENATION"
                        or len(candidate_concat_op.outputs) != 1
                        or str(candidate_concat_op.outputs[0]) != str(mul_data_input_name)
                    ):
                        continue
                    axis = int(candidate_concat_op.options.get("axis", 1))
                    if axis < 0:
                        axis += 4
                    if axis != 1:
                        continue
                    if str(mul_data_input_name) in model_outputs:
                        continue

                    mul_const_input_index = 1 - int(mul_data_input_index)
                    mul_const_name = str(mul_inputs[int(mul_const_input_index)])
                    mul_const_tensor = model_ir.tensors.get(str(mul_const_name), None)
                    if mul_const_tensor is None or mul_const_tensor.data is None:
                        continue

                    pre_add_const_input_index = int(1 - int(pre_add_input_index))
                    pre_add_const_name = str(side_name)
                    pre_mul_data_name = str(mul_data_input_name)
                    pre_mul_idx = int(candidate_mul_idx)
                    pre_mul_op = candidate_mul_op
                    pre_mul_data_input_index = int(mul_data_input_index)
                    pre_mul_const_input_index = int(mul_const_input_index)
                    pre_mul_const_name = str(mul_const_name)
                    concat_idx = int(candidate_concat_idx)
                    concat_op = candidate_concat_op
                    concat_out_name = str(mul_data_input_name)
                    break
                if concat_idx is not None:
                    break

            if (
                pre_add_const_input_index is None
                or pre_add_const_name is None
                or pre_mul_data_name is None
                or pre_mul_idx is None
                or pre_mul_op is None
                or pre_mul_data_input_index is None
                or pre_mul_const_input_index is None
                or pre_mul_const_name is None
                or concat_idx is None
                or concat_op is None
                or concat_out_name is None
            ):
                continue

            tail_users = [int(v) for v in consumers.get(str(post_out_name), [])]
            if len(tail_users) != 1:
                continue
            tail_add_idx = int(tail_users[0])
            tail_add_op = model_ir.operators[int(tail_add_idx)]
            if str(tail_add_op.op_type) != "ADD" or len(tail_add_op.inputs) != 2 or len(tail_add_op.outputs) != 1:
                continue
            tail_add_inputs = [str(v) for v in list(tail_add_op.inputs)]
            if str(post_out_name) == tail_add_inputs[0]:
                tail_add_const_name = str(tail_add_inputs[1])
            elif str(post_out_name) == tail_add_inputs[1]:
                tail_add_const_name = str(tail_add_inputs[0])
            else:
                continue
            tail_add_const_tensor = model_ir.tensors.get(str(tail_add_const_name), None)
            if tail_add_const_tensor is None or tail_add_const_tensor.data is None:
                continue
            tail_add_const_data = np.asarray(tail_add_const_tensor.data)
            if int(tail_add_const_data.size) != 1:
                tail_add_shape = [int(v) for v in list(tail_add_const_data.shape)]
                if len(tail_add_shape) != 4:
                    continue
                if not (
                    int(tail_add_shape[0]) == 1
                    and int(tail_add_shape[1]) == 1
                    and int(tail_add_shape[2]) == 1
                    and int(tail_add_shape[3]) > 0
                ):
                    continue

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue
            pre_indices: List[int] = []
            new_concat_inputs: List[str] = []
            valid_pre = True
            for concat_input_name in concat_inputs:
                pre_idx = producers.get(str(concat_input_name), None)
                if pre_idx is None:
                    valid_pre = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(concat_input_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                    or str(concat_input_name) in model_outputs
                ):
                    valid_pre = False
                    break
                if set(int(v) for v in consumers.get(str(concat_input_name), [])) != {int(concat_idx)}:
                    valid_pre = False
                    break
                source_name = str(pre_op.inputs[0])
                if model_ir.tensors.get(str(source_name), None) is None:
                    valid_pre = False
                    break
                pre_indices.append(int(pre_idx))
                new_concat_inputs.append(str(source_name))
            if not valid_pre:
                continue

            concat_out_tensor = model_ir.tensors.get(str(concat_out_name), None)
            target_shape_nhwc: Optional[List[int]] = None
            if concat_out_tensor is not None and len(list(concat_out_tensor.shape)) == 4:
                target_shape_nhwc = _permute_shape(list(concat_out_tensor.shape), perm_nchw_to_nhwc)
            pre_add_out_tensor = model_ir.tensors.get(str(pre_add_out_name), None)
            if target_shape_nhwc is None and pre_add_out_tensor is not None and len(list(pre_add_out_tensor.shape)) == 4:
                target_shape_nhwc = _permute_shape(list(pre_add_out_tensor.shape), perm_nchw_to_nhwc)

            chain_indices = {int(concat_idx), int(pre_mul_idx), int(pre_add_idx), int(post_idx), int(tail_add_idx)}
            chain_indices.update(int(v) for v in pre_indices)

            if _rewrite_binary_const_to_nhwc_if_needed(
                op=pre_mul_op,
                const_input_index=int(pre_mul_const_input_index),
                const_input_name=str(pre_mul_const_name),
                target_shape_nhwc=target_shape_nhwc,
                consumers=consumers,
                chain_indices=chain_indices,
            ) is None:
                continue
            if _rewrite_binary_const_to_nhwc_if_needed(
                op=pre_add_op,
                const_input_index=int(pre_add_const_input_index),
                const_input_name=str(pre_add_const_name),
                target_shape_nhwc=target_shape_nhwc,
                consumers=consumers,
                chain_indices=chain_indices,
            ) is None:
                continue

            concat_users = [int(v) for v in consumers.get(str(concat_out_name), [])]
            legacy_concat_users = [int(v) for v in concat_users if int(v) != int(pre_mul_idx)]
            canonical_concat_out_name = str(concat_out_name)
            if len(legacy_concat_users) > 0 or str(concat_out_name) in model_outputs:
                canonical_concat_out_name = _unique_tensor_name(f"{concat_out_name}_nhwc")
                old_concat_tensor = model_ir.tensors.get(str(concat_out_name), None)
                if old_concat_tensor is not None:
                    old_shape = [int(v) for v in list(old_concat_tensor.shape)]
                    old_sig = (
                        [int(v) for v in list(old_concat_tensor.shape_signature)]
                        if old_concat_tensor.shape_signature is not None
                        else [int(v) for v in list(old_concat_tensor.shape)]
                    )
                    new_shape = (
                        _permute_shape(old_shape, perm_nchw_to_nhwc)
                        if len(old_shape) == 4
                        else [int(v) for v in list(old_shape)]
                    )
                    new_sig = (
                        _permute_shape(old_sig, perm_nchw_to_nhwc)
                        if len(old_sig) == 4
                        else [int(v) for v in list(old_sig)]
                    )
                    model_ir.tensors[canonical_concat_out_name] = TensorIR(
                        name=canonical_concat_out_name,
                        dtype=str(old_concat_tensor.dtype),
                        shape=[int(v) for v in list(new_shape)],
                        shape_signature=[int(v) for v in list(new_sig)],
                        data=None,
                        is_variable=False,
                        quantization=_clone_quantization(old_concat_tensor.quantization),
                    )
                _set_operator_outputs(
                    model_ir=model_ir,
                    op=concat_op,
                    new_outputs=[str(canonical_concat_out_name)],
                )
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=pre_mul_op,
                    input_index=int(pre_mul_data_input_index),
                    new_input_name=str(canonical_concat_out_name),
                )
            else:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(concat_out_name), None),
                    perm_nchw_to_nhwc,
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in list(new_concat_inputs)],
            )
            if not isinstance(concat_op.options, dict):
                concat_op.options = {}
            concat_op.options["axis"] = 3

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(pre_mul_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(pre_add_out_name), None),
                perm_nchw_to_nhwc,
            )

            if str(tail_add_op.inputs[0]) == str(post_out_name):
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=tail_add_op,
                    input_index=0,
                    new_input_name=str(pre_add_out_name),
                )
            else:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=tail_add_op,
                    input_index=1,
                    new_input_name=str(pre_add_out_name),
                )

            remove_indices = set(int(v) for v in pre_indices)
            remove_indices.add(int(post_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            if len(legacy_concat_users) > 0 or str(concat_out_name) in model_outputs:
                adapter_perm_name = _find_or_create_nhwc_to_nchw_perm_tensor()
                model_ir.operators.append(
                    OperatorIR(
                        op_type="TRANSPOSE",
                        inputs=[str(canonical_concat_out_name), str(adapter_perm_name)],
                        outputs=[str(concat_out_name)],
                    )
                )

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_concat_mul_add_transpose_add_nhwc_bridge_chains": int(optimized)}


def _optimize_concat_mul_add_add_mean_reshape_tail_nhwc_bridge_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate DenseNet-like terminal NCHW bridge tails:

      leaf_i_nhwc --TRANSPOSE(0,3,1,2)--> leaf_i_nchw
      CONCAT(axis=1, leaf_i_nchw) -> c_nchw
      MUL(c_nchw, const) -> m_nchw
      ADD(m_nchw, const) -> a0_nchw
      ADD(a0_nchw, const) -> a1_nchw
      MEAN(a1_nchw, axes_nchw, keepDims=True) -> mean_nchw
      RESHAPE(mean_nchw, shape4) -> out_nhwc

    Rewrite:
      CONCAT(axis=3, leaf_i_nhwc) -> c_nhwc
      MUL(c_nhwc, const_nhwc) -> m_nhwc
      ADD(m_nhwc, const_nhwc) -> a0_nhwc
      ADD(a0_nhwc, const_nhwc) -> a1_nhwc
      MEAN(a1_nhwc, axes_nhwc, keepDims=True) -> mean_nhwc
      RESHAPE(mean_nhwc, shape4_nhwc) -> out_nhwc
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _rewrite_binary_const_to_nhwc_if_needed(
        *,
        op: OperatorIR,
        const_input_index: int,
        const_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        chain_indices: set[int],
        consumers: Dict[str, List[int]],
    ) -> Optional[str]:
        const_tensor = model_ir.tensors.get(str(const_input_name), None)
        if const_tensor is None or const_tensor.data is None:
            return None
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return str(const_input_name)

        target_shape = (
            [int(v) for v in list(target_shape_nhwc)]
            if _is_fully_known_positive_shape(target_shape_nhwc)
            else None
        )

        current_shape = [int(v) for v in list(const_data.shape)]
        if target_shape is not None and _broadcast_static_shapes(target_shape, current_shape) is not None:
            return str(const_input_name)

        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is None or _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                rotated = np.asarray(rotated_candidate)
        elif int(const_data.ndim) == 3:
            rotated_candidate = np.transpose(const_data, [1, 2, 0]).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is None or _broadcast_static_shapes(target_shape, rotated_shape) is not None:
                rotated = np.asarray(rotated_candidate)
        else:
            return None

        if rotated is None:
            return None

        const_users = [int(v) for v in consumers.get(str(const_input_name), [])]
        shared_outside = any(int(v) not in chain_indices for v in const_users)
        if shared_outside:
            rotated_name = _unique_tensor_name(f"{const_input_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=op,
                input_index=int(const_input_index),
                new_input_name=str(rotated_name),
            )
            return str(rotated_name)

        const_tensor.data = np.asarray(rotated)
        const_tensor.shape = [int(v) for v in list(rotated.shape)]
        const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return str(const_input_name)

    def _rewrite_mean_axes_to_nhwc_if_needed(
        *,
        mean_op: OperatorIR,
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> bool:
        if len(mean_op.inputs) < 2:
            return False
        axes_input_name = str(mean_op.inputs[1])
        axes_tensor = model_ir.tensors.get(axes_input_name, None)
        axes_vals = _read_const_ints_from_tensor(axes_tensor)
        if axes_tensor is None or axes_vals is None:
            return False
        remapped_axes: List[int] = []
        for axis in axes_vals:
            norm_axis = int(axis)
            if norm_axis < 0:
                norm_axis += 4
            if norm_axis < 0 or norm_axis >= 4:
                return False
            # Axis mapping is old-dim -> new-dim (inverse of transpose perm).
            remapped_axes.append(int(perm_nhwc_to_nchw[int(norm_axis)]))

        axes_users = [int(v) for v in consumers.get(axes_input_name, [])]
        shared_outside = any(int(v) not in chain_indices for v in axes_users)
        if shared_outside:
            new_axes_name = _unique_tensor_name(f"{axes_input_name}_nhwc")
            model_ir.tensors[new_axes_name] = TensorIR(
                name=new_axes_name,
                dtype="INT32",
                shape=[int(len(remapped_axes))],
                shape_signature=[int(len(remapped_axes))],
                data=np.asarray(remapped_axes, dtype=np.int32),
                is_variable=False,
                quantization=None,
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=mean_op,
                input_index=1,
                new_input_name=str(new_axes_name),
            )
            return True
        return bool(_write_const_ints_to_tensor(axes_tensor, remapped_axes))

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for reshape_idx, reshape_op in enumerate(model_ir.operators):
            if str(reshape_op.op_type) != "RESHAPE" or len(reshape_op.inputs) < 2 or len(reshape_op.outputs) != 1:
                continue

            mean_out_name = str(reshape_op.inputs[0])
            reshape_out_name = str(reshape_op.outputs[0])
            if mean_out_name in model_outputs:
                continue

            mean_idx = producers.get(str(mean_out_name), None)
            if mean_idx is None:
                continue
            mean_op = model_ir.operators[int(mean_idx)]
            if (
                str(mean_op.op_type) != "MEAN"
                or len(mean_op.inputs) < 2
                or len(mean_op.outputs) != 1
                or str(mean_op.outputs[0]) != str(mean_out_name)
                or not bool(mean_op.options.get("keepDims", False))
            ):
                continue
            if set(int(v) for v in consumers.get(str(mean_out_name), [])) != {int(reshape_idx)}:
                continue

            add1_out_name = str(mean_op.inputs[0])
            if add1_out_name in model_outputs:
                continue
            add1_idx = producers.get(str(add1_out_name), None)
            if add1_idx is None:
                continue
            add1_op = model_ir.operators[int(add1_idx)]
            if (
                str(add1_op.op_type) != "ADD"
                or len(add1_op.inputs) != 2
                or len(add1_op.outputs) != 1
                or str(add1_op.outputs[0]) != str(add1_out_name)
            ):
                continue
            if set(int(v) for v in consumers.get(str(add1_out_name), [])) != {int(mean_idx)}:
                continue

            add1_inputs = [str(v) for v in list(add1_op.inputs)]
            add0_idx: Optional[int] = None
            add1_const_index: Optional[int] = None
            add1_const_name: Optional[str] = None
            add0_out_name: Optional[str] = None
            for i, candidate_name in enumerate(add1_inputs):
                candidate_idx = producers.get(str(candidate_name), None)
                if candidate_idx is None:
                    continue
                candidate_op = model_ir.operators[int(candidate_idx)]
                if (
                    str(candidate_op.op_type) != "ADD"
                    or len(candidate_op.inputs) != 2
                    or len(candidate_op.outputs) != 1
                    or str(candidate_op.outputs[0]) != str(candidate_name)
                ):
                    continue
                add0_idx = int(candidate_idx)
                add1_const_index = int(1 - int(i))
                add1_const_name = str(add1_inputs[int(add1_const_index)])
                add0_out_name = str(candidate_name)
                break
            if add0_idx is None or add1_const_index is None or add1_const_name is None or add0_out_name is None:
                continue
            if set(int(v) for v in consumers.get(str(add0_out_name), [])) != {int(add1_idx)}:
                continue

            add0_op = model_ir.operators[int(add0_idx)]
            add0_inputs = [str(v) for v in list(add0_op.inputs)]
            mul_idx: Optional[int] = None
            add0_const_index: Optional[int] = None
            add0_const_name: Optional[str] = None
            mul_out_name: Optional[str] = None
            for i, candidate_name in enumerate(add0_inputs):
                candidate_idx = producers.get(str(candidate_name), None)
                if candidate_idx is None:
                    continue
                candidate_op = model_ir.operators[int(candidate_idx)]
                if (
                    str(candidate_op.op_type) != "MUL"
                    or len(candidate_op.inputs) != 2
                    or len(candidate_op.outputs) != 1
                    or str(candidate_op.outputs[0]) != str(candidate_name)
                ):
                    continue
                mul_idx = int(candidate_idx)
                add0_const_index = int(1 - int(i))
                add0_const_name = str(add0_inputs[int(add0_const_index)])
                mul_out_name = str(candidate_name)
                break
            if mul_idx is None or add0_const_index is None or add0_const_name is None or mul_out_name is None:
                continue
            if set(int(v) for v in consumers.get(str(mul_out_name), [])) != {int(add0_idx)}:
                continue

            mul_op = model_ir.operators[int(mul_idx)]
            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            concat_idx: Optional[int] = None
            concat_data_index: Optional[int] = None
            concat_out_name: Optional[str] = None
            mul_const_index: Optional[int] = None
            mul_const_name: Optional[str] = None
            for i, candidate_name in enumerate(mul_inputs):
                candidate_idx = producers.get(str(candidate_name), None)
                if candidate_idx is None:
                    continue
                candidate_op = model_ir.operators[int(candidate_idx)]
                if (
                    str(candidate_op.op_type) != "CONCATENATION"
                    or len(candidate_op.outputs) != 1
                    or str(candidate_op.outputs[0]) != str(candidate_name)
                ):
                    continue
                axis = int(candidate_op.options.get("axis", 1))
                if axis < 0:
                    axis += 4
                if axis != 1:
                    continue
                concat_idx = int(candidate_idx)
                concat_data_index = int(i)
                concat_out_name = str(candidate_name)
                mul_const_index = int(1 - int(i))
                mul_const_name = str(mul_inputs[int(mul_const_index)])
                break
            if (
                concat_idx is None
                or concat_data_index is None
                or concat_out_name is None
                or mul_const_index is None
                or mul_const_name is None
            ):
                continue
            if set(int(v) for v in consumers.get(str(concat_out_name), [])) != {int(mul_idx)}:
                continue
            if concat_out_name in model_outputs:
                continue

            concat_op = model_ir.operators[int(concat_idx)]
            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue

            pre_indices: List[int] = []
            new_concat_inputs: List[str] = []
            valid_pre = True
            for concat_input_name in concat_inputs:
                pre_idx = producers.get(str(concat_input_name), None)
                if pre_idx is None:
                    valid_pre = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(concat_input_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                ):
                    valid_pre = False
                    break
                if str(concat_input_name) in model_outputs:
                    valid_pre = False
                    break
                if set(int(v) for v in consumers.get(str(concat_input_name), [])) != {int(concat_idx)}:
                    valid_pre = False
                    break
                source_name = str(pre_op.inputs[0])
                if model_ir.tensors.get(str(source_name), None) is None:
                    valid_pre = False
                    break
                pre_indices.append(int(pre_idx))
                new_concat_inputs.append(str(source_name))
            if not valid_pre:
                continue

            target_shape_nhwc: Optional[List[int]] = None
            concat_out_tensor = model_ir.tensors.get(str(concat_out_name), None)
            if concat_out_tensor is not None and len(list(concat_out_tensor.shape)) == 4:
                target_shape_nhwc = _permute_shape(list(concat_out_tensor.shape), perm_nchw_to_nhwc)
            elif len(new_concat_inputs) > 0:
                first_input_tensor = model_ir.tensors.get(str(new_concat_inputs[0]), None)
                if first_input_tensor is not None and len(list(first_input_tensor.shape)) == 4:
                    inferred_shape = [int(v) for v in list(first_input_tensor.shape)]
                    channel_sum = 0
                    channel_ok = True
                    for input_name in new_concat_inputs:
                        input_tensor = model_ir.tensors.get(str(input_name), None)
                        if input_tensor is None or len(list(input_tensor.shape)) != 4:
                            channel_ok = False
                            break
                        shape = [int(v) for v in list(input_tensor.shape)]
                        if shape[0] != inferred_shape[0] or shape[1] != inferred_shape[1] or shape[2] != inferred_shape[2]:
                            channel_ok = False
                            break
                        channel_sum += int(shape[3])
                    if channel_ok:
                        target_shape_nhwc = [
                            int(inferred_shape[0]),
                            int(inferred_shape[1]),
                            int(inferred_shape[2]),
                            int(channel_sum),
                        ]

            chain_indices = {
                int(concat_idx),
                int(mul_idx),
                int(add0_idx),
                int(add1_idx),
                int(mean_idx),
                int(reshape_idx),
            }
            chain_indices.update(int(v) for v in pre_indices)

            if _rewrite_binary_const_to_nhwc_if_needed(
                op=mul_op,
                const_input_index=int(mul_const_index),
                const_input_name=str(mul_const_name),
                target_shape_nhwc=target_shape_nhwc,
                chain_indices=chain_indices,
                consumers=consumers,
            ) is None:
                continue
            if _rewrite_binary_const_to_nhwc_if_needed(
                op=add0_op,
                const_input_index=int(add0_const_index),
                const_input_name=str(add0_const_name),
                target_shape_nhwc=target_shape_nhwc,
                chain_indices=chain_indices,
                consumers=consumers,
            ) is None:
                continue
            if _rewrite_binary_const_to_nhwc_if_needed(
                op=add1_op,
                const_input_index=int(add1_const_index),
                const_input_name=str(add1_const_name),
                target_shape_nhwc=target_shape_nhwc,
                chain_indices=chain_indices,
                consumers=consumers,
            ) is None:
                continue
            if not _rewrite_mean_axes_to_nhwc_if_needed(
                mean_op=mean_op,
                consumers=consumers,
                chain_indices=chain_indices,
            ):
                continue

            # If RESHAPE shape is exactly old NCHW mean shape, rewrite to NHWC shape.
            reshape_shape_tensor = model_ir.tensors.get(str(reshape_op.inputs[1]), None)
            reshape_shape_vals = _read_const_ints_from_tensor(reshape_shape_tensor)
            mean_out_tensor = model_ir.tensors.get(str(mean_out_name), None)
            if (
                reshape_shape_tensor is not None
                and reshape_shape_vals is not None
                and mean_out_tensor is not None
                and len(reshape_shape_vals) == 4
                and len(list(mean_out_tensor.shape)) == 4
            ):
                mean_out_shape_nchw = [int(v) for v in list(mean_out_tensor.shape)]
                mean_out_shape_nhwc = _permute_shape(list(mean_out_shape_nchw), perm_nchw_to_nhwc)
                if [int(v) for v in list(reshape_shape_vals)] == mean_out_shape_nchw:
                    _write_const_ints_to_tensor(reshape_shape_tensor, [int(v) for v in list(mean_out_shape_nhwc)])

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=[str(v) for v in list(new_concat_inputs)],
            )
            if not isinstance(concat_op.options, dict):
                concat_op.options = {}
            concat_op.options["axis"] = 3

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(add0_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(add1_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mean_out_name), None),
                perm_nchw_to_nhwc,
            )

            for remove_idx in sorted(list(set(int(v) for v in pre_indices)), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_concat_mul_add_add_mean_reshape_tail_nhwc_bridge_chains": int(optimized)}


def _optimize_concat_tree_mul_add_transpose_nhwc_bridge_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate nested CONCAT(axis=1) trees wrapped by NHWC<->NCHW adapters:

      leaves: x_i_nhwc --TRANSPOSE(0,3,1,2)--> x_i_nchw
      concat tree over x_i_nchw (axis=1) -> c_nchw
      MUL(c_nchw, const_nchw) -> m_nchw
      TRANSPOSE(m_nchw, 0,2,3,1) -> m_nhwc
      ADD(m_nhwc, const_nhwc) -> y_nhwc

    Rewrite:
      concat tree over x_i_nhwc (axis=3) -> c_nhwc
      MUL(c_nhwc, const_nhwc) -> m_nhwc
      ADD(m_nhwc, const_nhwc) -> y_nhwc
      remove leaf pre/post TRANSPOSE operators.
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _remap_nchw_axis_to_nhwc(axis: int) -> Optional[int]:
        remapped_axis = int(axis)
        if remapped_axis < 0:
            remapped_axis += 4
        if remapped_axis < 0 or remapped_axis >= 4:
            return None
        try:
            return int(perm_nchw_to_nhwc.index(int(remapped_axis)))
        except ValueError:
            return None

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_op: OperatorIR,
        data_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_input_index: Optional[int] = None
        for i, input_name in enumerate(mul_inputs):
            if str(input_name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(input_name), None)
            if tensor is not None and tensor.data is not None:
                const_input_index = int(i)
                break
        if const_input_index is None:
            return False

        const_name = str(mul_inputs[int(const_input_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        target_shape = (
            [int(v) for v in list(target_shape_nhwc)]
            if _is_fully_known_positive_shape(target_shape_nhwc)
            else None
        )
        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            as_is_shape = [int(v) for v in list(const_data.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, as_is_shape) is not None:
                return True
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, rotated_shape) is None:
                return False
            rotated = np.asarray(rotated_candidate)
        else:
            if target_shape is None:
                return False
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(const_data.shape)]) is None:
                return False
            return True

        if rotated is None:
            return False
        const_users = [int(v) for v in consumers.get(const_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in const_users)
        if shared_outside_chain:
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_input_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return True

    def _collect_concat_tree(
        *,
        root_concat_idx: int,
        mul_idx: int,
        consumers: Dict[str, List[int]],
        producers: Dict[str, int],
        model_outputs: set[str],
    ) -> Optional[Tuple[set[int], set[int]]]:
        concat_indices: set[int] = set()
        leaf_pre_indices: set[int] = set()
        stack: List[Tuple[int, int]] = [(int(root_concat_idx), int(mul_idx))]

        while len(stack) > 0:
            concat_idx, expected_consumer_idx = stack.pop()
            if int(concat_idx) in concat_indices:
                continue
            concat_op = model_ir.operators[int(concat_idx)]
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.outputs) != 1:
                return None
            axis = int(concat_op.options.get("axis", 1))
            if _remap_nchw_axis_to_nhwc(axis) is None:
                return None
            concat_out_name = str(concat_op.outputs[0])
            if concat_out_name in model_outputs:
                return None
            if set(int(v) for v in consumers.get(concat_out_name, [])) != {int(expected_consumer_idx)}:
                return None
            concat_indices.add(int(concat_idx))

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                return None
            for concat_input_name in concat_inputs:
                producer_idx = producers.get(str(concat_input_name), None)
                if producer_idx is None:
                    return None
                producer_op = model_ir.operators[int(producer_idx)]
                if (
                    str(producer_op.op_type) == "CONCATENATION"
                    and len(producer_op.outputs) == 1
                    and str(producer_op.outputs[0]) == str(concat_input_name)
                ):
                    stack.append((int(producer_idx), int(concat_idx)))
                    continue
                if (
                    str(producer_op.op_type) == "TRANSPOSE"
                    and len(producer_op.inputs) >= 2
                    and len(producer_op.outputs) == 1
                    and str(producer_op.outputs[0]) == str(concat_input_name)
                    and _read_transpose_perm(model_ir, producer_op) == perm_nhwc_to_nchw
                    and str(concat_input_name) not in model_outputs
                    and set(int(v) for v in consumers.get(str(concat_input_name), [])) == {int(concat_idx)}
                ):
                    leaf_pre_indices.add(int(producer_idx))
                    continue
                return None

        if len(leaf_pre_indices) < 2:
            return None
        return concat_indices, leaf_pre_indices

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue

            mul_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if mul_out_name in model_outputs or post_out_name in model_outputs:
                continue

            mul_idx = producers.get(str(mul_out_name), None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(mul_op.outputs[0]) != str(mul_out_name):
                continue
            if set(int(v) for v in consumers.get(str(mul_out_name), [])) != {int(post_idx)}:
                continue

            add_users = [int(v) for v in consumers.get(str(post_out_name), [])]
            if len(add_users) != 1:
                continue
            add_idx = int(add_users[0])
            add_op = model_ir.operators[int(add_idx)]
            if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                continue
            add_inputs = [str(v) for v in list(add_op.inputs)]
            if str(post_out_name) not in add_inputs:
                continue
            add_side_name = add_inputs[1] if add_inputs[0] == str(post_out_name) else add_inputs[0]
            add_side_tensor = model_ir.tensors.get(str(add_side_name), None)
            if add_side_tensor is None or add_side_tensor.data is None:
                continue
            add_side_data = np.asarray(add_side_tensor.data)
            if int(add_side_data.size) != 1:
                side_shape = [int(v) for v in list(add_side_data.shape)]
                if len(side_shape) != 4:
                    continue
                if not (
                    int(side_shape[0]) == 1
                    and int(side_shape[1]) == 1
                    and int(side_shape[2]) == 1
                    and int(side_shape[3]) > 0
                ):
                    continue

            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            root_concat_idx: Optional[int] = None
            root_concat_out_name: Optional[str] = None
            for mul_input_name in mul_inputs:
                producer_idx = producers.get(str(mul_input_name), None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if str(producer_op.op_type) != "CONCATENATION" or len(producer_op.outputs) != 1:
                    continue
                axis = int(producer_op.options.get("axis", 1))
                if axis < 0:
                    axis += 4
                if axis != 1:
                    continue
                if str(producer_op.outputs[0]) != str(mul_input_name):
                    continue
                root_concat_idx = int(producer_idx)
                root_concat_out_name = str(mul_input_name)
                break
            if root_concat_idx is None or root_concat_out_name is None:
                continue
            if set(int(v) for v in consumers.get(str(root_concat_out_name), [])) != {int(mul_idx)}:
                continue
            if str(root_concat_out_name) in model_outputs:
                continue

            tree = _collect_concat_tree(
                root_concat_idx=int(root_concat_idx),
                mul_idx=int(mul_idx),
                consumers=consumers,
                producers=producers,
                model_outputs=model_outputs,
            )
            if tree is None:
                continue
            concat_indices, leaf_pre_indices = tree
            concat_axis_updates: Dict[int, int] = {}
            axis_map_ok = True
            for concat_idx in concat_indices:
                concat_op = model_ir.operators[int(concat_idx)]
                remapped_axis = _remap_nchw_axis_to_nhwc(int(concat_op.options.get("axis", 1)))
                if remapped_axis is None:
                    axis_map_ok = False
                    break
                concat_axis_updates[int(concat_idx)] = int(remapped_axis)
            if not axis_map_ok:
                continue

            chain_indices: set[int] = {
                int(post_idx),
                int(mul_idx),
                int(add_idx),
            }
            chain_indices.update(int(v) for v in concat_indices)
            chain_indices.update(int(v) for v in leaf_pre_indices)

            target_shape_nhwc: Optional[List[int]] = None
            root_concat_tensor = model_ir.tensors.get(str(root_concat_out_name), None)
            if root_concat_tensor is not None and len(list(root_concat_tensor.shape)) == 4:
                target_shape_nhwc = _permute_shape(list(root_concat_tensor.shape), perm_nchw_to_nhwc)
            if not _rewrite_mul_const_to_nhwc_if_needed(
                mul_op=mul_op,
                data_input_name=str(root_concat_out_name),
                target_shape_nhwc=target_shape_nhwc,
                consumers=consumers,
                chain_indices=chain_indices,
            ):
                continue

            for pre_idx in sorted(list(leaf_pre_indices)):
                pre_op = model_ir.operators[int(pre_idx)]
                pre_out_name = str(pre_op.outputs[0])
                pre_in_name = str(pre_op.inputs[0])
                _replace_tensor_inputs(model_ir, pre_out_name, pre_in_name)

            for concat_idx in sorted(list(concat_indices)):
                concat_op = model_ir.operators[int(concat_idx)]
                if not isinstance(concat_op.options, dict):
                    concat_op.options = {}
                concat_op.options["axis"] = int(concat_axis_updates[int(concat_idx)])
                if len(concat_op.outputs) == 1:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(concat_op.outputs[0]), None),
                        perm_nchw_to_nhwc,
                    )

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul_out_name), None),
                perm_nchw_to_nhwc,
            )

            if str(add_op.inputs[0]) == str(post_out_name):
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=add_op,
                    input_index=0,
                    new_input_name=str(mul_out_name),
                )
            else:
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=add_op,
                    input_index=1,
                    new_input_name=str(mul_out_name),
                )

            remove_indices = set(int(v) for v in leaf_pre_indices)
            remove_indices.add(int(post_idx))
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_concat_tree_mul_add_transpose_nhwc_bridge_chains": int(optimized)}


def _optimize_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Eliminate strict NCHW bridge chains:
      xi_nhwc -> TRANSPOSE(0,3,1,2) -> STRIDED_SLICE -> PAD
      (two or more branches) -> CONCAT(axis=1) -> MUL(const) -> TRANSPOSE(0,2,3,1) -> ADD(const)

    Rewrite to NHWC:
      xi_nhwc -> STRIDED_SLICE(NHWC constants) -> PAD(NHWC pads)
      -> CONCAT(axis=3) -> MUL(const_nhwc) -> ADD(const)
    """
    optimized = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _unique_tensor_name(base: str) -> str:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 1
        return candidate

    def _permute_axis_mask(mask: int, perm: List[int]) -> int:
        remapped = 0
        for new_axis, old_axis in enumerate(perm):
            if ((int(mask) >> int(old_axis)) & 1) != 0:
                remapped |= (1 << int(new_axis))
        return int(remapped)

    def _read_stridedslice_consts(
        op: OperatorIR,
    ) -> Optional[Tuple[List[int], List[int], List[int], int, int]]:
        if str(op.op_type) != "STRIDED_SLICE" or len(op.inputs) < 4:
            return None
        opts = dict(op.options) if isinstance(op.options, dict) else {}
        if any(int(opts.get(k, 0)) != 0 for k in ["ellipsisMask", "newAxisMask", "shrinkAxisMask"]):
            return None
        begin_mask = int(opts.get("beginMask", 0))
        end_mask = int(opts.get("endMask", 0))
        begin_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[1]), None))
        end_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[2]), None))
        stride_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[3]), None))
        if begin_vals is None or end_vals is None or stride_vals is None:
            return None
        if len(begin_vals) != 4 or len(end_vals) != 4 or len(stride_vals) != 4:
            return None
        return (
            [int(v) for v in list(begin_vals)],
            [int(v) for v in list(end_vals)],
            [int(v) for v in list(stride_vals)],
            int(begin_mask),
            int(end_mask),
        )

    def _write_or_clone_i32_vector_tensor(
        *,
        op: OperatorIR,
        input_index: int,
        values: List[int],
        chain_indices: set[int],
        consumers: Dict[str, List[int]],
        clone_suffix: str,
    ) -> bool:
        input_name = str(op.inputs[int(input_index)])
        tensor = model_ir.tensors.get(input_name, None)
        if tensor is None:
            return False
        users = [int(v) for v in consumers.get(input_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in users)
        if shared_outside_chain:
            clone_name = _unique_tensor_name(f"{input_name}_{clone_suffix}")
            arr = np.asarray(values, dtype=np.int32)
            model_ir.tensors[clone_name] = TensorIR(
                name=clone_name,
                dtype="INT32",
                shape=[int(arr.size)],
                shape_signature=[int(arr.size)],
                data=np.asarray(arr),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=op,
                input_index=int(input_index),
                new_input_name=clone_name,
            )
            return True
        _write_const_ints_to_tensor(tensor, values)
        return True

    def _write_or_clone_i32_4x2_tensor(
        *,
        op: OperatorIR,
        input_index: int,
        values_4x2: np.ndarray,
        chain_indices: set[int],
        consumers: Dict[str, List[int]],
        clone_suffix: str,
    ) -> bool:
        input_name = str(op.inputs[int(input_index)])
        tensor = model_ir.tensors.get(input_name, None)
        if tensor is None:
            return False
        users = [int(v) for v in consumers.get(input_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in users)
        if shared_outside_chain:
            clone_name = _unique_tensor_name(f"{input_name}_{clone_suffix}")
            model_ir.tensors[clone_name] = TensorIR(
                name=clone_name,
                dtype="INT32",
                shape=[4, 2],
                shape_signature=[4, 2],
                data=np.asarray(values_4x2, dtype=np.int32),
                is_variable=False,
                quantization=_clone_quantization(tensor.quantization),
            )
            _replace_operator_input_at(
                model_ir=model_ir,
                op=op,
                input_index=int(input_index),
                new_input_name=clone_name,
            )
            return True
        tensor.data = np.asarray(values_4x2, dtype=np.int32)
        tensor.shape = [4, 2]
        tensor.shape_signature = [4, 2]
        return True

    def _rewrite_mul_const_to_nhwc_if_needed(
        *,
        mul_op: OperatorIR,
        data_input_name: str,
        target_shape_nhwc: Optional[List[int]],
        consumers: Dict[str, List[int]],
        chain_indices: set[int],
    ) -> bool:
        mul_inputs = [str(v) for v in list(mul_op.inputs)]
        const_input_index: Optional[int] = None
        for i, input_name in enumerate(mul_inputs):
            if str(input_name) == str(data_input_name):
                continue
            tensor = model_ir.tensors.get(str(input_name), None)
            if tensor is not None and tensor.data is not None:
                const_input_index = int(i)
                break
        if const_input_index is None:
            return False

        const_name = str(mul_inputs[int(const_input_index)])
        const_tensor = model_ir.tensors.get(const_name, None)
        if const_tensor is None or const_tensor.data is None:
            return False
        const_data = np.asarray(const_tensor.data)
        if int(const_data.size) == 1:
            return True

        target_shape = (
            [int(v) for v in list(target_shape_nhwc)]
            if _is_fully_known_positive_shape(target_shape_nhwc)
            else None
        )
        rotated: Optional[np.ndarray] = None
        if int(const_data.ndim) == 4:
            as_is_shape = [int(v) for v in list(const_data.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, as_is_shape) is not None:
                return True
            rotated_candidate = np.transpose(const_data, perm_nchw_to_nhwc).astype(const_data.dtype, copy=False)
            rotated_shape = [int(v) for v in list(rotated_candidate.shape)]
            if target_shape is not None and _broadcast_static_shapes(target_shape, rotated_shape) is None:
                return False
            rotated = np.asarray(rotated_candidate)
        else:
            if target_shape is None:
                return False
            if _broadcast_static_shapes(target_shape, [int(v) for v in list(const_data.shape)]) is None:
                return False
            return True

        if rotated is None:
            return False
        const_users = [int(v) for v in consumers.get(const_name, [])]
        shared_outside_chain = any(int(v) not in chain_indices for v in const_users)
        if shared_outside_chain:
            rotated_name = _unique_tensor_name(f"{const_name}_nhwc")
            model_ir.tensors[rotated_name] = TensorIR(
                name=rotated_name,
                dtype=str(const_tensor.dtype),
                shape=[int(v) for v in list(rotated.shape)],
                shape_signature=[int(v) for v in list(rotated.shape)],
                data=np.asarray(rotated),
                is_variable=False,
                quantization=_clone_quantization(const_tensor.quantization),
            )
            mul_inputs[int(const_input_index)] = str(rotated_name)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mul_op,
                new_inputs=mul_inputs,
            )
        else:
            const_tensor.data = np.asarray(rotated)
            const_tensor.shape = [int(v) for v in list(rotated.shape)]
            const_tensor.shape_signature = [int(v) for v in list(rotated.shape)]
        return True

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for post_idx, post_op in enumerate(model_ir.operators):
            if str(post_op.op_type) != "TRANSPOSE" or len(post_op.inputs) < 2 or len(post_op.outputs) != 1:
                continue
            if _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc:
                continue
            mul_out_name = str(post_op.inputs[0])
            post_out_name = str(post_op.outputs[0])
            if mul_out_name in model_outputs or post_out_name in model_outputs:
                continue

            mul_idx = producers.get(str(mul_out_name), None)
            if mul_idx is None:
                continue
            mul_op = model_ir.operators[int(mul_idx)]
            if str(mul_op.op_type) != "MUL" or len(mul_op.inputs) != 2 or len(mul_op.outputs) != 1:
                continue
            if str(mul_op.outputs[0]) != str(mul_out_name):
                continue
            if set(int(v) for v in consumers.get(str(mul_out_name), [])) != {int(post_idx)}:
                continue

            add_users = [int(v) for v in consumers.get(str(post_out_name), [])]
            if len(add_users) == 0:
                continue
            add_indices: List[int] = []
            valid_add_tail = True
            for add_idx in add_users:
                add_op = model_ir.operators[int(add_idx)]
                if str(add_op.op_type) != "ADD" or len(add_op.inputs) != 2 or len(add_op.outputs) != 1:
                    valid_add_tail = False
                    break
                add_inputs = [str(v) for v in list(add_op.inputs)]
                if str(post_out_name) == add_inputs[0]:
                    side_name = str(add_inputs[1])
                elif str(post_out_name) == add_inputs[1]:
                    side_name = str(add_inputs[0])
                else:
                    valid_add_tail = False
                    break
                side_tensor = model_ir.tensors.get(side_name, None)
                if side_tensor is None or side_tensor.data is None:
                    valid_add_tail = False
                    break
                side_data = np.asarray(side_tensor.data)
                if int(side_data.size) != 1:
                    side_shape = [int(v) for v in list(side_data.shape)]
                    if len(side_shape) != 4:
                        valid_add_tail = False
                        break
                    if not (
                        int(side_shape[0]) == 1
                        and int(side_shape[1]) == 1
                        and int(side_shape[2]) == 1
                        and int(side_shape[3]) > 0
                    ):
                        valid_add_tail = False
                        break
                add_indices.append(int(add_idx))
            if not valid_add_tail:
                continue

            mul_inputs = [str(v) for v in list(mul_op.inputs)]
            concat_idx: Optional[int] = None
            concat_out_name: Optional[str] = None
            concat_data_input_name: Optional[str] = None
            for mul_input_name in mul_inputs:
                producer_idx = producers.get(str(mul_input_name), None)
                if producer_idx is None:
                    continue
                producer_op = model_ir.operators[int(producer_idx)]
                if str(producer_op.op_type) != "CONCATENATION" or len(producer_op.outputs) != 1:
                    continue
                axis = int(producer_op.options.get("axis", 1))
                if axis < 0:
                    axis += 4
                if axis != 1:
                    continue
                if str(producer_op.outputs[0]) != str(mul_input_name):
                    continue
                concat_idx = int(producer_idx)
                concat_out_name = str(mul_input_name)
                concat_data_input_name = str(mul_input_name)
                break
            if concat_idx is None or concat_out_name is None or concat_data_input_name is None:
                continue
            if set(int(v) for v in consumers.get(str(concat_out_name), [])) != {int(mul_idx)}:
                continue
            if concat_out_name in model_outputs:
                continue

            concat_op = model_ir.operators[int(concat_idx)]
            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            if len(concat_inputs) < 2:
                continue

            branch_plans: List[Dict[str, Any]] = []
            valid = True
            for concat_input_name in concat_inputs:
                pad_idx = producers.get(str(concat_input_name), None)
                if pad_idx is None:
                    valid = False
                    break
                pad_op = model_ir.operators[int(pad_idx)]
                if (
                    str(pad_op.op_type) not in {"PAD", "MIRROR_PAD"}
                    or len(pad_op.inputs) < 2
                    or len(pad_op.outputs) != 1
                    or str(pad_op.outputs[0]) != str(concat_input_name)
                ):
                    valid = False
                    break
                if set(int(v) for v in consumers.get(str(concat_input_name), [])) != {int(concat_idx)}:
                    valid = False
                    break
                pad_input_name = str(pad_op.inputs[0])
                slice_idx = producers.get(str(pad_input_name), None)
                if slice_idx is None:
                    valid = False
                    break
                slice_op = model_ir.operators[int(slice_idx)]
                if (
                    str(slice_op.op_type) != "STRIDED_SLICE"
                    or len(slice_op.inputs) < 4
                    or len(slice_op.outputs) != 1
                    or str(slice_op.outputs[0]) != str(pad_input_name)
                ):
                    valid = False
                    break
                if set(int(v) for v in consumers.get(str(pad_input_name), [])) != {int(pad_idx)}:
                    valid = False
                    break
                consts = _read_stridedslice_consts(slice_op)
                if consts is None:
                    valid = False
                    break
                begin_vals, end_vals, stride_vals, begin_mask, end_mask = consts
                pre_out_name = str(slice_op.inputs[0])
                pre_idx = producers.get(str(pre_out_name), None)
                if pre_idx is None:
                    valid = False
                    break
                pre_op = model_ir.operators[int(pre_idx)]
                if (
                    str(pre_op.op_type) != "TRANSPOSE"
                    or len(pre_op.inputs) < 2
                    or len(pre_op.outputs) != 1
                    or str(pre_op.outputs[0]) != str(pre_out_name)
                    or _read_transpose_perm(model_ir, pre_op) != perm_nhwc_to_nchw
                ):
                    valid = False
                    break
                if pre_out_name in model_outputs:
                    valid = False
                    break
                if set(int(v) for v in consumers.get(str(pre_out_name), [])) != {int(slice_idx)}:
                    valid = False
                    break
                pads_name = str(pad_op.inputs[1])
                pads_tensor = model_ir.tensors.get(str(pads_name), None)
                if pads_tensor is None or pads_tensor.data is None:
                    valid = False
                    break
                try:
                    pads_pairs = np.asarray(pads_tensor.data).reshape(4, 2)
                except Exception:
                    valid = False
                    break
                if int(pads_pairs.size) != 8:
                    valid = False
                    break
                branch_plans.append(
                    {
                        "pre_idx": int(pre_idx),
                        "pre_op": pre_op,
                        "slice_idx": int(slice_idx),
                        "slice_op": slice_op,
                        "pad_idx": int(pad_idx),
                        "pad_op": pad_op,
                        "begin_vals": [int(v) for v in list(begin_vals)],
                        "end_vals": [int(v) for v in list(end_vals)],
                        "stride_vals": [int(v) for v in list(stride_vals)],
                        "begin_mask": int(begin_mask),
                        "end_mask": int(end_mask),
                        "pads_pairs": np.asarray(pads_pairs, dtype=np.int32),
                    }
                )
            if not valid or len(branch_plans) < 2:
                continue

            chain_indices: set[int] = {int(post_idx), int(mul_idx), int(concat_idx)}
            chain_indices.update(int(v) for v in add_indices)
            for plan in branch_plans:
                chain_indices.add(int(plan["pre_idx"]))
                chain_indices.add(int(plan["slice_idx"]))
                chain_indices.add(int(plan["pad_idx"]))

            target_shape_nhwc: Optional[List[int]] = None
            concat_out_tensor = model_ir.tensors.get(str(concat_out_name), None)
            if concat_out_tensor is not None and len(list(concat_out_tensor.shape)) == 4:
                target_shape_nhwc = _permute_shape(list(concat_out_tensor.shape), perm_nchw_to_nhwc)

            if not _rewrite_mul_const_to_nhwc_if_needed(
                mul_op=mul_op,
                data_input_name=str(concat_data_input_name),
                target_shape_nhwc=target_shape_nhwc,
                consumers=consumers,
                chain_indices=chain_indices,
            ):
                continue

            apply_ok = True
            for plan in branch_plans:
                pre_input_name = str(plan["pre_op"].inputs[0])
                _replace_operator_input_at(
                    model_ir=model_ir,
                    op=plan["slice_op"],
                    input_index=0,
                    new_input_name=pre_input_name,
                )
                new_begin = [int(plan["begin_vals"][idx]) for idx in perm_nchw_to_nhwc]
                new_end = [int(plan["end_vals"][idx]) for idx in perm_nchw_to_nhwc]
                new_stride = [int(plan["stride_vals"][idx]) for idx in perm_nchw_to_nhwc]
                if not _write_or_clone_i32_vector_tensor(
                    op=plan["slice_op"],
                    input_index=1,
                    values=new_begin,
                    chain_indices=chain_indices,
                    consumers=consumers,
                    clone_suffix="nhwc_begin",
                ):
                    apply_ok = False
                    break
                if not _write_or_clone_i32_vector_tensor(
                    op=plan["slice_op"],
                    input_index=2,
                    values=new_end,
                    chain_indices=chain_indices,
                    consumers=consumers,
                    clone_suffix="nhwc_end",
                ):
                    apply_ok = False
                    break
                if not _write_or_clone_i32_vector_tensor(
                    op=plan["slice_op"],
                    input_index=3,
                    values=new_stride,
                    chain_indices=chain_indices,
                    consumers=consumers,
                    clone_suffix="nhwc_stride",
                ):
                    apply_ok = False
                    break
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(plan["slice_op"].outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
                slice_opts = dict(plan["slice_op"].options) if isinstance(plan["slice_op"].options, dict) else {}
                slice_opts["beginMask"] = int(_permute_axis_mask(int(plan["begin_mask"]), perm_nchw_to_nhwc))
                slice_opts["endMask"] = int(_permute_axis_mask(int(plan["end_mask"]), perm_nchw_to_nhwc))
                plan["slice_op"].options = slice_opts

                pads_nhwc = np.asarray(
                    [
                        plan["pads_pairs"][0],
                        plan["pads_pairs"][2],
                        plan["pads_pairs"][3],
                        plan["pads_pairs"][1],
                    ],
                    dtype=np.int32,
                )
                if not _write_or_clone_i32_4x2_tensor(
                    op=plan["pad_op"],
                    input_index=1,
                    values_4x2=np.asarray(pads_nhwc, dtype=np.int32),
                    chain_indices=chain_indices,
                    consumers=consumers,
                    clone_suffix="nhwc_pads",
                ):
                    apply_ok = False
                    break
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(plan["pad_op"].outputs[0]), None),
                    perm_nchw_to_nhwc,
                )
            if not apply_ok:
                continue

            concat_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat_out_name), None),
                perm_nchw_to_nhwc,
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul_out_name), None),
                perm_nchw_to_nhwc,
            )

            old_mul_tensor = model_ir.tensors.get(str(mul_out_name), None)
            _set_operator_outputs(
                model_ir=model_ir,
                op=mul_op,
                new_outputs=[str(post_out_name)],
            )
            post_out_tensor = model_ir.tensors.get(str(post_out_name), None)
            if old_mul_tensor is not None and post_out_tensor is not None:
                post_out_tensor.dtype = str(old_mul_tensor.dtype)
                post_out_tensor.quantization = _clone_quantization(old_mul_tensor.quantization)
                post_out_tensor.shape = [int(v) for v in list(old_mul_tensor.shape)]
                post_out_tensor.shape_signature = (
                    [int(v) for v in list(old_mul_tensor.shape_signature)]
                    if old_mul_tensor.shape_signature is not None
                    else [int(v) for v in list(old_mul_tensor.shape)]
                )

            remove_indices = {int(post_idx)}
            remove_indices.update(int(plan["pre_idx"]) for plan in branch_plans)
            for remove_idx in sorted(list(remove_indices), reverse=True):
                del model_ir.operators[int(remove_idx)]

            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains": int(optimized)}


def _optimize_boundary_input_transpose_stridedslice_qdq_concat_blocks(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NHWC<->NCHW boundary transpose round-trips around STRIDED_SLICE QDQ-CONCAT stems.

    Target pattern (strict):
      input(NHWC) --TRANSPOSE(0,3,1,2)--> input_onnx_ncx_internal
        -> STRIDED_SLICE* -> QUANTIZE -> DEQUANTIZE
        -> CONCAT(axis=1) -> QUANTIZE
        -> TRANSPOSE(0,2,3,1)*

    Rewrite:
    - Move STRIDED_SLICE to NHWC by permuting begin/end/stride constants.
    - Switch CONCAT axis from 1(NCHW-C) to 3(NHWC-C).
    - Remove trailing NCHW->NHWC transposes by writing QUANTIZE directly to
      their output tensor names.
    - Remove shared boundary NHWC->NCHW transpose.
    """
    removed_boundary = 0
    removed_post_transposes = 0
    rewritten_slices = 0
    rewritten_concat = 0

    perm_nchw_to_nhwc = [0, 2, 3, 1]

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        model_inputs = set(str(v) for v in model_ir.inputs)
        model_outputs = set(str(v) for v in model_ir.outputs)

        for pre_idx, pre_op in enumerate(model_ir.operators):
            if str(pre_op.op_type) != "TRANSPOSE" or len(pre_op.inputs) < 2 or len(pre_op.outputs) != 1:
                continue

            input_name = str(pre_op.inputs[0])
            internal_name = str(pre_op.outputs[0])
            if input_name not in model_inputs:
                continue
            if input_name in model_outputs or internal_name in model_outputs:
                continue
            if not str(internal_name).endswith("_onnx_ncx_internal"):
                continue
            if _read_transpose_perm(model_ir, pre_op) != [0, 3, 1, 2]:
                continue

            internal_users = [int(v) for v in consumers.get(internal_name, [])]
            if len(internal_users) == 0:
                continue

            slice_indices: List[int] = []
            slice_rewrites: List[Dict[str, Any]] = []
            q_indices: List[int] = []
            dq_indices: List[int] = []
            dq_output_names: List[str] = []
            concat_idx: Optional[int] = None

            valid = True
            for user_idx in sorted(internal_users):
                op = model_ir.operators[int(user_idx)]
                if str(op.op_type) != "STRIDED_SLICE" or len(op.inputs) < 4 or len(op.outputs) != 1:
                    valid = False
                    break
                if str(op.inputs[0]) != internal_name:
                    valid = False
                    break
                # Keep the rewrite strict: axis masks must be absent or zero.
                opts = dict(op.options) if isinstance(op.options, dict) else {}
                if any(int(opts.get(k, 0)) != 0 for k in ["beginMask", "endMask", "ellipsisMask", "newAxisMask", "shrinkAxisMask"]):
                    valid = False
                    break
                begin_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
                end_tensor = model_ir.tensors.get(str(op.inputs[2]), None)
                stride_tensor = model_ir.tensors.get(str(op.inputs[3]), None)
                begin_vals = _read_const_ints_from_tensor(begin_tensor)
                end_vals = _read_const_ints_from_tensor(end_tensor)
                stride_vals = _read_const_ints_from_tensor(stride_tensor)
                if (
                    begin_vals is None
                    or end_vals is None
                    or stride_vals is None
                    or len(begin_vals) != 4
                    or len(end_vals) != 4
                    or len(stride_vals) != 4
                ):
                    valid = False
                    break
                new_begin = [int(begin_vals[idx]) for idx in perm_nchw_to_nhwc]
                new_end = [int(end_vals[idx]) for idx in perm_nchw_to_nhwc]
                new_stride = [int(stride_vals[idx]) for idx in perm_nchw_to_nhwc]

                slice_out = str(op.outputs[0])
                if slice_out in model_outputs:
                    valid = False
                    break
                slice_users = [int(v) for v in consumers.get(slice_out, [])]
                if len(slice_users) != 1:
                    valid = False
                    break
                q_idx = int(slice_users[0])
                q_op = model_ir.operators[int(q_idx)]
                if str(q_op.op_type) != "QUANTIZE" or len(q_op.inputs) != 1 or len(q_op.outputs) != 1:
                    valid = False
                    break
                if str(q_op.inputs[0]) != slice_out:
                    valid = False
                    break

                q_out = str(q_op.outputs[0])
                q_users = [int(v) for v in consumers.get(q_out, [])]
                if len(q_users) != 1:
                    valid = False
                    break
                dq_idx = int(q_users[0])
                dq_op = model_ir.operators[int(dq_idx)]
                if str(dq_op.op_type) != "DEQUANTIZE" or len(dq_op.inputs) != 1 or len(dq_op.outputs) != 1:
                    valid = False
                    break
                if str(dq_op.inputs[0]) != q_out:
                    valid = False
                    break

                dq_out = str(dq_op.outputs[0])
                dq_users = [int(v) for v in consumers.get(dq_out, [])]
                if len(dq_users) != 1:
                    valid = False
                    break
                cand_concat_idx = int(dq_users[0])
                cand_concat = model_ir.operators[int(cand_concat_idx)]
                if str(cand_concat.op_type) != "CONCATENATION" or len(cand_concat.outputs) != 1:
                    valid = False
                    break
                axis = int(cand_concat.options.get("axis", 1))
                if axis < 0:
                    axis += 4
                if axis != 1:
                    valid = False
                    break

                if concat_idx is None:
                    concat_idx = int(cand_concat_idx)
                elif int(concat_idx) != int(cand_concat_idx):
                    valid = False
                    break

                slice_indices.append(int(user_idx))
                slice_rewrites.append(
                    {
                        "slice_idx": int(user_idx),
                        "begin_name": str(op.inputs[1]),
                        "end_name": str(op.inputs[2]),
                        "stride_name": str(op.inputs[3]),
                        "new_begin": [int(v) for v in list(new_begin)],
                        "new_end": [int(v) for v in list(new_end)],
                        "new_stride": [int(v) for v in list(new_stride)],
                    }
                )
                q_indices.append(int(q_idx))
                dq_indices.append(int(dq_idx))
                dq_output_names.append(str(dq_out))

            if not valid or concat_idx is None or len(slice_indices) == 0:
                continue

            concat_op = model_ir.operators[int(concat_idx)]
            if set(str(v) for v in list(concat_op.inputs)) != set(dq_output_names):
                continue

            concat_out = str(concat_op.outputs[0])
            if concat_out in model_outputs:
                continue
            concat_users = [int(v) for v in consumers.get(concat_out, [])]
            if len(concat_users) != 1:
                continue
            q_concat_idx = int(concat_users[0])
            q_concat_op = model_ir.operators[int(q_concat_idx)]
            if str(q_concat_op.op_type) != "QUANTIZE" or len(q_concat_op.inputs) != 1 or len(q_concat_op.outputs) != 1:
                continue
            if str(q_concat_op.inputs[0]) != concat_out:
                continue

            q_concat_out = str(q_concat_op.outputs[0])
            if q_concat_out in model_outputs:
                continue
            post_indices = [int(v) for v in consumers.get(q_concat_out, [])]
            if len(post_indices) == 0:
                continue
            post_output_names: List[str] = []
            valid_posts = True
            for post_idx in post_indices:
                post_op = model_ir.operators[int(post_idx)]
                if (
                    str(post_op.op_type) != "TRANSPOSE"
                    or len(post_op.inputs) < 2
                    or len(post_op.outputs) != 1
                    or str(post_op.inputs[0]) != q_concat_out
                    or _read_transpose_perm(model_ir, post_op) != [0, 2, 3, 1]
                ):
                    valid_posts = False
                    break
                post_out = str(post_op.outputs[0])
                if post_out in model_outputs:
                    valid_posts = False
                    break
                post_output_names.append(post_out)
            if not valid_posts:
                continue

            # Rewrite each STRIDED_SLICE from NCHW params to NHWC params.
            for rewrite in slice_rewrites:
                slice_idx = int(rewrite["slice_idx"])
                slice_op = model_ir.operators[int(slice_idx)]
                begin_name = str(rewrite["begin_name"])
                end_name = str(rewrite["end_name"])
                stride_name = str(rewrite["stride_name"])
                new_begin = [int(v) for v in list(rewrite["new_begin"])]
                new_end = [int(v) for v in list(rewrite["new_end"])]
                new_stride = [int(v) for v in list(rewrite["new_stride"])]

                begin_tensor = model_ir.tensors.get(begin_name, None)
                end_tensor = model_ir.tensors.get(end_name, None)
                stride_tensor = model_ir.tensors.get(stride_name, None)
                begin_vals = _read_const_ints_from_tensor(begin_tensor)
                end_vals = _read_const_ints_from_tensor(end_tensor)
                stride_vals = _read_const_ints_from_tensor(stride_tensor)
                if (
                    begin_vals is None
                    or end_vals is None
                    or stride_vals is None
                    or len(begin_vals) != 4
                    or len(end_vals) != 4
                    or len(stride_vals) != 4
                ):
                    valid = False
                    break

                if begin_vals != new_begin:
                    _write_const_ints_to_tensor(begin_tensor, new_begin)
                    if _read_const_ints_from_tensor(begin_tensor) != new_begin:
                        valid = False
                        break
                if end_vals != new_end:
                    _write_const_ints_to_tensor(end_tensor, new_end)
                    if _read_const_ints_from_tensor(end_tensor) != new_end:
                        valid = False
                        break
                if stride_vals != new_stride:
                    _write_const_ints_to_tensor(stride_tensor, new_stride)
                    if _read_const_ints_from_tensor(stride_tensor) != new_stride:
                        valid = False
                        break

                _set_operator_inputs(
                    model_ir=model_ir,
                    op=slice_op,
                    new_inputs=[input_name, begin_name, end_name, stride_name],
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(slice_op.outputs[0]), None),
                    [0, 2, 3, 1],
                )
                rewritten_slices += 1

            if not valid:
                continue

            # Propagate permuted metadata through Q/DQ/CONCAT/Q.
            for q_idx in q_indices:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(model_ir.operators[int(q_idx)].outputs[0]), None),
                    [0, 2, 3, 1],
                )
            for dq_idx in dq_indices:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(model_ir.operators[int(dq_idx)].outputs[0]), None),
                    [0, 2, 3, 1],
                )
            concat_op.options["axis"] = 3
            rewritten_concat += 1
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(concat_out, None),
                [0, 2, 3, 1],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(q_concat_out, None),
                [0, 2, 3, 1],
            )

            # Remove trailing transposes by writing quantized concat output directly.
            canonical_post_output = str(post_output_names[0])
            _set_operator_outputs(
                model_ir=model_ir,
                op=q_concat_op,
                new_outputs=[canonical_post_output],
            )
            for alias_name in post_output_names[1:]:
                _replace_tensor_inputs(model_ir, alias_name, canonical_post_output)

            canonical_tensor = model_ir.tensors.get(canonical_post_output, None)
            q_concat_tensor = model_ir.tensors.get(q_concat_out, None)
            if canonical_tensor is not None and q_concat_tensor is not None:
                canonical_tensor.dtype = str(q_concat_tensor.dtype)
                canonical_tensor.quantization = _clone_quantization(q_concat_tensor.quantization)
                canonical_tensor.shape = [int(v) for v in list(q_concat_tensor.shape)]
                canonical_tensor.shape_signature = (
                    [int(v) for v in list(q_concat_tensor.shape_signature)]
                    if q_concat_tensor.shape_signature is not None
                    else [int(v) for v in list(q_concat_tensor.shape)]
                )

            remove_indices = sorted(
                set([int(pre_idx)] + [int(v) for v in post_indices]),
                reverse=True,
            )
            for remove_idx in remove_indices:
                del model_ir.operators[int(remove_idx)]
                if int(remove_idx) in post_indices:
                    removed_post_transposes += 1
            removed_boundary += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {
        "removed_boundary_input_transpose_stridedslice_blocks": int(removed_boundary),
        "rewritten_boundary_stridedslices": int(rewritten_slices),
        "rewritten_boundary_qdq_concat_axis": int(rewritten_concat),
        "removed_boundary_post_transposes": int(removed_post_transposes),
    }


def _optimize_boundary_input_layout_transposes(model_ir: ModelIR) -> Dict[str, int]:
    """
    Elide synthetic input-boundary layout adapters inserted as:
      input --TRANSPOSE--> input_onnx_ncx_internal

    Safety:
    - Matches only synthetic wrapper outputs named '*_onnx_ncx_internal'.
    - Input tensor must be a model input and not a model output.
    - Never mutates model-input metadata (public I/O contract).
    - Rewrites only when input/internal tensor metadata already match exactly.
    """
    removed = 0

    while True:
        model_inputs = set(str(v) for v in model_ir.inputs)
        model_outputs = set(str(v) for v in model_ir.outputs)
        changed = False

        for op_idx, op in enumerate(model_ir.operators):
            if str(op.op_type) != "TRANSPOSE":
                continue
            if len(op.inputs) < 2 or len(op.outputs) != 1:
                continue

            input_name = str(op.inputs[0])
            output_name = str(op.outputs[0])
            if input_name not in model_inputs:
                continue
            if input_name in model_outputs or output_name in model_outputs:
                continue
            if not str(output_name).endswith("_onnx_ncx_internal"):
                continue
            # Ensure this is the boundary adapter form [N,C,*] internalization.
            perm = _read_transpose_perm(model_ir, op)
            if perm is None:
                continue
            rank = len(perm)
            if rank == 3 and perm != [0, 2, 1]:
                continue
            if rank == 4 and perm != [0, 3, 1, 2]:
                continue
            if rank == 5 and perm != [0, 4, 1, 2, 3]:
                continue
            if rank not in {3, 4, 5}:
                continue

            internal_tensor = model_ir.tensors.get(output_name, None)
            input_tensor = model_ir.tensors.get(input_name, None)
            if internal_tensor is None or input_tensor is None:
                continue

            input_shape = [int(v) for v in list(input_tensor.shape)]
            internal_shape = [int(v) for v in list(internal_tensor.shape)]
            if input_shape != internal_shape:
                continue
            input_signature = (
                [int(v) for v in list(input_tensor.shape_signature)]
                if input_tensor.shape_signature is not None
                else input_shape
            )
            internal_signature = (
                [int(v) for v in list(internal_tensor.shape_signature)]
                if internal_tensor.shape_signature is not None
                else internal_shape
            )
            if input_signature != internal_signature:
                continue
            if str(input_tensor.dtype) != str(internal_tensor.dtype):
                continue
            if input_tensor.quantization != internal_tensor.quantization:
                continue

            # Redirect all internal users to model input without changing
            # external input tensor metadata.
            _replace_tensor_inputs(model_ir, output_name, input_name)

            del model_ir.operators[int(op_idx)]
            removed += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"removed_boundary_input_layout_transpose": int(removed)}


def _collect_schema_ops_for_range(
    *,
    opset_min: int = 13,
    opset_max: int = 18,
) -> List[str]:
    schema_ops = set()
    for schema in onnx.defs.get_all_schemas_with_history():
        if schema.domain != "":
            continue
        if int(schema.since_version) < int(opset_min) or int(schema.since_version) > int(opset_max):
            continue
        schema_ops.add(str(schema.name))
    return sorted(schema_ops)


def _build_schema_policy_matrix(
    *,
    schema_ops: set,
    supported_registry_ops: set,
    custom_candidate_ops: set,
) -> List[Dict[str, Any]]:
    matrix: List[Dict[str, Any]] = []
    for op in sorted(list(schema_ops)):
        if op in supported_registry_ops:
            policy = "builtin_supported"
        elif op in custom_candidate_ops:
            policy = "custom_candidate"
        else:
            policy = "explicit_error"
        matrix.append(
            {
                "onnx_op": str(op),
                "policy": str(policy),
            }
        )
    return matrix


def build_op_coverage_report(
    *,
    onnx_graph: onnx.ModelProto,
    output_file_name: str,
    opset_min: int = 13,
    opset_max: int = 18,
    conversion_error: Optional[str] = None,
    allow_custom_ops: bool = False,
    custom_op_allowlist: Optional[List[str]] = None,
    disable_group_convolution: bool = False,
    preprocess_report: Optional[Dict[str, Any]] = None,
    output_nms_with_argmax: bool = False,
    switch_nms_version: str = "v4",
) -> Dict[str, Any]:
    onnx_graph = _infer_shapes_with_fallback(onnx_graph)

    shape_map, dtype_map = _extract_tensor_info(onnx_graph)
    constants: Dict[str, np.ndarray] = {}
    for ini in onnx_graph.graph.initializer:
        constants[ini.name] = np.asarray(numpy_helper.to_array(ini))
    tensor_consumer_count: Dict[str, int] = {}
    for node in onnx_graph.graph.node:
        for input_name in node.input:
            if str(input_name) == "":
                continue
            tensor_consumer_count[str(input_name)] = int(tensor_consumer_count.get(str(input_name), 0) + 1)
    graph_output_names = [str(o.name) for o in onnx_graph.graph.output]

    model_ir = ModelIR(name=output_file_name)
    model_ir.metadata["tensor_lineage_events"] = []
    ctx = LoweringContext(
        model_ir=model_ir,
        shape_map=shape_map,
        dtype_map=dtype_map,
        constants=constants,
        onnx_model=onnx_graph,
        allow_custom_ops=allow_custom_ops,
        custom_op_allowlist=custom_op_allowlist,
        disable_group_convolution=disable_group_convolution,
        tensor_consumer_count=tensor_consumer_count,
        graph_output_names=graph_output_names,
        output_nms_with_argmax=output_nms_with_argmax,
        switch_nms_version=switch_nms_version,
    )
    initializer_names = {ini.name for ini in onnx_graph.graph.initializer}
    for graph_input in onnx_graph.graph.input:
        if graph_input.name in initializer_names:
            continue
        ctx.ensure_tensor(graph_input.name)
        model_ir.inputs.append(graph_input.name)
    for name, value in constants.items():
        if name not in model_ir.tensors:
            ctx.add_const_tensor(name, value)

    node_reports: List[Dict[str, Any]] = []
    unsupported_nodes: List[Dict[str, Any]] = []
    custom_lowered_nodes: List[Dict[str, Any]] = []
    graph_unique_ops: set = set()
    for node in onnx_graph.graph.node:
        node_name = node.name if node.name else node.op_type
        graph_unique_ops.add(node.op_type)
        if node.op_type == "Constant":
            value_attr = None
            for attr in node.attribute:
                if attr.name == "value":
                    value_attr = attr
                    break
            if value_attr is not None and len(node.output) > 0:
                const_array = np.asarray(numpy_helper.to_array(value_attr.t))
                out_name = node.output[0]
                if out_name in model_ir.tensors:
                    t = model_ir.tensors[out_name]
                    t.data = const_array
                    t.dtype = tflite_dtype_from_numpy(const_array.dtype)
                    t.shape, t.shape_signature = normalize_onnx_shape(list(const_array.shape))
                    constants[out_name] = const_array
                else:
                    _added = ctx.add_const_tensor(out_name, const_array)
                    if _added != out_name:
                        model_ir.tensors[out_name] = model_ir.tensors.pop(_added)
                        model_ir.tensors[out_name].name = out_name
                        constants[out_name] = constants.pop(_added)
            node_reports.append(
                {
                    "node_name": node_name,
                    "onnx_op": "Constant",
                    "supported": True,
                    "reason_code": "handled_inline",
                    "message": "Constant node is handled inline in lowering pass.",
                }
            )
            continue

        wrapped = _NodeWrap(node)
        try:
            resolution = resolve_node_dispatch(wrapped, ctx)
            dispatch_mode = str(resolution.dispatch_mode)
            reason_code = resolution.reason_code
            message = resolution.message
            node_reports.append(
                {
                    "node_name": node_name,
                    "onnx_op": node.op_type,
                    "supported": True,
                    "dispatch_mode": dispatch_mode,
                    "reason_code": reason_code,
                    "message": message,
                }
            )
            if dispatch_mode == "custom":
                custom_lowered_nodes.append(
                    {
                        "node_name": node_name,
                        "onnx_op": node.op_type,
                        "reason_code": reason_code,
                        "message": message,
                    }
                )
        except NodeValidationError as ve:
            issue = ve.to_dict()
            issue["supported"] = False
            issue["dispatch_mode"] = "unsupported"
            node_reports.append(issue)
            unsupported_nodes.append(issue)
        except Exception as ex:
            issue = {
                "node_name": node_name,
                "onnx_op": node.op_type,
                "supported": False,
                "dispatch_mode": "unsupported",
                "reason_code": "validation_exception",
                "message": str(ex),
            }
            node_reports.append(issue)
            unsupported_nodes.append(issue)

    supported_registry_ops = set(get_supported_onnx_ops())
    custom_candidate_ops = set(get_custom_op_candidate_ops())
    schema_ops = set(_collect_schema_ops_for_range(opset_min=opset_min, opset_max=opset_max))
    schema_policy_matrix = _build_schema_policy_matrix(
        schema_ops=schema_ops,
        supported_registry_ops=supported_registry_ops,
        custom_candidate_ops=custom_candidate_ops,
    )
    schema_policy_counts: Dict[str, int] = {
        "builtin_supported": 0,
        "custom_candidate": 0,
        "explicit_error": 0,
    }
    for item in schema_policy_matrix:
        p = str(item["policy"])
        schema_policy_counts[p] = int(schema_policy_counts.get(p, 0) + 1)
    graph_ops = sorted(graph_unique_ops)
    supported_graph_ops = sorted(
        list({r["onnx_op"] for r in node_reports if r["supported"] is True})
    )
    unsupported_graph_ops = sorted(
        list({r["onnx_op"] for r in node_reports if r["supported"] is False})
    )
    reason_counts: Dict[str, int] = {}
    for issue in unsupported_nodes:
        reason = str(issue.get("reason_code", "unknown"))
        reason_counts[reason] = int(reason_counts.get(reason, 0) + 1)

    normalized_allowlist = (
        [str(v).strip() for v in custom_op_allowlist if str(v).strip() != ""]
        if custom_op_allowlist is not None
        else None
    )
    allowlist_set = (
        {str(v) for v in normalized_allowlist}
        if normalized_allowlist is not None
        else set()
    )
    candidate_ops_now_builtin_supported = sorted(
        list(custom_candidate_ops & supported_registry_ops)
    )
    allowlist_builtin_supported_ops = sorted(
        list(allowlist_set & supported_registry_ops)
    )
    allowlist_custom_candidate_ops = sorted(
        list(allowlist_set & custom_candidate_ops)
    )
    allowlist_unknown_ops = sorted(
        list(
            allowlist_set
            - supported_registry_ops
            - custom_candidate_ops
            - schema_ops
        )
    )

    total_nodes = len(node_reports)
    supported_nodes = len([r for r in node_reports if r["supported"] is True])
    coverage = float(supported_nodes / total_nodes) if total_nodes > 0 else 1.0
    report: Dict[str, Any] = {
        "schema_version": 1,
        "target_opset_min": int(opset_min),
        "target_opset_max": int(opset_max),
        "supported_onnx_ops_registry": sorted(list(supported_registry_ops)),
        "custom_op_candidate_ops": sorted(list(custom_candidate_ops)),
        "schema_onnx_ops_target_range": sorted(list(schema_ops)),
        "schema_policy_matrix": schema_policy_matrix,
        "schema_policy_counts": schema_policy_counts,
        "schema_unresolved_ops": [],
        "registry_missing_from_schema_range": sorted(list(schema_ops - supported_registry_ops)),
        "registry_extra_outside_schema_range": sorted(list(supported_registry_ops - schema_ops)),
        "graph_ops": graph_ops,
        "graph_supported_ops": supported_graph_ops,
        "graph_unsupported_ops": unsupported_graph_ops,
        "graph_custom_ops": sorted(list({r["onnx_op"] for r in custom_lowered_nodes})),
        "graph_node_reports": node_reports,
        "custom_lowered_nodes": custom_lowered_nodes,
        "unsupported_nodes": unsupported_nodes,
        "unsupported_reason_counts": reason_counts,
        "graph_summary": {
            "total_nodes": int(total_nodes),
            "supported_nodes": int(supported_nodes),
            "custom_lowered_nodes": int(len(custom_lowered_nodes)),
            "unsupported_nodes": int(total_nodes - supported_nodes),
            "coverage_ratio": float(coverage),
        },
        "custom_op_policy": {
            "allow_custom_ops": bool(allow_custom_ops),
            "custom_op_allowlist": (
                [str(v) for v in normalized_allowlist]
                if normalized_allowlist is not None
                else None
            ),
            "candidate_count": int(len(custom_candidate_ops)),
            "candidate_count_excluding_builtin_supported": int(
                len(custom_candidate_ops - supported_registry_ops)
            ),
            "candidate_ops_now_builtin_supported": candidate_ops_now_builtin_supported,
            "allowlist_builtin_supported_ops": allowlist_builtin_supported_ops,
            "allowlist_custom_candidate_ops": allowlist_custom_candidate_ops,
            "allowlist_unknown_ops": allowlist_unknown_ops,
        },
        "preprocess_report": (
            dict(preprocess_report)
            if isinstance(preprocess_report, dict)
            else {
                "schema_version": 1,
                "pipeline_version": 1,
                "registered_rule_ids": [],
                "enabled_rule_ids": [],
                "applied_rules": [],
                "summary": {
                    "registered_rule_count": 0,
                    "enabled_rule_count": 0,
                    "executed_rule_count": 0,
                    "changed_rule_count": 0,
                    "total_matched_nodes": 0,
                    "total_rewritten_nodes": 0,
                },
            }
        ),
        "conversion_error": conversion_error,
    }
    return report


def write_op_coverage_report(
    *,
    report: Dict[str, Any],
    output_report_path: str,
) -> str:
    import json
    import os

    os.makedirs(os.path.dirname(output_report_path) or ".", exist_ok=True)
    with open(output_report_path, "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    return output_report_path


def _trace_tensor_rewrite_history(
    *,
    original_name: str,
    lineage_events: List[Dict[str, Any]],
) -> Tuple[str, List[Dict[str, Any]], List[str]]:
    current_name = str(original_name)
    applied_events: List[Dict[str, Any]] = []
    name_chain: List[str] = [str(original_name)]
    for event in lineage_events:
        kind = str(event.get("kind", ""))
        if kind == "rename_tensor":
            old_name = str(event.get("old_name", ""))
            new_name = str(event.get("new_name", ""))
            if current_name == old_name and old_name != "" and new_name != "":
                current_name = new_name
                applied_events.append(
                    {
                        "event_index": int(event.get("event_index", -1)),
                        "kind": kind,
                        "from": old_name,
                        "to": new_name,
                    }
                )
                name_chain.append(current_name)
        elif kind == "replace_input":
            src_name = str(event.get("src_name", ""))
            dst_name = str(event.get("dst_name", ""))
            if current_name == src_name and src_name != "" and dst_name != "":
                current_name = dst_name
                applied_events.append(
                    {
                        "event_index": int(event.get("event_index", -1)),
                        "kind": kind,
                        "from": src_name,
                        "to": dst_name,
                    }
                )
                name_chain.append(current_name)
    return current_name, applied_events, name_chain


def _build_onnx_tensor_consumer_graph(
    onnx_graph: onnx.ModelProto,
) -> Tuple[Dict[str, List[int]], List[List[str]], List[str]]:
    tensor_consumers: Dict[str, List[int]] = {}
    node_outputs: List[List[str]] = []
    node_op_types: List[str] = []
    for node_index, node in enumerate(onnx_graph.graph.node):
        outputs = [str(v) for v in node.output if str(v) != ""]
        node_outputs.append(outputs)
        node_op_types.append(str(node.op_type))
        for input_name in node.input:
            input_name = str(input_name)
            if input_name == "":
                continue
            if input_name not in tensor_consumers:
                tensor_consumers[input_name] = []
            tensor_consumers[input_name].append(int(node_index))
    return tensor_consumers, node_outputs, node_op_types


def _infer_correspondence_via_downstream(
    *,
    records: List[Dict[str, Any]],
    onnx_graph: onnx.ModelProto,
) -> None:
    record_by_output = {
        str(record.get("onnx_output_name", "")): record
        for record in records
        if str(record.get("onnx_output_name", "")) != ""
    }
    tensor_consumers, node_outputs, node_op_types = _build_onnx_tensor_consumer_graph(onnx_graph)
    resolvable_statuses = {"direct", "rewritten", "inferred"}
    traceable_ops = {
        "Transpose",
        "Reshape",
        "QuantizeLinear",
        "DequantizeLinear",
        "Identity",
        "Cast",
        "Squeeze",
        "Unsqueeze",
        "Flatten",
    }
    max_infer_hops = 4

    changed = True
    while changed:
        changed = False
        for record in records:
            if str(record.get("status", "")) != "missing":
                continue
            start_name = str(record.get("onnx_output_name", ""))
            if start_name == "":
                continue

            queue: List[Tuple[str, int]] = [(start_name, 0)]
            visited: set = {start_name}
            inferred_target: Optional[Dict[str, Any]] = None
            inferred_from_output: Optional[str] = None
            inferred_hops: Optional[int] = None

            while len(queue) > 0 and inferred_target is None:
                tensor_name, hops = queue.pop(0)
                if int(hops) >= int(max_infer_hops):
                    continue
                consumer_indices = tensor_consumers.get(str(tensor_name), [])
                for consumer_index in consumer_indices:
                    if int(consumer_index) < 0 or int(consumer_index) >= len(node_outputs):
                        continue
                    if (
                        int(consumer_index) >= len(node_op_types)
                        or str(node_op_types[int(consumer_index)]) not in traceable_ops
                    ):
                        continue
                    for out_name in node_outputs[int(consumer_index)]:
                        out_name = str(out_name)
                        if out_name == "" or out_name in visited:
                            continue
                        visited.add(out_name)
                        downstream_record = record_by_output.get(out_name, None)
                        if (
                            downstream_record is not None
                            and str(downstream_record.get("status", "")) in resolvable_statuses
                            and bool(downstream_record.get("exists_in_final_model", False))
                        ):
                            inferred_target = downstream_record
                            inferred_from_output = out_name
                            inferred_hops = int(hops + 1)
                            break
                        queue.append((out_name, int(hops + 1)))
                    if inferred_target is not None:
                        break

            if inferred_target is None:
                continue

            resolved_name = str(inferred_target.get("resolved_tflite_tensor_name", ""))
            if resolved_name == "":
                continue

            record["status"] = "inferred"
            record["is_rewritten"] = True
            record["resolved_tflite_tensor_name"] = resolved_name
            record["exists_in_final_model"] = bool(
                inferred_target.get("exists_in_final_model", False)
            )
            record["producer_operator_index"] = inferred_target.get(
                "producer_operator_index", None
            )
            record["producer_op_type"] = inferred_target.get("producer_op_type", None)
            record["consumer_count"] = inferred_target.get("consumer_count", 0)
            record["is_graph_output"] = bool(inferred_target.get("is_graph_output", False))
            record["inferred_from_onnx_output_name"] = str(inferred_from_output)
            record["inferred_hops"] = int(inferred_hops) if inferred_hops is not None else None
            chain = record.get("rewrite_name_chain", [])
            if not isinstance(chain, list):
                chain = [str(start_name)]
            chain = [str(v) for v in chain]
            if resolved_name not in chain:
                chain.append(resolved_name)
            record["rewrite_name_chain"] = chain
            rewrite_events = record.get("rewrite_events", [])
            if not isinstance(rewrite_events, list):
                rewrite_events = []
            rewrite_events.append(
                {
                    "event_index": None,
                    "kind": "inferred_via_downstream",
                    "from": str(start_name),
                    "to": str(resolved_name),
                    "via_onnx_output": str(inferred_from_output),
                    "hops": int(inferred_hops) if inferred_hops is not None else None,
                }
            )
            record["rewrite_events"] = rewrite_events
            changed = True


def build_tensor_correspondence_report(
    *,
    onnx_graph: onnx.ModelProto,
    model_ir: ModelIR,
) -> Dict[str, Any]:
    lineage_events = model_ir.metadata.get("tensor_lineage_events", [])
    if not isinstance(lineage_events, list):
        lineage_events = []
    lineage_events = [
        dict(event) for event in lineage_events if isinstance(event, dict)
    ]

    producers = _build_tensor_producer_map(model_ir)
    consumers = _build_tensor_consumer_map(model_ir)
    model_output_set = set(str(v) for v in model_ir.outputs)
    model_tensor_set = set(str(v) for v in model_ir.tensors.keys())
    channel_last_hint_names = {
        str(v)
        for v in model_ir.metadata.get("assume_channel_last_layout_tensor_names", [])
        if str(v) != ""
    }

    def _is_channel_last_layout_name(tensor_name: str) -> bool:
        normalized_name = str(tensor_name).split(":")[0].lower()
        if normalized_name == "":
            return False
        return (
            "_nwc" in normalized_name
            or "_nhwc" in normalized_name
            or "_ndhwc" in normalized_name
        )

    records: List[Dict[str, Any]] = []
    total_onnx_outputs = 0

    for identifier, node in enumerate(onnx_graph.graph.node, start=1):
        node_name = str(node.name)
        node_op = str(node.op_type)
        for output_index, output_name in enumerate(node.output):
            output_name = str(output_name)
            if output_name == "":
                continue
            total_onnx_outputs += 1

            resolved_name, applied_events, name_chain = _trace_tensor_rewrite_history(
                original_name=output_name,
                lineage_events=lineage_events,
            )
            exists_in_final_model = resolved_name in model_tensor_set
            is_rewritten = str(resolved_name) != str(output_name)

            producer_index = producers.get(resolved_name, None)
            producer_op_type = None
            if producer_index is not None:
                producer_op_type = str(model_ir.operators[int(producer_index)].op_type)
            consumer_count = int(len(consumers.get(resolved_name, [])))
            is_graph_output = bool(resolved_name in model_output_set)

            if not is_rewritten and exists_in_final_model:
                status = "direct"
            elif is_rewritten and exists_in_final_model:
                status = "rewritten"
            elif is_rewritten and not exists_in_final_model:
                status = "rewritten_but_missing"
            else:
                status = "missing"

            records.append(
                {
                    "identifier": int(identifier),
                    "onnx_op_type": node_op,
                    "onnx_op_name": node_name,
                    "onnx_output_index": int(output_index),
                    "onnx_output_name": output_name,
                    "resolved_tflite_tensor_name": str(resolved_name),
                    "exists_in_final_model": bool(exists_in_final_model),
                    "is_graph_output": bool(is_graph_output),
                    "consumer_count": int(consumer_count),
                    "producer_operator_index": (
                        int(producer_index)
                        if producer_index is not None
                        else None
                    ),
                    "producer_op_type": producer_op_type,
                    "assume_channel_last_layout": bool(
                        resolved_name in channel_last_hint_names
                        or _is_channel_last_layout_name(resolved_name)
                    ),
                    "status": status,
                    "is_rewritten": bool(is_rewritten),
                    "rewrite_events": applied_events,
                    "rewrite_name_chain": name_chain,
                }
            )

    for record in records:
        if str(record.get("status", "")) != "missing":
            continue
        if str(record.get("onnx_op_type", "")) != "QLinearConv":
            continue
        onnx_op_name = str(record.get("onnx_op_name", ""))
        if onnx_op_name == "":
            continue
        candidate_name = f"{onnx_op_name}_output_nhwc"
        if candidate_name not in model_tensor_set:
            continue

        producer_index = producers.get(candidate_name, None)
        producer_op_type = None
        if producer_index is not None:
            producer_op_type = str(model_ir.operators[int(producer_index)].op_type)
        consumer_count = int(len(consumers.get(candidate_name, [])))
        is_graph_output = bool(candidate_name in model_output_set)

        record["status"] = "inferred"
        record["is_rewritten"] = True
        record["resolved_tflite_tensor_name"] = str(candidate_name)
        record["exists_in_final_model"] = True
        record["producer_operator_index"] = (
            int(producer_index) if producer_index is not None else None
        )
        record["producer_op_type"] = producer_op_type
        record["consumer_count"] = int(consumer_count)
        record["is_graph_output"] = bool(is_graph_output)
        record["assume_channel_last_layout"] = bool(
            candidate_name in channel_last_hint_names
            or _is_channel_last_layout_name(candidate_name)
        )
        chain = record.get("rewrite_name_chain", [])
        if not isinstance(chain, list):
            chain = [str(record.get("onnx_output_name", ""))]
        chain = [str(v) for v in chain]
        if candidate_name not in chain:
            chain.append(str(candidate_name))
        record["rewrite_name_chain"] = chain
        rewrite_events = record.get("rewrite_events", [])
        if not isinstance(rewrite_events, list):
            rewrite_events = []
        rewrite_events.append(
            {
                "event_index": None,
                "kind": "inferred_via_qlinearconv_output_nhwc",
                "from": str(record.get("onnx_output_name", "")),
                "to": str(candidate_name),
            }
        )
        record["rewrite_events"] = rewrite_events

    _infer_correspondence_via_downstream(
        records=records,
        onnx_graph=onnx_graph,
    )

    direct_count = int(sum(1 for r in records if str(r.get("status", "")) == "direct"))
    rewritten_count = int(sum(1 for r in records if str(r.get("status", "")) == "rewritten"))
    inferred_count = int(sum(1 for r in records if str(r.get("status", "")) == "inferred"))
    rewritten_and_missing_count = int(
        sum(1 for r in records if str(r.get("status", "")) == "rewritten_but_missing")
    )
    missing_count = int(
        sum(
            1
            for r in records
            if str(r.get("status", "")) in {"missing", "rewritten_but_missing"}
        )
    )

    report = {
        "schema_version": 1,
        "summary": {
            "total_onnx_node_outputs": int(total_onnx_outputs),
            "direct_count": int(direct_count),
            "rewritten_count": int(rewritten_count),
            "inferred_count": int(inferred_count),
            "missing_count": int(missing_count),
            "rewritten_but_missing_count": int(rewritten_and_missing_count),
            "lineage_event_count": int(len(lineage_events)),
        },
        "lineage_events": lineage_events,
        "records": records,
    }
    return report


def write_tensor_correspondence_report(
    *,
    report: Dict[str, Any],
    output_report_path: str,
) -> str:
    import json
    import os

    os.makedirs(os.path.dirname(output_report_path) or ".", exist_ok=True)
    with open(output_report_path, "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)
    return output_report_path


def lower_onnx_to_ir(
    onnx_graph: onnx.ModelProto,
    output_file_name: str,
    allow_custom_ops: bool = False,
    custom_op_allowlist: Optional[List[str]] = None,
    optimize_layout_transpose_chains: bool = True,
    transpose_inputs_to_nhwc: bool = False,
    keep_ncw_or_nchw_or_ncdhw_input_names: Optional[List[str]] = None,
    keep_nwc_or_nhwc_or_ndhwc_input_names: Optional[List[str]] = None,
    keep_shape_absolutely_input_names: Optional[List[str]] = None,
    disable_group_convolution: bool = False,
    output_nms_with_argmax: bool = False,
    switch_nms_version: str = "v4",
    mvn_epsilon: float = 1e-10,
    show_progress: bool = False,
    apply_safe_transpose_reduction_lite_on_no_layout_opt: bool = False,
    disable_suppression_flextranspose: bool = False,
    number_of_dimensions_after_flextranspose_compression: int = 6,
    disable_suppression_flexstridedslice: bool = False,
    number_of_dimensions_after_flexstridedslice_compression: int = 5,
    optimization_for_gpu_delegate: bool = False,
    replace_argmax_to_reducemax_and_indices_is_int64: bool = False,
    replace_argmax_to_reducemax_and_indices_is_float32: bool = False,
    replace_argmax_to_fused_argmax_and_indices_is_int64: bool = False,
    replace_argmax_to_fused_argmax_and_indices_is_float32: bool = False,
    fused_argmax_scale_ratio: float = 0.5,
    replace_to_pseudo_operators: Optional[List[str]] = None,
    protected_boundary_tensor_names: Optional[List[str]] = None,
) -> ModelIR:
    onnx_graph = _infer_shapes_with_fallback(onnx_graph)

    shape_map, dtype_map = _extract_tensor_info(onnx_graph)
    dynamic_boundary_tensors = _collect_dynamic_boundary_tensor_names(onnx_graph)
    onnx_boundary_signature_map = _build_onnx_boundary_shape_signature_map(
        onnx_graph=onnx_graph,
        shape_map=shape_map,
    )
    constants: Dict[str, np.ndarray] = {}
    for ini in onnx_graph.graph.initializer:
        constants[ini.name] = np.asarray(numpy_helper.to_array(ini))
    tensor_consumer_count: Dict[str, int] = {}
    for node in onnx_graph.graph.node:
        for input_name in node.input:
            if str(input_name) == "":
                continue
            tensor_consumer_count[str(input_name)] = int(
                tensor_consumer_count.get(str(input_name), 0) + 1
            )
    graph_output_names = [str(o.name) for o in onnx_graph.graph.output]

    model_ir = ModelIR(name=output_file_name)
    model_ir.metadata["tensor_lineage_events"] = []
    model_ir.metadata["onnx_dynamic_input_tensor_names"] = list(
        dynamic_boundary_tensors["inputs"]
    )
    model_ir.metadata["onnx_dynamic_output_tensor_names"] = list(
        dynamic_boundary_tensors["outputs"]
    )
    model_ir.metadata["onnx_boundary_shape_signature_map"] = dict(
        onnx_boundary_signature_map
    )
    model_ir.metadata["original_graph_output_names"] = [
        str(output.name)
        for output in onnx_graph.graph.output
        if str(output.name) != ""
    ]
    model_ir.metadata["protected_boundary_tensor_names"] = list(
        dict.fromkeys(
            [
                str(name)
                for name in (protected_boundary_tensor_names or [])
                if str(name).strip() != ""
            ]
        )
    )
    ctx = LoweringContext(
        model_ir=model_ir,
        shape_map=shape_map,
        dtype_map=dtype_map,
        constants=constants,
        onnx_model=onnx_graph,
        allow_custom_ops=allow_custom_ops,
        custom_op_allowlist=custom_op_allowlist,
        disable_group_convolution=disable_group_convolution,
        tensor_consumer_count=tensor_consumer_count,
        graph_output_names=graph_output_names,
        output_nms_with_argmax=output_nms_with_argmax,
        switch_nms_version=switch_nms_version,
        mvn_epsilon=mvn_epsilon,
        disable_suppression_flextranspose=disable_suppression_flextranspose,
        number_of_dimensions_after_flextranspose_compression=number_of_dimensions_after_flextranspose_compression,
        disable_suppression_flexstridedslice=disable_suppression_flexstridedslice,
        number_of_dimensions_after_flexstridedslice_compression=number_of_dimensions_after_flexstridedslice_compression,
        optimization_for_gpu_delegate=optimization_for_gpu_delegate,
        replace_argmax_to_reducemax_and_indices_is_int64=replace_argmax_to_reducemax_and_indices_is_int64,
        replace_argmax_to_reducemax_and_indices_is_float32=replace_argmax_to_reducemax_and_indices_is_float32,
        replace_argmax_to_fused_argmax_and_indices_is_int64=replace_argmax_to_fused_argmax_and_indices_is_int64,
        replace_argmax_to_fused_argmax_and_indices_is_float32=replace_argmax_to_fused_argmax_and_indices_is_float32,
        fused_argmax_scale_ratio=fused_argmax_scale_ratio,
        replace_to_pseudo_operators=replace_to_pseudo_operators,
    )

    keep_ncw_input_names = {
        str(v) for v in (keep_ncw_or_nchw_or_ncdhw_input_names or [])
    }
    keep_nwc_input_names = {
        str(v) for v in (keep_nwc_or_nhwc_or_ndhwc_input_names or [])
    }
    keep_shape_abs_input_names = {
        str(v) for v in (keep_shape_absolutely_input_names or [])
    }
    input_name_remap: Dict[str, str] = {}

    # Inputs
    initializer_names = {ini.name for ini in onnx_graph.graph.initializer}
    for graph_input in onnx_graph.graph.input:
        if graph_input.name in initializer_names:
            continue
        input_name = str(graph_input.name)
        ctx.ensure_tensor(input_name)
        input_tensor = model_ir.tensors[input_name]
        model_ir.inputs.append(input_name)

        if not transpose_inputs_to_nhwc:
            continue

        input_rank = len(list(input_tensor.shape))
        if input_rank not in [3, 4, 5]:
            continue

        keep_shape_abs = input_name in keep_shape_abs_input_names
        keep_ncw = input_name in keep_ncw_input_names
        keep_nwc = (
            input_name in keep_nwc_input_names
            and input_name not in keep_ncw_input_names
        )
        if keep_shape_abs or keep_ncw or keep_nwc:
            continue

        original_shape = list(input_tensor.shape)
        original_signature = (
            list(input_tensor.shape_signature)
            if input_tensor.shape_signature is not None
            else list(input_tensor.shape)
        )
        if input_rank == 3:
            perm_internal_to_external = [0, 2, 1]
            perm_external_to_internal = [0, 2, 1]
        elif input_rank == 4:
            perm_internal_to_external = [0, 2, 3, 1]
            perm_external_to_internal = [0, 3, 1, 2]
        else:
            perm_internal_to_external = [0, 2, 3, 4, 1]
            perm_external_to_internal = [0, 4, 1, 2, 3]

        external_shape = _permute_shape(original_shape, perm_internal_to_external)
        external_signature = _permute_shape(
            original_signature, perm_internal_to_external
        )
        if external_shape is None or external_signature is None:
            continue

        input_tensor.shape = list(external_shape)
        input_tensor.shape_signature = list(external_signature)

        internal_input_name = ctx.add_intermediate_tensor(
            f"{input_name}_onnx_ncx_internal",
            dtype=str(input_tensor.dtype),
            shape=original_shape,
        )
        internal_tensor = model_ir.tensors[internal_input_name]
        internal_tensor.shape_signature = list(original_signature)
        internal_tensor.quantization = _clone_quantization(input_tensor.quantization)

        perm_name = ctx.add_const_tensor(
            f"{internal_input_name}_perm",
            np.asarray(perm_external_to_internal, dtype=np.int32),
        )
        ctx.add_operator(
            OperatorIR(
                op_type="TRANSPOSE",
                inputs=[input_name, perm_name],
                outputs=[internal_input_name],
            )
        )
        input_name_remap[input_name] = internal_input_name
        # Keep ONNX consumer-count semantics for remapped synthetic input names.
        # Several builders use this count to safely elide inverse transpose pairs;
        # without this handoff, synthetic names appear to have zero consumers and
        # can incorrectly drop the boundary transpose required by other branches.
        ctx.tensor_consumer_count[internal_input_name] = int(
            ctx.tensor_consumer_count.get(input_name, 0)
        )

    # Initializers as tensors
    for name, value in constants.items():
        if name not in model_ir.tensors:
            ctx.add_const_tensor(name, value)

    # Nodes
    graph_nodes = list(onnx_graph.graph.node)
    progress_bar = _create_progress_bar(
        total=len(graph_nodes),
        desc="flatbuffer_direct lowering",
        enabled=bool(show_progress),
    )
    try:
        for node in graph_nodes:
            try:
                if node.op_type == "Constant":
                    output_name = node.output[0]
                    value_attr = None
                    for attr in node.attribute:
                        if attr.name == "value":
                            value_attr = attr
                            break
                    if value_attr is None:
                        raise NotImplementedError(f"Constant node without value is not supported. op={node.name}")
                    const_array = np.asarray(numpy_helper.to_array(value_attr.t))
                    if output_name in model_ir.tensors:
                        # Replace existing placeholder tensor data.
                        t = model_ir.tensors[output_name]
                        t.data = const_array
                        t.dtype = tflite_dtype_from_numpy(const_array.dtype)
                        t.shape, t.shape_signature = normalize_onnx_shape(list(const_array.shape))
                        constants[output_name] = const_array
                    else:
                        name = ctx.add_const_tensor(output_name, const_array)
                        if name != output_name:
                            # keep graph output name stable if collision happened.
                            model_ir.tensors[output_name] = model_ir.tensors.pop(name)
                            model_ir.tensors[output_name].name = output_name
                            constants[output_name] = constants.pop(name)
                    continue

                wrapped = _NodeWrap(
                    node,
                    input_name_remap=input_name_remap,
                )
                try:
                    dispatch_node(wrapped, ctx)
                except NodeValidationError as ve:
                    raise NotImplementedError(
                        f"flatbuffer_direct validation failed: "
                        f"op={ve.node_op} node={ve.node_name} "
                        f"reason_code={ve.reason_code} message={ve.message}"
                    ) from ve
            finally:
                if progress_bar is not None:
                    progress_bar.update(1)
    finally:
        if progress_bar is not None:
            progress_bar.close()

    post_progress_total = 6 if optimize_layout_transpose_chains else 4
    post_progress_step = 0
    post_progress_bar = _create_progress_bar(
        total=post_progress_total,
        desc="flatbuffer_direct post-lowering",
        enabled=bool(show_progress),
    )

    def _set_post_progress_desc(stage_label: str) -> None:
        if post_progress_bar is None:
            return
        post_progress_bar.set_description_str(
            f"flatbuffer_direct post-lowering [{post_progress_step + 1}/{post_progress_total}] {stage_label}"
        )

    def _advance_post_progress() -> None:
        nonlocal post_progress_step
        if post_progress_bar is None:
            return
        post_progress_bar.update(1)
        post_progress_step = int(post_progress_step + 1)

    _set_post_progress_desc("outputs")

    # Outputs
    for graph_output in onnx_graph.graph.output:
        ctx.ensure_tensor(graph_output.name)
        _append_model_outputs_preserving_order(
            model_ir,
            [str(graph_output.name)],
        )
    protected_boundary_outputs = [
        str(name)
        for name in _get_protected_boundary_tensor_names(model_ir)
        if str(name) in model_ir.tensors
    ]
    _append_model_outputs_preserving_order(
        model_ir,
        protected_boundary_outputs,
    )
    dynamic_boundary_signature_map: Dict[str, List[int]] = {}
    onnx_boundary_signature_map = model_ir.metadata.get(
        "onnx_boundary_shape_signature_map",
        {},
    )
    if not isinstance(onnx_boundary_signature_map, dict):
        onnx_boundary_signature_map = {}
    dynamic_boundary_name_set = set(
        list(model_ir.metadata.get("onnx_dynamic_input_tensor_names", []))
        + list(model_ir.metadata.get("onnx_dynamic_output_tensor_names", []))
    )
    for boundary_name in list(model_ir.inputs) + list(model_ir.outputs):
        if str(boundary_name) not in dynamic_boundary_name_set:
            continue
        boundary_tensor = model_ir.tensors.get(str(boundary_name), None)
        if boundary_tensor is None or boundary_tensor.shape_signature is None:
            continue
        snapshot_signature = [int(v) for v in list(boundary_tensor.shape_signature)]
        hinted_signature = onnx_boundary_signature_map.get(str(boundary_name), None)
        hinted_signature_aligned = _align_boundary_signature_to_current_shape(
            boundary_signature=(
                [int(v) for v in list(hinted_signature)]
                if isinstance(hinted_signature, list)
                else None
            ),
            current_shape=[int(v) for v in list(boundary_tensor.shape)],
        )
        if hinted_signature_aligned is not None:
            if len(hinted_signature_aligned) == len(snapshot_signature):
                for axis in range(len(snapshot_signature)):
                    if int(hinted_signature_aligned[axis]) > 0:
                        snapshot_signature[axis] = int(hinted_signature_aligned[axis])
        if any(int(v) < 0 for v in snapshot_signature):
            dynamic_boundary_signature_map[str(boundary_name)] = snapshot_signature
    model_ir.metadata["dynamic_boundary_shape_signature_map"] = (
        dynamic_boundary_signature_map
    )
    _advance_post_progress()

    if optimize_layout_transpose_chains:
        _set_post_progress_desc("layout optimize pass-set 1")
        has_qdq_ops = any(
            str(op.op_type) in {"QUANTIZE", "DEQUANTIZE"}
            for op in model_ir.operators
        )
        # NOTE:
        # Binary/fanout transpose rewrites can invalidate tensor shape/layout
        # assumptions in multi-branch int8 graphs. Keep them enabled on
        # non-QDQ graphs where the pass is effective and stable.
        enable_transpose_binary_bridge_optimizations = not has_qdq_ops
        # Duplicate transpose fanout dedup is safe on pure float graphs and
        # removes redundant NHWC/NCHW adapters in multi-branch heads.
        enable_duplicate_transpose_fanout_optimizations = not has_qdq_ops

        _optimize_layout_transpose_chains(model_ir)
        _optimize_transpose_quant_dequant_bridges(model_ir)
        _optimize_boundary_input_transpose_batchmatmul_chains(model_ir)
        _optimize_leading_input_transpose_passthrough_chains(model_ir)
        _optimize_asin_transpose_passthrough_chains(model_ir)
        _optimize_erf_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nchw_nhwc_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(model_ir)
        _optimize_hardswish_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_mul_transpose_passthrough_chains(model_ir)
        _optimize_swish_transpose_passthrough_chains(model_ir)
        _optimize_gelu_tanh_transpose_passthrough_chains(model_ir)
        _optimize_center_size_offset_terminal_transpose_chains(model_ir)
        _optimize_leakyrelu_transpose_passthrough_chains(model_ir)
        _optimize_prelu_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_concat_conv_nhwc_groups(model_ir)
        _optimize_transpose_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_pre_concat_ndhwc_chains(model_ir)
        _optimize_transpose_stridedslice_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(model_ir)
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir)
        _optimize_transpose_slice_logistic_concat_reshape_tail_nhwc_chains(model_ir)
        _optimize_shufflenet_transpose_shuffle_chains(model_ir)
        _optimize_shufflenet_reshape_transpose_shuffle_nhwc_chains(model_ir)
        _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(model_ir)
        _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mulconst_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_reshape_transpose_to_expanddims_nhwc_chains(model_ir)
        _optimize_transpose_reshape_transpose_to_flatten_hw_nhwc_chains(model_ir)
        _optimize_reshape_transpose_reshape_transpose_to_nhwc_reshape_chains(model_ir)
        _optimize_attention_qkv_reshape_transpose_reshape_to_reshape_transpose_chains(model_ir)
        _optimize_attention_gather_transpose_reshape_cleanup_chains(model_ir)
        _optimize_gather_axis0_singleton_to_reshape_input_chains(model_ir)
        _optimize_attention_preproj_reshape_to_batchmatmul_ranklift_chains(model_ir)
        _optimize_window_partition_reshape_transpose_to_space_to_depth_chains(model_ir)
        _optimize_window_reverse_reshape_transpose_to_depth_to_space_chains(model_ir)
        _optimize_transpose_pre_unary_squeeze_transpose_suffix_nhwc_chains(model_ir)
        _optimize_squeeze_unary_reshape_passthrough_chains(model_ir)
        _optimize_squeeze_reshape_identity_chains(model_ir)
        _optimize_fold_mul_add_mul_affine_chains(model_ir)
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mean_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir)
        _optimize_transpose_layernorm_stats_nhwc_propagation_chains(model_ir)
        _optimize_layernorm_stats_via_existing_post_transpose_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir)
        _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_conv_attention_nhwc_propagation_chains(model_ir)
        _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir)
        _optimize_sinet_mix_attention_double_logistic_nhwc_chains(model_ir)
        _optimize_transpose_sum_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_nested_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
        _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains(model_ir)
        _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir)
        _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir)
        _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir)
        _optimize_transpose_add_concat_const_suffix_nhwc_chains(model_ir)
        _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir)
        _optimize_transposeconv_output_nhwc_passthrough_chains(model_ir)
        _optimize_transposeconv_output_channel1_terminal_transpose_chains(model_ir)
        _optimize_transpose_unary_passthrough_chains(model_ir)
        _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
        _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_dequant_relu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
        _optimize_trailing_output_transpose_passthrough_chains(model_ir)
        _optimize_transpose_dequant_mul_add_prelu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_prelu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_prelu_transpose_bridges(model_ir)
        _optimize_dequant_prelu_quantize_chains(model_ir)
        _optimize_dequant_prelu_depthwise_quantize_chains(model_ir)
        _optimize_dequant_transposeconv_quantize_chains(model_ir)
        _optimize_dequant_reshape_quantize_chains(model_ir)
        _optimize_dequant_hardsigmoid_quantize_chains(model_ir)
        _optimize_dequant_maxpool_quantize_chains(model_ir)
        _optimize_dequant_softmax_quantize_chains(model_ir)
        _optimize_dequant_logistic_quantize_chains(model_ir)
        _canonicalize_softmax_transpose_chains(model_ir)
        _optimize_transpose_binary_symmetric_legacy_only_bridges_safe(model_ir)
        _optimize_transpose_binary_single_post_bridges_safe(model_ir)
        _optimize_transpose_binary_mixed_fanout_bridges_safe(model_ir)
        _optimize_transpose_binary_asymmetric_fanout_bridges(model_ir)
        _optimize_transpose_binary_full_post_fanout_bridges(model_ir)
        if enable_transpose_binary_bridge_optimizations:
            _optimize_transpose_binary_bridges(model_ir)
        if enable_duplicate_transpose_fanout_optimizations:
            _optimize_duplicate_transpose_fanout(model_ir)
        _optimize_duplicate_reshape_fanout(model_ir)
        # Binary bridge rewrites can introduce new transpose-(q|dq)-transpose patterns.
        _optimize_transpose_quant_dequant_bridges(model_ir)
        _optimize_boundary_input_transpose_batchmatmul_chains(model_ir)
        _optimize_leading_input_transpose_passthrough_chains(model_ir)
        _optimize_asin_transpose_passthrough_chains(model_ir)
        _optimize_erf_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nchw_nhwc_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(model_ir)
        _optimize_hardswish_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_mul_transpose_passthrough_chains(model_ir)
        _optimize_swish_transpose_passthrough_chains(model_ir)
        _optimize_gelu_tanh_transpose_passthrough_chains(model_ir)
        _optimize_center_size_offset_terminal_transpose_chains(model_ir)
        _optimize_leakyrelu_transpose_passthrough_chains(model_ir)
        _optimize_prelu_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_concat_conv_nhwc_groups(model_ir)
        _optimize_transpose_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_pre_concat_ndhwc_chains(model_ir)
        _optimize_transpose_stridedslice_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(model_ir)
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir)
        _optimize_transpose_slice_logistic_concat_reshape_tail_nhwc_chains(model_ir)
        _optimize_shufflenet_transpose_shuffle_chains(model_ir)
        _optimize_shufflenet_reshape_transpose_shuffle_nhwc_chains(model_ir)
        _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(model_ir)
        _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mulconst_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_reshape_transpose_to_expanddims_nhwc_chains(model_ir)
        _optimize_transpose_reshape_transpose_to_flatten_hw_nhwc_chains(model_ir)
        _optimize_reshape_transpose_reshape_transpose_to_nhwc_reshape_chains(model_ir)
        _optimize_attention_qkv_reshape_transpose_reshape_to_reshape_transpose_chains(model_ir)
        _optimize_attention_gather_transpose_reshape_cleanup_chains(model_ir)
        _optimize_gather_axis0_singleton_to_reshape_input_chains(model_ir)
        _optimize_attention_preproj_reshape_to_batchmatmul_ranklift_chains(model_ir)
        _optimize_window_partition_reshape_transpose_to_space_to_depth_chains(model_ir)
        _optimize_window_reverse_reshape_transpose_to_depth_to_space_chains(model_ir)
        _optimize_transpose_pre_unary_squeeze_transpose_suffix_nhwc_chains(model_ir)
        _optimize_squeeze_unary_reshape_passthrough_chains(model_ir)
        _optimize_squeeze_reshape_identity_chains(model_ir)
        _optimize_fold_mul_add_mul_affine_chains(model_ir)
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mean_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir)
        _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir)
        _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_conv_attention_nhwc_propagation_chains(model_ir)
        _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir)
        _optimize_sinet_mix_attention_double_logistic_nhwc_chains(model_ir)
        _optimize_transpose_sum_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_nested_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
        _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains(model_ir)
        _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir)
        _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir)
        _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir)
        _optimize_transpose_add_concat_const_suffix_nhwc_chains(model_ir)
        _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir)
        _optimize_transposeconv_output_nhwc_passthrough_chains(model_ir)
        _optimize_transposeconv_output_channel1_terminal_transpose_chains(model_ir)
        _optimize_transpose_unary_passthrough_chains(model_ir)
        _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
        _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_dequant_relu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
        _optimize_trailing_output_transpose_passthrough_chains(model_ir)
        _optimize_transpose_dequant_mul_add_prelu_quantize_bridges(model_ir)
        if enable_duplicate_transpose_fanout_optimizations:
            _optimize_duplicate_transpose_fanout(model_ir)
        _optimize_duplicate_reshape_fanout(model_ir)
        _optimize_transpose_dequant_prelu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_prelu_transpose_bridges(model_ir)
        _optimize_dequant_prelu_quantize_chains(model_ir)
        _optimize_dequant_prelu_depthwise_quantize_chains(model_ir)
        _optimize_dequant_transposeconv_quantize_chains(model_ir)
        _optimize_dequant_reshape_quantize_chains(model_ir)
        _optimize_dequant_hardsigmoid_quantize_chains(model_ir)
        _optimize_dequant_maxpool_quantize_chains(model_ir)
        _optimize_dequant_softmax_quantize_chains(model_ir)
        _optimize_dequant_logistic_quantize_chains(model_ir)
        _canonicalize_softmax_transpose_chains(model_ir)
        _optimize_transpose_binary_symmetric_legacy_only_bridges_safe(model_ir)
        _optimize_transpose_binary_single_post_bridges_safe(model_ir)
        _optimize_transpose_binary_mixed_fanout_bridges_safe(model_ir)
        _optimize_transpose_binary_asymmetric_fanout_bridges(model_ir)
        _optimize_transpose_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_dequantize_mean_quantize_bridges(model_ir)
        _optimize_transpose_mean_hardsigmoid_muladd_chains(model_ir)
        _optimize_nhwc_prefix_qlinear_silu_chains(model_ir)
        _optimize_nhwc_propagation_qlinear_concat_conv(model_ir)
        _optimize_concat_pre_quantize_dequantize(model_ir)
        _optimize_transpose_mean_maxpool_concat_conv_chains(model_ir)
        _optimize_transpose_quant_dequant_bridges(model_ir)
        _optimize_boundary_input_transpose_batchmatmul_chains(model_ir)
        _optimize_leading_input_transpose_passthrough_chains(model_ir)
        _optimize_asin_transpose_passthrough_chains(model_ir)
        _optimize_erf_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nchw_nhwc_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(model_ir)
        _optimize_hardswish_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_mul_transpose_passthrough_chains(model_ir)
        _optimize_swish_transpose_passthrough_chains(model_ir)
        _optimize_gelu_tanh_transpose_passthrough_chains(model_ir)
        _optimize_center_size_offset_terminal_transpose_chains(model_ir)
        _optimize_leakyrelu_transpose_passthrough_chains(model_ir)
        _optimize_prelu_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_concat_conv_nhwc_groups(model_ir)
        _optimize_transpose_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_pre_concat_ndhwc_chains(model_ir)
        _optimize_transpose_stridedslice_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(model_ir)
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir)
        _optimize_transpose_slice_logistic_concat_reshape_tail_nhwc_chains(model_ir)
        _optimize_shufflenet_transpose_shuffle_chains(model_ir)
        _optimize_shufflenet_reshape_transpose_shuffle_nhwc_chains(model_ir)
        _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(model_ir)
        _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mulconst_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_reshape_transpose_suffix_nhwc_chains(model_ir)
        _optimize_transpose_reshape_transpose_to_expanddims_nhwc_chains(model_ir)
        _optimize_transpose_reshape_transpose_to_flatten_hw_nhwc_chains(model_ir)
        _optimize_reshape_transpose_reshape_transpose_to_nhwc_reshape_chains(model_ir)
        _optimize_attention_qkv_reshape_transpose_reshape_to_reshape_transpose_chains(model_ir)
        _optimize_attention_gather_transpose_reshape_cleanup_chains(model_ir)
        _optimize_gather_axis0_singleton_to_reshape_input_chains(model_ir)
        _optimize_attention_preproj_reshape_to_batchmatmul_ranklift_chains(model_ir)
        _optimize_window_partition_reshape_transpose_to_space_to_depth_chains(model_ir)
        _optimize_window_reverse_reshape_transpose_to_depth_to_space_chains(model_ir)
        _optimize_transpose_pre_unary_squeeze_transpose_suffix_nhwc_chains(model_ir)
        _optimize_squeeze_unary_reshape_passthrough_chains(model_ir)
        _optimize_squeeze_reshape_identity_chains(model_ir)
        _optimize_transpose_instancenorm_prepost_nhwc_chains(model_ir)
        _optimize_squeeze_unary_reshape_passthrough_chains(model_ir)
        _optimize_squeeze_reshape_identity_chains(model_ir)
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mean_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir)
        _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir)
        _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_conv_attention_nhwc_propagation_chains(model_ir)
        _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir)
        _optimize_sinet_mix_attention_double_logistic_nhwc_chains(model_ir)
        _optimize_transpose_sum_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_nested_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
        _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains(model_ir)
        _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir)
        _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir)
        _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir)
        _optimize_transpose_add_concat_const_suffix_nhwc_chains(model_ir)
        _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir)
        _optimize_transposeconv_output_nhwc_passthrough_chains(model_ir)
        _optimize_transposeconv_output_channel1_terminal_transpose_chains(model_ir)
        _optimize_transpose_unary_passthrough_chains(model_ir)
        _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
        _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_dequant_relu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
        _optimize_trailing_output_transpose_passthrough_chains(model_ir)
        _optimize_transpose_dequant_mul_add_prelu_quantize_bridges(model_ir)
        if enable_duplicate_transpose_fanout_optimizations:
            _optimize_duplicate_transpose_fanout(model_ir)
        _optimize_duplicate_reshape_fanout(model_ir)
        _optimize_transpose_dequant_prelu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_prelu_transpose_bridges(model_ir)
        _optimize_dequant_prelu_quantize_chains(model_ir)
        _optimize_dequant_prelu_depthwise_quantize_chains(model_ir)
        _optimize_dequant_transposeconv_quantize_chains(model_ir)
        _optimize_dequant_reshape_quantize_chains(model_ir)
        _optimize_dequant_hardsigmoid_quantize_chains(model_ir)
        _optimize_dequant_maxpool_quantize_chains(model_ir)
        _optimize_dequant_softmax_quantize_chains(model_ir)
        _optimize_dequant_logistic_quantize_chains(model_ir)
        _canonicalize_softmax_transpose_chains(model_ir)
        _optimize_layout_transpose_chains(model_ir)
        _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
        _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_binary_symmetric_legacy_only_bridges_safe(model_ir)
        _optimize_transpose_binary_single_post_bridges_safe(model_ir)
        _optimize_transpose_binary_mixed_fanout_bridges_safe(model_ir)
        _optimize_transpose_binary_asymmetric_fanout_bridges(model_ir)
        _optimize_transpose_binary_full_post_fanout_bridges(model_ir)
        _advance_post_progress()
    _set_post_progress_desc("core cleanup passes")
    _optimize_fuse_pseudo_leakyrelu_chains(model_ir)
    _optimize_yolo_decode_mul_square_anchor_chains(model_ir)
    _optimize_fold_consecutive_mul_constants_chains(model_ir)
    _sanitize_terminal_transpose_before_dequantize(model_ir)
    _optimize_terminal_quantize_dequantize(model_ir)
    _optimize_fold_conv_mul_add_affine_chains(
        model_ir,
        enable_conv_add_only_fold=True,
    )
    _optimize_fuse_conv_activation_chains(model_ir)
    _resolve_dynamic_reshape_shapes(model_ir)
    _optimize_squeeze_unary_reshape_passthrough_chains(model_ir)
    _optimize_squeeze_reshape_identity_chains(model_ir)
    _prune_dead_operators(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    _advance_post_progress()
    if optimize_layout_transpose_chains:
        _set_post_progress_desc("layout recovery pass-set 2")
        # Final recovery sweep:
        # some transpose-binary patterns become shape-safe only after static
        # metadata reconciliation, so run bridge passes once more.
        _optimize_transpose_mean_hardsigmoid_muladd_chains(model_ir)
        _optimize_nhwc_prefix_qlinear_silu_chains(model_ir)
        _optimize_nhwc_propagation_qlinear_concat_conv(model_ir)
        _optimize_concat_pre_quantize_dequantize(model_ir)
        _optimize_transpose_mean_maxpool_concat_conv_chains(model_ir)
        _optimize_transpose_quant_dequant_bridges(model_ir)
        _optimize_boundary_input_transpose_batchmatmul_chains(model_ir)
        _optimize_leading_input_transpose_passthrough_chains(model_ir)
        _optimize_asin_transpose_passthrough_chains(model_ir)
        _optimize_erf_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nchw_nhwc_chains(model_ir)
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(model_ir)
        _optimize_hardswish_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_transpose_passthrough_chains(model_ir)
        _optimize_hardsigmoid_mul_transpose_passthrough_chains(model_ir)
        _optimize_swish_transpose_passthrough_chains(model_ir)
        _optimize_gelu_tanh_transpose_passthrough_chains(model_ir)
        _optimize_center_size_offset_terminal_transpose_chains(model_ir)
        _optimize_leakyrelu_transpose_passthrough_chains(model_ir)
        _optimize_prelu_transpose_passthrough_chains(model_ir)
        _optimize_transpose_elementwise_concat_conv_nhwc_groups(model_ir)
        _optimize_transpose_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_pre_concat_ndhwc_chains(model_ir)
        _optimize_transpose_stridedslice_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(model_ir)
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir)
        _optimize_transpose_slice_logistic_concat_reshape_tail_nhwc_chains(model_ir)
        _optimize_shufflenet_transpose_shuffle_chains(model_ir)
        _optimize_shufflenet_reshape_transpose_shuffle_nhwc_chains(model_ir)
        _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(model_ir)
        _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mean_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir)
        _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir)
        _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_conv_attention_nhwc_propagation_chains(model_ir)
        _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir)
        _optimize_sinet_mix_attention_double_logistic_nhwc_chains(model_ir)
        _optimize_transpose_sum_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_nested_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
        _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains(model_ir)
        _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir)
        _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir)
        _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir)
        _optimize_transpose_add_concat_const_suffix_nhwc_chains(model_ir)
        _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir)
        _optimize_transposeconv_output_nhwc_passthrough_chains(model_ir)
        _optimize_transposeconv_output_channel1_terminal_transpose_chains(model_ir)
        _optimize_transpose_unary_passthrough_chains(model_ir)
        _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
        _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_dequant_relu_quantize_bridges(model_ir)
        _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
        _optimize_trailing_output_transpose_passthrough_chains(model_ir)
        _optimize_transpose_dequant_mul_add_prelu_quantize_bridges(model_ir)
        _optimize_dequant_transposeconv_quantize_chains(model_ir)
        _optimize_dequant_hardsigmoid_quantize_chains(model_ir)
        _optimize_dequant_maxpool_quantize_chains(model_ir)
        _optimize_dequant_softmax_quantize_chains(model_ir)
        _optimize_dequant_logistic_quantize_chains(model_ir)
        _canonicalize_softmax_transpose_chains(model_ir)
        _optimize_transpose_binary_symmetric_legacy_only_bridges_safe(model_ir)
        _optimize_transpose_binary_single_post_bridges_safe(model_ir)
        _optimize_transpose_binary_mixed_fanout_bridges_safe(model_ir)
        _optimize_transpose_binary_asymmetric_fanout_bridges(model_ir)
        _optimize_transpose_binary_full_post_fanout_bridges(model_ir)
        # Binary bridge recovery can recreate pre/post transpose wrappers around CONCAT.
        _optimize_transpose_elementwise_concat_conv_nhwc_groups(model_ir)
        _optimize_transpose_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_pre_concat_ndhwc_chains(model_ir)
        _optimize_transpose_stridedslice_pre_concat_nhwc_chains(model_ir)
        _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(model_ir)
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir)
        _optimize_transpose_slice_logistic_concat_reshape_tail_nhwc_chains(model_ir)
        _optimize_shufflenet_transpose_shuffle_chains(model_ir)
        _optimize_shufflenet_reshape_transpose_shuffle_nhwc_chains(model_ir)
        _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(model_ir)
        _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
        _optimize_layout_transpose_chains(model_ir)
        _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
        _optimize_transpose_unary_binary_full_post_fanout_bridges(model_ir)
        _optimize_transpose_pre_add_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
        _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pre_unary_mul_add_transpose_fanout_nhwc_chains(model_ir)
        _optimize_transpose_mean_mul_add_const_prepost_nhwc_chains(model_ir)
        _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir)
        _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir)
        _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_conv_attention_nhwc_propagation_chains(model_ir)
        _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir)
        _optimize_transpose_sum_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_nested_weighted_add_swish_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_muladd_prepost_nhwc_chains(model_ir)
        _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
        _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
        _optimize_transpose_logistic_sub_muladd_dual_postconv_nhwc_chains(model_ir)
        _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir)
        _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir)
        _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir)
        _optimize_transpose_add_concat_const_suffix_nhwc_chains(model_ir)
        _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir)
        for _ in range(2):
            rewritten_instnorm = int(
                _optimize_transpose_instancenorm_prepost_nhwc_chains(model_ir).get(
                    "optimized_transpose_instancenorm_prepost_nhwc_chains", 0
                )
            )
            rewritten_instnorm_posttranspose_bias = int(
                _optimize_transpose_instancenorm_posttranspose_bias_add_nhwc_chains(model_ir).get(
                    "optimized_transpose_instancenorm_posttranspose_bias_add_nhwc_chains", 0
                )
            )
            rewritten_instnorm_pad = int(
                _optimize_transpose_instancenorm_pad_prepost_nhwc_chains(model_ir).get(
                    "optimized_transpose_instancenorm_pad_prepost_nhwc_chains", 0
                )
            )
            rewritten_flat_globalnorm_pad = int(
                _optimize_transpose_flatten_globalnorm_pad_prepost_nhwc_chains(model_ir).get(
                    "optimized_transpose_flatten_globalnorm_pad_prepost_nhwc_chains", 0
                )
            )
            rewritten_instnorm_residual = int(
                _optimize_transpose_instancenorm_residual_add_to_single_post_adapter_nhwc_chains(model_ir).get(
                    "optimized_transpose_instancenorm_residual_add_to_single_post_adapter_nhwc_chains", 0
                )
            )
            rewritten_instnorm_residual_tail = int(
                _optimize_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains(model_ir).get(
                    "optimized_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains", 0
                )
            )
            rewritten_instnorm_dualstats_residual = int(
                _optimize_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains(model_ir).get(
                    "optimized_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains", 0
                )
            )
            if (
                rewritten_instnorm
                + rewritten_instnorm_posttranspose_bias
                + rewritten_instnorm_pad
                + rewritten_flat_globalnorm_pad
                + rewritten_instnorm_residual
                + rewritten_instnorm_residual_tail
                + rewritten_instnorm_dualstats_residual
                <= 0
            ):
                break
        _optimize_squeeze_unary_reshape_passthrough_chains(model_ir)
        _optimize_squeeze_reshape_identity_chains(model_ir)
        _prune_dead_operators(model_ir)
        _reconcile_static_tensor_shapes(model_ir)
        _advance_post_progress()
    _set_post_progress_desc("terminal cleanup passes")
    # Recovery sweeps above can re-introduce terminal TRANSPOSE->DEQUANTIZE.
    # Run terminal sanitizers once more at the very end.
    _sanitize_terminal_transpose_before_dequantize(model_ir)
    _optimize_terminal_quantize_dequantize(model_ir)
    _optimize_fold_conv_mul_add_affine_chains(
        model_ir,
        enable_conv_add_only_fold=True,
    )
    _optimize_fuse_conv_activation_chains(model_ir)
    _optimize_transpose_pre_argmax_nhwc_terminal_chains(model_ir)
    _optimize_transpose_gather_transpose_nhwc_channel_chains(model_ir)
    _optimize_terminal_softmax_transpose_after_nhwc_propagation(model_ir)
    _optimize_boundary_input_transpose_mul_sum_reshape_nhwc_chains(model_ir)
    _optimize_boundary_input_transpose_channel_slice_blocks(model_ir)
    _optimize_internal_transpose_channel_slice_nhwc_propagation_chains(model_ir)
    _optimize_transpose_channel_slice_muladd_nhwc_bridge_chains(model_ir)
    _optimize_transpose_channel_slice_dual_add_bridges_strict(model_ir)
    _optimize_transpose_slice_muladd_conv_mergeadd_strict(model_ir)
    _optimize_transpose_slice_muladd_mergeadd_posttranspose_strict(model_ir)
    _optimize_transpose_pad_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_transpose_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_concat_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_transpose_add_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_add_mean_reshape_tail_nhwc_bridge_chains(model_ir)
    _optimize_concat_tree_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_nhwc_chains(model_ir)
    _optimize_layout_transpose_chains(model_ir)
    _optimize_boundary_input_transpose_stridedslice_qdq_concat_blocks(model_ir)
    _optimize_transpose_swish_residual_concat_closure_nhwc_chains(model_ir)
    _optimize_transpose_dequant_logistic_mul_quantize_bridges(model_ir)
    _optimize_transpose_swish_qdq_nhwc_islands(model_ir)
    # Late recovery passes can recreate Conv->InstNorm(NCHW)->Pad wrappers.
    _optimize_transpose_instancenorm_posttranspose_bias_add_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_pad_prepost_nhwc_chains(model_ir)
    _optimize_transpose_flatten_globalnorm_pad_prepost_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_residual_add_to_single_post_adapter_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains(model_ir)
    _optimize_transpose_dual_mul_concat_prepost_nhwc_chains(model_ir)
    _optimize_boundary_input_layout_transposes(model_ir)
    # Boundary input layout recovery can recreate input-head PAD wrappers.
    _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
    _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
    _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
    # Final transpose-only sweep after boundary rewrites. This removes
    # inverse/identity pairs introduced late by localized layout adapters.
    _optimize_layout_transpose_chains(model_ir)
    _optimize_transpose_gather_transpose_nhwc_channel_chains(model_ir)
    if optimize_layout_transpose_chains:
        # Boundary/layout recovery can still recreate NCHW wrappers around MEAN.
        # Run dedicated NHWC passthrough once more in the terminal stage.
        _optimize_transpose_mean_prepost_nhwc_passthrough_chains(model_ir)
        _optimize_transpose_pre_unary_mean_terminal_nhwc_chains(model_ir)
        _optimize_transpose_se_conv_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_batchmatmul_affine_transpose_input_chains(model_ir)
        _optimize_batchmatmul_reshape_se_nhwc_chains(model_ir)
        _optimize_batchmatmul_transpose_input_to_adj_flags(model_ir)
        _optimize_attention_qkv_gather_reshape_transpose_hoist_chains(model_ir)
        _optimize_attention_qkv_slice_replace_gather_reshape_chains(model_ir)
        _optimize_attention_qkv_slice_to_split_chains(model_ir)
        _optimize_attention_split_post_reshape_collapse_chains(model_ir)
        _optimize_attention_qkv_shared_pretranspose_slice_nchw_chains(model_ir)
        _optimize_attention_qkv_weighted_sum_bridge_to_nhwc_chains(model_ir)
        _optimize_split_conv_concat_transpose_bridge_to_single_post_nchw(model_ir)
        _optimize_layout_transpose_chains(model_ir)
        _optimize_singleton_channel_layout_transpose_to_reshape(model_ir)
        _optimize_singleton_layout_reshape_unary_passthrough_chains(model_ir)
        _optimize_consecutive_inverse_singleton_layout_reshapes(model_ir)
        _optimize_singleton_layout_reshape_maxpool_binary_cast_chains(model_ir)
        _optimize_singleton_nms_maxpool_nhwc_chains(model_ir)
        _optimize_flatten_concat_expanddims_to_nhwc_concat(model_ir)
        _optimize_consecutive_reshape_passthrough_chains(model_ir)
        _optimize_squeeze_reshape_identity_chains(model_ir)
        _optimize_singleton_spatial_nhwc_transpose_reshape_flatten(model_ir)
        _optimize_singleton_reshape_concat_post_transpose_nhwc_chains(model_ir)
        # Run OSNet-specific multi-branch gate rewrite at terminal stage so
        # earlier generic passes do not re-wrap rewritten NHWC tensors.
        _optimize_transpose_osnet_multi_gate_muladd_prepost_nhwc_chains(model_ir)
    _optimize_maximum_minimum_relu0to1_chains(model_ir)
    # MAXIMUM/MINIMUM -> RELU_0_TO_1 canonicalization can expose fresh
    # NHWC<->NCHW unary wrappers (e.g. HardSigmoid clamp). Re-run the
    # strict transpose-unary passthrough fold once in terminal stage.
    _optimize_transpose_unary_passthrough_chains(model_ir)
    _optimize_maximum_with_zero_input2_to_relu(model_ir)
    _optimize_sinet_shuffle_residual_transpose_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
    # This generic pass can recreate SiNet concat+resize transpose adapters.
    _optimize_sinet_concat_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_dual_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_concat_resize_affine_tail_concat_transpose_chains(model_ir)
    _optimize_sinet_softmax_mask_residual_nhwc_tail_chains(model_ir)
    _optimize_transpose_mul_add_const_prelu_prepost_nhwc_terminal_chains(model_ir)
    _optimize_transpose_hardswish_se_conv_hardsigmoid_mul_prepost_nhwc_chains(model_ir)
    _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
    # Terminal MUL/ADD/PRELU rewriting can recreate NCHW bridge wrappers.
    _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
    _optimize_sinet_concat_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_dual_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_concat_resize_affine_tail_concat_transpose_chains(model_ir)
    _optimize_sinet_softmax_mask_residual_nhwc_tail_chains(model_ir)
    # Apply singleton transpose->reshape rewrite regardless of layout-opt mode.
    # This is required for fallback relowering (optimize_layout_transpose_chains=False)
    # where channelwise [1,C,1,1] -> [1,1,1,C] adapters can remain as TRANSPOSE.
    _optimize_singleton_channel_layout_transpose_to_reshape(model_ir)
    _optimize_duplicate_reshape_fanout(model_ir)
    _optimize_singleton_layout_reshape_unary_passthrough_chains(model_ir)
    _optimize_consecutive_inverse_singleton_layout_reshapes(model_ir)
    _optimize_singleton_layout_reshape_maxpool_binary_cast_chains(model_ir)
    _optimize_singleton_nms_maxpool_nhwc_chains(model_ir)
    _optimize_flatten_concat_expanddims_to_nhwc_concat(model_ir)
    _optimize_consecutive_reshape_passthrough_chains(model_ir)
    _optimize_squeeze_reshape_identity_chains(model_ir)
    _optimize_singleton_spatial_nhwc_transpose_reshape_flatten(model_ir)
    _prune_dead_operators(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    _resolve_dynamic_reshape_shapes(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # Very late shape reconciliation can expose strict shuffle-residual patterns.
    # Re-run terminal transpose reducers once at absolute end.
    _optimize_sinet_shuffle_residual_transpose_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
    _optimize_sinet_concat_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_dual_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_concat_resize_affine_tail_concat_transpose_chains(model_ir)
    _optimize_sinet_softmax_mask_residual_nhwc_tail_chains(model_ir)
    _optimize_transpose_mul_add_const_prelu_prepost_nhwc_terminal_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
    _optimize_sinet_concat_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_dual_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_concat_resize_affine_tail_concat_transpose_chains(model_ir)
    _optimize_sinet_softmax_mask_residual_nhwc_tail_chains(model_ir)
    # Final cleanup for residual transpose bridges introduced in late SiNet blocks.
    _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
    _prune_dead_operators(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # Dead-op pruning can unblock strict locality guards in this rewrite.
    _optimize_transpose_pre_add_mul_add_prelu_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_mul_add_transpose_fanout_nhwc_chains(model_ir)
    # The pre_add/prelu rewrite above can recreate SiNet dual-resize adapters.
    _optimize_sinet_concat_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_dual_resize_affine_transpose_chains(model_ir)
    _optimize_sinet_concat_resize_affine_tail_concat_transpose_chains(model_ir)
    _optimize_sinet_softmax_mask_residual_nhwc_tail_chains(model_ir)
    # Late SiNet rewrites can recreate SA/PA MIRROR_PAD NCHW<->NHWC bridges.
    _optimize_transpose_csp_attention_nhwc_chains(model_ir)
    _optimize_transpose_sa_pa_mirrorpad_nhwc_propagation_chains(model_ir)
    _optimize_batchmatmul_affine_transpose_input_chains(model_ir)
    _optimize_batchmatmul_reshape_se_nhwc_chains(model_ir)
    _optimize_batchmatmul_transpose_input_to_adj_flags(model_ir)
    _optimize_attention_qkv_gather_reshape_transpose_hoist_chains(model_ir)
    _optimize_attention_qkv_slice_replace_gather_reshape_chains(model_ir)
    _optimize_attention_qkv_slice_to_split_chains(model_ir)
    _optimize_attention_split_post_reshape_collapse_chains(model_ir)
    _optimize_attention_qkv_shared_pretranspose_slice_nchw_chains(model_ir)
    _optimize_attention_qkv_weighted_sum_bridge_to_nhwc_chains(model_ir)
    _optimize_split_conv_concat_transpose_bridge_to_single_post_nchw(model_ir)
    _optimize_sinet_mix_attention_double_logistic_nhwc_chains(model_ir)
    _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
    _optimize_transpose_3d_leaky_logistic_muladd_ndhwc_chains(model_ir)
    _optimize_transpose_conv3d_leaky_mul_unsqueeze_ndhwc_chains(model_ir)
    _optimize_transpose_cost_volume_scatter_ndhwc_chains(model_ir)
    _optimize_fold_conv_mul_add_affine_chains(
        model_ir,
        enable_conv_add_only_fold=True,
    )
    _optimize_transpose_axis3_const_concat_bridge_nhwc_chains(model_ir)
    _optimize_transpose_pre_dequant_concat_quantize_post_nhwc_chains(model_ir)
    _optimize_transpose_layernorm_stats_nhwc_propagation_chains(model_ir)
    _optimize_layernorm_stats_via_existing_post_transpose_nhwc_chains(model_ir)
    _optimize_layout_transpose_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(model_ir)
    _optimize_transpose_reshape_transpose_to_expanddims_nhwc_chains(model_ir)
    _optimize_transpose_reshape_transpose_to_flatten_hw_nhwc_chains(model_ir)
    _optimize_reshape_transpose_reshape_transpose_to_nhwc_reshape_chains(model_ir)
    _optimize_nchw_channel_shuffle_reshape_transpose_reshape_to_gather(model_ir)
    _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
    _optimize_attention_qkv_reshape_transpose_reshape_to_reshape_transpose_chains(model_ir)
    _optimize_attention_gather_transpose_reshape_cleanup_chains(model_ir)
    _optimize_gather_axis0_singleton_to_reshape_input_chains(model_ir)
    _optimize_attention_preproj_reshape_to_batchmatmul_ranklift_chains(model_ir)
    _optimize_window_partition_reshape_transpose_to_space_to_depth_chains(model_ir)
    _optimize_window_reverse_reshape_transpose_to_depth_to_space_chains(model_ir)
    _prune_dead_operators(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # Late transpose/layout rewrites can invalidate previously resolved
    # RESHAPE constants. Re-resolve once at absolute end.
    _resolve_dynamic_reshape_shapes(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # Keep HARD_SWISH outputs shape-preserving at the final serialization boundary.
    _sanitize_hardswish_tensor_shapes(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # If HARD_SWISH metadata was corrected above, allow RESHAPE constants to
    # converge one last time from the repaired input signatures.
    _resolve_dynamic_reshape_shapes(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # Late transpose/reshape rewrites can re-expose strict single-path
    # CONV_2D/DEPTHWISE_CONV_2D -> RELU(RELU6) chains.
    # Run activation fusion once more at absolute end.
    _optimize_fuse_conv_activation_chains(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    # Final boundary cleanup after all late transpose/layout rewrites.
    _optimize_boundary_input_transpose_mul_sum_reshape_nhwc_chains(model_ir)
    _optimize_internal_transpose_channel_slice_nhwc_propagation_chains(model_ir)
    _optimize_transpose_channel_slice_muladd_nhwc_bridge_chains(model_ir)
    _optimize_transpose_channel_slice_dual_add_bridges_strict(model_ir)
    _optimize_transpose_slice_muladd_conv_mergeadd_strict(model_ir)
    _optimize_transpose_slice_muladd_mergeadd_posttranspose_strict(model_ir)
    _optimize_transpose_pad_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_transpose_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_concat_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_transpose_add_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_add_mean_reshape_tail_nhwc_bridge_chains(model_ir)
    _optimize_concat_tree_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains(model_ir)
    _optimize_transpose_pre_add_nhwc_chains(model_ir)
    _optimize_layout_transpose_chains(model_ir)
    _optimize_transpose_slice_prepost_nhwc_passthrough_chains(model_ir)
    # Keep pre-concat NHWC relayout at terminal stage as late strict rewrites
    # can recreate CONCAT(axis=1)+post-transpose wrappers.
    _optimize_transpose_pre_concat_nhwc_chains(model_ir)
    _optimize_transpose_split_mixed_pre_concat_to_single_post_adapter_nhwc_chains(model_ir)
    _optimize_transpose_input_chains_pre_concat_to_single_post_adapter(model_ir)
    _optimize_transpose_concat_unary_fanout_conv_nhwc_chains(model_ir)
    _optimize_transpose_shape_extract_nhwc_to_nchw_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_transpose_elementwise_roundtrip_nhwc_nchw_fanout_chains(model_ir)
    _optimize_singleton_layout_reshape_maxpool_binary_cast_chains(model_ir)
    _optimize_singleton_nms_maxpool_nhwc_chains(model_ir)
    # Boundary cleanup can recreate a terminal no-op RESHAPE.
    # Run one last pass to remove shape-preserving single reshapes.
    _optimize_consecutive_reshape_passthrough_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_convpool_output_transpose_nhwc_passthrough_chains(model_ir)
    elif apply_safe_transpose_reduction_lite_on_no_layout_opt:
        _apply_safe_transpose_reduction_lite(model_ir)
        # Keep strict, const-only NHWC<->NCHW affine bridge folding enabled
        # in no-layout fallback so simple TRANSPOSE->MUL->ADD->TRANSPOSE
        # chains are still reduced.
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
    _optimize_transpose_dequant_hardsigmoid_quantize_bridges(model_ir)
    _optimize_transpose_pre_dequant_concat_quantize_post_nhwc_chains(model_ir)
    # Some late specialized rewrites can still leave trivial
    # NHWC->NCHW->NHWC wrappers around unary activations.
    # Run one final strict unary transpose fold before serialization guards.
    _optimize_transpose_unary_passthrough_chains(model_ir)
    _optimize_transpose_unary_fanout_inverse_post_bridges(model_ir)
    # No-layout fallback relowering can still keep strict
    # TRANSPOSE->LOGISTIC->MUL->TRANSPOSE swish wrappers (e.g. MobileViT stem).
    _optimize_swish_transpose_passthrough_chains(model_ir)
    # Conv1D shim patterns can retain:
    # TRANSPOSE -> SQUEEZE -> UNARY -> EXPAND_DIMS -> TRANSPOSE.
    # Fold them to a single rank-4 UNARY op.
    _optimize_transpose_squeeze_unary_expanddims_transpose_nhwc_chains(model_ir)
    # Variant with intermediate rank-4 transpose:
    # TRANSPOSE -> UNARY -> TRANSPOSE -> RESHAPE -> EXPAND_DIMS -> TRANSPOSE.
    _optimize_transpose_unary_transpose_reshape_expanddims_transpose_nhwc_chains(model_ir)
    # Fanout variant:
    # Keep NCHW side branch and bypass only the NHWC wrapper branch.
    _optimize_transpose_squeeze_unary_expanddims_transpose_nhwc_fanout_bypass_chains(model_ir)
    # InstanceNorm(flat) bridge variant:
    # T -> SQUEEZE -> RESHAPE -> IN -> RESHAPE -> UNARY -> EXPAND -> T
    # can be rewritten in NHWC by swapping C/W at reshape boundary.
    _optimize_transpose_squeeze_instancenorm_unary_expanddims_transpose_nhwc_chains(model_ir)
    # tencoder residual-gated branch variant:
    # dual (post-conv T->SQUEEZE) branches merge via ADD then EXPAND->T before next CONV.
    _optimize_tencoder_add_expand_transpose_conv_nhwc_chains(model_ir)
    # Conv1D tail patterns may still contain:
    # TRANSPOSE -> SQUEEZE -> UNARY* -> BATCH_MATMUL.
    # Rewire to NHWC + adjX matmul when mathematically equivalent.
    _optimize_transpose_squeeze_unary_batchmatmul_nhwc_chains(model_ir)
    # Decoder linear->deconv tails can keep:
    # BATCH_MATMUL -> ADD -> EXPAND_DIMS -> TRANSPOSE -> TRANSPOSE_CONV.
    # Transpose BATCH_MATMUL/ADD layout so TRANSPOSE before deconv is removed.
    _optimize_decoder_batchmatmul_add_expand_transpose_to_nhwc_deconv_input(model_ir)
    # Decoder terminal tails can keep:
    # TRANSPOSE -> SQUEEZE -> MEAN -> SQUEEZE.
    # Remap squeeze/mean axes in NHWC and drop the transpose.
    _optimize_transpose_squeeze_mean_squeeze_terminal_nhwc_chains(model_ir)
    # Very late transpose/layout rewrites can recreate PAD-adjacent NCHW wrappers.
    _optimize_transpose_pad_prepost_nhwc_chains(model_ir)
    _optimize_transpose_unary_pad_prepost_to_single_adapter_nhwc_chains(model_ir)
    _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_posttranspose_bias_add_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains(model_ir)
    # Norm-subgraph fallback rewrites can introduce channelwise
    # [1,C,1,1]->[1,1,1,C] adapters as TRANSPOSE in no-layout mode.
    # Re-canonicalize them to RESHAPE at the very end.
    _optimize_singleton_channel_layout_transpose_to_reshape(model_ir)
    _optimize_duplicate_reshape_fanout(model_ir)
    _optimize_consecutive_reshape_passthrough_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_layout_transpose_chains(model_ir)
    _repair_rank4_channelwise_broadcast_constants_to_runtime_layout(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    _realign_dynamic_boundary_shape_signature_map(model_ir)
    # Keep final serialized metadata consistent for tools that render
    # shape_signature (e.g. Netron): HARD_SWISH is shape-preserving.
    _sanitize_hardswish_tensor_shapes(model_ir)
    # Final guardrail for runtime validity: ensure SQUEEZE axes target
    # singleton dimensions after all late layout/shape rewrites.
    _sanitize_squeeze_axes_with_static_input_shapes(model_ir)
    _sanitize_wrong_way_nchw_to_nhwc_transpose_before_conv(model_ir)
    _repair_rank4_binary_layout_mismatch_with_transpose_adapter(model_ir)
    _optimize_singleton_channel_layout_transpose_to_reshape(model_ir)
    _optimize_duplicate_reshape_fanout(model_ir)
    _optimize_consecutive_reshape_passthrough_chains(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    _sanitize_static_shape_signature_consistency(model_ir)
    _repair_rank4_binary_layout_mismatch_with_transpose_adapter(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    if optimize_layout_transpose_chains or apply_safe_transpose_reduction_lite_on_no_layout_opt:
        # Late binary-layout repairs can recreate PRELU bridge wrappers.
        # Run PRELU transpose passthrough once more before final safety checks,
        # including the safe-reduction fallback relowering path.
        _optimize_prelu_transpose_passthrough_chains(model_ir)
        _optimize_transpose_dual_pre_add_to_single_post_adapter_nhwc_chains(model_ir)
        _optimize_terminal_transpose_mul_add_reshape_fc_nhwc_chains(model_ir)
        if optimize_layout_transpose_chains:
            _optimize_terminal_transpose_prelu_reshape_batchmatmul_nhwc_chains(model_ir)
        # PRELU rewrite can recreate strict TRANSPOSE->MUL(const)->ADD(const)->TRANSPOSE
        # bridges on legacy branches. Fold them again in this final stage.
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        if optimize_layout_transpose_chains:
            _optimize_layout_transpose_chains(model_ir)
        _reconcile_static_tensor_shapes(model_ir)
    # Keep this at absolute end of optimization pipeline: several late
    # shape/layout repair passes can recreate the exact tail pattern.
    _optimize_transpose_instancenorm_posttranspose_bias_add_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_residual_mul_concat_conv_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_dualstats_residual_add_resize_nhwc_chains(model_ir)
    # Late bridge rewrites above can recreate strict
    # TRANSPOSE->MUL(const)->ADD(const)->TRANSPOSE fragments.
    _optimize_fold_mul_add_mul_affine_chains(model_ir)
    _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
    _optimize_concat_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_transpose_add_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_add_mean_reshape_tail_nhwc_bridge_chains(model_ir)
    _optimize_concat_tree_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_transpose_pre_add_nhwc_chains(model_ir)
    _optimize_transpose_channel_slice_dual_add_bridges_strict(model_ir)
    _optimize_transpose_slice_muladd_conv_mergeadd_strict(model_ir)
    _optimize_transpose_slice_muladd_mergeadd_posttranspose_strict(model_ir)
    _optimize_transpose_pad_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_transpose_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains(model_ir)
    # Strict slice/merge cleanup above can recreate simple affine bridge tails:
    # TRANSPOSE->MUL(const)->TRANSPOSE->ADD(const).
    _optimize_fold_mul_add_mul_affine_chains(model_ir)
    _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
    # Keep this after pre_add/slice/pad strict rewrites: those passes can
    # recreate CONCAT->MUL->TRANSPOSE->ADD NHWC bridge tails.
    _optimize_concat_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_transpose_add_nhwc_bridge_chains(model_ir)
    _optimize_concat_mul_add_add_mean_reshape_tail_nhwc_bridge_chains(model_ir)
    _optimize_concat_tree_mul_add_transpose_nhwc_bridge_chains(model_ir)
    _optimize_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains(model_ir)
    _optimize_transpose_concat_unary_fanout_conv_nhwc_chains(model_ir)
    _optimize_transpose_shape_extract_nhwc_to_nchw_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_layout_transpose_chains(model_ir)
    # Keep QKV bridge reductions at the terminal stage: some late strict
    # transpose/add/slice rewrites above can recreate this exact motif.
    _optimize_attention_qkv_shared_pretranspose_slice_nchw_chains(model_ir)
    _optimize_attention_qkv_weighted_sum_bridge_to_nhwc_chains(model_ir)
    _optimize_split_conv_concat_transpose_bridge_to_single_post_nchw(model_ir)
    _optimize_transpose_hardswish_se_conv_hardsigmoid_mul_prepost_nhwc_chains(model_ir)
    # Late affine/fusion cleanups can recreate
    # TRANSPOSE->(ADD/MUL hard-sigmoid-like)->MUL->TRANSPOSE wrappers.
    # Run strict hard-sigmoid transpose passthrough once more at terminal stage.
    _optimize_hardsigmoid_mul_transpose_passthrough_chains(model_ir)
    _optimize_hardsigmoid_transpose_passthrough_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_layout_transpose_chains(model_ir)
    # Absolute-end cleanup: late bridge rewrites can recreate strict
    # pre/post CONCAT transpose wrappers and SHAPE-extract transposes.
    _optimize_transpose_pre_concat_nhwc_chains(model_ir)
    _optimize_transpose_shape_extract_nhwc_to_nchw_chains(model_ir)
    if optimize_layout_transpose_chains:
        _optimize_layout_transpose_chains(model_ir)
    _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
    _optimize_constant_input_cast_chains(model_ir)
    _optimize_redundant_int64_to_int32_cast_chains(model_ir)
    _replace_expand_dims_and_squeeze_with_reshape(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    _advance_post_progress()

    _repair_unbound_nonconstant_operator_inputs_with_layout_transpose(model_ir)
    # The late unbound-input repair can inject strict
    # NHWC->NCHW->NHWC MUL/ADD wrappers (repair_perm tensors).
    # Fold them again before final shape/topology reconciliation.
    _optimize_transpose_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
    _optimize_constant_input_cast_chains(model_ir)
    _optimize_redundant_int64_to_int32_cast_chains(model_ir)
    _optimize_transpose_flatten_globalnorm_pad_prepost_nhwc_chains(model_ir)
    # Very late terminal bridge/transpose rewrites above can still stale out
    # RESHAPE constant inputs. Re-resolve once immediately before final sort.
    _resolve_dynamic_reshape_shapes(
        model_ir,
        prefer_runtime_inferable_from_onnx_raw=True,
    )
    _reconcile_static_tensor_shapes(model_ir)
    split_fallback_stats = _replace_unsupported_split_with_slice(model_ir)
    if int(split_fallback_stats.get("replaced_unsupported_split_with_slice", 0)) > 0:
        _reconcile_static_tensor_shapes(model_ir)

    # Safety fallback:
    # Some aggressive transpose/layout rewrites can leave dangling dynamic inputs
    # (tensor is consumed but has no producer and no embedded constant data),
    # which later appears as unexpected runtime-fed input in TFLite.
    # When detected, rebuild with transpose-chain optimization disabled to
    # preserve graph connectivity.
    unbound_inputs = _find_unbound_nonconstant_operator_inputs(model_ir)
    if optimize_layout_transpose_chains and len(unbound_inputs) > 0:
        _set_post_progress_desc("fallback relowering")
        _advance_post_progress()
        if post_progress_bar is not None:
            post_progress_bar.close()
        fallback_ir = lower_onnx_to_ir(
            onnx_graph=onnx_graph,
            output_file_name=output_file_name,
            allow_custom_ops=allow_custom_ops,
            custom_op_allowlist=custom_op_allowlist,
            optimize_layout_transpose_chains=False,
            transpose_inputs_to_nhwc=transpose_inputs_to_nhwc,
            keep_ncw_or_nchw_or_ncdhw_input_names=keep_ncw_or_nchw_or_ncdhw_input_names,
            keep_nwc_or_nhwc_or_ndhwc_input_names=keep_nwc_or_nhwc_or_ndhwc_input_names,
            keep_shape_absolutely_input_names=keep_shape_absolutely_input_names,
            disable_group_convolution=disable_group_convolution,
            output_nms_with_argmax=output_nms_with_argmax,
            switch_nms_version=switch_nms_version,
            show_progress=show_progress,
            apply_safe_transpose_reduction_lite_on_no_layout_opt=False,
            number_of_dimensions_after_flextranspose_compression=number_of_dimensions_after_flextranspose_compression,
            number_of_dimensions_after_flexstridedslice_compression=number_of_dimensions_after_flexstridedslice_compression,
            protected_boundary_tensor_names=protected_boundary_tensor_names,
        )
        fallback_norm_stats = _optimize_transpose_norm_subgraph_pad_prepost_nhwc_chains(fallback_ir)
        if int(fallback_norm_stats.get("optimized_transpose_norm_subgraph_pad_prepost_nhwc_chains", 0)) > 0:
            _repair_rank4_binary_layout_mismatch_with_transpose_adapter(fallback_ir)
            _optimize_singleton_channel_layout_transpose_to_reshape(fallback_ir)
            _optimize_duplicate_reshape_fanout(fallback_ir)
            _optimize_consecutive_reshape_passthrough_chains(fallback_ir)
            _reconcile_static_tensor_shapes(fallback_ir)
            _topologically_sort_operators(fallback_ir)
        fallback_ir.metadata["layout_optimize_fallback"] = {
            "reason": "dangling_dynamic_inputs_detected",
            "count": int(len(unbound_inputs)),
            "samples": [dict(v) for v in list(unbound_inputs[:8])],
        }
        return fallback_ir

    _set_post_progress_desc("topological sort")
    _topologically_sort_operators(model_ir)
    if apply_safe_transpose_reduction_lite_on_no_layout_opt:
        # In no-layout fallback path, some strict MUL/ADD affine bridges become
        # reducible only after final topological normalization.
        _optimize_transpose_se_fc_mul_prepost_nhwc_chains(model_ir)
        _optimize_transpose_mul_add_const_prepost_nhwc_chains(model_ir)
        _topologically_sort_operators(model_ir)
    # Final boundary-signature restore:
    # late static-shape reconciliations may overwrite graph-boundary dynamic
    # contracts (e.g. NMS selected_indices leading axis).
    _realign_dynamic_boundary_shape_signature_map(model_ir)
    _sanitize_static_shape_signature_consistency(model_ir)
    # Absolute-final guard: topological sort + signature sanitize can expose
    # one more strict TRANSPOSE->MUL(const)->TRANSPOSE->ADD(const) fragment.
    _optimize_transpose_mul_posttranspose_add_nhwc_chains(model_ir)
    _optimize_transpose_instancenorm_posttranspose_bias_add_nhwc_chains(model_ir)
    _optimize_transpose_flatten_globalnorm_pad_prepost_nhwc_chains(model_ir)
    _topologically_sort_operators(model_ir)
    _advance_post_progress()
    if post_progress_bar is not None:
        post_progress_bar.close()

    return model_ir
