diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py
new file mode 100644
index 00000000000..42660a3e31e
--- /dev/null
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py
@@ -0,0 +1,180 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Convert InstanceNormalization nodes to 4D-compatible form.
+
+Transform rule
+--------------
+Only 3D and 4D inputs are supported:
+
+* **4D input** ``[N, C, H, W]`` — no transformation needed; node is left untouched.
+* **3D input** ``[N, C, D]`` — the node is replaced by a four-node sequence:
+
+  .. code-block::
+
+      [N,C,D]  →  Reshape[N,C,D,1]  →  InstanceNorm  →  Reshape[1,N,C,D]  →  Squeeze(axis=0)  →  [N,C,D]
+
+  The intermediate tensor after the second Reshape is ``[1,N,C,D]`` — a fully 4D tensor
+  compatible with downstream DLA ops.  The final Squeeze restores the original 3D shape
+  so that downstream consumers see the expected rank.
+
+* **>4D input** — raises ``ValueError`` (unsupported).
+
+Numerical correctness
+---------------------
+InstanceNorm normalises each ``(N, C)`` slice over its spatial dimensions.
+For ``[N,C,D,1]`` the spatial product is ``D x 1 = D``, identical to the original
+``[N,C,D]`` case, so the transform is numerically lossless.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+
+from ...logging_config import logger
+from ..utils.graph_utils import (
+    GraphCache,
+    add_unique_initializers,
+    add_value_info,
+    batch_replace_nodes,
+)
+from .dla_constants_to_initializers import _add_legacy_initializer_inputs
+from .dla_indexing import _known_rank, _operator_nodes
+from .dla_model_boundaries import _require_adapter_opset
+
+
+def _apply_fix_instancenorm_channel_mismatch(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Replace every 3D InstanceNormalization with a 4D-compatible Reshape/InstanceNorm/Reshape/Squeeze sequence."""
+    # ── Shape inference ───────────────────────────────────────────────────────
+    graph = model.graph
+    cache = GraphCache(graph)
+    node_idx = {id(n): i for i, n in enumerate(graph.node)}
+    unique_init: set = set(cache.init_names)
+
+    replacements: dict[int, tuple[list, list]] = {}
+    cnt = 0
+
+    nodes = _operator_nodes(model, {"InstanceNormalization"})
+    for node in nodes:
+        rank = _known_rank(cache, node.input[0])
+        if rank is not None and rank not in (3, 4):
+            raise ValueError(
+                f"InstanceNormalization {node.name!r}: input rank {rank} is not supported"
+            )
+    if any(
+        (shape := cache.get_shape(node.input[0])) is not None and len(shape) == 3 and all(shape)
+        for node in nodes
+    ):
+        _require_adapter_opset(model)
+
+    for node in nodes:
+        in_name = node.input[0]
+        out_name = node.output[0]
+
+        in_shape = cache.get_shape(in_name)
+        if in_shape is None:
+            logger.warning(
+                "InstanceNormalization node %r: input shape unknown; skipping.", node.name
+            )
+            continue
+
+        rank = len(in_shape)
+
+        if rank == 4:
+            continue  # Already 4D — nothing to do.
+
+        if not all(in_shape):
+            continue
+
+        n, c, d = in_shape
+        pfx = node.name or f"instnorm_{cnt}"
+
+        # 1. Pre-Reshape: [N,C,D] → [N,C,D,1]  (add trailing unary spatial dim)
+        pre_shape_init = numpy_helper.from_array(
+            np.array([n, c, d, 1], dtype=np.int64),
+            name=cache.unique_name(f"{pfx}_pre_shape_{cnt}"),
+        )
+        add_unique_initializers(graph, unique_init, [pre_shape_init])
+        pre_out = cache.unique_name(f"{in_name}_trail4d_{cnt}")
+        pre_reshape = helper.make_node(
+            "Reshape",
+            inputs=[in_name, pre_shape_init.name],
+            outputs=[pre_out],
+            name=cache.unique_name(f"{pfx}_pre_reshape_{cnt}"),
+        )
+        in_dtype = cache.get_dtype(in_name) or onnx.TensorProto.FLOAT
+        add_value_info(graph, pre_out, in_dtype, [n, c, d, 1])
+
+        # 2. InstanceNorm on [N,C,D,1] → [N,C,D,1]
+        instrnorm_out = cache.unique_name(f"{out_name}_4d_{cnt}")
+        new_instrnorm = onnx.NodeProto()
+        new_instrnorm.CopyFrom(node)
+        new_instrnorm.input[0] = pre_out
+        new_instrnorm.output[0] = instrnorm_out
+        new_instrnorm.name = cache.unique_name(f"{pfx}_4d_{cnt}")
+        add_value_info(graph, instrnorm_out, in_dtype, [n, c, d, 1])
+
+        # 3. Post-Reshape: [N,C,D,1] → [1,N,C,D]  (move spatial dim to batch position)
+        post_shape_init = numpy_helper.from_array(
+            np.array([1, n, c, d], dtype=np.int64),
+            name=cache.unique_name(f"{pfx}_post_shape_{cnt}"),
+        )
+        add_unique_initializers(graph, unique_init, [post_shape_init])
+        post_out = cache.unique_name(f"{out_name}_1ncd_{cnt}")
+        post_reshape = helper.make_node(
+            "Reshape",
+            inputs=[instrnorm_out, post_shape_init.name],
+            outputs=[post_out],
+            name=cache.unique_name(f"{pfx}_post_reshape_{cnt}"),
+        )
+        add_value_info(graph, post_out, in_dtype, [1, n, c, d])
+
+        # 4. Squeeze axis=0: [1,N,C,D] → [N,C,D]  (restore original rank)
+        squeeze_axes_init = numpy_helper.from_array(
+            np.array([0], dtype=np.int64),
+            name=cache.unique_name(f"{pfx}_squeeze_axes_{cnt}"),
+        )
+        add_unique_initializers(graph, unique_init, [squeeze_axes_init])
+        squeeze = helper.make_node(
+            "Squeeze",
+            inputs=[post_out, squeeze_axes_init.name],
+            outputs=[out_name],
+            name=cache.unique_name(f"{pfx}_squeeze_{cnt}"),
+        )
+
+        replacements[node_idx[id(node)]] = (
+            [node],
+            [pre_reshape, new_instrnorm, post_reshape, squeeze],
+        )
+        logger.debug(
+            "InstanceNormalization %r [%s] → Reshape[%s,1]→InstanceNorm→Reshape[1,%s]→Squeeze",
+            node.name,
+            ",".join(str(d) for d in in_shape),
+            ",".join(str(d) for d in in_shape),
+            ",".join(str(d) for d in in_shape),
+        )
+        cnt += 1
+
+    # ── Apply replacements ─────────────────────────────────────────────────────
+    batch_replace_nodes(graph, replacements)
+    _add_legacy_initializer_inputs(model)
+
+    logger.debug(
+        "fix_instancenorm_channel_mismatch: transformed %d InstanceNormalization node(s).",
+        len(replacements),
+    )
+    return model
diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py
new file mode 100644
index 00000000000..01669428071
--- /dev/null
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py
@@ -0,0 +1,204 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Replace proven private TopK index tiling with a vector Gather."""
+
+from __future__ import annotations
+
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper
+
+from ...logging_config import logger
+from ..utils.graph_utils import (
+    GraphCache,
+    batch_replace_nodes,
+    get_node_attr_i,
+    resolve_reshape_shape,
+    set_node_attr_i,
+)
+from .dla_constants_to_initializers import _add_legacy_initializer_inputs
+from .dla_graph_cleanup import _subgraph_captures
+from .dla_indexing import _insert_support_nodes, _operator_nodes
+from .dla_rank_conversion import _runtime_input_names
+
+
+def _private_consumer(cache, name, operation, standard, protected, port=0):
+    """Match one standard consumer without changing a public or captured owner."""
+    users = cache.get_consumers(name)
+    if name in protected or len(users) != 1:
+        return None
+    node, index = users[0]
+    return (
+        node
+        if index == port
+        and id(node) in standard
+        and node.op_type == operation
+        and len(node.output) == 1
+        else None
+    )
+
+
+def _apply_topk(model: onnx.ModelProto, *, rank_converted=False) -> onnx.ModelProto:
+    """Plan exact private index rewrites before changing nodes, controls or metadata."""
+    if next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1) < 13:
+        return model
+    graph = model.graph
+    cache = GraphCache(graph)
+    standard_nodes = _operator_nodes(model, {"TopK", "Cast", "Reshape", "Tile", "GatherElements"})
+    standard = {id(node) for node in standard_nodes}
+    topks = [node for node in standard_nodes if node.op_type == "TopK"]
+    for node in topks:
+        if len(node.input) != 2 or len(node.output) != 2:
+            raise ValueError(f"[dla_topk] TopK {node.name!r}: expected two inputs and two outputs")
+    protected = cache.graph_output_names | _subgraph_captures(graph)
+    runtime = _runtime_input_names(model)
+    plans = []
+    for node in topks:
+        values, indices = node.output
+        if values in protected or cache.get_consumers(values):
+            continue
+        cast1 = _private_consumer(cache, indices, "Cast", standard, protected)
+        if cast1 is None or get_node_attr_i(cast1, "to", 0) not in {
+            TensorProto.INT32,
+            TensorProto.INT64,
+        }:
+            continue
+        reshape = _private_consumer(cache, cast1.output[0], "Reshape", standard, protected)
+        tile = (
+            _private_consumer(cache, reshape.output[0], "Tile", standard, protected)
+            if reshape
+            else None
+        )
+        if tile is None or len(reshape.input) != 2 or len(tile.input) != 2:
+            continue
+        cast2 = (
+            _private_consumer(cache, tile.output[0], "Cast", standard, protected)
+            if rank_converted
+            else None
+        )
+        if cast2 is not None and get_node_attr_i(cast2, "to", 0) not in {
+            TensorProto.INT32,
+            TensorProto.INT64,
+        }:
+            continue
+        gather = _private_consumer(
+            cache,
+            cast2.output[0] if cast2 else tile.output[0],
+            "GatherElements",
+            standard,
+            protected,
+            1,
+        )
+        if gather is None or len(gather.input) != 2:
+            continue
+        axis = 2 if rank_converted else 1
+        if get_node_attr_i(gather, "axis", 0) % 4 != axis:
+            continue
+        shape = cache.get_shape(node.input[0])
+        data_shape = cache.get_shape(gather.input[0])
+        if (
+            shape is None
+            or len(shape) != (4 if rank_converted else 2)
+            or any(dim != 1 for dim in shape[:-1])
+        ):
+            continue
+        n = shape[-1]
+        if (
+            not 1 < n <= np.iinfo(np.int32).max
+            or get_node_attr_i(node, "axis", -1) % len(shape) != len(shape) - 1
+        ):
+            continue
+        if data_shape is None or len(data_shape) != 4 or data_shape[axis] != n:
+            continue
+        controls = [node.input[1], reshape.input[1], tile.input[1]]
+        if any(name in runtime or not cache.is_init(name) for name in controls):
+            continue
+        k_array, target, repeats = (cache.get_init_array(name) for name in controls)
+        assert k_array is not None and target is not None and repeats is not None
+        if k_array.shape != (1,) or k_array.dtype != np.int64 or not 0 < int(k_array[0]) <= n:
+            continue
+        k = int(k_array[0])
+        original_indices = [*shape[:-1], k]
+        reshaped = resolve_reshape_shape(target, original_indices)
+        required_shape = [1, 1, 1, 1]
+        required_shape[axis] = k
+        required_repeats = list(data_shape)
+        required_repeats[axis] = 1
+        if reshaped != required_shape or repeats.tolist() != required_repeats:
+            continue
+        plans.append((node, cast1, reshape, tile, cast2, gather, n, k))
+
+    before = {}
+    removed = set()
+    for count, (node, cast1, reshape, tile, cast2, gather, n, k) in enumerate(plans):
+        prefix = node.name or f"topk_{count}"
+        shape_name = cache.unique_name(prefix + "_reshape_const")
+        reshape_out = cache.unique_name(prefix + "_reshape_out")
+        axes_name = cache.unique_name(prefix + "_sq_axes")
+        squeezed = cache.unique_name(prefix + "_sq_out")
+        graph.initializer.extend(
+            [
+                numpy_helper.from_array(np.asarray([1, n, 1, 1], np.int64), shape_name),
+                numpy_helper.from_array(np.asarray([0, 2, 3], np.int64), axes_name),
+            ]
+        )
+        graph.value_info.extend(
+            [
+                helper.make_tensor_value_info(
+                    reshape_out, cache.get_dtype(node.input[0]), [1, n, 1, 1]
+                ),
+                helper.make_tensor_value_info(squeezed, TensorProto.INT32, [k]),
+            ]
+        )
+        before[id(node)] = [
+            helper.make_node(
+                "Reshape",
+                [node.input[0], shape_name],
+                [reshape_out],
+                name=cache.unique_name(prefix + "_reshape"),
+            )
+        ]
+        before[id(gather)] = [
+            helper.make_node(
+                "Squeeze",
+                [cast1.output[0], axes_name],
+                [squeezed],
+                name=cache.unique_name(prefix + "_squeeze_indices"),
+            )
+        ]
+        node.input[0] = reshape_out
+        set_node_attr_i(node, "axis", 1)
+        for name in [*node.output, cast1.output[0]]:
+            cache.update_shape(name, [1, k, 1, 1])
+        set_node_attr_i(cast1, "to", TensorProto.INT32)
+        if cast1.output[0] in cache.vi_map:
+            cache.vi_map[cast1.output[0]].type.tensor_type.elem_type = TensorProto.INT32
+        gather.op_type = "Gather"
+        gather.input[1] = squeezed
+        removed.update(id(item) for item in (reshape, tile, cast2) if item is not None)
+    nodes = []
+    replacements = {}
+    for index, node in enumerate(graph.node):
+        if id(node) in removed:
+            replacements[index] = ([node], [])
+        else:
+            nodes.extend(before.get(id(node), []))
+            nodes.append(node)
+    batch_replace_nodes(graph, replacements)
+    _insert_support_nodes(graph, nodes)
+    _add_legacy_initializer_inputs(model)
+    logger.debug("topk: rewrote %d TopK chain(s).", len(plans))
+    return model
diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py
new file mode 100644
index 00000000000..6b45944d19a
--- /dev/null
+++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py
@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Prepare proven FLOAT consumer bindings while preserving integer Unsqueeze owners."""
+
+import onnx
+from onnx import TensorProto, helper
+
+from ..utils.graph_utils import GraphCache, add_value_info
+from .dla_indexing import _insert_support_nodes, _operator_nodes
+
+
+def _float_input(cache, node, port, source, version):
+    """Require a FLOAT schema binding instead of guessing which ports accept integers."""
+    if node.op_type == "Cast":
+        return False
+    try:
+        schema = onnx.defs.get_schema(node.op_type, version, "")
+    except onnx.defs.SchemaError:
+        return False
+    if not schema.inputs:
+        return False
+    formal = schema.inputs[min(port, len(schema.inputs) - 1)].type_str
+    if formal == "tensor(float)":
+        return True
+    constraints = {item.type_param_str: item.allowed_type_strs for item in schema.type_constraints}
+    if "tensor(float)" not in constraints.get(formal, []):
+        return False
+    bindings = [
+        cache.get_dtype(name)
+        for index, name in enumerate(node.input)
+        if name
+        and name != source
+        and schema.inputs[min(index, len(schema.inputs) - 1)].type_str == formal
+    ]
+    bindings.extend(
+        dtype
+        for index, name in enumerate(node.output)
+        if name
+        and schema.outputs[min(index, len(schema.outputs) - 1)].type_str == formal
+        and (dtype := cache.get_dtype(name)) is not None
+    )
+    return bool(bindings) and all(dtype == TensorProto.FLOAT for dtype in bindings)
+
+
+def _apply_unsqueeze(model):
+    """Add private FLOAT views only for consumers whose schema already requires FLOAT."""
+    graph = model.graph
+    cache = GraphCache(graph)
+    standard = {id(node) for node in _operator_nodes(model, {node.op_type for node in graph.node})}
+    version = next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1)
+    after = {}
+    for node in _operator_nodes(model, {"Unsqueeze"}):
+        if cache.get_dtype(node.input[0]) not in (TensorProto.INT32, TensorProto.INT64):
+            continue
+        consumers = [
+            (consumer, port)
+            for consumer, port in cache.get_consumers(node.output[0])
+            if id(consumer) in standard
+            and _float_input(cache, consumer, port, node.output[0], version)
+        ]
+        if not consumers:
+            continue
+        output = cache.unique_name(node.output[0] + "_float32")
+        cast = helper.make_node(
+            "Cast",
+            [node.output[0]],
+            [output],
+            to=TensorProto.FLOAT,
+            name=cache.unique_name(node.output[0] + "_cast_to_float32"),
+        )
+        add_value_info(graph, output, TensorProto.FLOAT, cache.get_shape(node.output[0]))
+        for consumer, port in consumers:
+            consumer.input[port] = output
+        after[id(node)] = cast
+    nodes = []
+    for node in graph.node:
+        nodes.append(node)
+        if id(node) in after:
+            nodes.append(after[id(node)])
+    _insert_support_nodes(graph, nodes)
+    return model
diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py
index 10c879867df..1b2b6388eb5 100644
--- a/modelopt/onnx/graph_surgery/make_dla_compatible.py
+++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py
@@ -71,6 +71,8 @@
     _apply_reshape_reduce_sum,
     _apply_reshape_transpose_reshape,
 )
+from .dla_transforms.dla_topk import _apply_topk
+from .dla_transforms.dla_unsqueeze import _apply_unsqueeze
 from .dla_transforms.dla_where import _apply_where
 from .utils.graph_utils import (
     DEFAULT_QDQ_REMOVE_QUANT_TYPES,
@@ -137,8 +139,10 @@ def _apply_constants_to_initializers(model, **options):
     ("UpdateAxis", _apply_update_axis),
     ("Where", _apply_where),
     ("Not", _apply_not),
+    ("Unsqueeze", _apply_unsqueeze),
     ("Greater", _apply_greater),
     ("ShapeInference", infer_shapes),
+    ("TopK", partial(_apply_topk, rank_converted=True)),
     ("RemoveUnusedInitializers", _remove_unused_initializers),
 )
 
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py
new file mode 100644
index 00000000000..685666f0662
--- /dev/null
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py
@@ -0,0 +1,453 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_fix_instancenorm_channel_mismatch`.
+
+"""ORT CPU parity tests for ``dla_fix_instancenorm_channel_mismatch``.
+
+Transform rule
+--------------
+For 3D input ``[N, c,D]``:
+
+  ``[N,C,D] → Reshape[N,C,D,1] → InstanceNorm → Reshape[1,N,C,D] → Squeeze(axis=0) → [N,C,D]``
+
+Numerical correctness holds because InstanceNorm normalises each (N,C) slice over
+its spatial dims; ``[N,C,D,1]`` has the same spatial product as ``[N,C,D]``.
+
+For 4D input the transform is a no-op (node unchanged, no extra nodes inserted).
+For >4D input a ``ValueError`` is raised.
+
+Coverage
+--------
+3D input [2,4,6]     → output shape (2,4,6)   with ORT parity
+4D input [2,4,8,6]   → no-op (0 nodes inserted, output unchanged)
+5D input [2,4,8,6,3] → raises ValueError
+"""
+
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto, helper, numpy_helper
+from onnx.reference import ReferenceEvaluator
+
+from modelopt.onnx.graph_surgery.dla_transforms.dla_fix_instancenorm_channel_mismatch import (
+    _apply_fix_instancenorm_channel_mismatch as _apply_fix_instancenorm,
+)
+
+ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks")
+
+_REPO_ROOT = Path(__file__).resolve().parents[5]
+_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "instancenorm_channel_mismatch_test_models"
+_OPSET_VERSION = 17
+_MAX_IR_VERSION_FOR_ORT = 9
+
+
+# ── Shared utilities ──────────────────────────────────────────────────────────
+
+
+def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto:
+    """Declare a typed native tensor."""
+    return helper.make_tensor_value_info(name, elem_type, shape)
+
+
+def _clamp_ir(model: onnx.ModelProto) -> None:
+    """Use a model IR supported by the CPU runtime."""
+    model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT)
+
+
+def _make_instancenorm_model(input_shape: list, epsilon: float = 1e-5) -> onnx.ModelProto:
+    """Build a minimal model with a single InstanceNormalization node."""
+    rank = len(input_shape)
+    assert rank >= 3, "InstanceNorm requires at least 3D input"
+    c = input_shape[1]
+
+    scale_data = np.ones(c, dtype=np.float32)
+    bias_data = np.zeros(c, dtype=np.float32)
+    scale_init = numpy_helper.from_array(scale_data, name="scale")
+    bias_init = numpy_helper.from_array(bias_data, name="bias")
+
+    node = helper.make_node(
+        "InstanceNormalization",
+        inputs=["X", "scale", "bias"],
+        outputs=["Y"],
+        name="instnorm",
+        epsilon=epsilon,
+    )
+    graph = helper.make_graph(
+        [node],
+        "instnorm_graph",
+        [_vi("X", TensorProto.FLOAT, input_shape)],
+        [_vi("Y", TensorProto.FLOAT, input_shape)],
+        initializer=[scale_init, bias_init],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)])
+    _clamp_ir(model)
+    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
+    onnx.checker.check_model(model)
+    return model
+
+
+def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict:
+    """Execute all public outputs on the CPU runtime."""
+    _clamp_ir(model)
+    sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
+    return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output}
+
+
+def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
+    """Compare finite tensor directions in double precision."""
+    a64, b64 = a.astype(np.float64).ravel(), b.astype(np.float64).ravel()
+    na, nb = float(np.linalg.norm(a64)), float(np.linalg.norm(b64))
+    if na == 0.0 and nb == 0.0:
+        return 1.0
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    return float(np.dot(a64, b64) / (na * nb))
+
+
+def _assert_values_match(ref: dict, got: dict) -> None:
+    """Compare restored values after checking element counts."""
+    assert ref.keys() == got.keys()
+    for name in ref:
+        a, b = ref[name].ravel(), got[name].ravel()
+        assert a.shape == b.shape, f"{name}: element count {a.shape} vs {b.shape}"
+        if a.dtype.kind == "f":
+            cos = _cosine_similarity(a, b)
+            assert cos >= 1.0 - 1e-5, f"{name}: cosine {cos:.8f}"
+            np.testing.assert_allclose(a, b, rtol=1e-4, atol=1e-4, err_msg=name)
+
+
+_RNG = np.random.default_rng(2024)
+
+
+# ── 3D input [2,4,6] → output shape (2,4,6) ──────────────────────────────────
+
+
+def test_instancenorm_3d_parity():
+    """3D input [2,4,6]: transform produces same numerics and restores original shape."""
+    shape = [2, 4, 6]
+    model = _make_instancenorm_model(shape)
+    feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)}
+
+    ref = _run_ort(model, feeds)
+
+    tr = copy.deepcopy(model)
+    tr = _apply_fix_instancenorm(tr)
+    _clamp_ir(tr)
+    got = _run_ort(tr, feeds)
+
+    # Output shape must be restored to original 3D
+    assert got["Y"].shape == tuple(shape), f"Expected {tuple(shape)}, got {got['Y'].shape}"
+
+    # Values must match
+    _assert_values_match(ref, got)
+
+    # The original InstanceNorm node must be gone
+    assert not any(
+        n.op_type == "InstanceNormalization" and n.name == "instnorm" for n in tr.graph.node
+    ), "Original InstanceNorm must be replaced"
+
+    # Exactly one new InstanceNorm, two Reshape nodes, and one Squeeze
+    in_nodes = [n for n in tr.graph.node if n.op_type == "InstanceNormalization"]
+    reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"]
+    squeeze_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"]
+    assert len(in_nodes) == 1, f"Expected 1 InstanceNorm, got {len(in_nodes)}"
+    assert len(reshape_nodes) == 2, f"Expected 2 Reshape nodes, got {len(reshape_nodes)}"
+    assert len(squeeze_nodes) == 1, f"Expected 1 Squeeze node, got {len(squeeze_nodes)}"
+
+
+# ── 3D input with non-trivial scale/bias ─────────────────────────────────────
+
+
+def test_instancenorm_3d_nontrivial_scale_bias():
+    """3D input [3,8,16] with random scale/bias: parity and shape preserved."""
+    shape = [3, 8, 16]
+    c = shape[1]
+    scale_data = _RNG.standard_normal(c).astype(np.float32) * 2.0 + 1.0
+    bias_data = _RNG.standard_normal(c).astype(np.float32) * 0.5
+
+    scale_init = numpy_helper.from_array(scale_data, name="scale")
+    bias_init = numpy_helper.from_array(bias_data, name="bias")
+    node = helper.make_node(
+        "InstanceNormalization",
+        inputs=["X", "scale", "bias"],
+        outputs=["Y"],
+        name="instnorm",
+        epsilon=1e-3,
+    )
+    graph = helper.make_graph(
+        [node],
+        "instnorm_nontrivial",
+        [_vi("X", TensorProto.FLOAT, shape)],
+        [_vi("Y", TensorProto.FLOAT, shape)],
+        initializer=[scale_init, bias_init],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)])
+    _clamp_ir(model)
+    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
+    onnx.checker.check_model(model)
+
+    feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)}
+    ref = _run_ort(model, feeds)
+
+    tr = copy.deepcopy(model)
+    tr = _apply_fix_instancenorm(tr)
+    _clamp_ir(tr)
+    got = _run_ort(tr, feeds)
+
+    assert got["Y"].shape == tuple(shape)
+    _assert_values_match(ref, got)
+
+
+# ── 4D input: no-op ───────────────────────────────────────────────────────────
+
+
+def test_instancenorm_4d_noop():
+    """4D input [2,4,8,6]: transform must not modify the graph; ORT outputs unchanged."""
+    shape = [2, 4, 8, 6]
+    model = _make_instancenorm_model(shape)
+    original_node_count = len(model.graph.node)
+    feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)}
+
+    ref = _run_ort(model, feeds)
+
+    tr = copy.deepcopy(model)
+    tr = _apply_fix_instancenorm(tr)
+    _clamp_ir(tr)
+    got = _run_ort(tr, feeds)
+
+    assert len(tr.graph.node) == original_node_count, (
+        f"4D model should be unchanged: {original_node_count} → {len(tr.graph.node)} nodes"
+    )
+    assert any(n.op_type == "InstanceNormalization" for n in tr.graph.node)
+    assert got["Y"].shape == tuple(shape)
+    _assert_values_match(ref, got)
+
+
+# ── >4D input: raises ValueError ─────────────────────────────────────────────
+
+
+def test_instancenorm_5d_raises():
+    """5D input must raise ValueError."""
+    shape = [2, 4, 8, 6, 3]
+    c = shape[1]
+    scale_init = numpy_helper.from_array(np.ones(c, dtype=np.float32), name="scale")
+    bias_init = numpy_helper.from_array(np.zeros(c, dtype=np.float32), name="bias")
+    node = helper.make_node(
+        "InstanceNormalization",
+        inputs=["X", "scale", "bias"],
+        outputs=["Y"],
+        name="instnorm",
+    )
+    graph = helper.make_graph(
+        [node],
+        "instnorm_5d",
+        [_vi("X", TensorProto.FLOAT, shape)],
+        [_vi("Y", TensorProto.FLOAT, shape)],
+        initializer=[scale_init, bias_init],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)])
+    _clamp_ir(model)
+    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
+
+    with pytest.raises(ValueError, match=r"rank.*5.*not supported|not supported"):
+        _apply_fix_instancenorm(copy.deepcopy(model))
+
+
+# ── Multiple InstanceNorm nodes in one graph ──────────────────────────────────
+
+
+def test_instancenorm_multiple_nodes():
+    """Graph with two InstanceNorm nodes (one 3D, one 4D): only 3D one is transformed."""
+    c = 4
+
+    scale3 = numpy_helper.from_array(np.ones(c, dtype=np.float32), name="scale3")
+    bias3 = numpy_helper.from_array(np.zeros(c, dtype=np.float32), name="bias3")
+    scale4 = numpy_helper.from_array(np.ones(c, dtype=np.float32), name="scale4")
+    bias4 = numpy_helper.from_array(np.zeros(c, dtype=np.float32), name="bias4")
+
+    node3d = helper.make_node(
+        "InstanceNormalization",
+        inputs=["X3", "scale3", "bias3"],
+        outputs=["Y3"],
+        name="instnorm_3d",
+    )
+    node4d = helper.make_node(
+        "InstanceNormalization",
+        inputs=["X4", "scale4", "bias4"],
+        outputs=["Y4"],
+        name="instnorm_4d",
+    )
+    graph = helper.make_graph(
+        [node3d, node4d],
+        "multi_instnorm",
+        [
+            _vi("X3", TensorProto.FLOAT, [2, c, 6]),
+            _vi("X4", TensorProto.FLOAT, [2, c, 8, 6]),
+        ],
+        [
+            _vi("Y3", TensorProto.FLOAT, [2, c, 6]),
+            _vi("Y4", TensorProto.FLOAT, [2, c, 8, 6]),
+        ],
+        initializer=[scale3, bias3, scale4, bias4],
+    )
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)])
+    _clamp_ir(model)
+    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
+    onnx.checker.check_model(model)
+
+    feeds = {
+        "X3": _RNG.standard_normal([2, c, 6]).astype(np.float32),
+        "X4": _RNG.standard_normal([2, c, 8, 6]).astype(np.float32),
+    }
+    ref = _run_ort(model, feeds)
+
+    tr = copy.deepcopy(model)
+    tr = _apply_fix_instancenorm(tr)
+    _clamp_ir(tr)
+    got = _run_ort(tr, feeds)
+
+    # Y3 shape must remain (2,4,6)
+    assert got["Y3"].shape == (2, c, 6)
+    # Y4 shape must remain (2,4,8,6)
+    assert got["Y4"].shape == (2, c, 8, 6)
+    _assert_values_match(ref, got)
+
+    # Exactly 2 InstanceNorm nodes remain (1 original 4D + 1 new for 3D)
+    in_nodes = [n for n in tr.graph.node if n.op_type == "InstanceNormalization"]
+    assert len(in_nodes) == 2, f"Expected 2 InstanceNorm nodes, got {len(in_nodes)}"
+
+
+@pytest.mark.parametrize("dtype", [TensorProto.FLOAT, TensorProto.FLOAT16])
+@pytest.mark.parametrize("legacy", [False, True])
+def test_instancenorm_preserves_dtype_epsilon_and_allocated_names(dtype, legacy):
+    """Keep typed metadata, attributes, legacy declarations and existing name owners."""
+    model = _make_instancenorm_model([2, 4, 6], epsilon=0.125)
+    numpy_dtype = np.float16 if dtype == TensorProto.FLOAT16 else np.float32
+    for tensor in model.graph.initializer:
+        tensor.CopyFrom(
+            numpy_helper.from_array(numpy_helper.to_array(tensor).astype(numpy_dtype), tensor.name)
+        )
+    for value in [*model.graph.input, *model.graph.output]:
+        value.type.tensor_type.elem_type = dtype
+    model.graph.initializer.append(
+        numpy_helper.from_array(np.asarray([99], np.int64), "instnorm_pre_shape_0")
+    )
+    model.graph.node.insert(
+        0,
+        helper.make_node(
+            "Constant",
+            [],
+            ["X_trail4d_0"],
+            name="instnorm_pre_reshape_0",
+            value=numpy_helper.from_array(np.asarray([42], np.float32)),
+        ),
+    )
+    constant = model.graph.node[0]
+    constant_bytes = constant.SerializeToString()
+    if legacy:
+        model.ir_version = 3
+        model.graph.input.extend(
+            _vi(t.name, t.data_type, list(t.dims)) for t in model.graph.initializer
+        )
+    feed = {"X": np.arange(48, dtype=numpy_dtype).reshape(2, 4, 6)}
+    onnx.checker.check_model(model, full_check=True)
+    expected = ReferenceEvaluator(model).run(None, feed)[0]
+    changed = _apply_fix_instancenorm(model)
+    assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant
+    assert constant.SerializeToString() == constant_bytes
+    onnx.checker.check_model(changed, full_check=True)
+    actual = ReferenceEvaluator(changed).run(None, feed)[0]
+    assert actual.dtype == numpy_dtype and actual.shape == expected.shape
+    np.testing.assert_allclose(actual, expected, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize("symbolic", [False, True])
+def test_instancenorm_late_unsupported_rank_leaves_caller_unchanged(symbolic):
+    """Reject a later five-dimensional operation before adapting an earlier valid one."""
+    model = _make_instancenorm_model([2, 4, 6])
+    model.graph.input.append(_vi("x5", TensorProto.FLOAT, [2, 4, 2, 3, 4]))
+    model.graph.output.append(_vi("y5", TensorProto.FLOAT, [2, 4, 2, 3, 4]))
+    model.graph.node.append(
+        helper.make_node("InstanceNormalization", ["x5", "scale", "bias"], ["y5"])
+    )
+    if symbolic:
+        model.graph.input[-1].type.tensor_type.shape.dim[0].dim_param = "batch"
+        model.graph.output[-1].type.tensor_type.shape.dim[0].dim_param = "batch"
+    onnx.checker.check_model(model, full_check=True)
+    feeds = {
+        "X": np.arange(48, dtype=np.float32).reshape(2, 4, 6),
+        "x5": np.arange(192, dtype=np.float32).reshape(2, 4, 2, 3, 4),
+    }
+    _run_ort(copy.deepcopy(model), feeds)
+    before = model.SerializeToString()
+    with pytest.raises(ValueError, match="rank"):
+        _apply_fix_instancenorm(model)
+    assert model.SerializeToString() == before
+
+
+@pytest.mark.parametrize("version", [11, 12])
+def test_instancenorm_legacy_opset_refuses_before_mutation(version):
+    """Refuse unsupported axes-input adapters before allocating controls."""
+    model = _make_instancenorm_model([2, 4, 6])
+    model.opset_import[0].version = version
+    before = model.SerializeToString()
+    with pytest.raises(ValueError, match="opset"):
+        _apply_fix_instancenorm(model)
+    assert model.SerializeToString() == before
+
+
+def test_opaque_instancenorm_function_is_unchanged():
+    """Do not replace model-local operators by interpreting their standard-op names."""
+    model = _make_instancenorm_model([2, 4, 6])
+    model.graph.node[0].domain = "opaque"
+    model.functions.append(
+        helper.make_function(
+            "opaque",
+            "InstanceNormalization",
+            ["x", "s", "b"],
+            ["y"],
+            [helper.make_node("Neg", ["x"], ["y"])],
+            [helper.make_opsetid("", 17)],
+            attributes=["epsilon"],
+        )
+    )
+    model.opset_import.append(helper.make_opsetid("opaque", 1))
+    before = model.SerializeToString()
+    feeds = {"X": np.arange(48, dtype=np.float32).reshape(2, 4, 6)}
+    expected = ReferenceEvaluator(model).run(None, feeds)[0]
+    changed = _apply_fix_instancenorm(model)
+    assert changed.SerializeToString() == before
+    np.testing.assert_array_equal(ReferenceEvaluator(changed).run(None, feeds)[0], expected)
+
+
+def test_instancenorm_zero_batch_stays_native():
+    """Preserve zero dimensions instead of applying copied-zero Reshape semantics."""
+    model = _make_instancenorm_model([0, 4, 6])
+    before = model.SerializeToString()
+    changed = _apply_fix_instancenorm(model)
+    assert changed.SerializeToString() == before
+    onnx.checker.check_model(changed, full_check=True)
+    assert ReferenceEvaluator(changed).run(None, {"X": np.empty((0, 4, 6), np.float32)})[
+        0
+    ].shape == (0, 4, 6)
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py
new file mode 100644
index 00000000000..08907464789
--- /dev/null
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py
@@ -0,0 +1,336 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_topk`.
+
+"""Execute private TopK index rewrites and preserve incompatible owners and layouts."""
+
+import copy
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto, helper, numpy_helper
+from onnx.reference import ReferenceEvaluator
+
+from modelopt.onnx.graph_surgery.dla_transforms.dla_topk import _apply_topk
+
+ort = pytest.importorskip("onnxruntime")
+
+
+def _vi(name, dtype, shape):
+    """Declare a typed tensor shape."""
+    return helper.make_tensor_value_info(name, dtype, shape)
+
+
+def _make_topk_chain_model(n=16, k=3, repeats=4, *, rank_converted=False):
+    """Build an executable index chain with a fixed rank-four gather layout."""
+    x_shape = [1, 1, 1, n] if rank_converted else [1, n]
+    rs_shape = [1, 1, k, 1] if rank_converted else [1, k, 1, 1]
+    tile_repeats = [1, 1, 1, repeats] if rank_converted else [1, 1, repeats, 1]
+    data_shape = [1, 1, n, repeats] if rank_converted else [1, n, repeats, 1]
+    out_shape = [1, 1, k, repeats] if rank_converted else [1, k, repeats, 1]
+    nodes = [
+        helper.make_node(
+            "TopK", ["X", "k_val"], ["values", "indices"], name="topk", axis=len(x_shape) - 1
+        ),
+        helper.make_node("Cast", ["indices"], ["cast1"], name="cast1", to=TensorProto.INT64),
+        helper.make_node("Reshape", ["cast1", "target"], ["reshaped"], name="reshape"),
+        helper.make_node("Tile", ["reshaped", "repeats"], ["tiled"], name="tile"),
+    ]
+    if rank_converted:
+        nodes.append(
+            helper.make_node("Cast", ["tiled"], ["cast2"], name="cast2", to=TensorProto.INT64)
+        )
+    nodes.append(
+        helper.make_node(
+            "GatherElements",
+            ["data", "cast2" if rank_converted else "tiled"],
+            ["Y"],
+            name="gather",
+            axis=2 if rank_converted else 1,
+        )
+    )
+    model = helper.make_model(
+        helper.make_graph(
+            nodes,
+            "private_indices",
+            [_vi("X", TensorProto.FLOAT, x_shape), _vi("data", TensorProto.FLOAT, data_shape)],
+            [_vi("Y", TensorProto.FLOAT, out_shape)],
+            [
+                numpy_helper.from_array(np.asarray([k], np.int64), "k_val"),
+                numpy_helper.from_array(np.asarray(rs_shape, np.int64), "target"),
+                numpy_helper.from_array(np.asarray(tile_repeats, np.int64), "repeats"),
+            ],
+        ),
+        opset_imports=[helper.make_opsetid("", 17)],
+    )
+    model.ir_version = 9
+    return onnx.shape_inference.infer_shapes(model, strict_mode=True)
+
+
+def _feeds(model):
+    """Use distinct scores and gather values to expose changed index selection."""
+    result = {}
+    for value in model.graph.input:
+        shape = [dim.dim_value for dim in value.type.tensor_type.shape.dim]
+        result[value.name] = np.arange(np.prod(shape), dtype=np.float32).reshape(shape)
+    return result
+
+
+def _execute(model, feeds, *, reference=False):
+    """Fully check and run the real graph on CPU or the opaque-function evaluator."""
+    onnx.checker.check_model(model, full_check=True)
+    if reference:
+        return ReferenceEvaluator(model).run(None, feeds)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+    return ort.InferenceSession(
+        model.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, feeds)
+
+
+@pytest.mark.parametrize("rank_converted", [False, True])
+@pytest.mark.parametrize("legacy", [False, True])
+def test_topk_chain_preserves_checked_outputs(rank_converted, legacy):
+    """Execute both supported layouts and declare controls for legacy model IR."""
+    model = _make_topk_chain_model(rank_converted=rank_converted)
+    if legacy:
+        model.ir_version = 3
+        model.graph.input.extend(
+            _vi(t.name, t.data_type, list(t.dims)) for t in model.graph.initializer
+        )
+    feeds = {name: value for name, value in _feeds(model).items() if name in {"X", "data"}}
+    expected = _execute(model, feeds)
+    transformed = _apply_topk(copy.deepcopy(model), rank_converted=rank_converted)
+    assert not any(node.op_type == "Tile" for node in transformed.graph.node)
+    assert any(node.op_type == "Gather" for node in transformed.graph.node)
+    for actual, original in zip(_execute(transformed, feeds), expected):
+        np.testing.assert_array_equal(actual, original)
+
+
+@pytest.mark.parametrize("owner", ["values", "indices", "cast1", "reshaped", "tiled", "cast2"])
+def test_public_topk_owners_remain_native(owner):
+    """Keep public values and middleware owners without changing their ranks or types."""
+    model = _make_topk_chain_model(rank_converted=True)
+    model.graph.output.append(
+        copy.deepcopy(next(v for v in model.graph.value_info if v.name == owner))
+    )
+    before = model.SerializeToString()
+    feeds = _feeds(model)
+    expected = _execute(model, feeds)
+    transformed = _apply_topk(copy.deepcopy(model), rank_converted=True)
+    assert transformed.SerializeToString() == before
+    for actual, original in zip(_execute(transformed, feeds), expected):
+        np.testing.assert_array_equal(actual, original)
+
+
+@pytest.mark.parametrize("owner", ["cast1", "reshaped"])
+def test_captured_topk_owners_remain_native(owner):
+    """Keep index producers that nested graphs capture outside the direct consumer map."""
+    model = _make_topk_chain_model(rank_converted=True)
+    info = copy.deepcopy(next(v for v in model.graph.value_info if v.name == owner))
+    info.name = "captured"
+    branch = helper.make_graph(
+        [helper.make_node("Identity", [owner], ["captured"])], "capture", [], [info]
+    )
+    model.graph.input.append(_vi("choose", TensorProto.BOOL, []))
+    model.graph.node.append(
+        helper.make_node("If", ["choose"], ["capture_out"], then_branch=branch, else_branch=branch)
+    )
+    info = copy.deepcopy(info)
+    info.name = "capture_out"
+    model.graph.output.append(info)
+    feeds = _feeds(model)
+    feeds["choose"] = np.asarray(True)
+    expected = _execute(model, feeds)
+    before = model.SerializeToString()
+    transformed = _apply_topk(copy.deepcopy(model), rank_converted=True)
+    assert transformed.SerializeToString() == before
+    for actual, original in zip(_execute(transformed, feeds), expected):
+        np.testing.assert_array_equal(actual, original)
+
+
+@pytest.mark.parametrize("rank_converted", [False, True])
+def test_gather_elements_partial_non_axis_extent_is_not_expanded(rank_converted):
+    """Decline layouts where Gather would select more non-axis rows than GatherElements."""
+    model = _make_topk_chain_model(rank_converted=rank_converted)
+    model.graph.input[1].type.tensor_type.shape.dim[1 if rank_converted else 0].dim_value = 2
+    feeds = _feeds(model)
+    expected = _execute(model, feeds)
+    before = model.SerializeToString()
+    transformed = _apply_topk(copy.deepcopy(model), rank_converted=rank_converted)
+    assert transformed.SerializeToString() == before
+    np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0])
+
+
+def test_tile_orientation_must_preserve_each_selected_index():
+    """Decline same-shaped tiling that interleaves indices along the feature axis."""
+    model = _make_topk_chain_model(k=2, rank_converted=True)
+    for tensor in model.graph.initializer:
+        if tensor.name == "target":
+            tensor.CopyFrom(numpy_helper.from_array(np.asarray([1, 1, 1, 2], np.int64), "target"))
+        if tensor.name == "repeats":
+            tensor.CopyFrom(numpy_helper.from_array(np.asarray([1, 1, 2, 2], np.int64), "repeats"))
+    next(v for v in model.graph.value_info if v.name == "reshaped").CopyFrom(
+        _vi("reshaped", TensorProto.INT64, [1, 1, 1, 2])
+    )
+    feeds = _feeds(model)
+    expected = _execute(model, feeds)
+    before = model.SerializeToString()
+    transformed = _apply_topk(copy.deepcopy(model), rank_converted=True)
+    assert transformed.SerializeToString() == before
+    np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0])
+
+
+def test_topk_float_cast_rounding_is_not_removed():
+    """Preserve an intentional FLOAT16 rounding step before the final integer cast."""
+    model = _make_topk_chain_model(n=4096, k=1, rank_converted=True)
+    cast = next(node for node in model.graph.node if node.name == "cast1")
+    next(attr for attr in cast.attribute if attr.name == "to").i = TensorProto.FLOAT16
+    for value in model.graph.value_info:
+        if value.name in {"cast1", "reshaped", "tiled"}:
+            value.type.tensor_type.elem_type = TensorProto.FLOAT16
+    feeds = _feeds(model)
+    feeds["X"].fill(0)
+    feeds["X"].reshape(-1)[2049] = 100
+    expected = _execute(model, feeds)
+    before = model.SerializeToString()
+    transformed = _apply_topk(copy.deepcopy(model), rank_converted=True)
+    assert transformed.SerializeToString() == before
+    np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0])
+
+
+@pytest.mark.parametrize("site", ["TopK", "Cast"])
+def test_opaque_index_functions_remain_native(site):
+    """Leave model-local operators intact even when their names match the pattern."""
+    model = _make_topk_chain_model()
+    node = next(node for node in model.graph.node if node.op_type == site)
+    node.domain = "custom"
+    body = (
+        [helper.make_node("TopK", ["a", "b"], ["v", "i"], axis=1)]
+        if site == "TopK"
+        else [helper.make_node("Identity", ["a"], ["v"])]
+    )
+    function = helper.make_function(
+        "custom",
+        site,
+        ["a", "b"] if site == "TopK" else ["a"],
+        ["v", "i"] if site == "TopK" else ["v"],
+        body,
+        [helper.make_opsetid("", 17)],
+        attributes=["axis"] if site == "TopK" else ["to"],
+    )
+    model.functions.append(function)
+    model.opset_import.append(helper.make_opsetid("custom", 1))
+    feeds = _feeds(model)
+    expected = _execute(model, feeds, reference=True)
+    before = model.SerializeToString()
+    transformed = _apply_topk(copy.deepcopy(model))
+    assert transformed.SerializeToString() == before
+    np.testing.assert_array_equal(_execute(transformed, feeds, reference=True)[0], expected[0])
+
+
+@pytest.mark.parametrize("version", [11, 12])
+def test_legacy_adapter_schemas_remain_native(version):
+    """Avoid inserting input-based Squeeze controls into older operator sets."""
+    model = _make_topk_chain_model()
+    model.opset_import[0].version = version
+    feeds = _feeds(model)
+    expected = _execute(model, feeds)
+    before = model.SerializeToString()
+    transformed = _apply_topk(copy.deepcopy(model))
+    assert transformed.SerializeToString() == before
+    np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0])
+
+
+def test_topk_without_index_chain_is_unchanged():
+    """Keep an ordinary TopK value path that has no index rewrite pattern."""
+    model = helper.make_model(
+        helper.make_graph(
+            [
+                helper.make_node("TopK", ["X", "K"], ["values", "indices"], axis=3),
+                helper.make_node("Relu", ["values"], ["Y"]),
+            ],
+            "unmatched_topk",
+            [_vi("X", TensorProto.FLOAT, [1, 1, 1, 16])],
+            [_vi("Y", TensorProto.FLOAT, [1, 1, 1, 3])],
+            [numpy_helper.from_array(np.asarray([3], np.int64), "K")],
+        ),
+        opset_imports=[helper.make_opsetid("", 17)],
+    )
+    model.ir_version = 9
+    feeds = _feeds(model)
+    expected = _execute(model, feeds)
+    original = model.SerializeToString()
+    _apply_topk(model)
+    assert model.SerializeToString() == original
+    np.testing.assert_array_equal(_execute(model, feeds)[0], expected[0])
+
+
+@pytest.mark.parametrize("control", ["k_val", "target", "repeats"])
+def test_topk_runtime_initializer_defaults_remain_native(control):
+    """Do not freeze overridable structural controls while collapsing the index chain."""
+    model = _make_topk_chain_model(rank_converted=True)
+    tensor = next(item for item in model.graph.initializer if item.name == control)
+    model.graph.input.append(_vi(control, tensor.data_type, list(tensor.dims)))
+    feeds = {name: value for name, value in _feeds(model).items() if name in {"X", "data"}}
+    expected = _execute(model, feeds)
+    before = model.SerializeToString()
+    changed = _apply_topk(copy.deepcopy(model), rank_converted=True)
+    assert changed.SerializeToString() == before
+    np.testing.assert_array_equal(_execute(changed, feeds)[0], expected[0])
+
+
+def test_topk_allocates_distinct_names_and_preserves_constant_identity():
+    """Reuse the shared allocator and batch updates without copying unrelated payloads."""
+    model = _make_topk_chain_model(rank_converted=True)
+    constant = helper.make_node(
+        "Constant",
+        [],
+        ["topk_sq_out"],
+        name="topk_reshape",
+        value=numpy_helper.from_array(np.arange(16, dtype=np.float32)),
+    )
+    model.graph.node.insert(0, constant)
+    constant = model.graph.node[0]
+    original = constant.SerializeToString()
+    model.graph.initializer.append(
+        numpy_helper.from_array(np.asarray([42], np.int64), "topk_reshape_const")
+    )
+    feeds = _feeds(model)
+    expected = _execute(model, feeds)
+    changed = _apply_topk(model, rank_converted=True)
+    assert next(node for node in changed.graph.node if node.op_type == "Constant") is constant
+    assert constant.SerializeToString() == original
+    definitions = [name for node in changed.graph.node for name in node.output] + [
+        item.name for item in changed.graph.initializer
+    ]
+    assert len(definitions) == len(set(definitions))
+    np.testing.assert_array_equal(_execute(changed, feeds)[0], expected[0])
+
+
+def test_malformed_late_topk_refuses_before_any_mutation():
+    """Validate every candidate's arity before changing an earlier valid chain."""
+    model = _make_topk_chain_model(rank_converted=True)
+    model.graph.node.append(helper.make_node("TopK", ["X"], ["bad_values", "bad_indices"]))
+    original = model.SerializeToString()
+    with pytest.raises(ValueError, match="expected"):
+        _apply_topk(model, rank_converted=True)
+    assert model.SerializeToString() == original
diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_unsqueeze.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_unsqueeze.py
new file mode 100644
index 00000000000..baf46c9b508
--- /dev/null
+++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_unsqueeze.py
@@ -0,0 +1,244 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Keep native integer Unsqueeze contracts across data and structural consumers."""
+
+import copy
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto as T
+from onnx import helper as h
+from onnx import numpy_helper as n
+from onnx.reference import ReferenceEvaluator
+
+from modelopt.onnx.graph_surgery.dla_transforms.dla_unsqueeze import _apply_unsqueeze
+
+ort = pytest.importorskip("onnxruntime")
+
+
+def _model(operation):
+    """Build a checked native integer graph with values too large for FLOAT precision."""
+    inputs = [h.make_tensor_value_info("x", T.INT64, [2])]
+    initializers = [n.from_array(np.asarray([0], np.int64), "axes")]
+    nodes = [h.make_node("Unsqueeze", ["x", "axes"], ["expanded"], name="expand")]
+    dtype, shape = T.INT64, [1, 2]
+    if operation == "Cast":
+        nodes.append(h.make_node("Cast", ["expanded"], ["y"], to=T.FLOAT))
+        dtype = T.FLOAT
+    elif operation == "Double":
+        nodes.append(h.make_node("Cast", ["expanded"], ["y"], to=T.DOUBLE))
+        dtype = T.DOUBLE
+    elif operation == "Add":
+        initializers.append(n.from_array(np.asarray([[1, 2]], np.int64), "other"))
+        nodes.append(h.make_node("Add", ["expanded", "other"], ["y"]))
+    elif operation == "Shape":
+        nodes.append(h.make_node("Shape", ["expanded"], ["y"]))
+        shape = [2]
+    elif operation == "GatherND":
+        initializers.append(n.from_array(np.arange(8, dtype=np.float32), "data"))
+        nodes.append(h.make_node("GatherND", ["data", "expanded"], ["y"]))
+        inputs[0] = h.make_tensor_value_info("x", T.INT64, [1])
+        dtype, shape = T.FLOAT, [1]
+    else:
+        nodes.append(h.make_node("Identity", ["expanded"], ["y"]))
+    model = h.make_model(
+        h.make_graph(
+            nodes,
+            "integer_unsqueeze",
+            inputs,
+            [h.make_tensor_value_info("y", dtype, shape)],
+            initializers,
+        ),
+        opset_imports=[h.make_opsetid("", 17)],
+    )
+    model.ir_version = 9
+    return onnx.shape_inference.infer_shapes(model, strict_mode=True)
+
+
+@pytest.mark.parametrize("operation", ["Identity", "Add", "Shape", "Double", "Cast", "GatherND"])
+def test_unsqueeze_preserves_native_consumer_type_contracts(operation):
+    """Keep integer owners and controls, and let explicit casts perform their own conversion."""
+    model = _model(operation)
+    data = (
+        np.asarray([2**24 + 1, 2**24 + 3], np.int64)
+        if operation != "GatherND"
+        else np.asarray([3], np.int64)
+    )
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    onnx.checker.check_model(model, full_check=True)
+    expected = ort.InferenceSession(
+        model.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, {"x": data})[0]
+    before = model.SerializeToString()
+    changed = _apply_unsqueeze(copy.deepcopy(model))
+    assert changed.SerializeToString() == before
+    onnx.checker.check_model(changed, full_check=True)
+    actual = ort.InferenceSession(
+        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, {"x": data})[0]
+    np.testing.assert_array_equal(actual, expected)
+
+
+def test_opaque_unsqueeze_producer_stays_native():
+    """Preserve an opaque integer function despite its standard-looking operator name."""
+    model = _model("Identity")
+    model.graph.node[0].domain = "opaque"
+    model.functions.append(
+        h.make_function(
+            "opaque",
+            "Unsqueeze",
+            ["x", "axes"],
+            ["out"],
+            [h.make_node("Unsqueeze", ["x", "axes"], ["out"])],
+            [h.make_opsetid("", 17)],
+        )
+    )
+    model.opset_import.append(h.make_opsetid("opaque", 1))
+    data = np.asarray([2**24 + 1, 2**24 + 3], np.int64)
+    expected = ReferenceEvaluator(model).run(None, {"x": data})[0]
+    before = model.SerializeToString()
+    changed = _apply_unsqueeze(model)
+    assert changed.SerializeToString() == before
+    onnx.checker.check_model(changed, full_check=True)
+    np.testing.assert_array_equal(ReferenceEvaluator(changed).run(None, {"x": data})[0], expected)
+
+
+@pytest.mark.parametrize("operation", ["ConstantOfShape", "ScatterND"])
+def test_unsqueeze_keeps_structural_indices_for_additional_schemas(operation):
+    """Derive integer-only ports from schemas beyond the previous manual list."""
+    input_shape = [] if operation == "ConstantOfShape" else [1]
+    indices = (
+        np.asarray(3, np.int64) if operation == "ConstantOfShape" else np.asarray([2], np.int64)
+    )
+    initializers = [n.from_array(np.asarray([0], np.int64), "axes")]
+    nodes = [h.make_node("Unsqueeze", ["x", "axes"], ["expanded"])]
+    if operation == "ConstantOfShape":
+        nodes.append(h.make_node(operation, ["expanded"], ["y"]))
+        output_shape = [3]
+    else:
+        initializers.extend(
+            [
+                n.from_array(np.arange(4, dtype=np.float32), "data"),
+                n.from_array(np.asarray([99], np.float32), "updates"),
+            ]
+        )
+        nodes.append(h.make_node(operation, ["data", "expanded", "updates"], ["y"]))
+        output_shape = [4]
+    model = h.make_model(
+        h.make_graph(
+            nodes,
+            "integer_controls",
+            [h.make_tensor_value_info("x", T.INT64, input_shape)],
+            [h.make_tensor_value_info("y", T.FLOAT, output_shape)],
+            initializers,
+        ),
+        opset_imports=[h.make_opsetid("", 17)],
+    )
+    model.ir_version = 9
+    model = onnx.shape_inference.infer_shapes(model, strict_mode=True)
+    before = model.SerializeToString()
+    expected = ReferenceEvaluator(model).run(None, {"x": indices})[0]
+    changed = _apply_unsqueeze(model)
+    assert changed.SerializeToString() == before
+    onnx.checker.check_model(changed, full_check=True)
+    np.testing.assert_array_equal(
+        ReferenceEvaluator(changed).run(None, {"x": indices})[0], expected
+    )
+
+
+def test_unsqueeze_does_not_rewire_opaque_consumer():
+    """Keep opaque integer consumers independent of standard schema lookup."""
+    model = _model("Identity")
+    model.graph.node[1].domain = "opaque"
+    model.functions.append(
+        h.make_function(
+            "opaque",
+            "Identity",
+            ["x"],
+            ["out"],
+            [h.make_node("Neg", ["x"], ["out"])],
+            [h.make_opsetid("", 17)],
+        )
+    )
+    model.opset_import.append(h.make_opsetid("opaque", 1))
+    before = model.SerializeToString()
+    data = np.asarray([2**24 + 1, 2**24 + 3], np.int64)
+    expected = ReferenceEvaluator(model).run(None, {"x": data})[0]
+    changed = _apply_unsqueeze(model)
+    assert changed.SerializeToString() == before
+    onnx.checker.check_model(changed, full_check=True)
+    np.testing.assert_array_equal(ReferenceEvaluator(changed).run(None, {"x": data})[0], expected)
+
+
+@pytest.mark.parametrize("operation", ["Add", "Resize"])
+def test_unsqueeze_repairs_prepared_float_bindings_with_private_views(operation):
+    """Repair an intermediate FLOAT binding against an executable explicit-cast reference."""
+    scalar = operation == "Resize"
+    inputs = [h.make_tensor_value_info("x", T.INT64, [] if scalar else [2])]
+    initializers = [n.from_array(np.asarray([0], np.int64), "axes")]
+    nodes = [
+        h.make_node("Unsqueeze", ["x", "axes"], ["expanded"]),
+        h.make_node("Cast", ["expanded"], ["prepared_float"], to=T.FLOAT),
+    ]
+    if scalar:
+        inputs.append(h.make_tensor_value_info("data", T.FLOAT, [2]))
+        nodes.append(h.make_node("Resize", ["data", "", "prepared_float"], ["y"]))
+        out_shape = [4]
+        feeds = {"x": np.asarray(2, np.int64), "data": np.asarray([3, 7], np.float32)}
+    else:
+        initializers.append(n.from_array(np.asarray([[1, 2]], np.float32), "other"))
+        nodes.append(h.make_node("Add", ["prepared_float", "other"], ["y"]))
+        out_shape = [1, 2]
+        feeds = {"x": np.asarray([2**24 + 1, 2**24 + 3], np.int64)}
+    reference = h.make_model(
+        h.make_graph(
+            nodes,
+            "prepared_float_binding",
+            inputs,
+            [
+                h.make_tensor_value_info("y", T.FLOAT, out_shape),
+                h.make_tensor_value_info("expanded", T.INT64, [1] if scalar else [1, 2]),
+            ],
+            initializers,
+        ),
+        opset_imports=[h.make_opsetid("", 17)],
+    )
+    reference.ir_version = 9
+    reference = onnx.shape_inference.infer_shapes(reference, strict_mode=True)
+    onnx.checker.check_model(reference, full_check=True)
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    expected = ort.InferenceSession(
+        reference.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, feeds)
+    prepared = copy.deepcopy(reference)
+    del prepared.graph.node[1]
+    consumer = prepared.graph.node[1]
+    for port, name in enumerate(consumer.input):
+        if name == "prepared_float":
+            consumer.input[port] = "expanded"
+    # Earlier preparation stages can leave integer views on an already FLOAT-bound port.
+    changed = _apply_unsqueeze(prepared)
+    onnx.checker.check_model(changed, full_check=True)
+    casts = [node for node in changed.graph.node if node.op_type == "Cast"]
+    assert len(casts) == 1 and consumer.input[2 if scalar else 0] == casts[0].output[0]
+    actual = ort.InferenceSession(
+        changed.SerializeToString(), options, providers=["CPUExecutionProvider"]
+    ).run(None, feeds)
+    for got, wanted in zip(actual, expected):
+        np.testing.assert_array_equal(got, wanted)