from __future__ import annotations

import copy
import re
import threading
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,
    infer_model_ir_logical_layouts,
    is_channel_first_logical_layout,
    is_channel_last_logical_layout,
    normalize_logical_layout,
    normalize_onnx_shape,
    validate_model_ir_layout_annotations,
)
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 _rewrite_constant_divisors_to_multiplicative_reciprocals(
    model_ir: ModelIR,
) -> Dict[str, int]:
    integer_output_dtypes = {
        "INT8",
        "INT16",
        "INT32",
        "INT64",
        "UINT8",
        "UINT16",
        "UINT32",
        "UINT64",
    }
    changed_count = 0
    existing_tensor_names = set(str(name) for name in model_ir.tensors.keys())
    consumer_index: Dict[str, List[OperatorIR]] = {}
    for candidate_op in model_ir.operators:
        for input_name in candidate_op.inputs:
            consumer_index.setdefault(str(input_name), []).append(candidate_op)

    def _unique_tensor_name(base_name: str) -> str:
        candidate = str(base_name)
        suffix = 0
        while candidate in existing_tensor_names:
            suffix += 1
            candidate = f"{base_name}_{suffix}"
        existing_tensor_names.add(candidate)
        return candidate

    rewritten_ops: List[OperatorIR] = []
    for op in model_ir.operators:
        if str(op.op_type) != "DIV" or len(op.inputs) != 2 or len(op.outputs) != 1:
            rewritten_ops.append(op)
            continue

        lhs_name = str(op.inputs[0])
        rhs_name = str(op.inputs[1])
        output_name = str(op.outputs[0])
        rhs_tensor = model_ir.tensors.get(rhs_name, None)
        output_tensor = model_ir.tensors.get(output_name, None)
        lhs_tensor = model_ir.tensors.get(lhs_name, None)
        if (
            rhs_tensor is None
            or not isinstance(rhs_tensor.data, np.ndarray)
            or output_tensor is None
        ):
            rewritten_ops.append(op)
            continue

        output_dtype = str(output_tensor.dtype).upper()
        if output_dtype in integer_output_dtypes:
            rewritten_ops.append(op)
            continue
        direct_consumers = consumer_index.get(output_name, [])
        if any(
            str(consumer.op_type) == "CAST"
            and str(consumer.options.get("outDataType", "")).upper() in integer_output_dtypes
            for consumer in direct_consumers
        ):
            rewritten_ops.append(op)
            continue

        calc_dtype = "FLOAT16" if output_dtype == "FLOAT16" else "FLOAT32"
        np_calc_dtype = np.float16 if calc_dtype == "FLOAT16" else np.float32
        reciprocal_name = _unique_tensor_name(f"{output_name}_div_reciprocal")
        model_ir.tensors[reciprocal_name] = TensorIR(
            name=reciprocal_name,
            dtype=calc_dtype,
            shape=[int(v) for v in list(np.asarray(rhs_tensor.data, dtype=np_calc_dtype).shape)],
            shape_signature=[int(v) for v in list(np.asarray(rhs_tensor.data, dtype=np_calc_dtype).shape)],
            data=np.asarray(
                np.reciprocal(np.asarray(rhs_tensor.data, dtype=np_calc_dtype)),
                dtype=np_calc_dtype,
            ),
            logical_layout=normalize_logical_layout(rhs_tensor.logical_layout),
        )

        mul_lhs_name = lhs_name
        lhs_dtype = str(lhs_tensor.dtype).upper() if lhs_tensor is not None else "FLOAT32"
        if lhs_dtype != calc_dtype:
            lhs_shape = (
                [int(v) for v in list(lhs_tensor.shape)]
                if lhs_tensor is not None
                else [int(v) for v in list(output_tensor.shape)]
            )
            lhs_shape_signature = (
                [int(v) for v in list(lhs_tensor.shape_signature)]
                if lhs_tensor is not None and lhs_tensor.shape_signature is not None
                else list(lhs_shape)
            )
            lhs_cast_name = _unique_tensor_name(f"{output_name}_div_lhs_cast")
            model_ir.tensors[lhs_cast_name] = TensorIR(
                name=lhs_cast_name,
                dtype=calc_dtype,
                shape=lhs_shape,
                shape_signature=lhs_shape_signature,
                quantization=copy.deepcopy(lhs_tensor.quantization) if lhs_tensor is not None else None,
                logical_layout=normalize_logical_layout(
                    lhs_tensor.logical_layout if lhs_tensor is not None else output_tensor.logical_layout
                ),
            )
            rewritten_ops.append(
                OperatorIR(
                    op_type="CAST",
                    inputs=[lhs_name],
                    outputs=[lhs_cast_name],
                    options={
                        "inDataType": lhs_dtype,
                        "outDataType": calc_dtype,
                    },
                )
            )
            mul_lhs_name = lhs_cast_name

        mul_out_name = output_name
        if output_dtype != calc_dtype:
            mul_out_name = _unique_tensor_name(f"{output_name}_div_mul_out")
            model_ir.tensors[mul_out_name] = TensorIR(
                name=mul_out_name,
                dtype=calc_dtype,
                shape=[int(v) for v in list(output_tensor.shape)],
                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)]
                ),
                quantization=None,
                logical_layout=normalize_logical_layout(output_tensor.logical_layout),
            )

        rewritten_ops.append(
            OperatorIR(
                op_type="MUL",
                inputs=[mul_lhs_name, reciprocal_name],
                outputs=[mul_out_name],
                options=dict(op.options),
            )
        )

        if mul_out_name != output_name:
            rewritten_ops.append(
                OperatorIR(
                    op_type="CAST",
                    inputs=[mul_out_name],
                    outputs=[output_name],
                    options={
                        "inDataType": calc_dtype,
                        "outDataType": output_dtype,
                    },
                )
            )
        changed_count += 1

    if changed_count > 0:
        model_ir.operators = rewritten_ops
    return {
        "rewritten_constant_div_to_mul": int(changed_count),
    }


class _ProgressSpinner:
    def __init__(self, progress_bar: Any) -> None:
        self._progress_bar = progress_bar
        self._stop_event = threading.Event()
        self._thread: Optional[threading.Thread] = None

    def start(self) -> None:
        self.stop()
        if self._progress_bar is None:
            return
        self._stop_event = threading.Event()
        self._progress_bar.set_postfix_str("|", refresh=True)
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        thread = self._thread
        if thread is not None:
            self._stop_event.set()
            thread.join(timeout=0.5)
        self._thread = None
        if self._progress_bar is not None:
            self._progress_bar.set_postfix_str("", refresh=True)

    def _run(self) -> None:
        frames = ["|", "/", "-", "\\"]
        frame_index = 0
        while not self._stop_event.wait(0.1):
            if self._progress_bar is None:
                return
            frame_index = (frame_index + 1) % len(frames)
            self._progress_bar.set_postfix_str(frames[frame_index], refresh=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 _evaluate_constant_pool2d(
    *,
    input_data: Any,
    op_type: str,
    padding: str,
    stride_h: int,
    stride_w: int,
    filter_h: int,
    filter_w: int,
) -> Optional[np.ndarray]:
    try:
        x = np.asarray(input_data)
    except Exception:
        return None
    if x.ndim != 4:
        return None
    if int(stride_h) <= 0 or int(stride_w) <= 0 or int(filter_h) <= 0 or int(filter_w) <= 0:
        return None

    x_work = np.asarray(x)
    n_dim, in_h, in_w, channels = [int(v) for v in list(x_work.shape)]
    if min(n_dim, in_h, in_w, channels) <= 0:
        return None

    padding_mode = str(padding).upper()
    if padding_mode not in {"SAME", "VALID"}:
        return None

    if padding_mode == "SAME":
        out_h = int(np.ceil(float(in_h) / float(stride_h)))
        out_w = int(np.ceil(float(in_w) / float(stride_w)))
        pad_along_h = max((int(out_h) - 1) * int(stride_h) + int(filter_h) - int(in_h), 0)
        pad_along_w = max((int(out_w) - 1) * int(stride_w) + int(filter_w) - int(in_w), 0)
        pad_top = int(pad_along_h // 2)
        pad_bottom = int(pad_along_h - pad_top)
        pad_left = int(pad_along_w // 2)
        pad_right = int(pad_along_w - pad_left)
    else:
        out_h = ((int(in_h) - int(filter_h)) // int(stride_h)) + 1
        out_w = ((int(in_w) - int(filter_w)) // int(stride_w)) + 1
        pad_top = pad_bottom = pad_left = pad_right = 0

    if int(out_h) <= 0 or int(out_w) <= 0:
        return None

    pool_type = str(op_type).upper()
    if pool_type == "AVERAGE_POOL_2D":
        x_eval = np.asarray(x_work, dtype=np.float64)
        pad_value = 0.0
        out_dtype = np.float64
        reduce_fn = np.mean
    elif pool_type == "MAX_POOL_2D":
        if x_work.dtype == np.bool_:
            x_eval = np.asarray(x_work, dtype=np.bool_)
            pad_value = False
            out_dtype = np.bool_
        elif np.issubdtype(x_work.dtype, np.integer):
            x_eval = np.asarray(x_work)
            pad_value = np.iinfo(x_eval.dtype).min
            out_dtype = x_eval.dtype
        else:
            x_eval = np.asarray(x_work, dtype=np.float64)
            pad_value = -np.inf
            out_dtype = np.float64
        reduce_fn = np.max
    else:
        return None

    x_eval = np.pad(
        x_eval,
        ((0, 0), (pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
        mode="constant",
        constant_values=pad_value,
    )
    out = np.empty((int(n_dim), int(out_h), int(out_w), int(channels)), dtype=out_dtype)
    for out_y in range(int(out_h)):
        start_y = int(out_y) * int(stride_h)
        end_y = start_y + int(filter_h)
        for out_x in range(int(out_w)):
            start_x = int(out_x) * int(stride_w)
            end_x = start_x + int(filter_w)
            window = x_eval[:, start_y:end_y, start_x:end_x, :]
            out[:, out_y, out_x, :] = reduce_fn(window, axis=(1, 2))
    return np.asarray(out)


def _evaluate_constant_pad(
    *,
    input_data: Any,
    pads_data: Any,
    pad_value_data: Optional[Any] = None,
) -> Optional[np.ndarray]:
    try:
        x = np.asarray(input_data)
        pads = np.asarray(pads_data, dtype=np.int64)
    except Exception:
        return None
    if pads.ndim != 2 or pads.shape[1] != 2 or pads.shape[0] != x.ndim:
        return None
    if np.any(pads < 0):
        return None

    if pad_value_data is None:
        if x.dtype == np.bool_:
            pad_value: Any = False
        else:
            pad_value = 0
    else:
        try:
            pad_value_arr = np.asarray(pad_value_data)
        except Exception:
            return None
        if pad_value_arr.size != 1:
            return None
        pad_value = pad_value_arr.reshape(-1)[0].item()

    try:
        return np.pad(
            x,
            [(int(v[0]), int(v[1])) for v in pads.tolist()],
            mode="constant",
            constant_values=pad_value,
        )
    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_orphan_recurrent_step_tensors(model_ir: ModelIR) -> Dict[str, int]:
    """
    Repair orphan recurrent step aliases left behind by late graph-output rewrites.

    Some recurrent lowerings materialize both:
    - an internal step tensor such as `*_h_step_9`, and
    - a user-visible graph output such as `hidden_out`
      from the same RESHAPE(input, `*_h_step_shape_9`).

    Late cleanup can collapse the producing RESHAPE to the graph output name
    while leaving internal CONCAT users still pointing at the orphaned
    `*_h_step_*` tensor. That later surfaces as `Input tensor N lacks data` in
    TFLite because the orphan tensor has no producer and no constant buffer.
    """
    repaired = 0
    producer_index = _build_tensor_producer_map(model_ir)
    consumers = _build_tensor_consumer_map(model_ir)
    model_inputs = {str(v) for v in list(model_ir.inputs)}
    model_outputs = {str(v) for v in list(model_ir.outputs)}

    for tensor_name in list(model_ir.tensors.keys()):
        tensor_name = str(tensor_name)
        if tensor_name in producer_index or tensor_name in model_inputs:
            continue
        match = re.match(r"^(.+_(?:h|c)_step_)(\d+)$", tensor_name)
        if match is None:
            continue
        shape_tensor_name = f"{match.group(1)}shape_{match.group(2)}"
        replacement_name: Optional[str] = None
        for op in model_ir.operators:
            if str(op.op_type) != "RESHAPE" or len(op.inputs) < 2 or len(op.outputs) != 1:
                continue
            if str(op.inputs[1]) != shape_tensor_name:
                continue
            candidate_name = str(op.outputs[0])
            if candidate_name == tensor_name:
                replacement_name = None
                break
            replacement_name = candidate_name
            break
        if replacement_name is None:
            continue
        for consumer_idx in consumers.get(tensor_name, []):
            consumer = model_ir.operators[int(consumer_idx)]
            consumer.inputs = [
                replacement_name if str(input_name) == tensor_name else str(input_name)
                for input_name in consumer.inputs
            ]
        if tensor_name not in model_outputs:
            model_ir.tensors.pop(tensor_name, None)
        repaired += 1
        producer_index = _build_tensor_producer_map(model_ir)
        consumers = _build_tensor_consumer_map(model_ir)

    return {"repaired_orphan_recurrent_step_tensors": int(repaired)}


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_resize_add_concat_affine_conv_spp_nhwc_chains,
        _optimize_transpose_input_chains_pre_concat_to_single_post_adapter,
        _optimize_transpose_pre_add_nhwc_chains,
        _optimize_transpose_mean_mul_reshape_add_conv_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],
        allow_zero: Optional[bool] = None,
    ) -> List[int]:
        sanitized = [int(v) for v in list(template)]
        for idx, dim in enumerate(list(sanitized)):
            if int(dim) == 0:
                if allow_zero is False:
                    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

        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
                and any(int(dim) <= 0 for dim in input_signature)
            ):
                # 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,
                    allow_zero=(
                        bool(op.options.get("allowZero"))
                        if "allowZero" in op.options
                        else None
                    ),
                )
            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,
                        allow_zero=(
                            bool(op.options.get("allowZero"))
                            if "allowZero" in op.options
                            else None
                        ),
                    ),
                    input_signature=signature_for_resolve,
                    allow_zero=(
                        bool(op.options.get("allowZero"))
                        if "allowZero" in op.options
                        else None
                    ),
                )
        else:
            if (
                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-inferable `-1` templates for shape-tensor driven
                # reshapes when the source rank is still dynamic. Concretizing
                # them from placeholder static metadata can freeze valid dynamic
                # extents to `1` and break runtime element counts.
                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,
                        allow_zero=(
                            bool(op.options.get("allowZero"))
                            if "allowZero" in op.options
                            else None
                        ),
                    ),
                    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,
                allow_zero=(
                    bool(op.options.get("allowZero"))
                    if "allowZero" in op.options
                    else None
                ),
            )
            if (
                sum(1 for dim in fallback_shape if int(dim) == -1) <= 1
                and all(int(dim) >= -1 for dim in fallback_shape)
                and not any(int(dim) == 0 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 _rewrite_dynamic_rank1_unsqueeze_reshape_shape_inputs(
    model_ir: ModelIR,
) -> Dict[str, int]:
    def _unique_tensor_name(base: str) -> str:
        if str(base) not in model_ir.tensors:
            return str(base)
        index = 1
        while f"{base}_{index}" in model_ir.tensors:
            index += 1
        return f"{base}_{index}"

    rewritten = 0
    rewritten_ops: List[OperatorIR] = []
    for op in model_ir.operators:
        if str(op.op_type) != "RESHAPE" or len(op.inputs) < 2 or len(op.outputs) != 1:
            rewritten_ops.append(op)
            continue

        input_name = str(op.inputs[0])
        shape_name = str(op.inputs[1])
        output_name = str(op.outputs[0])
        input_tensor = model_ir.tensors.get(input_name, None)
        output_tensor = model_ir.tensors.get(output_name, None)
        shape_tensor = model_ir.tensors.get(shape_name, None)
        if input_tensor is None or output_tensor is None or shape_tensor is None:
            rewritten_ops.append(op)
            continue
        if shape_tensor.data is None:
            rewritten_ops.append(op)
            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_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)]
        )
        try:
            shape_values = [
                int(v) for v in np.asarray(shape_tensor.data).reshape(-1).tolist()
            ]
        except Exception:
            rewritten_ops.append(op)
            continue

        if len(input_signature) != 1:
            rewritten_ops.append(op)
            continue
        if output_signature != [-1, 1] and output_signature != [1, -1]:
            rewritten_ops.append(op)
            continue
        if (
            shape_values not in ([-1, 1], [1, -1], [1, 1])
            and list(op.options.get("newShape", [])) not in ([], [1, 1])
        ):
            rewritten_ops.append(op)
            continue

        runtime_shape_name = _unique_tensor_name(f"{output_name}_runtime_shape")
        runtime_shape_tensor = TensorIR(
            name=runtime_shape_name,
            dtype="INT32",
            shape=[1],
            shape_signature=[1],
        )
        model_ir.tensors[runtime_shape_name] = runtime_shape_tensor
        rewritten_ops.append(
            OperatorIR(
                op_type="SHAPE",
                inputs=[input_name],
                outputs=[runtime_shape_name],
                options={"outType": "INT32"},
            )
        )

        one_name = _unique_tensor_name(f"{output_name}_unsqueeze_runtime_one")
        model_ir.tensors[one_name] = TensorIR(
            name=one_name,
            dtype="INT32",
            shape=[1],
            shape_signature=[1],
            data=np.asarray([1], dtype=np.int32),
        )
        merged_shape_name = _unique_tensor_name(f"{output_name}_unsqueeze_runtime_shape")
        model_ir.tensors[merged_shape_name] = TensorIR(
            name=merged_shape_name,
            dtype="INT32",
            shape=[2],
            shape_signature=[2],
        )
        concat_inputs = [runtime_shape_name, one_name]
        if output_signature == [1, -1]:
            concat_inputs = [one_name, runtime_shape_name]
        rewritten_ops.append(
            OperatorIR(
                op_type="CONCATENATION",
                inputs=concat_inputs,
                outputs=[merged_shape_name],
                options={
                    "axis": 0,
                    "fusedActivationFunction": "NONE",
                },
            )
        )

        op.inputs[1] = merged_shape_name
        op.options["newShape"] = []
        rewritten_ops.append(op)
        rewritten += 1

    if rewritten > 0:
        model_ir.operators = rewritten_ops
    return {"rewritten_dynamic_rank1_unsqueeze_reshape_shape_inputs": int(rewritten)}


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)
                output_tensor = model_ir.tensors.get(outputs[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):
                    out_signature = (
                        [int(v) for v in list(output_tensor.shape_signature)]
                        if output_tensor is not None
                        and output_tensor.shape_signature is not None
                        and len(list(output_tensor.shape_signature)) == len(resolved)
                        else [int(v) for v in list(resolved)]
                    )
                    if len(input_signature) == len(resolved) and input_tensor.shape is not None:
                        input_shape = [int(v) for v in list(input_tensor.shape)]
                        for axis, dim in enumerate(input_signature):
                            if (
                                int(dim) < 0
                                and axis < len(input_shape)
                                and int(input_shape[axis]) == int(resolved[axis])
                            ):
                                out_signature[axis] = -1
                    changed |= _update_tensor_shape(outputs[0], resolved, out_signature)
                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_resize_add_concat_affine_conv_spp_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Collapse SPP-style NCHW resize/add/concat/affine bridges back to NHWC.

    Strict target:
      base_nhwc --T--> base_nchw
      resize_i_nhwc --T--> resize_i_nchw
      ADD(base_nchw, resize_i_nchw) -> add_i_nchw (4 branches)
      CONCAT(axis=1, add_i_nchw...) -> cat0_nchw
      MUL(cat0_nchw, const_nchw) -> mul0_nchw --T--> mul0_nhwc
      ADD(mul0_nhwc, bias_nhwc) -> ... -> CONV -> conv0_nhwc --T--> conv0_nchw
      CONCAT(axis=1, base_nchw, conv0_nchw) -> cat1_nchw
      MUL(cat1_nchw, const_nchw) -> mul1_nchw --T--> mul1_nhwc
      ADD(mul1_nhwc, bias_nhwc) -> ... -> CONV

    Rewrite:
      Keep both concat islands in NHWC, rewrite CONCAT axis to 3,
      permute channelwise MUL constants to NHWC, and remove the transposes.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _channelwise_const_name(binary_op: OperatorIR, tensor_name_to_skip: str) -> Optional[str]:
        for input_name in [str(v) for v in list(binary_op.inputs)]:
            if str(input_name) == str(tensor_name_to_skip):
                continue
            tensor = model_ir.tensors.get(str(input_name), None)
            if tensor is None or tensor.data is None:
                continue
            const_arr = np.asarray(tensor.data)
            if const_arr.ndim == 4:
                return str(input_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.inputs) != 4 or len(concat0_op.outputs) != 1:
                continue
            axis0 = int(concat0_op.options.get("axis", 1))
            if axis0 < 0:
                axis0 += 4
            if axis0 != 1:
                continue

            concat0_out_name = str(concat0_op.outputs[0])
            if concat0_out_name in model_outputs:
                continue

            branch_add_indices: List[int] = []
            base_nchw_name: Optional[str] = None
            resize_post_indices: List[int] = []
            resize_nhwc_names: List[str] = []
            valid = True

            for add_out_name in [str(v) for v in list(concat0_op.inputs)]:
                add_idx = producers.get(add_out_name, None)
                if add_idx is None:
                    valid = 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
                    or str(add_op.outputs[0]) != str(add_out_name)
                    or set(int(v) for v in consumers.get(str(add_out_name), [])) != {int(concat0_idx)}
                ):
                    valid = False
                    break

                add_inputs = [str(v) for v in list(add_op.inputs)]
                transpose_inputs: List[Tuple[int, str, str]] = []
                for input_name in add_inputs:
                    pre_idx = producers.get(str(input_name), None)
                    if pre_idx is None:
                        continue
                    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(pre_op.outputs[0]) not in model_outputs
                    ):
                        transpose_inputs.append((int(pre_idx), str(pre_op.inputs[0]), str(pre_op.outputs[0])))
                if len(transpose_inputs) != 2:
                    valid = False
                    break

                branch_add_indices.append(int(add_idx))
                if base_nchw_name is None:
                    names = [name for _, _, name in transpose_inputs]
                    for candidate_name in names:
                        if all(
                            candidate_name in [str(v) for v in list(model_ir.operators[int(producers[str(other_out)])].inputs)]
                            for other_out in list(concat0_op.inputs)
                            if producers.get(str(other_out), None) is not None
                        ):
                            base_nchw_name = str(candidate_name)
                            break
                if base_nchw_name is None:
                    valid = False
                    break

                branch_info = None
                for pre_idx, nhwc_name, nchw_name in transpose_inputs:
                    if str(nchw_name) == str(base_nchw_name):
                        continue
                    branch_prod_idx = producers.get(str(nhwc_name), None)
                    if branch_prod_idx is None:
                        continue
                    branch_prod = model_ir.operators[int(branch_prod_idx)]
                    if str(branch_prod.op_type) == "RESIZE_BILINEAR":
                        branch_info = (int(pre_idx), str(nhwc_name), str(nchw_name))
                        break
                if branch_info is None:
                    valid = False
                    break
                resize_post_indices.append(int(branch_info[0]))
                resize_nhwc_names.append(str(branch_info[1]))

            if not valid or base_nchw_name is None:
                continue

            base_pre_idx = producers.get(str(base_nchw_name), None)
            if base_pre_idx is None:
                continue
            base_pre_op = model_ir.operators[int(base_pre_idx)]
            if (
                str(base_pre_op.op_type) != "TRANSPOSE"
                or len(base_pre_op.inputs) < 2
                or len(base_pre_op.outputs) != 1
                or _read_transpose_perm(model_ir, base_pre_op) != perm_nhwc_to_nchw
            ):
                continue
            base_nhwc_name = str(base_pre_op.inputs[0])

            concat0_users = [int(v) for v in consumers.get(str(concat0_out_name), [])]
            if len(concat0_users) != 1:
                continue
            mul0_idx = int(concat0_users[0])
            mul0_op = model_ir.operators[int(mul0_idx)]
            if str(mul0_op.op_type) != "MUL" or len(mul0_op.inputs) != 2 or len(mul0_op.outputs) != 1:
                continue
            mul0_out_name = str(mul0_op.outputs[0])
            mul0_const_name = _channelwise_const_name(mul0_op, concat0_out_name)
            if mul0_const_name is None:
                continue

            mul0_users = [int(v) for v in consumers.get(str(mul0_out_name), [])]
            if len(mul0_users) != 1:
                continue
            post0_idx = int(mul0_users[0])
            post0_op = model_ir.operators[int(post0_idx)]
            if (
                str(post0_op.op_type) != "TRANSPOSE"
                or len(post0_op.inputs) < 2
                or len(post0_op.outputs) != 1
                or str(post0_op.inputs[0]) != str(mul0_out_name)
                or _read_transpose_perm(model_ir, post0_op) != perm_nchw_to_nhwc
            ):
                continue
            post0_out_name = str(post0_op.outputs[0])

            post0_users = [int(v) for v in consumers.get(str(post0_out_name), [])]
            if len(post0_users) != 1:
                continue
            add0_idx = int(post0_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_out_name = str(add0_op.outputs[0])

            add0_users = [int(v) for v in consumers.get(str(add0_out_name), [])]
            if len(add0_users) != 1:
                continue
            conv0_idx = int(add0_users[0])
            conv0_op = model_ir.operators[int(conv0_idx)]
            if str(conv0_op.op_type) not in {"CONV_2D", "DEPTHWISE_CONV_2D"}:
                continue
            conv0_out_name = str(conv0_op.outputs[0])

            conv0_users = [int(v) for v in consumers.get(str(conv0_out_name), [])]
            if len(conv0_users) != 1:
                continue
            conv0_post_idx = int(conv0_users[0])
            conv0_post_op = model_ir.operators[int(conv0_post_idx)]
            if (
                str(conv0_post_op.op_type) != "TRANSPOSE"
                or len(conv0_post_op.inputs) < 2
                or len(conv0_post_op.outputs) != 1
                or str(conv0_post_op.inputs[0]) != str(conv0_out_name)
                or _read_transpose_perm(model_ir, conv0_post_op) != perm_nhwc_to_nchw
            ):
                continue
            conv0_nchw_name = str(conv0_post_op.outputs[0])

            conv0_nchw_users = [int(v) for v in consumers.get(str(conv0_nchw_name), [])]
            base_nchw_users = [int(v) for v in consumers.get(str(base_nchw_name), [])]
            if len(conv0_nchw_users) != 1 or int(conv0_nchw_users[0]) not in set(int(v) for v in base_nchw_users):
                continue

            concat1_idx = int(conv0_nchw_users[0])
            concat1_op = model_ir.operators[int(concat1_idx)]
            if (
                str(concat1_op.op_type) != "CONCATENATION"
                or len(concat1_op.inputs) != 2
                or len(concat1_op.outputs) != 1
            ):
                continue
            axis1 = int(concat1_op.options.get("axis", 1))
            if axis1 < 0:
                axis1 += 4
            if axis1 != 1:
                continue
            concat1_inputs = [str(v) for v in list(concat1_op.inputs)]
            if set(concat1_inputs) != {str(base_nchw_name), str(conv0_nchw_name)}:
                continue

            concat1_out_name = str(concat1_op.outputs[0])
            concat1_users = [int(v) for v in consumers.get(str(concat1_out_name), [])]
            if len(concat1_users) != 1:
                continue
            mul1_idx = int(concat1_users[0])
            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])
            mul1_const_name = _channelwise_const_name(mul1_op, concat1_out_name)
            if mul1_const_name is None:
                continue

            mul1_users = [int(v) for v in consumers.get(str(mul1_out_name), [])]
            if len(mul1_users) != 1:
                continue
            post1_idx = int(mul1_users[0])
            post1_op = model_ir.operators[int(post1_idx)]
            if (
                str(post1_op.op_type) != "TRANSPOSE"
                or len(post1_op.inputs) < 2
                or len(post1_op.outputs) != 1
                or str(post1_op.inputs[0]) != str(mul1_out_name)
                or _read_transpose_perm(model_ir, post1_op) != perm_nchw_to_nhwc
            ):
                continue
            post1_out_name = str(post1_op.outputs[0])

            post1_users = [int(v) for v in consumers.get(str(post1_out_name), [])]
            if len(post1_users) != 1:
                continue
            add1_idx = int(post1_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_out_name = str(add1_op.outputs[0])

            add1_users = [int(v) for v in consumers.get(str(add1_out_name), [])]
            if len(add1_users) != 1:
                continue
            conv1_op = model_ir.operators[int(add1_users[0])]
            if str(conv1_op.op_type) not in {"CONV_2D", "DEPTHWISE_CONV_2D"}:
                continue

            for add_idx, resize_nhwc_name in zip(branch_add_indices, resize_nhwc_names):
                add_op = model_ir.operators[int(add_idx)]
                add_inputs = [
                    str(base_nhwc_name) if str(input_name) == str(base_nchw_name)
                    else str(resize_nhwc_name) if str(input_name) != str(base_nchw_name)
                    else str(input_name)
                    for input_name in list(add_op.inputs)
                ]
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=add_op,
                    new_inputs=add_inputs,
                )
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(str(add_op.outputs[0]), None),
                    perm_nchw_to_nhwc,
                )

            concat0_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat0_out_name), None),
                perm_nchw_to_nhwc,
            )

            mul0_const_tensor = model_ir.tensors.get(str(mul0_const_name), None)
            if mul0_const_tensor is not None and mul0_const_tensor.data is not None:
                const_arr = np.asarray(mul0_const_tensor.data)
                if const_arr.ndim == 4:
                    const_arr = np.transpose(const_arr, axes=perm_nchw_to_nhwc)
                    mul0_const_tensor.data = np.asarray(const_arr)
                    mul0_const_tensor.shape = [int(v) for v in list(const_arr.shape)]
                    mul0_const_tensor.shape_signature = [int(v) for v in list(const_arr.shape)]
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul0_out_name), None),
                perm_nchw_to_nhwc,
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=add0_op,
                new_inputs=[
                    str(mul0_out_name) if str(input_name) == str(post0_out_name) else str(input_name)
                    for input_name in list(add0_op.inputs)
                ],
            )

            _set_operator_inputs(
                model_ir=model_ir,
                op=concat1_op,
                new_inputs=[
                    str(base_nhwc_name) if str(input_name) == str(base_nchw_name)
                    else str(conv0_out_name) if str(input_name) == str(conv0_nchw_name)
                    else str(input_name)
                    for input_name in list(concat1_op.inputs)
                ],
            )
            concat1_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(concat1_out_name), None),
                perm_nchw_to_nhwc,
            )

            mul1_const_tensor = model_ir.tensors.get(str(mul1_const_name), None)
            if mul1_const_tensor is not None and mul1_const_tensor.data is not None:
                const_arr = np.asarray(mul1_const_tensor.data)
                if const_arr.ndim == 4:
                    const_arr = np.transpose(const_arr, axes=perm_nchw_to_nhwc)
                    mul1_const_tensor.data = np.asarray(const_arr)
                    mul1_const_tensor.shape = [int(v) for v in list(const_arr.shape)]
                    mul1_const_tensor.shape_signature = [int(v) for v in list(const_arr.shape)]
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mul1_out_name), None),
                perm_nchw_to_nhwc,
            )
            _set_operator_inputs(
                model_ir=model_ir,
                op=add1_op,
                new_inputs=[
                    str(mul1_out_name) if str(input_name) == str(post1_out_name) else str(input_name)
                    for input_name in list(add1_op.inputs)
                ],
            )

            remove_indices = {
                int(base_pre_idx),
                int(post0_idx),
                int(conv0_post_idx),
                int(post1_idx),
                *[int(v) for v in resize_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_resize_add_concat_affine_conv_spp_nhwc_chains": int(rewritten)}


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_mean_mul_reshape_add_conv_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around
      TRANSPOSE -> MEAN -> MUL(const) -> RESHAPE -> ADD(const) -> CONV
    by keeping the branch in 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_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
            mean_idx = int(pre_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]) != pre_output_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
            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 or sorted(normalized_axes) != [2, 3]:
                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
            mul_idx = int(mean_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_out_name = str(mul_op.outputs[0])
            if mul_out_name in model_outputs:
                continue

            mul_const_name: Optional[str] = None
            for input_name in [str(v) for v in list(mul_op.inputs)]:
                if str(input_name) == str(mean_out_name):
                    continue
                tensor = model_ir.tensors.get(str(input_name), None)
                if tensor is None or tensor.data is None:
                    continue
                const_arr = np.asarray(tensor.data)
                if const_arr.ndim == 4:
                    mul_const_name = str(input_name)
                    break
            if mul_const_name is None:
                continue

            mul_users = [int(v) for v in consumers.get(mul_out_name, [])]
            if len(mul_users) != 1:
                continue
            reshape_idx = int(mul_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:
                continue
            if str(reshape_op.inputs[0]) != mul_out_name:
                continue
            reshape_out_name = str(reshape_op.outputs[0])
            if reshape_out_name in model_outputs:
                continue

            reshape_out_tensor = model_ir.tensors.get(reshape_out_name, None)
            if reshape_out_tensor is None or len(list(reshape_out_tensor.shape)) != 4:
                continue
            reshape_shape = [int(v) for v in list(reshape_out_tensor.shape)]
            if any(int(v) <= 0 for v in reshape_shape) or reshape_shape[1:3] != [1, 1]:
                continue

            reshape_users = [int(v) for v in consumers.get(reshape_out_name, [])]
            if len(reshape_users) != 1:
                continue
            add_idx = int(reshape_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_out_name = str(add_op.outputs[0])
            if add_out_name in model_outputs:
                continue

            add_const_ok = False
            for input_name in [str(v) for v in list(add_op.inputs)]:
                if str(input_name) == str(reshape_out_name):
                    continue
                tensor = model_ir.tensors.get(str(input_name), None)
                if tensor is None or tensor.data is None:
                    continue
                add_const_ok = True
                break
            if not add_const_ok:
                continue

            add_users = [int(v) for v in consumers.get(add_out_name, [])]
            if len(add_users) != 1:
                continue
            conv_op = model_ir.operators[int(add_users[0])]
            if str(conv_op.op_type) not in {"CONV_2D", "DEPTHWISE_CONV_2D"}:
                continue

            mapped_axes = [int(perm_nhwc_to_nchw[int(axis)]) for axis in normalized_axes]
            _write_const_ints_to_tensor(mean_axes_tensor, mapped_axes)
            _set_operator_inputs(
                model_ir=model_ir,
                op=mean_op,
                new_inputs=[str(pre_input_name), str(mean_axes_name)],
            )

            mul_const_tensor = model_ir.tensors.get(str(mul_const_name), None)
            if mul_const_tensor is not None and mul_const_tensor.data is not None:
                mul_const_arr = np.asarray(mul_const_tensor.data)
                if mul_const_arr.ndim == 4:
                    mul_const_arr = np.transpose(mul_const_arr, axes=perm_nchw_to_nhwc)
                    mul_const_tensor.data = np.asarray(mul_const_arr)
                    mul_const_tensor.shape = [int(v) for v in list(mul_const_arr.shape)]
                    mul_const_tensor.shape_signature = [int(v) for v in list(mul_const_arr.shape)]

            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(mean_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,
            )

            add_inputs = [
                str(mul_out_name) if str(input_name) == str(reshape_out_name) else str(input_name)
                for input_name in list(add_op.inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=add_op,
                new_inputs=add_inputs,
            )

            del model_ir.operators[int(reshape_idx)]
            del model_ir.operators[int(pre_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_transpose_mean_mul_reshape_add_conv_nhwc_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_logistic_sub_mul_postadd_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Eliminate NCHW round-trips around logistic/sub-gated dual-MUL branches whose MUL
    outputs are immediately transposed back to NHWC before layout-agnostic consumers.

    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) -> ma_nchw --T(0,2,3,1)--> ma_nhwc
      MUL(sub_nchw, b_nchw) -> mb_nchw --T(0,2,3,1)--> mb_nhwc
      ma_nhwc / mb_nhwc then feed NHWC-only consumers (e.g. ADD -> CONV)

    Rewrite:
      - Bypass all three pre-transposes to NHWC.
      - Keep LOGISTIC/SUB/MUL in NHWC.
      - Remove both post-transposes and make the MULs emit NHWC directly.
    """
    rewritten = 0
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    perm_nchw_to_nhwc = [0, 2, 3, 1]

    def _collect_inverse_posts(tensor_name: str) -> Optional[Tuple[List[int], List[str]]]:
        out_users = [int(v) for v in consumers.get(str(tensor_name), [])]
        if len(out_users) == 0:
            return None
        post_indices_local: List[int] = []
        post_output_names_local: List[str] = []
        for user_idx in 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]) != str(tensor_name)
                or _read_transpose_perm(model_ir, user_op) != perm_nchw_to_nhwc
                or str(user_op.outputs[0]) in model_outputs
            ):
                return None
            post_indices_local.append(int(user_idx))
            post_output_names_local.append(str(user_op.outputs[0]))
        return post_indices_local, post_output_names_local

    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
                or 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 or 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])
            if int(mul_sub_idx) == int(mul_sig_idx):
                continue

            mul_sig_op = model_ir.operators[int(mul_sig_idx)]
            mul_sub_op = model_ir.operators[int(mul_sub_idx)]
            if (
                str(mul_sig_op.op_type) != "MUL"
                or len(mul_sig_op.inputs) != 2
                or len(mul_sig_op.outputs) != 1
                or str(mul_sub_op.op_type) != "MUL"
                or len(mul_sub_op.inputs) != 2
                or len(mul_sub_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
                gate_tensor_set = set(str(v) for v in gate_tensor_candidates)
                for input_name in list(mul_op.inputs):
                    if str(input_name) in gate_tensor_set:
                        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

            if set(int(v) for v in consumers.get(str(gate_pre_out_name), [])) != {int(logistic_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(sig_pre_out_name), [])) != {int(mul_sig_idx)}:
                continue
            if set(int(v) for v in consumers.get(str(sub_pre_out_name), [])) != {int(mul_sub_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

            sig_post_info = _collect_inverse_posts(mul_sig_out_name)
            sub_post_info = _collect_inverse_posts(mul_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

            _set_operator_inputs(
                model_ir=model_ir,
                op=logistic_op,
                new_inputs=[str(gate_pre_in_name)],
            )

            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)]
                replaced = False
                for input_idx, input_name in enumerate(list(mul_inputs)):
                    if str(input_name) == str(pre_out_name):
                        mul_inputs[input_idx] = str(pre_in_name)
                        replaced = True
                if not replaced:
                    break
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=mul_op,
                    new_inputs=mul_inputs,
                )
            else:
                for tensor_name in [
                    str(logistic_out_name),
                    str(sub_out_name),
                    str(mul_sig_out_name),
                    str(mul_sub_out_name),
                ]:
                    _permute_tensor_metadata_if_rank_matches(
                        model_ir.tensors.get(str(tensor_name), None),
                        perm_nchw_to_nhwc,
                    )

                for mul_op, old_mul_out_name, post_output_names in [
                    (mul_sig_op, mul_sig_out_name, sig_post_output_names),
                    (mul_sub_op, mul_sub_out_name, sub_post_output_names),
                ]:
                    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],
                    )
                    for alias_name in post_output_names[1:]:
                        _replace_tensor_inputs(model_ir, str(alias_name), canonical_post_output_name)

                    old_mul_out_tensor = model_ir.tensors.get(str(old_mul_out_name), None)
                    canonical_tensor = model_ir.tensors.get(str(canonical_post_output_name), None)
                    if old_mul_out_tensor is not None and canonical_tensor is not None:
                        canonical_tensor.dtype = str(old_mul_out_tensor.dtype)
                        canonical_tensor.quantization = _clone_quantization(old_mul_out_tensor.quantization)
                        canonical_tensor.shape = [int(v) for v in list(old_mul_out_tensor.shape)]
                        canonical_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(
                            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_mul_postadd_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_relu_split_all_outputs_to_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove redundant NCHW adapters around RELU -> SPLIT branches that all
    immediately return to NHWC.

    Target motif:
      src_nhwc --T(0,3,1,2)--> src_nchw --RELU--> relu_nchw
      relu_nchw --SPLIT(axis=1)--> s*_nchw
      each s*_nchw --T(0,2,3,1)--> s*_nhwc -> consumer

    Rewrite:
      src_nhwc --RELU--> relu_nhwc
      relu_nhwc --SPLIT(axis=3)--> s*_nhwc -> consumer

    This removes the pre-RELU transpose and the post-SPLIT branch 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:
        candidate = str(base)
        suffix = 1
        while candidate in model_ir.tensors:
            candidate = f"{base}_{suffix}"
            suffix += 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 or len(split_op.outputs) < 2:
                continue

            split_axis_name = str(split_op.inputs[0])
            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) != 1:
                continue
            split_axis = int(split_axis_vals[0])
            if split_axis < 0:
                split_axis += 4
            if split_axis != 1:
                continue

            split_input_name = str(split_op.inputs[1])
            if split_input_name in model_outputs:
                continue
            unary_idx = producers.get(split_input_name, None)
            if unary_idx is None:
                continue
            unary_op = model_ir.operators[int(unary_idx)]
            if (
                str(unary_op.op_type) != "RELU"
                or len(unary_op.inputs) != 1
                or len(unary_op.outputs) != 1
                or str(unary_op.outputs[0]) != split_input_name
            ):
                continue

            unary_input_name = str(unary_op.inputs[0])
            if unary_input_name in model_outputs:
                continue
            pre_t_idx = producers.get(unary_input_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]) != unary_input_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

            remove_op_ids: set[int] = {int(id(pre_t_op))}
            branch_rewires: List[Tuple[str, str]] = []
            valid = True
            for split_output_name in [str(v) for v in list(split_op.outputs)]:
                if split_output_name in model_outputs:
                    valid = False
                    break
                split_users = [int(v) for v in consumers.get(split_output_name, [])]
                if len(split_users) != 1:
                    valid = False
                    break
                post_t_idx = int(split_users[0])
                post_t_op = model_ir.operators[int(post_t_idx)]
                if (
                    str(post_t_op.op_type) != "TRANSPOSE"
                    or len(post_t_op.inputs) < 2
                    or len(post_t_op.outputs) != 1
                    or str(post_t_op.inputs[0]) != split_output_name
                    or str(post_t_op.outputs[0]) in model_outputs
                    or _read_transpose_perm(model_ir, post_t_op) != perm_nchw_to_nhwc
                ):
                    valid = False
                    break
                branch_rewires.append((str(post_t_op.outputs[0]), str(split_output_name)))
                remove_op_ids.add(int(id(post_t_op)))

            if not valid:
                continue

            split_axis_tensor = model_ir.tensors.get(split_axis_name, None)
            if split_axis_tensor is None:
                continue
            split_axis_users = [int(v) for v in consumers.get(split_axis_name, [])]
            if set(split_axis_users) == {int(split_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),
                )
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=split_op,
                    new_inputs=[split_axis_nhwc_name, split_input_name],
                )

            _set_operator_inputs(
                model_ir=model_ir,
                op=unary_op,
                new_inputs=[src_nhwc_name],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(split_input_name, None),
                perm_nchw_to_nhwc,
            )
            for split_output_name in [str(v) for v in list(split_op.outputs)]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(split_output_name, None),
                    perm_nchw_to_nhwc,
                )

            for post_output_name, replacement_name in list(branch_rewires):
                _replace_tensor_inputs(model_ir, str(post_output_name), str(replacement_name))

            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_relu_split_all_outputs_to_nhwc_chains": int(optimized)}


def _optimize_transpose_relu_split_conv_relu_concat_posttranspose_to_nhwc_chains(
    model_ir: ModelIR,
) -> Dict[str, int]:
    """
    Remove NCHW bridge transposes around:

      NHWC -> T -> RELU -> SPLIT(axis=1)
      branch -> T -> CONV -> T -> RELU
      CONCAT(axis=1, [branch, keep]) -> T -> NHWC consumer

    Rewrite the entire fragment to NHWC and switch split/concat channel axes to 3.
    """
    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

    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 concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or len(concat_op.inputs) != 2 or len(concat_op.outputs) != 1:
                continue
            concat_axis = int(concat_op.options.get("axis", 0))
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue

            concat_output_name = str(concat_op.outputs[0])
            if concat_output_name in model_outputs:
                continue
            concat_output_users = [int(v) for v in consumers.get(concat_output_name, [])]
            if len(concat_output_users) != 1:
                continue
            post_concat_t_idx = int(concat_output_users[0])
            post_concat_t = model_ir.operators[int(post_concat_t_idx)]
            if (
                str(post_concat_t.op_type) != "TRANSPOSE"
                or len(post_concat_t.inputs) < 2
                or len(post_concat_t.outputs) != 1
                or str(post_concat_t.inputs[0]) != concat_output_name
                or str(post_concat_t.outputs[0]) in model_outputs
                or _read_transpose_perm(model_ir, post_concat_t) != perm_nchw_to_nhwc
            ):
                continue

            branch_input_name = ""
            keep_input_name = ""
            split_op: Optional[OperatorIR] = None
            split_idx: Optional[int] = None
            pre_relu_op: Optional[OperatorIR] = None
            pre_relu_idx: Optional[int] = None
            pre_t_op: Optional[OperatorIR] = None
            pre_t_idx: Optional[int] = None
            pre_branch_t_op: Optional[OperatorIR] = None
            pre_branch_t_idx: Optional[int] = None
            conv_op: Optional[OperatorIR] = None
            post_branch_t_op: Optional[OperatorIR] = None
            post_branch_t_idx: Optional[int] = None
            post_relu_op: Optional[OperatorIR] = None

            for candidate_branch_input in [str(v) for v in list(concat_op.inputs)]:
                candidate_keep_input = next(
                    (str(v) for v in list(concat_op.inputs) if str(v) != str(candidate_branch_input)),
                    "",
                )
                if candidate_keep_input == "":
                    continue

                post_relu_candidate_idx = producers.get(str(candidate_branch_input), None)
                if post_relu_candidate_idx is None:
                    continue
                post_relu_candidate = model_ir.operators[int(post_relu_candidate_idx)]
                if (
                    str(post_relu_candidate.op_type) != "RELU"
                    or len(post_relu_candidate.inputs) != 1
                    or len(post_relu_candidate.outputs) != 1
                    or str(post_relu_candidate.outputs[0]) != str(candidate_branch_input)
                ):
                    continue

                post_branch_t_output_name = str(post_relu_candidate.inputs[0])
                post_branch_t_candidate_idx = producers.get(post_branch_t_output_name, None)
                if post_branch_t_candidate_idx is None:
                    continue
                post_branch_t_candidate = model_ir.operators[int(post_branch_t_candidate_idx)]
                if (
                    str(post_branch_t_candidate.op_type) != "TRANSPOSE"
                    or len(post_branch_t_candidate.inputs) < 2
                    or len(post_branch_t_candidate.outputs) != 1
                    or str(post_branch_t_candidate.outputs[0]) != post_branch_t_output_name
                    or _read_transpose_perm(model_ir, post_branch_t_candidate) != perm_nhwc_to_nchw
                    or set(int(v) for v in consumers.get(post_branch_t_output_name, [])) != {int(post_relu_candidate_idx)}
                ):
                    continue

                conv_output_name = str(post_branch_t_candidate.inputs[0])
                conv_candidate_idx = producers.get(conv_output_name, None)
                if conv_candidate_idx is None:
                    continue
                conv_candidate = model_ir.operators[int(conv_candidate_idx)]
                if (
                    str(conv_candidate.op_type) != "CONV_2D"
                    or len(conv_candidate.inputs) < 1
                    or len(conv_candidate.outputs) != 1
                    or str(conv_candidate.outputs[0]) != conv_output_name
                ):
                    continue

                conv_input_name = str(conv_candidate.inputs[0])
                pre_branch_t_candidate_idx = producers.get(conv_input_name, None)
                if pre_branch_t_candidate_idx is None:
                    continue
                pre_branch_t_candidate = model_ir.operators[int(pre_branch_t_candidate_idx)]
                if (
                    str(pre_branch_t_candidate.op_type) != "TRANSPOSE"
                    or len(pre_branch_t_candidate.inputs) < 2
                    or len(pre_branch_t_candidate.outputs) != 1
                    or str(pre_branch_t_candidate.outputs[0]) != conv_input_name
                    or _read_transpose_perm(model_ir, pre_branch_t_candidate) != perm_nchw_to_nhwc
                    or set(int(v) for v in consumers.get(conv_input_name, [])) != {int(conv_candidate_idx)}
                ):
                    continue

                branch_split_output_name = str(pre_branch_t_candidate.inputs[0])
                split_candidate_idx = producers.get(branch_split_output_name, None)
                if split_candidate_idx is None:
                    continue
                split_candidate = model_ir.operators[int(split_candidate_idx)]
                if (
                    str(split_candidate.op_type) != "SPLIT"
                    or len(split_candidate.inputs) < 2
                    or branch_split_output_name not in [str(v) for v in list(split_candidate.outputs)]
                    or str(candidate_keep_input) not in [str(v) for v in list(split_candidate.outputs)]
                ):
                    continue

                split_axis_name = str(split_candidate.inputs[0])
                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) != 1:
                    continue
                split_axis = int(split_axis_vals[0])
                if split_axis < 0:
                    split_axis += 4
                if split_axis != 1:
                    continue

                if set(int(v) for v in consumers.get(str(candidate_keep_input), [])) != {int(concat_idx)}:
                    continue
                if set(int(v) for v in consumers.get(branch_split_output_name, [])) != {int(pre_branch_t_candidate_idx)}:
                    continue

                split_input_name = str(split_candidate.inputs[1])
                pre_relu_candidate_idx = producers.get(split_input_name, None)
                if pre_relu_candidate_idx is None:
                    continue
                pre_relu_candidate = model_ir.operators[int(pre_relu_candidate_idx)]
                if (
                    str(pre_relu_candidate.op_type) != "RELU"
                    or len(pre_relu_candidate.inputs) != 1
                    or len(pre_relu_candidate.outputs) != 1
                    or str(pre_relu_candidate.outputs[0]) != split_input_name
                ):
                    continue

                pre_t_output_name = str(pre_relu_candidate.inputs[0])
                pre_t_candidate_idx = producers.get(pre_t_output_name, None)
                if pre_t_candidate_idx is None:
                    continue
                pre_t_candidate = model_ir.operators[int(pre_t_candidate_idx)]
                if (
                    str(pre_t_candidate.op_type) != "TRANSPOSE"
                    or len(pre_t_candidate.inputs) < 2
                    or len(pre_t_candidate.outputs) != 1
                    or str(pre_t_candidate.outputs[0]) != pre_t_output_name
                    or _read_transpose_perm(model_ir, pre_t_candidate) != perm_nhwc_to_nchw
                    or set(int(v) for v in consumers.get(pre_t_output_name, [])) != {int(pre_relu_candidate_idx)}
                ):
                    continue

                if str(pre_t_candidate.inputs[0]) in model_outputs:
                    continue

                branch_input_name = str(candidate_branch_input)
                keep_input_name = str(candidate_keep_input)
                split_op = split_candidate
                split_idx = int(split_candidate_idx)
                pre_relu_op = pre_relu_candidate
                pre_relu_idx = int(pre_relu_candidate_idx)
                pre_t_op = pre_t_candidate
                pre_t_idx = int(pre_t_candidate_idx)
                pre_branch_t_op = pre_branch_t_candidate
                pre_branch_t_idx = int(pre_branch_t_candidate_idx)
                conv_op = conv_candidate
                post_branch_t_op = post_branch_t_candidate
                post_branch_t_idx = int(post_branch_t_candidate_idx)
                post_relu_op = post_relu_candidate
                break

            if (
                split_op is None
                or split_idx is None
                or pre_relu_op is None
                or pre_relu_idx is None
                or pre_t_op is None
                or pre_t_idx is None
                or pre_branch_t_op is None
                or pre_branch_t_idx is None
                or conv_op is None
                or post_branch_t_op is None
                or post_branch_t_idx is None
                or post_relu_op is None
                or branch_input_name == ""
                or keep_input_name == ""
            ):
                continue

            split_axis_name = str(split_op.inputs[0])
            split_axis_tensor = model_ir.tensors.get(split_axis_name, None)
            if split_axis_tensor is None:
                continue
            split_axis_users = [int(v) for v in consumers.get(split_axis_name, [])]
            if set(split_axis_users) == {int(split_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),
                )
                _set_operator_inputs(
                    model_ir=model_ir,
                    op=split_op,
                    new_inputs=[split_axis_nhwc_name, str(split_op.inputs[1])],
                )

            src_nhwc_name = str(pre_t_op.inputs[0])
            _set_operator_inputs(
                model_ir=model_ir,
                op=pre_relu_op,
                new_inputs=[src_nhwc_name],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(pre_relu_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )
            for split_output_name in [str(v) for v in list(split_op.outputs)]:
                _permute_tensor_metadata_if_rank_matches(
                    model_ir.tensors.get(split_output_name, None),
                    perm_nchw_to_nhwc,
                )

            branch_split_output_name = str(pre_branch_t_op.inputs[0])
            conv_inputs = [str(v) for v in list(conv_op.inputs)]
            conv_inputs[0] = branch_split_output_name
            _set_operator_inputs(
                model_ir=model_ir,
                op=conv_op,
                new_inputs=conv_inputs,
            )

            conv_output_name = str(conv_op.outputs[0])
            _set_operator_inputs(
                model_ir=model_ir,
                op=post_relu_op,
                new_inputs=[conv_output_name],
            )
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(str(post_relu_op.outputs[0]), None),
                perm_nchw_to_nhwc,
            )

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            replaced_concat_inputs = [
                str(branch_input_name) if str(v) == str(branch_input_name) else str(keep_input_name)
                for v in list(concat_inputs)
            ]
            _set_operator_inputs(
                model_ir=model_ir,
                op=concat_op,
                new_inputs=replaced_concat_inputs,
            )
            concat_op.options["axis"] = 3
            _permute_tensor_metadata_if_rank_matches(
                model_ir.tensors.get(concat_output_name, None),
                perm_nchw_to_nhwc,
            )

            post_concat_output_name = str(post_concat_t.outputs[0])
            _replace_tensor_inputs(model_ir, post_concat_output_name, concat_output_name)

            remove_ids = {
                int(id(pre_t_op)),
                int(id(pre_branch_t_op)),
                int(id(post_branch_t_op)),
                int(id(post_concat_t)),
            }
            remove_indices = sorted(
                [
                    int(op_idx)
                    for op_idx, op in enumerate(model_ir.operators)
                    if int(id(op)) in remove_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_relu_split_conv_relu_concat_posttranspose_to_nhwc_chains": int(optimized)}


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_mixed_mean_reducemax_concat_mirrorpad_nhwc_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Repair partial NHWC propagation around SiNet/DEA attention SA branches.

    Target broken sketch:
      x_nhwc --T(0,3,1,2)--> x_nchw
      mean_nhwc = MEAN(x_nhwc, axes=[3], keepDims=1)
      max_nchw  = REDUCE_MAX(x_nchw, axes=[1], keepDims=1)
      CONCAT(mean_nhwc, max_nchw, axis=1) -> MIRROR_PAD(nchw pads) -> T(0,2,3,1) -> CONV

    Rewrite:
      - move REDUCE_MAX to x_nhwc and remap axes to NHWC
      - rewrite CONCAT axis to NHWC channel axis
      - rewrite MIRROR_PAD pairs to NHWC ordering
      - remove the redundant MIRROR_PAD output transpose
    """
    optimized = 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]:
        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

    while True:
        changed = False
        consumers = _build_tensor_consumer_map(model_ir)
        producers = _build_tensor_producer_map(model_ir)
        model_outputs = {str(v) for v in list(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) != 2
                or len(concat_op.outputs) != 1
            ):
                continue
            concat_axis = int(concat_op.options.get("axis", 0))
            if concat_axis < 0:
                concat_axis += 4
            if concat_axis != 1:
                continue

            concat_input_names = [str(v) for v in list(concat_op.inputs)]
            concat_output_name = str(concat_op.outputs[0])
            if concat_output_name in model_outputs:
                continue

            input_producers = [producers.get(name, None) for name in concat_input_names]
            if any(idx is None for idx in input_producers):
                continue
            lhs_op = model_ir.operators[int(input_producers[0])]
            rhs_op = model_ir.operators[int(input_producers[1])]
            producer_pair = {str(lhs_op.op_type), str(rhs_op.op_type)}
            if producer_pair != {"MEAN", "REDUCE_MAX"}:
                continue

            mean_op = lhs_op if str(lhs_op.op_type) == "MEAN" else rhs_op
            max_op = rhs_op if str(rhs_op.op_type) == "REDUCE_MAX" else lhs_op
            mean_out_name = str(mean_op.outputs[0])
            max_out_name = str(max_op.outputs[0])
            mean_tensor = model_ir.tensors.get(mean_out_name, None)
            max_tensor = model_ir.tensors.get(max_out_name, None)
            mean_input_tensor = model_ir.tensors.get(str(mean_op.inputs[0]), None)
            max_input_tensor = model_ir.tensors.get(str(max_op.inputs[0]), None)
            if (
                mean_tensor is None
                or max_tensor is None
                or mean_input_tensor is None
                or max_input_tensor is None
                or len(list(mean_tensor.shape)) != 4
                or len(list(max_tensor.shape)) != 4
            ):
                continue

            mean_input_name = str(mean_op.inputs[0])
            max_input_name = str(max_op.inputs[0])
            source_nhwc_name: Optional[str] = None
            pre_idx = producers.get(max_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]) == max_input_name
                    and str(pre_op.inputs[0]) == mean_input_name
                    and _read_transpose_perm(model_ir, pre_op) == perm_nhwc_to_nchw
                ):
                    source_nhwc_name = str(mean_input_name)
            if source_nhwc_name is None:
                pre_idx = producers.get(mean_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]) != mean_input_name
                    or str(pre_op.inputs[0]) != max_input_name
                    or _read_transpose_perm(model_ir, pre_op) != perm_nchw_to_nhwc
                ):
                    continue
                source_nhwc_name = str(mean_input_name)
            if not bool(mean_op.options.get("keepDims", False)) or not bool(max_op.options.get("keepDims", False)):
                continue

            mirror_users = [int(v) for v in consumers.get(concat_output_name, [])]
            if len(mirror_users) != 1:
                continue
            mirror_idx = int(mirror_users[0])
            mirror_op = model_ir.operators[int(mirror_idx)]
            if (
                str(mirror_op.op_type) != "MIRROR_PAD"
                or len(mirror_op.inputs) < 2
                or len(mirror_op.outputs) != 1
                or str(mirror_op.inputs[0]) != concat_output_name
            ):
                continue

            mirror_output_name = str(mirror_op.outputs[0])
            if mirror_output_name in model_outputs:
                continue
            transpose_users = [int(v) for v in consumers.get(mirror_output_name, [])]
            if len(transpose_users) != 1:
                continue
            post_idx = int(transpose_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]) != mirror_output_name
                or _read_transpose_perm(model_ir, post_op) != perm_nchw_to_nhwc
            ):
                continue
            conv_input_name = str(post_op.outputs[0])
            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) != "CONV_2D"
                or len(conv_op.inputs) < 1
                or str(conv_op.inputs[0]) != conv_input_name
            ):
                continue

            if not _rewrite_reduce_axes_nchw_to_nhwc(max_op):
                continue
            if not _rewrite_pad_pairs_nchw_to_nhwc(mirror_op):
                continue

            _replace_operator_input_at(
                model_ir=model_ir,
                op=max_op,
                input_index=0,
                new_input_name=str(source_nhwc_name),
            )
            concat_op.options["axis"] = 3
            _replace_operator_input_at(
                model_ir=model_ir,
                op=conv_op,
                input_index=0,
                new_input_name=str(mirror_output_name),
            )

            for tensor_name in [str(max_out_name), str(concat_output_name), str(mirror_output_name)]:
                tensor = model_ir.tensors.get(str(tensor_name), None)
                _permute_tensor_metadata_if_rank_matches(tensor, perm_nchw_to_nhwc)
                if tensor is not None and len(list(tensor.shape)) == 4:
                    tensor.logical_layout = "NHWC"
            max_input_tensor.logical_layout = "NCHW"
            mean_input_tensor.logical_layout = "NHWC"

            del model_ir.operators[int(post_idx)]
            optimized += 1
            changed = True
            break

        if not changed:
            break

    _prune_unused_tensors(model_ir)
    return {"optimized_mixed_mean_reducemax_concat_mirrorpad_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 moved layout axes
        # are still provably singleton in signatures. Rank-4 layout bridges keep
        # batch at axis 0, so a dynamic batch dimension there remains safe.
        if (
            len(input_signature) != rank
            or len(output_signature) != rank
            or any(int(v) < 0 for idx, v in enumerate(input_signature) if int(idx) != 0)
            or any(int(v) < 0 for idx, v in enumerate(output_signature) if int(idx) != 0)
        ):
            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_constant_input_pool_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold constant-input pool operators by materializing their outputs as constant tensors.

    Rewrite:
      CONST(in) -> (AVERAGE_POOL_2D|MAX_POOL_2D) -> out
    Into:
      CONST(out_pooled_data) and remove the pool op.
    """
    rewritten = 0

    while True:
        changed = False
        for pool_idx, pool_op in enumerate(model_ir.operators):
            op_type = str(pool_op.op_type)
            if op_type not in {"AVERAGE_POOL_2D", "MAX_POOL_2D"}:
                continue
            if len(pool_op.inputs) != 1 or len(pool_op.outputs) != 1:
                continue

            in_name = str(pool_op.inputs[0])
            out_name = str(pool_op.outputs[0])
            if in_name == "" or out_name == "" or 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
            if not _is_fully_known_positive_shape(in_tensor.shape):
                continue

            pooled = _evaluate_constant_pool2d(
                input_data=in_tensor.data,
                op_type=op_type,
                padding=str(pool_op.options.get("padding", "SAME")),
                stride_h=int(pool_op.options.get("strideH", 1)),
                stride_w=int(pool_op.options.get("strideW", 1)),
                filter_h=int(pool_op.options.get("filterHeight", 1)),
                filter_w=int(pool_op.options.get("filterWidth", 1)),
            )
            if pooled is None:
                continue

            out_dtype = str(out_tensor.dtype).upper()
            casted = _cast_const_array_to_tflite_dtype(pooled, out_dtype)
            if casted is None:
                continue

            out_tensor.data = casted
            out_tensor.dtype = str(out_dtype)
            out_tensor.shape = [int(v) for v in list(casted.shape)]
            out_tensor.shape_signature = [int(v) for v in list(casted.shape)]
            _append_tensor_lineage_event(
                model_ir=model_ir,
                event={
                    "kind": "fold_constant_pool",
                    "op_type": str(op_type),
                    "input_name": str(in_name),
                    "output_name": str(out_name),
                },
            )

            del model_ir.operators[int(pool_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_constant_input_pool_chains": int(rewritten)}


def _optimize_constant_input_pad_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Fold constant-input PAD/PADV2 operators by materializing their outputs as constant tensors.

    Rewrite:
      CONST(in) + CONST(pads) [+ CONST(pad_value)] -> (PAD|PADV2) -> out
    Into:
      CONST(out_padded_data) and remove the pad op.
    """
    rewritten = 0

    while True:
        changed = False
        for pad_idx, pad_op in enumerate(model_ir.operators):
            op_type = str(pad_op.op_type)
            if op_type not in {"PAD", "PADV2"}:
                continue
            if len(pad_op.outputs) != 1:
                continue
            if len(pad_op.inputs) not in {2, 3}:
                continue

            in_name = str(pad_op.inputs[0])
            pads_name = str(pad_op.inputs[1])
            pad_value_name = str(pad_op.inputs[2]) if len(pad_op.inputs) == 3 else None
            out_name = str(pad_op.outputs[0])
            if in_name == "" or pads_name == "" or out_name == "" or in_name == out_name:
                continue

            in_tensor = model_ir.tensors.get(in_name, None)
            pads_tensor = model_ir.tensors.get(pads_name, None)
            out_tensor = model_ir.tensors.get(out_name, None)
            pad_value_tensor = model_ir.tensors.get(pad_value_name, None) if pad_value_name is not None else None
            if in_tensor is None or pads_tensor is None or out_tensor is None:
                continue
            if bool(in_tensor.is_variable) or in_tensor.data is None:
                continue
            if bool(pads_tensor.is_variable) or pads_tensor.data is None:
                continue
            if pad_value_name is not None and (pad_value_tensor is None or bool(pad_value_tensor.is_variable) or pad_value_tensor.data is None):
                continue

            padded = _evaluate_constant_pad(
                input_data=in_tensor.data,
                pads_data=pads_tensor.data,
                pad_value_data=None if pad_value_tensor is None else pad_value_tensor.data,
            )
            if padded is None:
                continue

            out_dtype = str(out_tensor.dtype).upper()
            casted = _cast_const_array_to_tflite_dtype(padded, out_dtype)
            if casted is None:
                continue

            out_tensor.data = casted
            out_tensor.dtype = str(out_dtype)
            out_tensor.shape = [int(v) for v in list(casted.shape)]
            out_tensor.shape_signature = [int(v) for v in list(casted.shape)]
            _append_tensor_lineage_event(
                model_ir=model_ir,
                event={
                    "kind": "fold_constant_pad",
                    "op_type": str(op_type),
                    "input_name": str(in_name),
                    "pads_name": str(pads_name),
                    "output_name": str(out_name),
                },
            )

            del model_ir.operators[int(pad_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_constant_input_pad_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_redundant_int32_to_int64_passthrough_cast_chains(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove redundant widening aliases:
      producer(INT32/UINT32) -> CAST(*32 -> *64 alias) -> CAST(*64 -> final)

    Keep the public alias tensor name on the producer output, but preserve the
    32-bit work dtype so downstream non-64-bit CAST consumers read the actual
    runtime dtype. This avoids unnecessary INT64 tensors in shape/index unary
    subgraphs such as Shape->Gather->Sign->Cast(float).
    """
    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 widen_idx, widen_op in enumerate(model_ir.operators):
            if str(widen_op.op_type) != "CAST" or len(widen_op.inputs) != 1 or len(widen_op.outputs) != 1:
                continue

            work_name = str(widen_op.inputs[0])
            alias_name = str(widen_op.outputs[0])
            if alias_name in model_outputs:
                continue

            widen_in_dtype = str(widen_op.options.get("inDataType", "")).upper()
            widen_out_dtype = str(widen_op.options.get("outDataType", "")).upper()
            if (widen_in_dtype, widen_out_dtype) not in {
                ("INT32", "INT64"),
                ("UINT32", "UINT64"),
            }:
                continue

            work_users = [int(v) for v in consumers.get(work_name, [])]
            if len(work_users) != 1 or int(work_users[0]) != int(widen_idx):
                continue

            alias_users = [int(v) for v in consumers.get(alias_name, [])]
            if len(alias_users) <= 0:
                continue
            if not all(str(model_ir.operators[int(user_idx)].op_type) == "CAST" for user_idx in alias_users):
                continue

            producer_idx = producers.get(work_name, None)
            if producer_idx is None or int(producer_idx) == int(widen_idx):
                continue
            producer_op = model_ir.operators[int(producer_idx)]
            if len(producer_op.outputs) != 1 or str(producer_op.outputs[0]) != work_name:
                continue

            work_tensor = model_ir.tensors.get(work_name, None)
            alias_tensor = model_ir.tensors.get(alias_name, None)
            if alias_tensor is None:
                continue

            _set_operator_outputs(
                model_ir=model_ir,
                op=producer_op,
                new_outputs=[alias_name],
            )
            alias_tensor.dtype = widen_in_dtype
            if work_tensor is not None:
                alias_tensor.shape = [int(v) for v in list(work_tensor.shape)]
                alias_tensor.shape_signature = (
                    [int(v) for v in list(work_tensor.shape_signature)]
                    if work_tensor.shape_signature is not None
                    else [int(v) for v in list(work_tensor.shape)]
                )
                alias_tensor.quantization = _clone_quantization(work_tensor.quantization)

            for alias_user_idx in alias_users:
                alias_user = model_ir.operators[int(alias_user_idx)]
                alias_user.options["inDataType"] = widen_in_dtype

            del model_ir.operators[int(widen_idx)]
            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_redundant_int32_to_int64_passthrough_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 {}
        )
        new_shape = options.get("newShape", None)
        if new_shape is not None:
            try:
                concrete_values = [int(v) for v in np.asarray(new_shape).reshape(-1).tolist()]
            except Exception:
                concrete_values = []
            if concrete_values and all(int(v) > 0 for v in concrete_values):
                return False
        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
        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_singleton_gate_conv_concat_nhwc_bridge_blocks(model_ir: ModelIR) -> Dict[str, int]:
    """
    Remove singleton-channel NCHW adapters around a localized gate->conv->concat block.

    Target:
      split_nchw --RESHAPE--> split_nhwc
      clip_nhwc --RESHAPE--> clip_nchw
      clip_nchw * split_nchw -> mul_gate
      scalar - clip_nchw -> sub_gate
      (optional) input_nhwc --TRANSPOSE--> input_nchw
      clip_nchw * input_nchw -> mul_rgb
      conv_out_nhwc --RESHAPE--> conv_out_nchw --LOGISTIC--> sig_nchw
      sub_gate * sig_nchw -> gated
      mul_gate + gated -> fused --UNARY--> clip3_nchw
      clip3_nchw --RESHAPE--> clip3_nhwc
      CONCAT(axis=3, [..., split_nhwc, ..., clip3_nhwc]) -> y_nhwc

    Rewrite:
      - Keep the whole gate path in NHWC.
      - Remove singleton layout RESHAPEs and optional input TRANSPOSE.
      - Rewire CONCAT to consume the NHWC tensors directly.
    """
    rewritten = 0
    perm_nchw_to_nhwc = [0, 2, 3, 1]
    perm_nhwc_to_nchw = [0, 3, 1, 2]
    unary_ops = {"RELU", "RELU6", "RELU_0_TO_1", "LOGISTIC", "HARD_SWISH", "LEAKY_RELU", "TANH"}

    def _is_singleton_layout_reshape(op: OperatorIR) -> bool:
        if str(op.op_type) != "RESHAPE" or len(op.inputs) < 1 or len(op.outputs) != 1:
            return False
        in_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
        out_tensor = model_ir.tensors.get(str(op.outputs[0]), None)
        if (
            in_tensor is None
            or out_tensor is None
            or not _is_fully_known_positive_shape(in_tensor.shape)
            or not _is_fully_known_positive_shape(out_tensor.shape)
        ):
            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 _permute_shape(in_shape, perm_nchw_to_nhwc) == out_shape or _permute_shape(in_shape, perm_nhwc_to_nchw) == out_shape

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

        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or int(concat_op.options.get("axis", 3)) != 3:
                continue
            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            clip3_adapter_idx = None
            clip3_adapter_op = None
            clip3_nchw_name = ""
            clip3_nhwc_name = ""
            clip3_unary_idx = None
            for name in concat_inputs:
                prod_idx = producers.get(name, None)
                if prod_idx is None or not _is_singleton_layout_reshape(model_ir.operators[int(prod_idx)]):
                    continue
                candidate_op = model_ir.operators[int(prod_idx)]
                candidate_source_name = str(candidate_op.inputs[0])
                candidate_unary_idx = producers.get(candidate_source_name, None)
                if candidate_unary_idx is None or str(model_ir.operators[int(candidate_unary_idx)].op_type) not in unary_ops:
                    continue
                add_candidate_idx = producers.get(str(model_ir.operators[int(candidate_unary_idx)].inputs[0]), None)
                if add_candidate_idx is None or str(model_ir.operators[int(add_candidate_idx)].op_type) != "ADD":
                    continue
                clip3_adapter_idx = int(prod_idx)
                clip3_adapter_op = candidate_op
                clip3_nchw_name = str(candidate_source_name)
                clip3_nhwc_name = str(candidate_op.outputs[0])
                clip3_unary_idx = int(candidate_unary_idx)
                break
            if clip3_adapter_idx is None or clip3_adapter_op is None or clip3_unary_idx is None:
                continue
            add_idx = producers.get(str(model_ir.operators[int(clip3_unary_idx)].inputs[0]), None)
            if add_idx is None or str(model_ir.operators[int(add_idx)].op_type) != "ADD":
                continue
            add_inputs = [str(v) for v in list(model_ir.operators[int(add_idx)].inputs)]
            mul_indices = [producers.get(name, None) for name in add_inputs]
            if any(idx is None or str(model_ir.operators[int(idx)].op_type) != "MUL" for idx in mul_indices):
                continue

            clip_adapter_idx = None
            clip_nhwc_name = ""
            split_nchw_name = ""
            sig_idx = None
            conv_reshape_idx = None
            conv_nhwc_name = ""
            sub_idx = None
            gate_mul_idx = None
            sig_mul_idx = None
            for mul_idx in [int(v) for v in mul_indices if v is not None]:
                mul_op = model_ir.operators[int(mul_idx)]
                mul_inputs = [str(v) for v in list(mul_op.inputs)]
                maybe_sig_idx = next((producers.get(name, None) for name in mul_inputs if producers.get(name, None) is not None and str(model_ir.operators[int(producers.get(name))].op_type) == "LOGISTIC"), None)
                if maybe_sig_idx is not None:
                    sig_idx = int(maybe_sig_idx)
                    other_name = next(name for name in mul_inputs if name != str(model_ir.operators[int(sig_idx)].outputs[0]))
                    sub_idx = producers.get(other_name, None)
                    if sub_idx is None or str(model_ir.operators[int(sub_idx)].op_type) != "SUB":
                        sig_idx = None
                        continue
                    clip_adapter_out = str(model_ir.operators[int(sub_idx)].inputs[1])
                    clip_adapter_idx = producers.get(clip_adapter_out, None)
                    if clip_adapter_idx is None or not _is_singleton_layout_reshape(model_ir.operators[int(clip_adapter_idx)]):
                        sig_idx = None
                        continue
                    clip_nhwc_name = str(model_ir.operators[int(clip_adapter_idx)].inputs[0])
                    conv_nchw_name = str(model_ir.operators[int(sig_idx)].inputs[0])
                    conv_reshape_idx = producers.get(conv_nchw_name, None)
                    if conv_reshape_idx is None or not _is_singleton_layout_reshape(model_ir.operators[int(conv_reshape_idx)]):
                        sig_idx = None
                        continue
                    conv_nhwc_name = str(model_ir.operators[int(conv_reshape_idx)].inputs[0])
                    sig_mul_idx = int(mul_idx)
                else:
                    clip_side = next((name for name in mul_inputs if producers.get(name, None) is not None and _is_singleton_layout_reshape(model_ir.operators[int(producers.get(name))])), None)
                    if clip_side is None:
                        continue
                    clip_adapter_idx = int(producers.get(clip_side))
                    clip_nhwc_name = str(model_ir.operators[int(clip_adapter_idx)].inputs[0])
                    split_nchw_name = next(name for name in mul_inputs if name != clip_side)
                    gate_mul_idx = int(mul_idx)
            if None in {clip_adapter_idx, sig_idx, conv_reshape_idx, sub_idx, gate_mul_idx, sig_mul_idx}:
                continue

            split_reshape_idx = producers.get(next((name for name in concat_inputs if name != clip3_nhwc_name and producers.get(name, None) is not None and _is_singleton_layout_reshape(model_ir.operators[int(producers.get(name))]) and str(model_ir.operators[int(producers.get(name))].inputs[0]) == split_nchw_name), ""), None)

            gate_clip_out = str(model_ir.operators[int(clip_adapter_idx)].outputs[0])
            gate_mul_op = model_ir.operators[int(gate_mul_idx)]
            sub_op = model_ir.operators[int(sub_idx)]
            sig_op = model_ir.operators[int(sig_idx)]
            sig_mul_op = model_ir.operators[int(sig_mul_idx)]

            _set_operator_inputs(model_ir=model_ir, op=gate_mul_op, new_inputs=[clip_nhwc_name if str(v) == gate_clip_out else str(v) for v in list(gate_mul_op.inputs)])
            _set_operator_inputs(model_ir=model_ir, op=sub_op, new_inputs=[clip_nhwc_name if str(v) == gate_clip_out else str(v) for v in list(sub_op.inputs)])
            _set_operator_inputs(model_ir=model_ir, op=sig_op, new_inputs=[conv_nhwc_name])

            rgb_mul_idx = next((int(user_idx) for user_idx in consumers.get(gate_clip_out, []) if int(user_idx) not in {int(gate_mul_idx), int(sub_idx)} and str(model_ir.operators[int(user_idx)].op_type) == "MUL"), None)
            input_transpose_idx = None
            if rgb_mul_idx is not None:
                rgb_mul_op = model_ir.operators[int(rgb_mul_idx)]
                other_name = next(name for name in list(rgb_mul_op.inputs) if name != gate_clip_out)
                input_transpose_idx = producers.get(other_name, None)
                if input_transpose_idx is not None and str(model_ir.operators[int(input_transpose_idx)].op_type) == "TRANSPOSE" and _read_transpose_perm(model_ir, model_ir.operators[int(input_transpose_idx)]) == perm_nhwc_to_nchw:
                    _set_operator_inputs(model_ir=model_ir, op=rgb_mul_op, new_inputs=[str(model_ir.operators[int(input_transpose_idx)].inputs[0]) if str(v) == other_name else clip_nhwc_name if str(v) == gate_clip_out else str(v) for v in list(rgb_mul_op.inputs)])
                    _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(rgb_mul_op.outputs[0]), None), perm_nchw_to_nhwc)
                else:
                    input_transpose_idx = None
                    _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(other_name, None), perm_nchw_to_nhwc)
                    _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(rgb_mul_op.outputs[0]), None), perm_nchw_to_nhwc)
                    _set_operator_inputs(
                        model_ir=model_ir,
                        op=rgb_mul_op,
                        new_inputs=[
                            clip_nhwc_name if str(v) == gate_clip_out else str(v)
                            for v in list(rgb_mul_op.inputs)
                        ],
                    )

            for tensor_name in [
                str(model_ir.tensors[split_nchw_name].name),
                str(gate_mul_op.outputs[0]),
                str(sub_op.outputs[0]),
                str(sig_op.outputs[0]),
                str(sig_mul_op.outputs[0]),
                str(model_ir.operators[int(add_idx)].outputs[0]),
                str(model_ir.operators[int(clip3_unary_idx)].outputs[0]),
            ]:
                _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(tensor_name, None), perm_nchw_to_nhwc)

            for user_idx in consumers.get(str(clip3_nhwc_name), []):
                user_op = model_ir.operators[int(user_idx)]
                _set_operator_inputs(model_ir=model_ir, op=user_op, new_inputs=[clip3_nchw_name if str(v) == clip3_nhwc_name else str(v) for v in list(user_op.inputs)])
            if split_reshape_idx is not None:
                split_reshape_out = str(model_ir.operators[int(split_reshape_idx)].outputs[0])
                for user_idx in consumers.get(split_reshape_out, []):
                    user_op = model_ir.operators[int(user_idx)]
                    _set_operator_inputs(model_ir=model_ir, op=user_op, new_inputs=[split_nchw_name if str(v) == split_reshape_out else str(v) for v in list(user_op.inputs)])

            remove_indices = {int(clip_adapter_idx), int(conv_reshape_idx), int(clip3_adapter_idx)}
            if split_reshape_idx is not None:
                remove_indices.add(int(split_reshape_idx))
            if input_transpose_idx is not None:
                remove_indices.add(int(input_transpose_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

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_singleton_gate_conv_concat_nhwc_bridge_blocks": int(rewritten)}


def _optimize_transpose_unary_split_concat_single_post_nchw(model_ir: ModelIR) -> Dict[str, int]:
    """
    Rewrite TRANSPOSE->UNARY->SPLIT->(optional UNARY)->CONCAT blocks to NHWC
    while preserving a single post-CONCAT NCHW contract.
    """
    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", "LOGISTIC", "HARD_SWISH", "LEAKY_RELU", "TANH"}

    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_singleton_layout_reshape(op: OperatorIR) -> bool:
        if str(op.op_type) != "RESHAPE" or len(op.inputs) < 1 or len(op.outputs) != 1:
            return False
        in_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
        out_tensor = model_ir.tensors.get(str(op.outputs[0]), None)
        if (
            in_tensor is None
            or out_tensor is None
            or not _is_fully_known_positive_shape(in_tensor.shape)
            or not _is_fully_known_positive_shape(out_tensor.shape)
        ):
            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 _permute_shape(in_shape, perm_nchw_to_nhwc) == out_shape or _permute_shape(in_shape, perm_nhwc_to_nchw) == out_shape

    while True:
        changed = False
        producers = _build_tensor_producer_map(model_ir)
        consumers = _build_tensor_consumer_map(model_ir)
        for concat_idx, concat_op in enumerate(model_ir.operators):
            if str(concat_op.op_type) != "CONCATENATION" or int(concat_op.options.get("axis", 1)) != 1 or len(concat_op.outputs) != 1:
                continue
            concat_output_name = str(concat_op.outputs[0])

            concat_inputs = [str(v) for v in list(concat_op.inputs)]
            branch_sources: List[str] = []
            branch_outputs: List[str] = []
            branch_unary_indices: List[int] = []
            split_idx: Optional[int] = None
            pre_unary_idx: Optional[int] = None
            pre_t_idx: Optional[int] = None
            ext_input_name = ""

            for input_name in concat_inputs:
                prod_idx = producers.get(input_name, None)
                source_name = str(input_name)
                branch_unary_idx = -1
                if prod_idx is not None and str(model_ir.operators[int(prod_idx)].op_type) in unary_ops:
                    branch_unary_idx = int(prod_idx)
                    source_name = str(model_ir.operators[int(prod_idx)].inputs[0])
                source_prod_idx = producers.get(source_name, None)
                if source_prod_idx is None or str(model_ir.operators[int(source_prod_idx)].op_type) != "SPLIT":
                    continue
                candidate_split_idx = int(source_prod_idx)
                candidate_split_op = model_ir.operators[int(candidate_split_idx)]
                axis_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(candidate_split_op.inputs[0]), None))
                if axis_vals is None or len(axis_vals) != 1 or int(axis_vals[0]) != 1:
                    continue
                candidate_pre_unary_idx = producers.get(str(candidate_split_op.inputs[1]), None)
                if candidate_pre_unary_idx is None or str(model_ir.operators[int(candidate_pre_unary_idx)].op_type) not in unary_ops:
                    continue
                candidate_pre_t_idx = producers.get(str(model_ir.operators[int(candidate_pre_unary_idx)].inputs[0]), None)
                if candidate_pre_t_idx is None:
                    continue
                candidate_pre_t_op = model_ir.operators[int(candidate_pre_t_idx)]
                if str(candidate_pre_t_op.op_type) != "TRANSPOSE" or _read_transpose_perm(model_ir, candidate_pre_t_op) != perm_nhwc_to_nchw:
                    continue
                if split_idx is None:
                    split_idx = int(candidate_split_idx)
                    pre_unary_idx = int(candidate_pre_unary_idx)
                    pre_t_idx = int(candidate_pre_t_idx)
                elif int(split_idx) != int(candidate_split_idx):
                    continue
                branch_sources.append(str(source_name))
                branch_outputs.append(str(input_name))
                branch_unary_indices.append(int(branch_unary_idx))

            if split_idx is None or pre_unary_idx is None or pre_t_idx is None:
                continue
            split_op = model_ir.operators[int(split_idx)]
            if len(branch_sources) != len(list(split_op.outputs)):
                continue
            ext_input_name = next((name for name in concat_inputs if name not in branch_outputs), "")
            if ext_input_name == "":
                continue
            ext_prod_idx = producers.get(ext_input_name, None)
            if ext_prod_idx is not None and _is_singleton_layout_reshape(model_ir.operators[int(ext_prod_idx)]):
                ext_input_name = str(model_ir.operators[int(ext_prod_idx)].inputs[0])

            pre_t_op = model_ir.operators[int(pre_t_idx)]
            source_nhwc_name = str(pre_t_op.inputs[0])
            split_axis_tensor = model_ir.tensors.get(str(split_op.inputs[0]), None)
            if split_axis_tensor is None:
                continue
            if set(int(v) for v in consumers.get(str(split_op.inputs[0]), [])) == {int(split_idx)}:
                _write_const_ints_to_tensor(split_axis_tensor, [3])
            else:
                axis_name = _unique_tensor_name(f"{str(split_op.inputs[0])}_nhwc")
                model_ir.tensors[axis_name] = TensorIR(name=axis_name, dtype=str(split_axis_tensor.dtype), shape=[1], shape_signature=[1], data=np.asarray([3], dtype=np.int32), is_variable=False, quantization=_clone_quantization(split_axis_tensor.quantization))
                _set_operator_inputs(model_ir=model_ir, op=split_op, new_inputs=[axis_name, str(split_op.inputs[1])])

            pre_unary_op = model_ir.operators[int(pre_unary_idx)]
            _set_operator_inputs(model_ir=model_ir, op=pre_unary_op, new_inputs=[source_nhwc_name])
            _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(pre_unary_op.outputs[0]), None), perm_nchw_to_nhwc)
            for split_output_name in list(split_op.outputs):
                _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(split_output_name), None), perm_nchw_to_nhwc)
            for unary_idx in branch_unary_indices:
                if unary_idx >= 0:
                    _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(model_ir.operators[int(unary_idx)].outputs[0]), None), perm_nchw_to_nhwc)

            new_concat_inputs: List[str] = []
            for input_name in concat_inputs:
                if input_name == ext_input_name or (ext_prod_idx is not None and producers.get(input_name, None) == ext_prod_idx):
                    new_concat_inputs.append(ext_input_name)
                else:
                    new_concat_inputs.append(input_name)
            concat_op.options["axis"] = 3
            _set_operator_inputs(model_ir=model_ir, op=concat_op, new_inputs=new_concat_inputs)

            concat_tensor = model_ir.tensors.get(concat_output_name, None)
            if concat_tensor is None or not _is_fully_known_positive_shape(concat_tensor.shape):
                continue
            nhwc_shape = _permute_shape([int(v) for v in list(concat_tensor.shape)], perm_nchw_to_nhwc)
            nhwc_sig = _permute_shape([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)], perm_nchw_to_nhwc)
            if nhwc_shape is None or nhwc_sig is None:
                continue
            concat_nhwc_name = _unique_tensor_name(f"{concat_output_name}_nhwc")
            model_ir.tensors[concat_nhwc_name] = TensorIR(name=concat_nhwc_name, dtype=str(concat_tensor.dtype), shape=[int(v) for v in list(nhwc_shape)], shape_signature=[int(v) for v in list(nhwc_sig)], data=None, is_variable=False, quantization=_clone_quantization(concat_tensor.quantization))
            _set_operator_outputs(model_ir=model_ir, op=concat_op, new_outputs=[concat_nhwc_name])
            perm_name = _unique_tensor_name(f"{concat_output_name}_perm")
            model_ir.tensors[perm_name] = TensorIR(name=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(concat_idx) + 1, OperatorIR(op_type="TRANSPOSE", inputs=[concat_nhwc_name, perm_name], outputs=[concat_output_name], options={}))
            del model_ir.operators[int(pre_t_idx)]

            rewritten += 1
            changed = True
            break

        if not changed:
            break

    if rewritten > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_unary_split_concat_single_post_nchw": int(rewritten)}


def _optimize_transpose_split_channelwise_tail_to_single_post_nchw(model_ir: ModelIR) -> Dict[str, int]:
    """
    Propagate NHWC through a local channelwise tail that currently starts with
    a NHWC->NCHW transpose feeding a channel SPLIT.
    """
    optimized = 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", "LOGISTIC", "HARD_SWISH", "LEAKY_RELU", "TANH", "NEG", "EXP"}
    binary_ops = {"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_probably_nhwc_tensor(name: str) -> bool:
        tensor = model_ir.tensors.get(str(name), None)
        if tensor is None or not _is_fully_known_positive_shape(tensor.shape):
            return False
        shape = [int(v) for v in list(tensor.shape)]
        if len(shape) != 4:
            return False
        if int(shape[1]) <= 4 and int(shape[2]) > 4 and int(shape[3]) > 4:
            return False
        return True

    def _rewrite_channel_slice_axis1_to_axis3(op: OperatorIR) -> bool:
        if str(op.op_type) != "SLICE" or len(op.inputs) < 3:
            return False
        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
        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 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])
            split_users = [int(v) for v in consumers.get(pre_out_name, [])]
            if len(split_users) != 1:
                continue
            split_idx = int(split_users[0])
            split_op = model_ir.operators[int(split_idx)]
            if str(split_op.op_type) != "SPLIT" or len(split_op.inputs) < 2:
                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) != 1 or int(axis_vals[0]) != 1:
                continue

            _replace_operator_input_at(model_ir=model_ir, op=split_op, input_index=1, new_input_name=pre_in_name)
            axis_tensor = model_ir.tensors.get(str(split_op.inputs[0]), None)
            if axis_tensor is not None:
                if set(int(v) for v in consumers.get(str(split_op.inputs[0]), [])) == {int(split_idx)}:
                    _write_const_ints_to_tensor(axis_tensor, [3])

            nhwc_tensors = {str(pre_in_name)}
            for out_name in list(split_op.outputs):
                _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(out_name), None), perm_nchw_to_nhwc)
                nhwc_tensors.add(str(out_name))

            propagated = True
            while propagated:
                propagated = False
                for op in model_ir.operators:
                    if len(op.outputs) != 1:
                        continue
                    out_name = str(op.outputs[0])
                    if out_name in nhwc_tensors:
                        continue
                    op_type = str(op.op_type)
                    if op_type in unary_ops and len(op.inputs) == 1 and 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)
                        propagated = True
                        continue
                    if op_type in binary_ops and len(op.inputs) == 2:
                        tensor_inputs = [str(v) for v in list(op.inputs) if model_ir.tensors.get(str(v), None) is not None and model_ir.tensors[str(v)].data is None]
                        if len(tensor_inputs) == 0:
                            continue
                        if all(name in nhwc_tensors or _is_probably_nhwc_tensor(name) for name in tensor_inputs):
                            _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(out_name, None), perm_nchw_to_nhwc)
                            nhwc_tensors.add(out_name)
                            propagated = True
                        continue
                    if op_type == "CONCATENATION" and len(op.inputs) > 0 and all(str(v) in nhwc_tensors for v in list(op.inputs)):
                        axis = int(op.options.get("axis", -1))
                        if axis in {1, -3}:
                            op.options["axis"] = 3
                        _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(out_name, None), perm_nchw_to_nhwc)
                        nhwc_tensors.add(out_name)
                        propagated = True
                        continue
                    if op_type == "SPLIT" and len(op.inputs) >= 2 and str(op.inputs[1]) in nhwc_tensors:
                        axis_tensor_local = model_ir.tensors.get(str(op.inputs[0]), None)
                        axis_vals_local = _read_const_ints_from_tensor(axis_tensor_local)
                        if axis_vals_local is None or len(axis_vals_local) != 1 or int(axis_vals_local[0]) not in {1, 3}:
                            continue
                        if int(axis_vals_local[0]) == 1:
                            axis_users = [int(v) for v in consumers.get(str(op.inputs[0]), [])]
                            if set(axis_users) == {int(next(idx for idx, candidate in enumerate(model_ir.operators) if candidate is op))}:
                                _write_const_ints_to_tensor(axis_tensor_local, [3])
                            else:
                                axis_name = _unique_tensor_name(f"{str(op.inputs[0])}_nhwc")
                                model_ir.tensors[axis_name] = TensorIR(
                                    name=axis_name,
                                    dtype=str(axis_tensor_local.dtype),
                                    shape=[1],
                                    shape_signature=[1],
                                    data=np.asarray([3], dtype=np.int32),
                                    is_variable=False,
                                    quantization=_clone_quantization(axis_tensor_local.quantization),
                                )
                                _set_operator_inputs(
                                    model_ir=model_ir,
                                    op=op,
                                    new_inputs=[axis_name, str(op.inputs[1])],
                                )
                        for split_out_name in list(op.outputs):
                            _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(str(split_out_name), None), perm_nchw_to_nhwc)
                            nhwc_tensors.add(str(split_out_name))
                        propagated = True
                        continue
                    if op_type == "SLICE" and len(op.inputs) >= 3 and str(op.inputs[0]) in nhwc_tensors:
                        if _rewrite_channel_slice_axis1_to_axis3(op):
                            _permute_tensor_metadata_if_rank_matches(model_ir.tensors.get(out_name, None), perm_nchw_to_nhwc)
                            nhwc_tensors.add(out_name)
                            propagated = True

            output_name = str(model_ir.outputs[0]) if len(model_ir.outputs) == 1 else ""
            if output_name != "" and output_name in nhwc_tensors:
                prod_idx = producers.get(output_name, None)
                if prod_idx is not None:
                    prod_op = model_ir.operators[int(prod_idx)]
                    output_nhwc_name = _unique_tensor_name(f"{output_name}_nhwc")
                    out_tensor = model_ir.tensors.get(output_name, None)
                    if out_tensor is not None and _is_fully_known_positive_shape(out_tensor.shape):
                        nhwc_shape = [int(v) for v in list(out_tensor.shape)]
                        nhwc_sig = [int(v) for v in list(out_tensor.shape_signature)] if out_tensor.shape_signature is not None else list(nhwc_shape)
                        model_ir.tensors[output_nhwc_name] = TensorIR(name=output_nhwc_name, dtype=str(out_tensor.dtype), shape=list(nhwc_shape), shape_signature=list(nhwc_sig), data=None, is_variable=False, quantization=_clone_quantization(out_tensor.quantization))
                        _set_operator_outputs(model_ir=model_ir, op=prod_op, new_outputs=[output_nhwc_name])
                        perm_name = _unique_tensor_name(f"{output_name}_perm")
                        model_ir.tensors[perm_name] = TensorIR(name=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(prod_idx) + 1, OperatorIR(op_type="TRANSPOSE", inputs=[output_nhwc_name, perm_name], outputs=[output_name], options={}))

            del model_ir.operators[int(pre_idx)]
            optimized += 1
            changed = True
            break

        if not changed:
            break

    if optimized > 0:
        _prune_unused_tensors(model_ir)
    return {"optimized_transpose_split_channelwise_tail_to_single_post_nchw": int(optimized)}


def _sanitize_probable_nhwc_axis_sensitive_ops(model_ir: ModelIR) -> Dict[str, int]:
    """
    Repair late axis-sensitive metadata when tensors are clearly NHWC but still
    carry NCHW-era axis constants or derived shapes.
    """
    rewritten = 0
    inserted_terminal_transposes = 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_probable_nhwc_shape(shape: Optional[List[int]]) -> bool:
        if shape is None or len(shape) != 4:
            return False
        values = [int(v) for v in list(shape)]
        return (
            all(int(v) > 0 for v in values)
            and int(values[1]) > 4
            and int(values[2]) > 4
            and 0 < int(values[3]) <= 8
        )

    def _set_axis_input(op: OperatorIR, op_idx: int, new_axis: int, consumers: Dict[str, List[int]]) -> None:
        axis_name = str(op.inputs[0])
        axis_tensor = model_ir.tensors.get(axis_name, None)
        if axis_tensor is None:
            return
        axis_users = [int(v) for v in consumers.get(axis_name, [])]
        if set(axis_users) == {int(op_idx)}:
            _write_const_ints_to_tensor(axis_tensor, [int(new_axis)])
            return
        clone_name = _unique_tensor_name(f"{axis_name}_nhwc")
        model_ir.tensors[clone_name] = TensorIR(
            name=clone_name,
            dtype=str(axis_tensor.dtype),
            shape=[1],
            shape_signature=[1],
            data=np.asarray([int(new_axis)], dtype=np.int32),
            is_variable=False,
            quantization=_clone_quantization(axis_tensor.quantization),
        )
        _set_operator_inputs(
            model_ir=model_ir,
            op=op,
            new_inputs=[clone_name] + [str(v) for v in list(op.inputs[1:])],
        )

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

        for op_idx, op in enumerate(model_ir.operators):
            op_type = str(op.op_type)

            if op_type == "SPLIT" and len(op.inputs) >= 2:
                input_name = str(op.inputs[1])
                input_tensor = model_ir.tensors.get(input_name, None)
                axis_vals = _read_const_ints_from_tensor(model_ir.tensors.get(str(op.inputs[0]), None))
                if (
                    input_tensor is not None
                    and _is_probable_nhwc_shape(input_tensor.shape)
                    and axis_vals is not None
                    and len(axis_vals) == 1
                    and int(axis_vals[0]) == 1
                ):
                    input_shape = [int(v) for v in list(input_tensor.shape)]
                    num_splits = int(op.options.get("numSplits", len(list(op.outputs))))
                    if num_splits > 0 and int(input_shape[3]) % int(num_splits) == 0:
                        _set_axis_input(op, int(op_idx), 3, consumers)
                        c_each = int(input_shape[3]) // int(num_splits)
                        for output_name in list(op.outputs):
                            output_tensor = model_ir.tensors.get(str(output_name), None)
                            if output_tensor is not None:
                                output_tensor.shape = [input_shape[0], input_shape[1], input_shape[2], c_each]
                                output_tensor.shape_signature = [input_shape[0], input_shape[1], input_shape[2], c_each]
                        rewritten += 1
                        changed = True
                continue

            if op_type == "CONCATENATION" and len(op.inputs) > 0:
                axis = int(op.options.get("axis", -1))
                input_tensors = [model_ir.tensors.get(str(name), None) for name in list(op.inputs)]
                if (
                    axis == 1
                    and all(t is not None and _is_probable_nhwc_shape(t.shape) for t in input_tensors)
                ):
                    base = [int(v) for v in list(input_tensors[0].shape)]
                    if all(
                        [int(v) for v in list(t.shape)][:3] == base[:3]
                        for t in input_tensors[1:]
                    ):
                        total_c = sum(int(list(t.shape)[3]) for t in input_tensors if t is not None)
                        op.options["axis"] = 3
                        out_tensor = model_ir.tensors.get(str(op.outputs[0]), None)
                        if out_tensor is not None:
                            out_tensor.shape = [base[0], base[1], base[2], total_c]
                            out_tensor.shape_signature = [base[0], base[1], base[2], total_c]
                        rewritten += 1
                        changed = True
                continue

            if op_type == "SLICE" and len(op.inputs) >= 3:
                input_name = str(op.inputs[0])
                input_tensor = model_ir.tensors.get(input_name, None)
                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 (
                    input_tensor is not None
                    and _is_probable_nhwc_shape(input_tensor.shape)
                    and begin_vals is not None
                    and size_vals is not None
                    and len(begin_vals) == 4
                    and len(size_vals) == 4
                    and int(size_vals[1]) == 1
                    and int(begin_vals[2]) == 0
                    and int(begin_vals[3]) == 0
                ):
                    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)
                    input_shape = [int(v) for v in list(input_tensor.shape)]
                    out_tensor = model_ir.tensors.get(str(op.outputs[0]), None)
                    if out_tensor is not None:
                        out_tensor.shape = [input_shape[0], input_shape[1], input_shape[2], int(new_size[3])]
                        out_tensor.shape_signature = [input_shape[0], input_shape[1], input_shape[2], int(new_size[3])]
                    rewritten += 1
                    changed = True
                continue

            if op_type in {"RELU", "RELU6", "RELU_0_TO_1", "LOGISTIC", "HARD_SWISH", "LEAKY_RELU", "TANH", "NEG", "EXP"} and len(op.inputs) == 1:
                input_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
                out_tensor = model_ir.tensors.get(str(op.outputs[0]), None) if len(op.outputs) == 1 else None
                if input_tensor is not None and out_tensor is not None and _is_probable_nhwc_shape(input_tensor.shape):
                    out_tensor.shape = [int(v) for v in list(input_tensor.shape)]
                    out_tensor.shape_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)]
                continue

            if op_type in {"ADD", "SUB", "MUL", "DIV", "MAXIMUM", "MINIMUM"} and len(op.inputs) == 2 and len(op.outputs) == 1:
                a_tensor = model_ir.tensors.get(str(op.inputs[0]), None)
                b_tensor = model_ir.tensors.get(str(op.inputs[1]), None)
                out_tensor = model_ir.tensors.get(str(op.outputs[0]), None)
                if out_tensor is None:
                    continue
                shapes = []
                for tensor in [a_tensor, b_tensor]:
                    if tensor is not None and tensor.data is None and tensor.shape is not None:
                        shapes.append([int(v) for v in list(tensor.shape)])
                if len(shapes) == 0:
                    continue
                if all(_is_probable_nhwc_shape(shape) or shape in ([1], []) for shape in shapes):
                    inferred = shapes[0]
                    for shape in shapes[1:]:
                        candidate = _broadcast_static_shapes(inferred, shape)
                        if candidate is not None:
                            inferred = [int(v) for v in list(candidate)]
                    if len(inferred) == 4:
                        out_tensor.shape = list(inferred)
                        out_tensor.shape_signature = list(inferred)
                continue

        if not changed:
            break

    if rewritten > 0:
        producers = _build_tensor_producer_map(model_ir)
        for output_name in list(model_ir.outputs):
            output_name = str(output_name)
            output_tensor = model_ir.tensors.get(output_name, None)
            if output_tensor is None or not _is_probable_nhwc_shape(output_tensor.shape):
                continue
            prod_idx = producers.get(output_name, None)
            if prod_idx is None:
                continue
            output_nhwc_name = _unique_tensor_name(f"{output_name}_nhwc")
            nhwc_shape = [int(v) for v in list(output_tensor.shape)]
            nhwc_sig = [int(v) for v in list(output_tensor.shape_signature)] if output_tensor.shape_signature is not None else list(nhwc_shape)
            nchw_shape = [nhwc_shape[0], nhwc_shape[3], nhwc_shape[1], nhwc_shape[2]]
            nchw_sig = [nhwc_sig[0], nhwc_sig[3], nhwc_sig[1], nhwc_sig[2]]
            model_ir.tensors[output_nhwc_name] = TensorIR(
                name=output_nhwc_name,
                dtype=str(output_tensor.dtype),
                shape=list(nhwc_shape),
                shape_signature=list(nhwc_sig),
                data=None,
                is_variable=False,
                quantization=_clone_quantization(output_tensor.quantization),
            )
            _set_operator_outputs(
                model_ir=model_ir,
                op=model_ir.operators[int(prod_idx)],
                new_outputs=[output_nhwc_name],
            )
            perm_name = _unique_tensor_name(f"{output_name}_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,
            )
            output_tensor.shape = list(nchw_shape)
            output_tensor.shape_signature = list(nchw_sig)
            model_ir.operators.insert(
                int(prod_idx) + 1,
                OperatorIR(
                    op_type="TRANSPOSE",
                    inputs=[output_nhwc_name, perm_name],
                    outputs=[output_name],
                    options={},
                ),
            )
            inserted_terminal_transposes += 1

    return {
        "sanitized_probable_nhwc_axis_sensitive_ops": int(rewritten),
        "inserted_probable_nhwc_terminal_transposes": int(inserted_terminal_transposes),
    }


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)
        data_tensor = model_ir.tensors.get(name, None)
        if data_tensor is not None:
            logical_layout = str(getattr(data_tensor, "logical_layout", "") or "").upper()
            if logical_layout in {"NHWC", "NCHW"}:
                return logical_layout
        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 _repair_decomposed_instance_normalization_layouts(
    model_ir: ModelIR,
) -> Dict[str, int]:
    repaired = 0
    consumers = _build_tensor_consumer_map(model_ir)

    def _single_consumer_index(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 _set_shape_metadata(tensor_name: str, shape: List[int]) -> bool:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None:
            return False
        changed = False
        normalized_shape = [int(v) for v in list(shape)]
        if [int(v) for v in list(tensor.shape)] != normalized_shape:
            tensor.shape = list(normalized_shape)
            changed = True
        if tensor.shape_signature is not None and len(list(tensor.shape_signature)) == len(normalized_shape):
            if [int(v) for v in list(tensor.shape_signature)] != normalized_shape:
                tensor.shape_signature = list(normalized_shape)
                changed = True
        return changed

    def _reshape_channelwise_const(tensor_name: str, broadcast_shape: List[int]) -> bool:
        tensor = model_ir.tensors.get(str(tensor_name), None)
        if tensor is None or tensor.data is None:
            return False
        data = np.asarray(tensor.data)
        if int(data.size) <= 1:
            return False
        if int(np.prod(np.asarray(broadcast_shape), dtype=np.int64)) != int(data.size):
            return False
        reshaped = np.asarray(data.reshape(broadcast_shape), dtype=data.dtype)
        changed = False
        if list(data.shape) != list(reshaped.shape) or not np.array_equal(data, reshaped):
            tensor.data = reshaped
            changed = True
        if [int(v) for v in list(tensor.shape)] != [int(v) for v in list(reshaped.shape)]:
            tensor.shape = [int(v) for v in list(reshaped.shape)]
            changed = True
        if tensor.shape_signature is not None and len(list(tensor.shape_signature)) == int(reshaped.ndim):
            if [int(v) for v in list(tensor.shape_signature)] != [int(v) for v in list(reshaped.shape)]:
                tensor.shape_signature = [int(v) for v in list(reshaped.shape)]
                changed = True
        return changed

    for mean1_idx, mean1_op in enumerate(list(model_ir.operators)):
        if str(mean1_op.op_type) != "MEAN" or len(mean1_op.inputs) != 2 or len(mean1_op.outputs) != 1:
            continue
        x_name = str(mean1_op.inputs[0])
        x_tensor = model_ir.tensors.get(x_name, None)
        if x_tensor is None:
            continue
        x_shape = [int(v) for v in list(x_tensor.shape)]
        rank = len(x_shape)
        if rank not in {3, 4, 5}:
            continue
        x_layout = normalize_logical_layout(getattr(x_tensor, "logical_layout", None))
        if is_channel_first_logical_layout(x_layout):
            channel_axis = 1
        elif is_channel_last_logical_layout(x_layout):
            channel_axis = rank - 1
        else:
            continue
        desired_axes = [int(axis) for axis in range(1, rank) if int(axis) != int(channel_axis)]
        axes_name = str(mean1_op.inputs[1])
        axes_values = _read_const_ints_from_tensor(model_ir.tensors.get(axes_name, None))
        if axes_values is None:
            continue

        mean1_out_name = str(mean1_op.outputs[0])
        sub_idx = None
        for user_idx in consumers.get(x_name, []):
            user_op = model_ir.operators[int(user_idx)]
            if (
                str(user_op.op_type) == "SUB"
                and len(user_op.inputs) == 2
                and len(user_op.outputs) == 1
                and x_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
                and mean1_out_name in {str(user_op.inputs[0]), str(user_op.inputs[1])}
            ):
                sub_idx = int(user_idx)
                break
        if sub_idx is None:
            continue
        centered_name = str(model_ir.operators[int(sub_idx)].outputs[0])

        mul_square_idx = None
        mul_norm_idx = None
        for user_idx in consumers.get(centered_name, []):
            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

        squared_name = str(model_ir.operators[int(mul_square_idx)].outputs[0])
        mean2_idx = _single_consumer_index(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
        mean2_axes_name = str(mean2_op.inputs[1])
        mean2_axes_values = _read_const_ints_from_tensor(model_ir.tensors.get(mean2_axes_name, None))
        if mean2_axes_values is None:
            continue

        mean2_out_name = str(mean2_op.outputs[0])
        add_eps_idx = _single_consumer_index(mean2_out_name)
        if add_eps_idx is None:
            continue
        add_eps_out_name = str(model_ir.operators[int(add_eps_idx)].outputs[0])
        sqrt_idx = _single_consumer_index(add_eps_out_name)
        if sqrt_idx is None:
            continue
        sqrt_out_name = str(model_ir.operators[int(sqrt_idx)].outputs[0])
        div_idx = _single_consumer_index(sqrt_out_name)
        if div_idx is None:
            continue
        div_out_name = str(model_ir.operators[int(div_idx)].outputs[0])
        mul_norm_op = model_ir.operators[int(mul_norm_idx)]
        if div_out_name not in {str(v) for v in list(mul_norm_op.inputs)}:
            continue
        normalized_name = str(mul_norm_op.outputs[0])

        mul_scale_idx = _single_consumer_index(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 {str(v) for v in list(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])
        desired_reduced_shape = list(x_shape)
        for axis in desired_axes:
            desired_reduced_shape[int(axis)] = 1
        desired_broadcast_shape = [1 for _ in range(rank)]
        desired_broadcast_shape[int(channel_axis)] = int(x_shape[int(channel_axis)])

        bias_const_name = None
        bias_broadcast_shape = None
        post_scaled_name = scaled_name
        add_bias_idx = _single_consumer_index(scaled_name)
        if add_bias_idx is not None:
            add_bias_op = model_ir.operators[int(add_bias_idx)]
            if (
                str(add_bias_op.op_type) == "ADD"
                and len(add_bias_op.inputs) == 2
                and len(add_bias_op.outputs) == 1
                and scaled_name in {str(v) for v in list(add_bias_op.inputs)}
            ):
                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])
                )
                bias_broadcast_shape = list(desired_broadcast_shape)
            else:
                add_bias_idx = None
        if add_bias_idx is None:
            post_idx = _single_consumer_index(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:
                continue
            post_perm = _read_transpose_perm(model_ir, post_op)
            if post_perm is None or len(post_perm) != rank:
                continue
            post_scaled_name = str(post_op.outputs[0])
            add_bias_idx = _single_consumer_index(post_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
                or post_scaled_name not in {str(v) for v in list(add_bias_op.inputs)}
            ):
                continue
            bias_const_name = (
                str(add_bias_op.inputs[0])
                if str(add_bias_op.inputs[1]) == post_scaled_name
                else str(add_bias_op.inputs[1])
            )
            bias_broadcast_shape = [1 for _ in range(rank)]
            inverse_post_perm = [0 for _ in range(rank)]
            for new_axis, old_axis in enumerate(post_perm):
                inverse_post_perm[int(old_axis)] = int(new_axis)
            post_channel_axis = int(inverse_post_perm[int(channel_axis)])
            bias_broadcast_shape[int(post_channel_axis)] = int(x_shape[int(channel_axis)])
        if bias_const_name is None or bias_broadcast_shape is None:
            continue

        changed = False
        if [int(v) for v in list(axes_values)] != desired_axes:
            changed = _write_const_ints_to_tensor(model_ir.tensors.get(axes_name, None), desired_axes) or changed
        if [int(v) for v in list(mean2_axes_values)] != desired_axes:
            changed = _write_const_ints_to_tensor(model_ir.tensors.get(mean2_axes_name, None), desired_axes) or changed

        for reduced_name in [mean1_out_name, mean2_out_name, add_eps_out_name, sqrt_out_name, div_out_name]:
            changed = _set_shape_metadata(reduced_name, desired_reduced_shape) or changed
        for full_name in [centered_name, squared_name, normalized_name, scaled_name]:
            changed = _set_shape_metadata(full_name, x_shape) or changed
        changed = _reshape_channelwise_const(scale_const_name, desired_broadcast_shape) or changed
        changed = _reshape_channelwise_const(bias_const_name, bias_broadcast_shape) or changed
        if post_scaled_name != scaled_name:
            post_scaled_tensor = model_ir.tensors.get(post_scaled_name, None)
            if post_scaled_tensor is not None:
                post_shape = [int(v) for v in list(post_scaled_tensor.shape)]
                if len(post_shape) == rank:
                    changed = _set_shape_metadata(post_scaled_name, post_shape) or changed

        if changed:
            repaired += 1

    return {"repaired_decomposed_instance_normalization_layouts": 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)
        consumers = _build_tensor_consumer_map(model_ir)
        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
            consumer_op_types = {
                str(model_ir.operators[int(consumer_idx)].op_type)
                for consumer_idx in consumers.get(str(output_name), [])
                if int(consumer_idx) != int(op_idx)
            }
            if len(consumer_op_types & {"GATHER", "GATHER_ND", "SLICE", "STRIDED_SLICE"}) > 0:
                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),
    )
    lowering_progress_spinner = _ProgressSpinner(progress_bar)
    try:
        lowering_progress_spinner.start()
        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:
            lowering_progress_spinner.stop()
            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),
    )
    post_progress_spinner = _ProgressSpinner(post_progress_bar)

    def _set_post_progress_desc(stage_label: str) -> None:
        if post_progress_bar is None:
            return
        post_progress_spinner.start()
        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_spinner.stop()
        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_resize_add_concat_affine_conv_spp_nhwc_chains(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_mean_mul_reshape_add_conv_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_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_mixed_mean_reducemax_concat_mirrorpad_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_logistic_sub_mul_postadd_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_resize_add_concat_affine_conv_spp_nhwc_chains(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_mean_mul_reshape_add_conv_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_mixed_mean_reducemax_concat_mirrorpad_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_logistic_sub_mul_postadd_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_resize_add_concat_affine_conv_spp_nhwc_chains(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_mean_mul_reshape_add_conv_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_mixed_mean_reducemax_concat_mirrorpad_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_logistic_sub_mul_postadd_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_resize_add_concat_affine_conv_spp_nhwc_chains(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_mean_mul_reshape_add_conv_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_mixed_mean_reducemax_concat_mirrorpad_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_logistic_sub_mul_postadd_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_resize_add_concat_affine_conv_spp_nhwc_chains(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_mean_mul_reshape_add_conv_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_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_logistic_sub_mul_postadd_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_singleton_gate_conv_concat_nhwc_bridge_blocks(model_ir)
    _optimize_transpose_unary_split_concat_single_post_nchw(model_ir)
    _optimize_transpose_split_channelwise_tail_to_single_post_nchw(model_ir)
    _sanitize_probable_nhwc_axis_sensitive_ops(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_mean_mul_reshape_add_conv_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_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_transpose_relu_split_all_outputs_to_nhwc_chains(model_ir)
    _optimize_transpose_relu_split_conv_relu_concat_posttranspose_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_mixed_mean_reducemax_concat_mirrorpad_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_singleton_gate_conv_concat_nhwc_bridge_blocks(model_ir)
    _optimize_transpose_unary_split_concat_single_post_nchw(model_ir)
    _optimize_transpose_split_channelwise_tail_to_single_post_nchw(model_ir)
    _sanitize_probable_nhwc_axis_sensitive_ops(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_relu_split_all_outputs_to_nhwc_chains(model_ir)
    _optimize_transpose_relu_split_conv_relu_concat_posttranspose_to_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_singleton_gate_conv_concat_nhwc_bridge_blocks(model_ir)
    _optimize_transpose_unary_split_concat_single_post_nchw(model_ir)
    _optimize_transpose_split_channelwise_tail_to_single_post_nchw(model_ir)
    _sanitize_probable_nhwc_axis_sensitive_ops(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_singleton_gate_conv_concat_nhwc_bridge_blocks(model_ir)
    _optimize_transpose_unary_split_concat_single_post_nchw(model_ir)
    _optimize_transpose_split_channelwise_tail_to_single_post_nchw(model_ir)
    _sanitize_probable_nhwc_axis_sensitive_ops(model_ir)
    _optimize_transpose_stridedslice_pad_concat_mul_add_posttranspose_nhwc_chains(model_ir)
    _optimize_transpose_resize_add_concat_affine_conv_spp_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_mean_mul_reshape_add_conv_nhwc_chains(model_ir)
    _optimize_transpose_resize_add_concat_affine_conv_spp_nhwc_chains(model_ir)
    _optimize_transpose_gather_transpose_axis_remap_nhwc_chains(model_ir)
    _optimize_constant_input_pad_chains(model_ir)
    _optimize_constant_input_pool_chains(model_ir)
    _optimize_constant_input_cast_chains(model_ir)
    _optimize_redundant_int32_to_int64_passthrough_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_orphan_recurrent_step_tensors(model_ir)
    _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_pad_chains(model_ir)
    _optimize_constant_input_pool_chains(model_ir)
    _optimize_constant_input_cast_chains(model_ir)
    _optimize_redundant_int32_to_int64_passthrough_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,
    )
    _rewrite_dynamic_rank1_unsqueeze_reshape_shape_inputs(model_ir)
    _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)
        _rewrite_dynamic_rank1_unsqueeze_reshape_shape_inputs(fallback_ir)
        _topologically_sort_operators(fallback_ir)
        infer_model_ir_logical_layouts(fallback_ir)
        fallback_broadcast_repair_stats = _repair_rank4_channelwise_broadcast_constants_to_runtime_layout(
            fallback_ir
        )
        if int(fallback_broadcast_repair_stats.get("repaired_rank4_channelwise_broadcast_constants", 0)) > 0:
            _reconcile_static_tensor_shapes(fallback_ir)
            _topologically_sort_operators(fallback_ir)
            infer_model_ir_logical_layouts(fallback_ir)
        _rewrite_constant_divisors_to_multiplicative_reciprocals(fallback_ir)
        _optimize_fold_consecutive_mul_constants_chains(fallback_ir)
        _topologically_sort_operators(fallback_ir)
        fallback_layout_problems = validate_model_ir_layout_annotations(fallback_ir)
        if len(fallback_layout_problems) > 0:
            fallback_ir.metadata["logical_layout_validation_errors"] = list(fallback_layout_problems)
        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

    _rewrite_constant_divisors_to_multiplicative_reciprocals(model_ir)
    _optimize_fold_consecutive_mul_constants_chains(model_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)
    # Some late boundary/layout repairs can still recreate the DEA/SiNet
    # mixed NHWC/NCHW SA branch around REDUCE_MAX->CONCAT->MIRROR_PAD.
    _optimize_mixed_mean_reducemax_concat_mirrorpad_nhwc_chains(model_ir)
    _rewrite_dynamic_rank1_unsqueeze_reshape_shape_inputs(model_ir)
    _topologically_sort_operators(model_ir)
    infer_model_ir_logical_layouts(model_ir)
    final_instancenorm_repair_stats = _repair_decomposed_instance_normalization_layouts(model_ir)
    if int(final_instancenorm_repair_stats.get("repaired_decomposed_instance_normalization_layouts", 0)) > 0:
        _reconcile_static_tensor_shapes(model_ir)
        _topologically_sort_operators(model_ir)
        infer_model_ir_logical_layouts(model_ir)
    final_broadcast_repair_stats = _repair_rank4_channelwise_broadcast_constants_to_runtime_layout(model_ir)
    if int(final_broadcast_repair_stats.get("repaired_rank4_channelwise_broadcast_constants", 0)) > 0:
        _reconcile_static_tensor_shapes(model_ir)
        _topologically_sort_operators(model_ir)
        infer_model_ir_logical_layouts(model_ir)
    # Absolute-final reshape cleanup:
    # very late repair/reconciliation passes above can still recreate trivial
    # singleton-growth RESHAPE chains (e.g. 2D->3D->4D Conv1D input shims).
    _optimize_consecutive_reshape_passthrough_chains(model_ir)
    _reconcile_static_tensor_shapes(model_ir)
    layout_problems = validate_model_ir_layout_annotations(model_ir)
    if len(layout_problems) > 0:
        model_ir.metadata["logical_layout_validation_errors"] = list(layout_problems)
    _advance_post_progress()
    if post_progress_bar is not None:
        post_progress_spinner.stop()
        post_progress_bar.close()

    return model_ir
