diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_convert_ops_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_convert_ops_to_4d.py new file mode 100644 index 00000000000..9bdede4a13b --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_convert_ops_to_4d.py @@ -0,0 +1,1808 @@ +# 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 non-4D operations to work with 4D tensors in a single graph pass. + +Each eligible op gets: + - Unsqueeze node(s) inserted before non-4D data inputs + - Axis / perm attribute adjustments to match the promoted 4D inputs + - Squeeze node(s) inserted after outputs to restore the original rank + +Constant initializer inputs are expanded in-place (prepend 1s) instead of +adding Unsqueeze nodes. + +A subsequent ``dla-remove-intermediary-squeeze-and-unsqueeze`` pass removes +adjacent Squeeze→Unsqueeze pairs, yielding fully-4D intermediate tensors. + +Ops **skipped** (handled by separate transforms or structurally special): + LSTM, MatMul, Gemm — dedicated transforms exist + Squeeze, Unsqueeze — structural; modifying them would corrupt the graph + Constant, Identity, Shape — no axes / passthrough +""" + +from __future__ import annotations + +import math +from itertools import combinations + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + batch_replace_nodes, + get_node_attr_i, + pad4d, + 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 _operator_nodes +from .dla_rank_conversion import _runtime_input_names + + +class _InitializerNames(set[str]): + """Track inserted controls while reusing the shared graph name allocator.""" + + __slots__ = ("_cache",) + + def __init__(self, cache: GraphCache) -> None: + """Seed existing initializer names and retain the shared allocator.""" + super().__init__(cache.init_names) + self._cache = cache + + def unique_name(self, preferred: str) -> str: + """Reserve a name across values, controls and nodes.""" + return self._cache.unique_name(preferred) + + +class _CacheShapeProxy: + """Dict-like wrapper around ``GraphCache.get_shape`` for handler compatibility.""" + + __slots__ = ("_cache", "_captured", "_index_views", "_runtime_inputs", "_standard_nodes") + + def __init__(self, cache: GraphCache, standard_nodes, captured, runtime_inputs) -> None: + """Retain native shape lookups and precomputed operator ownership guards.""" + self._cache = cache + self._standard_nodes = {id(node) for node in standard_nodes} + self._captured = captured + self._runtime_inputs = runtime_inputs + self._index_views = {} + + def get(self, name, default=None): + """Return a native shape or the caller default.""" + r = self._cache.get_shape(name) + return r if r is not None else default + + def __contains__(self, name): + """Check whether native shape metadata is available.""" + return self._cache.get_shape(name) is not None + + def __getitem__(self, name): + """Return the known native shape or raise KeyError.""" + r = self._cache.get_shape(name) + if r is None: + raise KeyError(name) + return r + + +# --------------------------------------------------------------------------- +# Ops that this transform must not touch +# --------------------------------------------------------------------------- +_OP_SKIP = frozenset( + { + "LSTM", + "MatMul", + "Gemm", + "Squeeze", + "Unsqueeze", + "Constant", + "Identity", + "Shape", + "GroupQueryAttention", + "Resize", + } +) + + +# --------------------------------------------------------------------------- +# Small utilities +# --------------------------------------------------------------------------- + + +def _find_squeeze_axes(shape_4d: list, target_shape: list) -> list[int]: + """Return the axes to remove from ``shape_4d`` to obtain ``target_shape``. + + Handles cases where the extra unary dim is at a non-leading position. + Example: shape_4d=[1,1,3,1], target=[1,1,3] → need axis=3, not 0. + + Falls back to leading axes when no exact match is found. + """ + n_remove = len(shape_4d) - len(target_shape) + if n_remove <= 0: + return [] + target = list(target_shape) + for axes_combo in combinations(range(len(shape_4d)), n_remove): + remaining = [shape_4d[i] for i in range(len(shape_4d)) if i not in set(axes_combo)] + if remaining == target: + return sorted(axes_combo) + return list(range(n_remove)) # fallback: leading axes + + +def _axis_delta(rank: int) -> int: + """Amount to add to a non-negative axis index when promoting rank-R → 4D.""" + return max(0, 4 - rank) + + +_get_attr_i = get_node_attr_i +_set_attr_i = set_node_attr_i + + +def _make_unsqueeze_to_4d( + input_name: str, + rank: int, + prefix: str, + unique_init: _InitializerNames, + graph, + c: int, +): + """Return (out_name, unsq_node). If rank==4 return (input_name, None).""" + if rank == 4: + return input_name, None + axes = list(range(4 - rank)) + axes_init = numpy_helper.from_array( + np.array(axes, dtype=np.int64), name=unique_init.unique_name(f"{prefix}_unsq_axes_{c}") + ) + add_unique_initializers(graph, unique_init, [axes_init]) + out_name = unique_init.unique_name(f"{prefix}_unsq4d_{c}") + node = helper.make_node( + "Unsqueeze", + inputs=[input_name, axes_init.name], + outputs=[out_name], + name=unique_init.unique_name(f"{prefix}_unsq4d_node_{c}"), + ) + return out_name, node + + +def _make_squeeze_to_rank( + input_name: str, + orig_rank: int, + output_name: str, + prefix: str, + unique_init: _InitializerNames, + graph, + c: int, + *, + shape_4d: list | None = None, + orig_out_shape: list | None = None, +): + """Return sq_node. If orig_rank==4 return None. + + When ``shape_4d`` and ``orig_out_shape`` are both supplied the squeeze axes + are determined by :func:`_find_squeeze_axes`, which handles the case where + the extra unary dim is at a non-leading position (e.g. shape_4d=[1,1,3,1] + with orig_out_shape=[1,1,3] → squeeze axis=3, not 0). + """ + if orig_rank == 4: + return None + if shape_4d is not None and orig_out_shape is not None: + axes = _find_squeeze_axes(shape_4d, orig_out_shape) + else: + axes = list(range(4 - orig_rank)) + axes_init = numpy_helper.from_array( + np.array(axes, dtype=np.int64), name=unique_init.unique_name(f"{prefix}_sq_axes_{c}") + ) + add_unique_initializers(graph, unique_init, [axes_init]) + node = helper.make_node( + "Squeeze", + inputs=[input_name, axes_init.name], + outputs=[output_name], + name=unique_init.unique_name(f"{prefix}_sq_node_{c}"), + ) + return node + + +def _expand_init_dims(init, cache) -> None: + """Prepend 1-dims to a static initializer to make it 4D. No-op if already ≥ 4D.""" + if len(init.dims) < 4: + shape = pad4d(list(init.dims)) + del init.dims[:] + init.dims.extend(shape) + info = cache.vi_map.get(init.name) + if info is not None and info.type.tensor_type.HasField("shape"): + info.type.tensor_type.shape.CopyFrom( + helper.make_tensor_value_info( + init.name, init.data_type, shape + ).type.tensor_type.shape + ) + + +# --------------------------------------------------------------------------- +# Per-op handlers +# Signature: +# (node, graph, shape_map, init_names, unique_init, +# inits_to_add, cnt_ref) -> list[NodeProto] | None +# +# Return value: +# None - node needs no change (skip) +# list - full replacement chain: [pre_nodes..., node_or_new_node, post_nodes...] +# The original `node` object (possibly modified in-place) is included +# when it should be kept; Flatten replaces it with a new Reshape node. +# +# cnt_ref is a list[int] used as a mutable counter (avoids nonlocal boilerplate). +# --------------------------------------------------------------------------- + + +def _handle_transpose(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Adapt supported permutations and restore the original output rank.""" + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + if rank == 4 or (rank == 5 and 1 not in orig_shape): + return None + + pfx = node.name or f"tr_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + # ── Adjust perm ────────────────────────────────────────────────────────── + perm_updated = False + for attr in node.attribute: + if attr.name == "perm": + perm = list(attr.ints) + if rank == 2: + new_perm = [0, 1, perm[0] + 2, perm[1] + 2] + elif rank == 3: + new_perm = [0, perm[0] + 1, perm[1] + 1, perm[2] + 1] + elif rank == 5: + # General algorithm: find rightmost unary dim in the OUTPUT shape, + # drop that permutation entry, shift remaining values. + current_shape = [orig_shape[i] for i in perm] + idx_to_remove = len(current_shape) - current_shape[::-1].index(1) - 1 + removed_orig_dim = perm[idx_to_remove] + new_perm_raw = [p for i, p in enumerate(perm) if i != idx_to_remove] + new_perm = [p if p < removed_orig_dim else p - 1 for p in new_perm_raw] + else: + return None # unsupported rank + del attr.ints[:] + attr.ints.extend(int(x) for x in new_perm) + perm_updated = True + break + + if not perm_updated: + return None + + if rank == 5: + # 5D collapse: squeeze the unary input dim → 4D Transpose → unsqueeze output back to 5D. + # Round-tripping through the unary dim preserves the output shape, so ORT parity holds + # and a subsequent dla-remove-intermediary-squeeze-and-unsqueeze pass can clean up. + c_5d = cnt_ref[0] + cnt_ref[0] += 1 + c_5d2 = cnt_ref[0] + cnt_ref[0] += 1 + + # Squeeze input at removed_orig_dim + sq_axis_init = numpy_helper.from_array( + np.array([removed_orig_dim], dtype=np.int64), + name=unique_init.unique_name(f"{pfx}_5d_sq_axes_{c_5d}"), + ) + add_unique_initializers(graph, unique_init, [sq_axis_init]) + sq_out = unique_init.unique_name(f"{pfx}_5d_sq4d_{c_5d}") + sq_node = helper.make_node( + "Squeeze", + inputs=[in_name, sq_axis_init.name], + outputs=[sq_out], + name=unique_init.unique_name(f"{pfx}_5d_sq_{c_5d}"), + ) + node.input[0] = sq_out + + # Unsqueeze Transpose output at idx_to_remove to restore the 5D output shape + unsq_axis_init = numpy_helper.from_array( + np.array([idx_to_remove], dtype=np.int64), + name=unique_init.unique_name(f"{pfx}_5d_unsq_axes_{c_5d2}"), + ) + add_unique_initializers(graph, unique_init, [unsq_axis_init]) + orig_out = node.output[0] + tr_out_4d = unique_init.unique_name(f"{pfx}_5d_tr4d_{c_5d}") + unsq_node = helper.make_node( + "Unsqueeze", + inputs=[tr_out_4d, unsq_axis_init.name], + outputs=[orig_out], + name=unique_init.unique_name(f"{pfx}_5d_unsq_{c_5d2}"), + ) + node.output[0] = tr_out_4d + return [sq_node, node, unsq_node] + + pre: list = [] + post: list = [] + + # ── Unsqueeze input ─────────────────────────────────────────────────────── + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # ── Squeeze output ──────────────────────────────────────────────────────── + c2 = cnt_ref[0] + cnt_ref[0] += 1 + orig_out = node.output[0] + temp_out = unique_init.unique_name(f"{pfx}_tr_4d_{c}") + sq = _make_squeeze_to_rank(temp_out, rank, orig_out, pfx, unique_init, graph, c2) + if sq: + node.output[0] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +def _handle_reshape(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Reshape: modify target shape initializer to 4D and Squeeze output back to original rank. + + The Reshape input is left as-is (Reshape accepts any-rank input). + A Squeeze is added after to restore the original output rank so downstream + shape_map lookups remain consistent with runtime shapes. + Leading-singleton 5D targets compute in 4D and then restore their native rank. + """ + if len(node.input) < 2 or node.input[1] not in init_names: + return None # dynamic shape — skip + + shape_init = shape_map._cache.get_init(node.input[1]) + if shape_init is None: + return None + + shape_arr = numpy_helper.to_array(shape_init) + + # Determine the concrete output shape. + # Priority 1: use the output shape from shape_map (populated by shape inference). + # Priority 2: compute from the Reshape's shape initializer, resolving -1/0 + # using the input shape via resolve_reshape_shape. + out_name = node.output[0] if node.output else None + out_shape_known = shape_map.get(out_name) if out_name else None + + if out_shape_known is not None: + resolved = list(out_shape_known) + else: + in_name = node.input[0] + input_shape = shape_map.get(in_name) + if input_shape is not None: + try: + resolved = resolve_reshape_shape(shape_arr, input_shape) + except ValueError: + resolved = [int(d) for d in shape_arr.tolist()] + else: + resolved = [int(d) for d in shape_arr.tolist()] + + if 0 in resolved: + return None + target_dims = len(resolved) + + if target_dims == 4: + return None # already 4D + + if target_dims == 2: + orig_out_rank = 2 + new_shape = np.array([1, 1, *resolved], dtype=np.int64) + elif target_dims == 3: + orig_out_rank = 3 + new_shape = np.array([1, *resolved], dtype=np.int64) + elif target_dims == 5 and resolved[0] == 1: + c_r = cnt_ref[0] + cnt_ref[0] += 1 + new_name = unique_init.unique_name(f"{shape_init.name}_4d_{c_r}") + inits_to_add.append(numpy_helper.from_array(np.asarray(resolved[1:], np.int64), new_name)) + axes = numpy_helper.from_array( + np.asarray([0], np.int64), unique_init.unique_name(f"{new_name}_axes") + ) + add_unique_initializers(graph, unique_init, [axes]) + original = node.output[0] + temporary = unique_init.unique_name(original + "_reshape_4d") + node.input[1] = new_name + node.output[0] = temporary + restore = helper.make_node( + "Unsqueeze", + [temporary, axes.name], + [original], + name=unique_init.unique_name(original + "_restore_5d"), + ) + return [node, restore] + else: + return None # unsupported + + c_r = cnt_ref[0] + cnt_ref[0] += 1 + new_shape_name = unique_init.unique_name(f"{shape_init.name}_4d_{c_r}") + # Add new init; leave old for any other consumers. + inits_to_add.append(numpy_helper.from_array(new_shape, new_shape_name)) + node.input[1] = new_shape_name + + # Squeeze output back to original rank so downstream shape_map stays consistent + pfx = node.name or f"rs_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + orig_out = node.output[0] + temp_out = unique_init.unique_name(f"{pfx}_rs_4d_{c}") + sq = _make_squeeze_to_rank(temp_out, orig_out_rank, orig_out, pfx, unique_init, graph, c) + if sq: + node.output[0] = temp_out + return [node, sq] + return [node] + + +def _handle_slice(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Adapt fixed slicing axes around leading input promotion.""" + if len(node.input) > 3 and node.input[3] and node.input[3] not in init_names: + return None + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + if rank == 4: + return None + + pfx = node.name or f"sl_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + post: list = [] + if len(node.input) < 4 or not node.input[3]: + starts = shape_map._cache.get_init(node.input[1]) + if starts is None: + return None + axes_name = unique_init.unique_name(f"{pfx}_default_axes_{c}") + inits_to_add.append( + numpy_helper.from_array( + np.arange(np.prod(starts.dims), dtype=np.int64) + 4 - rank, axes_name + ) + ) + while len(node.input) < 4: + node.input.append("") + node.input[3] = axes_name + + # ── Update axes (input[3]) ──────────────────────────────────────────────── + if len(node.input) > 3 and node.input[3] and node.input[3] in init_names: + init = shape_map._cache.get_init(node.input[3]) + if init is not None: + axes_arr = numpy_helper.to_array(init) + delta = _axis_delta(rank) + new_axes = np.array( + [int(a) + delta if int(a) >= 0 else int(a) for a in axes_arr], + dtype=np.int64, + ) + c_a = cnt_ref[0] + cnt_ref[0] += 1 + new_axes_name = unique_init.unique_name(f"{init.name}_4d_{c_a}") + # Add new init with unique name; leave old init for other consumers. + inits_to_add.append(numpy_helper.from_array(new_axes, new_axes_name)) + node.input[3] = new_axes_name + + # ── Unsqueeze data input ────────────────────────────────────────────────── + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # ── Squeeze outputs ─────────────────────────────────────────────────────── + for oi, orig_out in enumerate(node.output): + if not orig_out: + continue + ci = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_sl_4d_{c}_o{oi}") + sq = _make_squeeze_to_rank(temp_out, rank, orig_out, pfx, unique_init, graph, ci) + if sq: + node.output[oi] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +def _handle_argmax(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Prepare reduction axes while restoring native ArgMax outputs.""" + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + + pfx = node.name or f"argmax_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + post: list = [] + + delta = _axis_delta(rank) + orig_axis = _get_attr_i(node, "axis", 0) + adjusted_axis = orig_axis + delta if orig_axis >= 0 else orig_axis + _set_attr_i(node, "axis", adjusted_axis) + + orig_keepdims = _get_attr_i(node, "keepdims", 1) + _set_attr_i(node, "keepdims", 1) + + # ── Unsqueeze input ─────────────────────────────────────────────────────── + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # ── Squeeze outputs ─────────────────────────────────────────────────────── + c2 = cnt_ref[0] + cnt_ref[0] += 1 + orig_out = node.output[0] # ArgMax always has exactly one output + temp_out = unique_init.unique_name(f"{pfx}_argmax_4d_{c}") + + if orig_keepdims == 0: + # squeeze: leading (4-rank) axes + the reduced axis (held at 1 due to keepdims=1) + leading = list(range(4 - rank)) + sq_axes = sorted({*leading, adjusted_axis}) + axes_init = numpy_helper.from_array( + np.array(sq_axes, dtype=np.int64), + name=unique_init.unique_name(f"{pfx}_argmax_sq_axes_{c2}"), + ) + add_unique_initializers(graph, unique_init, [axes_init]) + sq = helper.make_node( + "Squeeze", + inputs=[temp_out, axes_init.name], + outputs=[orig_out], + name=unique_init.unique_name(f"{pfx}_argmax_sq_{c2}"), + ) + node.output[0] = temp_out + post.append(sq) + else: + sq = _make_squeeze_to_rank(temp_out, rank, orig_out, pfx, unique_init, graph, c2) + if sq: + node.output[0] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +def _handle_reduce(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Covers ReduceSum / ReduceMean / ReduceMax / ReduceMin.""" + if len(node.input) > 1 and node.input[1] and node.input[1] not in init_names: + return None + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + + pfx = node.name or f"reduce_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + post: list = [] + + delta = _axis_delta(rank) + orig_keepdims = _get_attr_i(node, "keepdims", 1) + _set_attr_i(node, "keepdims", 1) + + # ── Adjust axes ─────────────────────────────────────────────────────────── + adjusted_axes: list[int] = [] + + if len(node.input) > 1 and node.input[1] and node.input[1] in init_names: + # axes as input[1] initializer (opset 13+) + init = shape_map._cache.get_init(node.input[1]) + if init is not None: + orig_axes = list(numpy_helper.to_array(init).astype(int)) + adjusted_axes = [a + delta if a >= 0 else a for a in orig_axes] + c_a = cnt_ref[0] + cnt_ref[0] += 1 + new_axes_name = unique_init.unique_name(f"{init.name}_4d_{c_a}") + # Add new init; leave old for any other consumers. + inits_to_add.append( + numpy_helper.from_array(np.array(adjusted_axes, dtype=np.int64), new_axes_name) + ) + node.input[1] = new_axes_name + else: + for attr in node.attribute: + if attr.name == "axes": + orig_axes = list(attr.ints) + adjusted_axes = [a + delta if a >= 0 else a for a in orig_axes] + del attr.ints[:] + attr.ints.extend(int(x) for x in adjusted_axes) + break + + # ── Unsqueeze input ─────────────────────────────────────────────────────── + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # ── Squeeze outputs (Reduce always has one output) ──────────────────────── + c2 = cnt_ref[0] + cnt_ref[0] += 1 + orig_out = node.output[0] + temp_out = unique_init.unique_name(f"{pfx}_reduce_4d_{c}") + + if orig_keepdims == 0: + leading = list(range(4 - rank)) + noop = _get_attr_i(node, "noop_with_empty_axes", 0) + reduced = ( + [a % 4 for a in adjusted_axes] if adjusted_axes else [] if noop else list(range(4)) + ) + sq_axes = sorted(set(leading + reduced)) + axes_init = numpy_helper.from_array( + np.array(sq_axes, dtype=np.int64), + name=unique_init.unique_name(f"{pfx}_reduce_sq_axes_{c2}"), + ) + add_unique_initializers(graph, unique_init, [axes_init]) + sq = helper.make_node( + "Squeeze", + inputs=[temp_out, axes_init.name], + outputs=[orig_out], + name=unique_init.unique_name(f"{pfx}_reduce_sq_{c2}"), + ) + node.output[0] = temp_out + post.append(sq) + else: + sq = _make_squeeze_to_rank(temp_out, rank, orig_out, pfx, unique_init, graph, c2) + if sq: + node.output[0] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +def _handle_gather(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Promote Gather to 4D output. + + Strategy + -------- + Gather output rank = data_rank - 1 + indices_rank. + After unsqueezing data to 4D, we squeeze indices to rank-1 so that: + output_rank = 4 - 1 + 1 = 4 (no output squeeze needed). + + Indices are also cast to INT32 (DLA requirement). + Fixed indices use a shared private flattened view; runtime defaults stay dynamic. + For dynamic indices: Cast node followed by Squeeze node(s). + """ + in_name = node.input[0] + idx_name = node.input[1] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + idx_shape = shape_map.get(idx_name) + if idx_shape is None or sum(dim != 1 for dim in idx_shape) > 1: + return None + native_axis = _get_attr_i(node, "axis", 0) % rank + native_output = [*orig_shape[:native_axis], *idx_shape, *orig_shape[native_axis + 1 :]] + promoted = pad4d(orig_shape) + axis_4d = native_axis + _axis_delta(rank) + converted_output = [*promoted[:axis_4d], math.prod(idx_shape), *promoted[axis_4d + 1 :]] + if len(native_output) > 4: + return None + restore_axes = _find_squeeze_axes(converted_output, native_output) + if ( + any(converted_output[axis] != 1 for axis in restore_axes) + or [dim for axis, dim in enumerate(converted_output) if axis not in restore_axes] + != native_output + ): + return None + + pfx = node.name or f"gather_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + is_init = ( + idx_name in init_names + and idx_name not in shape_map._runtime_inputs + and not _has_unloaded_external(shape_map._cache, [idx_name]) + ) + + # ── 1. Cast indices to INT32 and flatten to 1-D ─────────────────────────── + if is_init: + if idx_name not in shape_map._index_views: + flat = ( + numpy_helper.to_array(shape_map._cache.get_init(idx_name)) + .reshape(-1) + .astype(np.int32, copy=False) + ) + private = unique_init.unique_name(f"{pfx}_idx_i32_{c}") + inits_to_add.append(numpy_helper.from_array(flat, private)) + shape_map._index_views[idx_name] = private + node.input[1] = shape_map._index_views[idx_name] + else: + # Check if the producer of idx_name is already a Cast node — reuse it. + cache = shape_map._cache + producer = cache.get_producer(idx_name) + producer_cast = ( + producer + if producer is not None + and producer.op_type == "Cast" + and producer.output + and producer.output[0] == idx_name + else None + ) + exclusive = ( + producer_cast is not None + and len(cache.get_consumers(idx_name)) == 1 + and idx_name not in cache.graph_output_names | shape_map._captured + and id(producer_cast) in shape_map._standard_nodes + ) + if exclusive: + _set_attr_i(producer_cast, "to", TensorProto.INT32) + value = cache.vi_map.get(idx_name) + if value is not None: + value.type.tensor_type.elem_type = TensorProto.INT32 + cast_name = idx_name + elif shape_map._cache.get_dtype(idx_name) != TensorProto.INT32: + # Insert a new Cast to INT32 + c_cast = cnt_ref[0] + cnt_ref[0] += 1 + cast_out = unique_init.unique_name(f"{pfx}_idx_i32_{c_cast}") + pre.append( + helper.make_node( + "Cast", + inputs=[idx_name], + outputs=[cast_out], + to=TensorProto.INT32, + name=unique_init.unique_name(f"{pfx}_idx_cast_{c_cast}"), + ) + ) + cast_name = cast_out + else: + cast_name = idx_name + + # Unsqueeze indices to 4D first, then Squeeze all unit dims to rank-1. + # This ensures the Squeeze axes are stable after upstream 4D promotion. + idx_shape = shape_map.get(idx_name) + idx_rank = len(idx_shape) if idx_shape is not None else None + + if idx_rank is not None and (idx_rank not in (1, 4)): + # Unsqueeze to 4D + c_unsq = cnt_ref[0] + cnt_ref[0] += 1 + unsq_axes = list(range(4 - idx_rank)) + unsq_axes_init = numpy_helper.from_array( + np.array(unsq_axes, dtype=np.int64), + unique_init.unique_name(f"{pfx}_idx_unsq_axes_{c_unsq}"), + ) + add_unique_initializers(graph, unique_init, [unsq_axes_init]) + unsq_out = unique_init.unique_name(f"{pfx}_idx_4d_{c_unsq}") + pre.append( + helper.make_node( + "Unsqueeze", + inputs=[cast_name, unsq_axes_init.name], + outputs=[unsq_out], + name=unique_init.unique_name(f"{pfx}_idx_unsq_{c_unsq}"), + ) + ) + cast_name = unsq_out + idx_shape_4d = pad4d(idx_shape) + else: + idx_shape_4d = list(idx_shape) if idx_shape is not None else None + + # Squeeze all unit dims from 4D shape to get rank-1 + if idx_shape_4d is not None and len(idx_shape_4d) > 1: + sq_axes = [i for i in range(len(idx_shape_4d)) if idx_shape_4d[i] == 1] + if len(sq_axes) == len(idx_shape_4d): + sq_axes = sq_axes[:-1] # keep at least one dim + if sq_axes: + c_sq = cnt_ref[0] + cnt_ref[0] += 1 + sq_axes_init = numpy_helper.from_array( + np.array(sq_axes, dtype=np.int64), + unique_init.unique_name(f"{pfx}_idx_sq_axes_{c_sq}"), + ) + add_unique_initializers(graph, unique_init, [sq_axes_init]) + sq_out = unique_init.unique_name(f"{pfx}_idx_1d_{c_sq}") + pre.append( + helper.make_node( + "Squeeze", + inputs=[cast_name, sq_axes_init.name], + outputs=[sq_out], + name=unique_init.unique_name(f"{pfx}_idx_sq_{c_sq}"), + ) + ) + node.input[1] = sq_out + else: + node.input[1] = cast_name + else: + node.input[1] = cast_name + + # ── 2. Unsqueeze data to 4D + adjust axis ──────────────────────────────── + adjusted_axis = _get_attr_i(node, "axis", 0) + if rank != 4: + delta = _axis_delta(rank) + orig_axis = _get_attr_i(node, "axis", 0) # ONNX Gather default: 0 + if orig_axis >= 0: + adjusted_axis = orig_axis + delta + else: + adjusted_axis = orig_axis # negative axes don't shift + _set_attr_i(node, "axis", adjusted_axis) + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.insert(0, unsq) + + # ── 3. Squeeze 4D output back to original output shape ──────────────────── + # After promoting data to 4D and flattening indices to 1D, the Gather output + # is 4D: data_4d[:axis] + [indices_len] + data_4d[axis+1:]. + # The extra unary dims may NOT be at the leading positions, so we must use + # _find_squeeze_axes with the actual 4D output shape rather than falling back + # to leading-axis removal. + orig_out = node.output[0] + orig_out_shape = native_output + post: list = [] + if len(orig_out_shape) < 4: + c_sq = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_gather_4d_{c_sq}") + + sq = _make_squeeze_to_rank( + temp_out, + len(orig_out_shape), + orig_out, + pfx, + unique_init, + graph, + c_sq, + shape_4d=converted_output, + orig_out_shape=list(orig_out_shape), + ) + if sq: + node.output[0] = temp_out + post.append(sq) + + if not pre and not post: + return None + return [*pre, node, *post] + + +def _handle_gatherelements(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Promote GatherElements to 4D output. + + GatherElements requires equal data and index ranks and produces output + with the same shape as indices. After unsqueezing data to 4D, indices must + also be unsqueezed to 4D so the ranks match → output is naturally 4D. + + Indices are cast to INT32 (DLA requirement). + Index Cast and Unsqueeze views preserve fixed storage and runtime defaults. + """ + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + if rank == 4: + return None + + pfx = node.name or f"gathel_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + delta = _axis_delta(rank) + + # ── 1. Adjust axis ──────────────────────────────────────────────────────── + orig_axis = _get_attr_i(node, "axis", 0) # ONNX GatherElements default: 0 + if orig_axis >= 0: + _set_attr_i(node, "axis", orig_axis + delta) + # negative axes don't need adjustment + + # ── 2. Unsqueeze data to 4D ─────────────────────────────────────────────── + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # ── 3. Cast indices to INT32 and unsqueeze to 4D (must match data shape) ── + if len(node.input) > 1 and node.input[1]: + idx_name = node.input[1] + # Only cast to INT32 if indices are not already INT32 + ge_need_cast = shape_map._cache.get_dtype(idx_name) != TensorProto.INT32 + + if ge_need_cast: + c_cast = cnt_ref[0] + cnt_ref[0] += 1 + cast_out = unique_init.unique_name(f"{pfx}_idx_i32_{c_cast}") + pre.append( + helper.make_node( + "Cast", + inputs=[idx_name], + outputs=[cast_out], + to=TensorProto.INT32, + name=unique_init.unique_name(f"{pfx}_idx_cast_{c_cast}"), + ) + ) + idx_after_cast = cast_out + else: + idx_after_cast = idx_name + + # Unsqueeze indices to 4D to match data + if rank < 4: + c_usq = cnt_ref[0] + cnt_ref[0] += 1 + idx_4d, idx_unsq = _make_unsqueeze_to_4d( + idx_after_cast, + rank, + unique_init.unique_name(f"{pfx}_idx"), + unique_init, + graph, + c_usq, + ) + node.input[1] = idx_4d + if idx_unsq: + pre.append(idx_unsq) + else: + node.input[1] = idx_after_cast + + # ── 4. Squeeze outputs back to original output shapes ───────────────────── + post: list = [] + for oi, orig_out in enumerate(node.output): + if not orig_out: + continue + if rank < 4: + c_sq = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_gathel_4d_{c_sq}_o{oi}") + sq = _make_squeeze_to_rank( + temp_out, + rank, + orig_out, + pfx, + unique_init, + graph, + c_sq, + ) + if sq: + node.output[oi] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +def _handle_expand(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Promote fixed Expand targets and restore native output shapes.""" + if len(node.input) < 2 or node.input[1] not in init_names: + return None + in_name = node.input[0] + orig_in_shape = shape_map.get(in_name) # may be None if input shape is unknown + + native_output = None + if orig_in_shape is not None: + target = shape_map._cache.get_init_array(node.input[1]) + native_output = list(np.broadcast_shapes(orig_in_shape, target.tolist())) + if native_output is None and any(shape_map.get(name) is None for name in node.output if name): + return None + pfx = node.name or f"exp_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + + # ── Update shape initializer (input[1]) to 4D ──────────────────────────── + if len(node.input) > 1 and node.input[1] and node.input[1] in init_names: + init = shape_map._cache.get_init(node.input[1]) + if init is not None: + arr = numpy_helper.to_array(init) + if len(arr) < 4: + new_arr = np.array(pad4d(arr.tolist()), dtype=arr.dtype) + c_t = cnt_ref[0] + cnt_ref[0] += 1 + new_tile_name = unique_init.unique_name(f"{init.name}_4d_{c_t}") + inits_to_add.append(numpy_helper.from_array(new_arr, new_tile_name)) + node.input[1] = new_tile_name + + # ── Unsqueeze data input to 4D (skip for scalar / unknown input shape) ─── + in_rank = len(orig_in_shape) if orig_in_shape is not None else None + if in_rank is not None and in_rank > 0: + out4d, unsq = _make_unsqueeze_to_4d(in_name, in_rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # ── Squeeze outputs based on each output's original shape ──────────────── + # The transformed Expand now produces 4D. Use each output's known original + # shape (not the input rank) to determine squeeze axes. This handles: + # * ops where input rank != output rank + # * ops where input shape is unknown but output shape is known + post: list = [] + for oi, orig_out in enumerate(node.output): + if not orig_out: + continue + orig_out_shape = shape_map.get(orig_out) + if orig_out_shape is None: + orig_out_shape = native_output + if orig_out_shape is None or len(orig_out_shape) >= 4: + continue + ci = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_exp_4d_{c}_o{oi}") + sq = _make_squeeze_to_rank( + temp_out, len(orig_out_shape), orig_out, pfx, unique_init, graph, ci + ) + if sq: + node.output[oi] = temp_out + post.append(sq) + return [*pre, node, *post] + + +def _handle_tile(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Promote fixed Tile repeats and restore native output shapes.""" + if len(node.input) < 2 or node.input[1] not in init_names: + return None + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + if rank == 4: + return None + + pfx = node.name or f"tile_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + + # ── Update repeats initializer (input[1]) if 3D ─────────────────────────── + if len(node.input) > 1 and node.input[1] and node.input[1] in init_names: + init = shape_map._cache.get_init(node.input[1]) + if init is not None: + arr = numpy_helper.to_array(init) + if len(arr) < 4: + new_arr = np.asarray(pad4d(arr.tolist()), dtype=arr.dtype) + name = unique_init.unique_name(f"{init.name}_tile4d_{c}") + inits_to_add.append(numpy_helper.from_array(new_arr, name)) + node.input[1] = name + + # ── Unsqueeze data input ────────────────────────────────────────────────── + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + # Squeeze outputs back to original rank + post: list = [] + for oi, orig_out in enumerate(node.output): + if not orig_out: + continue + ci = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_tile_4d_{c}_o{oi}") + sq = _make_squeeze_to_rank(temp_out, rank, orig_out, pfx, unique_init, graph, ci) + if sq: + node.output[oi] = temp_out + post.append(sq) + return [*pre, node, *post] + + +def _handle_lpnorm(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Adapt normalization axes around leading input promotion.""" + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + if rank == 4: + return None + + pfx = node.name or f"lpnorm_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + pre: list = [] + post: list = [] + + axis = _get_attr_i(node, "axis", -1) + _set_attr_i(node, "axis", axis + 4 - rank if axis >= 0 else axis) + + out4d, unsq = _make_unsqueeze_to_4d(in_name, rank, pfx, unique_init, graph, c) + node.input[0] = out4d + if unsq: + pre.append(unsq) + + for oi, orig_out in enumerate(node.output): + if not orig_out: + continue + ci = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_lpnorm_4d_{c}_o{oi}") + sq = _make_squeeze_to_rank(temp_out, rank, orig_out, pfx, unique_init, graph, ci) + if sq: + node.output[oi] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +def _axis_control_input_indices(node): + """Identify axis-operation controls whose dimensions must stay unchanged.""" + if node.op_type == "DequantizeLinear": + return {1, 2} + return {1} if node.op_type in {"Split", "TopK"} else set() + + +def _handle_axis_op(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Softmax, LogSoftmax, Concat, Split, TopK, DequantizeLinear.""" + skip_input_indices = _axis_control_input_indices(node) + data_inputs = [inp for i, inp in enumerate(node.input) if inp and i not in skip_input_indices] + if not data_inputs or any(shape_map.get(inp) is None for inp in data_inputs): + return None + first_inp_name = data_inputs[0] + orig_rank = len(shape_map.get(first_inp_name)) + if orig_rank == 4: + return None + + pfx = node.name or f"axop_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + # Expected 4D shape (shape-preserving assumption) for robust squeeze-axis finding + first_in_shape = shape_map.get(first_inp_name) if first_inp_name else None + expected_4d_shape = pad4d(first_in_shape) if first_in_shape else None + + pre: list = [] + post: list = [] + + delta = _axis_delta(orig_rank) + if not any(attr.name == "axis" for attr in node.attribute): + defaults = { + "Softmax": -1, + "LogSoftmax": -1, + "Split": 0, + "TopK": -1, + "QuantizeLinear": 1, + "DequantizeLinear": 1, + } + if node.op_type in defaults: + node.attribute.append(helper.make_attribute("axis", defaults[node.op_type])) + + # ── Adjust axis attribute ───────────────────────────────────────────────── + for attr in node.attribute: + if attr.name == "axis": + if attr.i >= 0: + attr.i += delta + break + + # ── Unsqueeze tensor inputs; expand initializer inputs ──────────────────── + for i, inp in enumerate(node.input): + if not inp or i in skip_input_indices: + continue + if inp in init_names: + _expand_init_dims(shape_map._cache.get_init(inp), shape_map._cache) + else: + s = shape_map.get(inp) + if s is None or len(s) == 4: + continue + ci = cnt_ref[0] + cnt_ref[0] += 1 + inp_4d, unsq = _make_unsqueeze_to_4d( + inp, len(s), unique_init.unique_name(f"{pfx}_in{i}"), unique_init, graph, ci + ) + node.input[i] = inp_4d + if unsq: + pre.append(unsq) + + # ── Squeeze outputs ─────────────────────────────────────────────────────── + for oi, out in enumerate(node.output): + if not out: + continue + ci = cnt_ref[0] + cnt_ref[0] += 1 + temp = unique_init.unique_name(f"{pfx}_axop_4d_{c}_o{oi}") + out_shape = shape_map.get(out) + sq = _make_squeeze_to_rank( + temp, + orig_rank, + out, + pfx, + unique_init, + graph, + ci, + shape_4d=expected_4d_shape, + orig_out_shape=out_shape, + ) + if sq: + node.output[oi] = temp + post.append(sq) + + return [*pre, node, *post] + + +def _handle_flatten(node, graph, shape_map, _init_names, unique_init, _inits_to_add, cnt_ref): + """Flatten → Unsqueeze(to 4D) → Reshape([1,1,outer,inner]) → Squeeze([outer,inner]). + + Same logic as the former ``dla_flatten_to_reshape`` transform, now inlined. + Input already 4D: Unsqueeze step is skipped. + """ + in_name = node.input[0] + out_name = node.output[0] + + in_shape = shape_map.get(in_name) + if in_shape is None: + logger.debug("Flatten %r: input shape unknown; skipping.", node.name) + return None + + rank = len(in_shape) + raw_axis = next((a.i for a in node.attribute if a.name == "axis"), 1) + axis = int(raw_axis) + if axis < 0: + axis = rank + axis + axis = max(0, min(axis, rank)) + + outer = int(np.prod(in_shape[:axis])) if axis > 0 else 1 + + pfx = node.name or f"fl_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + chain: list = [] + + # ── Step 1: Unsqueeze input to 4D (skip if already 4D) ─────────────────── + if rank < 4: + n_new = 4 - rank + unsq_axes_init = numpy_helper.from_array( + np.arange(n_new, dtype=np.int64), name=unique_init.unique_name(f"{pfx}_unsq_axes_{c}") + ) + add_unique_initializers(graph, unique_init, [unsq_axes_init]) + in4d = unique_init.unique_name(f"{pfx}_in4d_{c}") + chain.append( + helper.make_node( + "Unsqueeze", + inputs=[in_name, unsq_axes_init.name], + outputs=[in4d], + name=unique_init.unique_name(f"{pfx}_unsq_{c}"), + ) + ) + else: + in4d = in_name + + # ── Step 2: Reshape to [1, 1, outer, -1] (-1 infers inner automatically) ─ + shape_init = numpy_helper.from_array( + np.array([1, 1, outer, -1], dtype=np.int64), + name=unique_init.unique_name(f"{pfx}_reshape_shape_{c}"), + ) + add_unique_initializers(graph, unique_init, [shape_init]) + flat4d = unique_init.unique_name(f"{pfx}_flat4d_{c}") + chain.append( + helper.make_node( + "Reshape", + inputs=[in4d, shape_init.name], + outputs=[flat4d], + name=unique_init.unique_name(f"{pfx}_reshape_{c}"), + ) + ) + + # ── Step 3: Squeeze [1,1,outer,-1] → [outer,-1] ────────────────────────── + sq_axes_init = numpy_helper.from_array( + np.array([0, 1], dtype=np.int64), name=unique_init.unique_name(f"{pfx}_sq_axes_{c}") + ) + add_unique_initializers(graph, unique_init, [sq_axes_init]) + chain.append( + helper.make_node( + "Squeeze", + inputs=[flat4d, sq_axes_init.name], + outputs=[out_name], + name=unique_init.unique_name(f"{pfx}_sq_{c}"), + ) + ) + + logger.debug( + "Flatten %r [axis=%d]: %s → Unsqueeze→Reshape[1,1,%d,-1]→Squeeze", + node.name, + axis, + in_shape, + outer, + ) + return chain # replaces the Flatten node entirely + + +def _handle_instancenorm(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """3D InstanceNorm: add trailing dim [N,C,D]→[N,C,D,1] before; remove after.""" + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None: + return None + rank = len(orig_shape) + if rank != 3: + return None # 4D already fine; other ranks not handled + + pfx = node.name or f"in_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + c2 = cnt_ref[0] + cnt_ref[0] += 1 + + # Unsqueeze trailing dim: [N,C,D] → [N,C,D,1] + unsq_axes = numpy_helper.from_array( + np.array([3], dtype=np.int64), name=unique_init.unique_name(f"{pfx}_in_unsq_axes_{c}") + ) + add_unique_initializers(graph, unique_init, [unsq_axes]) + in4d = unique_init.unique_name(f"{pfx}_in_unsq4d_{c}") + unsq = helper.make_node( + "Unsqueeze", + inputs=[in_name, unsq_axes.name], + outputs=[in4d], + name=unique_init.unique_name(f"{pfx}_in_unsq_{c}"), + ) + node.input[0] = in4d + + # Squeeze trailing dim from output: [N,C,D,1] → [N,C,D] + sq_axes = numpy_helper.from_array( + np.array([3], dtype=np.int64), name=unique_init.unique_name(f"{pfx}_in_sq_axes_{c2}") + ) + add_unique_initializers(graph, unique_init, [sq_axes]) + out4d = unique_init.unique_name(f"{pfx}_in_4d_out_{c}") + sq = helper.make_node( + "Squeeze", + inputs=[out4d, sq_axes.name], + outputs=[node.output[0]], + name=unique_init.unique_name(f"{pfx}_in_sq_{c2}"), + ) + node.output[0] = out4d + + return [unsq, node, sq] + + +def _handle_layernorm(node, graph, shape_map, _init_names, unique_init, _inits_to_add, cnt_ref): + """LayerNormalization: unsqueeze data input to 4D, keep scale/bias unchanged, update axis. + + LayerNormalization inputs: + input[0]: data — promoted to 4D via Unsqueeze; output squeezed back to orig rank. + input[1]: scale — 1D, shape matches the normalised dims; left untouched. + input[2]: bias — 1D, optional; left untouched. + + The ``axis`` attribute is shifted by ``4 - orig_rank`` so it still addresses + the correct dimension after the data tensor is promoted. + """ + in_name = node.input[0] + orig_shape = shape_map.get(in_name) + if orig_shape is None or len(orig_shape) == 4: + return None + + orig_rank = len(orig_shape) + delta = 4 - orig_rank + pfx = node.name or f"ln_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + # ── Unsqueeze data input to 4D ──────────────────────────────────────────── + in4d, unsq = _make_unsqueeze_to_4d(in_name, orig_rank, pfx, unique_init, graph, c) + node.input[0] = in4d + + # ── Update axis attribute ───────────────────────────────────────────────── + # LayerNorm default axis is -1 (last dim) — negative stays the same. + orig_axis = _get_attr_i(node, "axis", -1) + if orig_axis >= 0: + _set_attr_i(node, "axis", orig_axis + delta) + + pre = [unsq] if unsq else [] + post = [] + native_axis = orig_axis % orig_rank + for index, original in enumerate(list(node.output)): + if not original: + continue + native_shape = shape_map.get(original) + if native_shape is None: + native_shape = ( + orig_shape + if index == 0 + else orig_shape[:native_axis] + [1] * (orig_rank - native_axis) + ) + c2 = cnt_ref[0] + cnt_ref[0] += 1 + temporary = unique_init.unique_name(f"{pfx}_ln_4d_{c}_{index}") + squeeze = _make_squeeze_to_rank( + temporary, + orig_rank, + original, + pfx, + unique_init, + graph, + c2, + shape_4d=pad4d(native_shape), + orig_out_shape=native_shape, + ) + if squeeze: + node.output[index] = temporary + post.append(squeeze) + return [*pre, node, *post] + + +def _handle_generic(node, graph, shape_map, init_names, unique_init, inits_to_add, cnt_ref): + """Default handler for unregistered ops (Add, Sub, Mul, Relu, Cast, etc.).""" + if node.op_type not in { + "Abs", + "Neg", + "Relu", + "Sigmoid", + "Tanh", + "Sqrt", + "Exp", + "Log", + "Erf", + "Cast", + "Clip", + "Add", + "Sub", + "Mul", + "Div", + "Pow", + "Greater", + "Less", + "Equal", + "And", + "Or", + "Xor", + "Where", + "Max", + "Min", + "PRelu", + "Sum", + "Conv", + "ConvTranspose", + }: + return None + # Determine orig_rank from first non-initializer input + first_inp_name = None + orig_rank = None + for inp in node.input: + if inp and inp not in init_names: + s = shape_map.get(inp) + if s is not None and len(s) > 0: + first_inp_name = inp + orig_rank = len(s) + break + if orig_rank is None: + return None + if node.op_type in {"Conv", "ConvTranspose"} and orig_rank != 4: + return None + broadcast_output = None + if node.op_type not in {"Conv", "ConvTranspose"}: + shapes = [shape_map.get(name) for name in node.input if name] + if any(shape is None for shape in shapes): + return None + broadcast_output = list(np.broadcast_shapes(*shapes)) + orig_rank = len(broadcast_output) + + pfx = node.name or f"gen_{cnt_ref[0]}" + c = cnt_ref[0] + cnt_ref[0] += 1 + + # Pre-compute expected 4D shape (shape-preserving assumption: 4D_out = 4D_in). + # This is used by _make_squeeze_to_rank to find the correct squeeze axis even + # when the unary dim is not at the leading position (e.g. [1,1,3,1] → [1,1,3]). + first_in_shape = shape_map.get(first_inp_name) if first_inp_name else None + expected_4d_shape = ( + pad4d(broadcast_output if broadcast_output is not None else first_in_shape) + if first_in_shape + else None + ) + + pre: list = [] + post: list = [] + + # Conv / ConvTranspose bias (input[2]) must stay 1-D per the ONNX spec. + conv_bias_indices: set[int] = {2} if node.op_type in ("Conv", "ConvTranspose") else set() + + for i, inp in enumerate(node.input): + if not inp: + continue + if i in conv_bias_indices: + continue # Conv/ConvTranspose bias — leave untouched + if inp in init_names: + init = shape_map._cache.get_init(inp) + if len(init.dims) != 0: + _expand_init_dims(init, shape_map._cache) + else: + s = shape_map.get(inp) + if s is None: + continue + if len(s) == 0: + continue # scalar dynamic tensor — leave untouched + # Convert all dynamic inputs to 4D — even already-4D inputs go through + # _make_unsqueeze_to_4d which is a no-op (returns the same name) for rank==4. + ci = cnt_ref[0] + cnt_ref[0] += 1 + inp_4d, unsq = _make_unsqueeze_to_4d( + inp, len(s), unique_init.unique_name(f"{pfx}_in{i}"), unique_init, graph, ci + ) + node.input[i] = inp_4d + if unsq: + pre.append(unsq) + + for oi, orig_out in enumerate(node.output): + if not orig_out: + continue + orig_out_shape = shape_map.get(orig_out, broadcast_output) + ci = cnt_ref[0] + cnt_ref[0] += 1 + temp_out = unique_init.unique_name(f"{pfx}_gen_4d_{c}_o{oi}") + sq = _make_squeeze_to_rank( + temp_out, + orig_rank, + orig_out, + pfx, + unique_init, + graph, + ci, + shape_4d=expected_4d_shape, + orig_out_shape=orig_out_shape, + ) + if sq: + node.output[oi] = temp_out + post.append(sq) + + return [*pre, node, *post] + + +# --------------------------------------------------------------------------- +# DequantizeLinear — static initializer data path +# --------------------------------------------------------------------------- + + +def _handle_dequantize_linear( + node, graph, shape_map, init_names, unique_init, _inits_to_add, cnt_ref +): + """Promote fixed nonblocked quantized data and restore its native output rank.""" + if not node.input or node.input[0] not in init_names: + return None + data_init = shape_map._cache.get_init(node.input[0]) + scale_shape = shape_map.get(node.input[1]) + if ( + data_init is None + or len(data_init.dims) >= 4 + or _get_attr_i(node, "block_size", 0) + or scale_shape is None + or len(scale_shape) > 1 + ): + return None + orig_rank = len(data_init.dims) + c = cnt_ref[0] + cnt_ref[0] += 1 + original = node.output[0] + temporary = shape_map._cache.unique_name(original + "_dq_4d") + squeeze = _make_squeeze_to_rank( + temporary, orig_rank, original, node.name or "dq", unique_init, graph, c + ) + _expand_init_dims(data_init, shape_map._cache) + if scale_shape: + axis = _get_attr_i(node, "axis", 1) + if axis >= 0: + _set_attr_i(node, "axis", axis + 4 - orig_rank) + node.output[0] = temporary + return [node, squeeze] + + +# --------------------------------------------------------------------------- +# Dispatch table +# --------------------------------------------------------------------------- + +_OP_HANDLERS = { + "Reshape": _handle_reshape, + "Transpose": _handle_transpose, + "Slice": _handle_slice, + "ArgMax": _handle_argmax, + "ReduceSum": _handle_reduce, + "ReduceMean": _handle_reduce, + "ReduceMax": _handle_reduce, + "ReduceMin": _handle_reduce, + "Gather": _handle_gather, + "GatherElements": _handle_gatherelements, + "Expand": _handle_expand, + "Tile": _handle_tile, + "LpNormalization": _handle_lpnorm, + "Softmax": _handle_axis_op, + "LogSoftmax": _handle_axis_op, + "Concat": _handle_axis_op, + "Split": _handle_axis_op, + "TopK": _handle_axis_op, + "DequantizeLinear": _handle_dequantize_linear, + "Flatten": _handle_flatten, + "InstanceNormalization": _handle_instancenorm, + "LayerNormalization": _handle_layernorm, +} + + +# --------------------------------------------------------------------------- +# Post-pass: split squeeze nodes whose output is both a graph output and +# consumed by internal nodes +# --------------------------------------------------------------------------- + + +def _split_graph_output_squeezes(model: onnx.ModelProto) -> None: + """Give internal consumers private native views without changing opaque operators or public bindings.""" + graph = model.graph + cache = GraphCache(graph) + outputs = {value.name: value for value in graph.output} + positions = {id(node): index for index, node in enumerate(graph.node)} + replacements = {} + for node in _operator_nodes(model, {"Squeeze"}): + original = node.output[0] + consumers = cache.get_consumers(original) + if original not in outputs or not consumers: + continue + internal = cache.unique_name(original + "_sq") + copied = onnx.NodeProto() + copied.CopyFrom(node) + copied.output[0] = internal + copied.name = cache.unique_name((node.name or original) + "_internal") + for consumer, port in consumers: + consumer.input[port] = internal + metadata = onnx.ValueInfoProto() + metadata.CopyFrom(outputs[original]) + metadata.name = internal + graph.value_info.append(metadata) + replacements[positions[id(node)]] = ([node], [node, copied]) + batch_replace_nodes(graph, replacements) + + +# --------------------------------------------------------------------------- +# Main transform +# --------------------------------------------------------------------------- + + +def _has_unloaded_external(cache, names): + """Identify controls whose payload cannot be read without the caller's model directory.""" + return any( + (tensor := cache.get_init(name)) is not None + and onnx.external_data_helper.uses_external_data(tensor) + and not tensor.HasField("raw_data") + for name in names + if name + ) + + +def _apply_convert_ops_to_4d(model: onnx.ModelProto) -> onnx.ModelProto: + """Wrap eligible non-4D operations while restoring their native output contracts.""" + if next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1) < 13: + return model + graph = model.graph + + # ── Build graph cache for O(1) lookups ──────────────────────────────────── + cache = GraphCache(graph) + unique_init = _InitializerNames(cache) + standard_nodes = _operator_nodes(model, {node.op_type for node in graph.node}) + captured = _subgraph_captures(graph) + shape_map = _CacheShapeProxy(cache, standard_nodes, captured, _runtime_input_names(model)) + init_names = cache.init_names + + node_idx = {id(n): i for i, n in enumerate(graph.node)} + inits_to_add: list = [] + replacements_batch: dict[int, tuple[list, list]] = {} + cnt_ref = [0] + + # ── Process each node ───────────────────────────────────────────────────── + for node in standard_nodes: + if node.op_type in _OP_SKIP: + continue + if ( + node.op_type in {"Reshape", "Slice", "Expand", "Tile"} + or node.op_type.startswith("Reduce") + ) and _has_unloaded_external(cache, node.input[1:]): + continue + if ( + any( + (shape := cache.get_shape(name)) is not None and len(shape) > 4 + for name in node.input + ) + and node.op_type != "Transpose" + ): + continue + if any( + name in cache.vi_map and cache.get_shape(name) is None + for name in [*node.input, *node.output] + if name + ): + continue + handler = _OP_HANDLERS.get(node.op_type, _handle_generic) + guarded_inputs = list(node.input) + if handler is _handle_dequantize_linear: + guarded_inputs = guarded_inputs[:1] + elif handler is _handle_axis_op: + controls = _axis_control_input_indices(node) + guarded_inputs = [ + name for index, name in enumerate(node.input) if index not in controls + ] + elif handler is _handle_generic: + guarded_inputs = [ + name + for index, name in enumerate(node.input) + if not (node.op_type in {"Conv", "ConvTranspose"} and index == 2) + and not (name in init_names and not cache.init_map[name].dims) + ] + if any( + name in init_names + and len(cache.init_map[name].dims) != 4 + and ( + name in cache.graph_input_names + or ( + name in guarded_inputs + and handler in {_handle_axis_op, _handle_generic, _handle_dequantize_linear} + and ( + len(cache.get_consumers(name)) > 1 + or name in cache.graph_output_names | captured + ) + ) + ) + for name in node.input + if name + ): + continue + + # Quick check: does any input OR output have rank ≠ 4? + # We need to transform if any tensor is non-4D so the graph is fully 4D. + any_non4d = False + for inp in node.input: + if inp: + s = shape_map.get(inp) + if s is not None and len(s) != 4: + any_non4d = True + break + if not any_non4d: + for out in node.output: + if out: + s = shape_map.get(out) + if s is not None and len(s) != 4: + any_non4d = True + break + if not any_non4d: + continue + + # Warn and skip 5D tensors without a unary dim (only Transpose handles 5D) + has_bad_5d = False + if node.op_type != "Transpose": + for inp in node.input: + if inp: + s = shape_map.get(inp) + if s is not None and len(s) == 5 and 1 not in s: + logger.warning( + "Node %r (%s): input %r has 5D shape %s with no unary dim; skipping.", + node.name, + node.op_type, + inp, + s, + ) + has_bad_5d = True + break + if has_bad_5d: + continue + + result = handler( + node, + graph, + shape_map, + init_names, + unique_init, + inits_to_add, + cnt_ref, + ) + if result is not None: + replacements_batch[node_idx[id(node)]] = ([node], result) + + # ── Apply all node replacements ─────────────────────────────────────────── + batch_replace_nodes(graph, replacements_batch) + + add_unique_initializers(graph, unique_init, inits_to_add) + + # ── Split squeeze nodes shared between graph outputs and internal nodes ──── + _split_graph_output_squeezes(model) + _add_legacy_initializer_inputs(model) + + logger.debug("convert_ops_to_4d: transformed %d node(s).", len(replacements_batch)) + return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py new file mode 100644 index 00000000000..4558c962581 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_convert_ops_to_4d.py @@ -0,0 +1,2157 @@ +# 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_convert_ops_to_4d`. + +"""ORT CPU parity tests for ``dla_convert_ops_to_4d``. + +Each test: + 1. Builds a minimal ONNX model with a non-4D op. + 2. Runs ORT to get reference outputs. + 3. Applies ``_apply_convert_ops_to_4d`` (via deep-copy). + 4. Runs ORT on the transformed model. + 5. Asserts cosine similarity ≥ 1-1e-5 and allclose(rtol=1e-4, atol=1e-4). + +Coverage +-------- +Reshape (2D→4D, 3D→4D) +Transpose (2D, 3D) +Slice (2D data) +ArgMax (2D keepdims=0, 3D keepdims=1) +ReduceSum (2D keepdims=0, 3D keepdims=1) +Gather (2D data with static indices) +GatherElements (2D data + 2D static indices) +Expand (3D input with static shape) +Tile (3D input with static repeats) +LpNormalization (2D) +Softmax (2D) +LogSoftmax (2D) +Concat (two 2D inputs) +Split (3D input) +InstanceNormalization (3D: [N,C,D]) +Relu — generic handler (2D) +4D input — should not be transformed +MatMul — in _OP_SKIP, must not be transformed +""" + +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_convert_ops_to_4d import ( + _apply_convert_ops_to_4d, +) + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +@pytest.mark.parametrize("public_control", [False, True]) +def test_shared_reshape_control_uses_private_targets(public_control): + """Non-mutating handlers can safely convert multiple readers of one fixed target.""" + target = _init("target", np.array([3, 2], np.int64)) + outputs = [_vi("y", TensorProto.FLOAT, [3, 2]), _vi("z", TensorProto.FLOAT, [3, 2])] + if public_control: + outputs.append(_vi("target", TensorProto.INT64, [2])) + model = _make_model( + [ + helper.make_node("Reshape", ["x", "target"], ["y"]), + helper.make_node("Reshape", ["x", "target"], ["z"]), + ], + "shared_reshape_target", + [_vi("x", TensorProto.FLOAT, [2, 3])], + outputs, + [target], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected, actual, transformed = _apply_and_run(model, feeds) + _assert_match(expected, actual) + controls = {node.input[1] for node in transformed.graph.node if node.op_type == "Reshape"} + assert len(controls) == 2 and "target" not in controls + values = { + tensor.name: numpy_helper.to_array(tensor) for tensor in transformed.graph.initializer + } + for name in controls: + np.testing.assert_array_equal(values[name], [1, 1, 3, 2]) + np.testing.assert_array_equal(values["target"], [3, 2]) + + +# ── Shared utilities ──────────────────────────────────────────────────────────────────────────────── + + +def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + """Build a named fixed tensor for a synthetic graph.""" + return numpy_helper.from_array(arr, name=name) + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Declare tensor metadata for a synthetic graph.""" + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + """Use the model IR supported by the installed CPU runtime.""" + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _make_model(nodes, name, inputs, outputs, inits=None, *, infer_symbolic_shapes=True): + """Build and check a model with optional symbolic shape inference.""" + graph = helper.make_graph(nodes, name, inputs, outputs, initializer=inits or []) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + if infer_symbolic_shapes: + try: + # Shape inference is optional in partial runtime installations. + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + except ImportError: + pass + else: + model = SymbolicShapeInference.infer_shapes(model, auto_merge=True) or model + 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} + + +@pytest.mark.parametrize("operation", ["Concat", "TopK"]) +def test_axis_operator_with_unknown_data_rank_stays_native(operation): + """Missing intermediate data metadata must not select a control rank or partially promote data.""" + nodes = [helper.make_node("Identity", ["x"], ["hidden"])] + inputs = [_vi("x", TensorProto.FLOAT, [2, 3])] + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + if operation == "Concat": + inputs.append(_vi("other", TensorProto.FLOAT, [2, 3])) + feeds["other"] = -feeds["x"] + nodes.append(helper.make_node("Concat", ["other", "hidden"], ["y"], axis=1)) + outputs = [_vi("y", TensorProto.FLOAT, [2, 6])] + else: + inputs.append(_vi("k", TensorProto.INT64, [1])) + feeds["k"] = np.array([2], np.int64) + nodes.append(helper.make_node("TopK", ["hidden", "k"], ["y", "indices"], axis=1)) + outputs = [_vi("y", TensorProto.FLOAT, [2, 2]), _vi("indices", TensorProto.INT64, [2, 2])] + model = _make_model(nodes, "unknown_axis_data", inputs, outputs, infer_symbolic_shapes=False) + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + target = next(n for n in converted.graph.node if n.op_type == operation) + assert target.SerializeToString() == model.graph.node[-1].SerializeToString() + onnx.checker.check_model(converted, full_check=True) + for name, actual in _run_ort(converted, feeds).items(): + np.testing.assert_array_equal(actual, expected[name]) + + +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_match(ref: dict, got: dict) -> None: + """Check public output shapes and numerical parity.""" + assert ref.keys() == got.keys() + for name in ref: + a, b = ref[name], got[name] + assert a.shape == b.shape, f"{name}: shape {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) + else: + np.testing.assert_array_equal(a, b, err_msg=name) + + +def _assert_values_match(ref: dict, got: dict) -> None: + """Element-wise comparison ignoring shape (for ops whose output rank changes).""" + 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) + else: + np.testing.assert_array_equal(a, b, err_msg=name) + + +def _apply_and_run(model, feeds): + """Execute native and converted graphs for parity checks.""" + ref = _run_ort(model, feeds) + tr = copy.deepcopy(model) + tr = _apply_convert_ops_to_4d(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + return ref, got, tr + + +_RNG = np.random.default_rng(2024) + + +# ── Reshape ───────────────────────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize(("shape", "output_shape"), [([8, 16], [16, 8]), ([3, 8, 16], [3, 128])]) +def test_reshape_to_4d(shape, output_shape): + """Preserve native Reshape outputs after private target promotion.""" + model = _make_model( + [helper.make_node("Reshape", ["X", "shape"], ["Y"], name="rs")], + "reshape", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, output_shape)], + [_init("shape", np.array(output_shape, dtype=np.int64))], + ) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + ref, got, _ = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +def test_transpose_2d(): + """Transpose [8,16] perm=[1,0] → [16,8]: perm adjusted to [0,1,3,2].""" + model = _make_model( + [helper.make_node("Transpose", ["X"], ["Y"], name="tr", perm=[1, 0])], + "tr2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [16, 8])], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # Verify perm was adjusted + tr_node = next( + n + for n in tr.graph.node + if n.op_type == "Transpose" + and any(a.name == "perm" and list(a.ints) == [0, 1, 3, 2] for a in n.attribute) + ) + assert tr_node is not None, "Expected perm [0,1,3,2] in transformed Transpose" + + +def test_transpose_3d(): + """Transpose [3,8,16] perm=[2,0,1] → [16,3,8]: perm adjusted to [0,3,1,2].""" + model = _make_model( + [helper.make_node("Transpose", ["X"], ["Y"], name="tr", perm=[2, 0, 1])], + "tr3d", + [_vi("X", TensorProto.FLOAT, [3, 8, 16])], + [_vi("Y", TensorProto.FLOAT, [16, 3, 8])], + ) + feeds = {"X": _RNG.standard_normal((3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── Slice ─────────────────────────────────────────────────────────────────────────────────────────── + + +def test_slice_2d(): + """Slice data[8,16] along axis=0 (starts=2,ends=6) → [4,16]; axis updated to 2.""" + starts = _init("starts", np.array([2], dtype=np.int64)) + ends = _init("ends", np.array([6], dtype=np.int64)) + axes = _init("axes", np.array([0], dtype=np.int64)) + model = _make_model( + [helper.make_node("Slice", ["X", "starts", "ends", "axes"], ["Y"], name="sl")], + "sl2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [4, 16])], + [starts, ends, axes], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── ArgMax ────────────────────────────────────────────────────────────────────────────────────────── + + +def test_argmax_2d_keepdims0(): + """ArgMax [8,16] axis=1 keepdims=0 → [8]; axis adjusted to 3, keepdims forced=1, Squeeze added.""" + model = _make_model( + [helper.make_node("ArgMax", ["X"], ["Y"], name="am", axis=1, keepdims=0)], + "am2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.INT64, [8])], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # keepdims must be forced to 1 on the transformed node + am_node = next(n for n in tr.graph.node if n.op_type == "ArgMax") + kd = next((int(a.i) for a in am_node.attribute if a.name == "keepdims"), 1) + assert kd == 1, f"ArgMax keepdims must be 1 after transform, got {kd}" + # a Squeeze node must be inserted to restore the original [8] shape + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be added for keepdims=0" + ) + + +def test_argmax_3d_keepdims1(): + """ArgMax [3,8,16] axis=2 keepdims=1 → [3,8,1]; axis adjusted to 3, keepdims unchanged.""" + model = _make_model( + [helper.make_node("ArgMax", ["X"], ["Y"], name="am", axis=2, keepdims=1)], + "am3d", + [_vi("X", TensorProto.FLOAT, [3, 8, 16])], + [_vi("Y", TensorProto.INT64, [3, 8, 1])], + ) + feeds = {"X": _RNG.standard_normal((3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # keepdims must remain 1 on the transformed node + am_node = next(n for n in tr.graph.node if n.op_type == "ArgMax") + kd = next((int(a.i) for a in am_node.attribute if a.name == "keepdims"), 1) + assert kd == 1, f"ArgMax keepdims must be 1 after transform, got {kd}" + + +# ── ReduceSum ─────────────────────────────────────────────────────────────────────────────────────── + + +def test_reducesum_2d_keepdims0(): + """ReduceSum [8,16] axes=[1] keepdims=0 → [8]; axis adjusted to 3, Squeeze covers adjusted axis.""" + axes_init = _init("axes_rd", np.array([1], dtype=np.int64)) + model = _make_model( + [helper.make_node("ReduceSum", ["X", "axes_rd"], ["Y"], name="rd", keepdims=0)], + "rd2d_kd0", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [8])], + [axes_init], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # keepdims must be forced to 1 on the transformed node + rd_node = next(n for n in tr.graph.node if n.op_type == "ReduceSum") + kd = next((int(a.i) for a in rd_node.attribute if a.name == "keepdims"), 1) + assert kd == 1, f"ReduceSum keepdims must be 1 after transform, got {kd}" + # a Squeeze node must be inserted to restore the original [8] shape + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be added for keepdims=0" + ) + + +def test_reducesum_3d_keepdims1(): + """ReduceSum [3,8,16] axes=[1] keepdims=1 → [3,1,16]; axis adjusted to 2, leading Squeeze only.""" + axes_init = _init("axes_rd3", np.array([1], dtype=np.int64)) + model = _make_model( + [helper.make_node("ReduceSum", ["X", "axes_rd3"], ["Y"], name="rd", keepdims=1)], + "rd3d_kd1", + [_vi("X", TensorProto.FLOAT, [3, 8, 16])], + [_vi("Y", TensorProto.FLOAT, [3, 1, 16])], + [axes_init], + ) + feeds = {"X": _RNG.standard_normal((3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # keepdims must remain 1 on the transformed node + rd_node = next(n for n in tr.graph.node if n.op_type == "ReduceSum") + kd = next((int(a.i) for a in rd_node.attribute if a.name == "keepdims"), 1) + assert kd == 1, f"ReduceSum keepdims must be 1 after transform, got {kd}" + + +def test_reducesum_4d_keepdims0(): + """ReduceSum [1,8,1,16] axes=[3] keepdims=0 → [1,8,1]; keepdims forced to 1, Squeeze added for axis 3.""" + axes_init = _init("axes_rd4", np.array([3], dtype=np.int64)) + model = _make_model( + [helper.make_node("ReduceSum", ["X", "axes_rd4"], ["Y"], name="rd", keepdims=0)], + "rd4d_kd0", + [_vi("X", TensorProto.FLOAT, [1, 8, 1, 16])], + [_vi("Y", TensorProto.FLOAT, [1, 8, 1])], + [axes_init], + ) + feeds = {"X": _RNG.standard_normal((1, 8, 1, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # keepdims must be forced to 1 even for 4D input + rd_node = next(n for n in tr.graph.node if n.op_type == "ReduceSum") + kd = next((int(a.i) for a in rd_node.attribute if a.name == "keepdims"), 1) + assert kd == 1, f"ReduceSum keepdims must be 1 after transform, got {kd}" + # Squeeze must be added to remove the reduced axis (axis 3) + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be added for keepdims=0" + ) + + +def test_argmax_4d_keepdims0(): + """ArgMax [1,8,1,16] axis=3 keepdims=0 → [1,8,1]; keepdims forced to 1, Squeeze added for axis 3.""" + model = _make_model( + [helper.make_node("ArgMax", ["X"], ["Y"], name="am", axis=3, keepdims=0)], + "am4d_kd0", + [_vi("X", TensorProto.FLOAT, [1, 8, 1, 16])], + [_vi("Y", TensorProto.INT64, [1, 8, 1])], + ) + feeds = {"X": _RNG.standard_normal((1, 8, 1, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # keepdims must be forced to 1 even for 4D input + am_node = next(n for n in tr.graph.node if n.op_type == "ArgMax") + kd = next((int(a.i) for a in am_node.attribute if a.name == "keepdims"), 1) + assert kd == 1, f"ArgMax keepdims must be 1 after transform, got {kd}" + # Squeeze must be added to remove the reduced axis (axis 3) + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be added for keepdims=0" + ) + + +# ── Gather ────────────────────────────────────────────────────────────────────────────────────────── + + +def test_gather_2d_data(): + """Gather data[8,16] axis=0 with 1D indices[2] → [2,16]; data unsqueezed to 4D, output squeezed back.""" + # Use 1D indices (shape [2]) → output shape [2,16] + idx = _init("idx", np.array([1, 3], dtype=np.int32)) + model = _make_model( + [helper.make_node("Gather", ["X", "idx"], ["Y"], name="ga", axis=0)], + "ga2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [2, 16])], + [idx], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze must restore output shape" + + +def test_gather_3d_axis1_scalar_index(): + """Gather [1900,1,2] axis=1 scalar-index=0 → [1900,2]. + + The extra unary dim in the 4D Gather output is NOT at the leading position + ([1,1900,1,2] after unsqueeze → axis shifts to 2 → output [1,1900,2]) + so the squeeze axes must be [0,?] determined by _find_squeeze_axes, not just [0,1]. + """ + # scalar index 0 → output drops axis=1 (size 1) → [1900,2] + idx = _init("idx3d", np.array(0, dtype=np.int32)) + model = _make_model( + [helper.make_node("Gather", ["X", "idx3d"], ["Y"], name="ga3d", axis=1)], + "ga3d_ax1", + [_vi("X", TensorProto.FLOAT, [1900, 1, 2])], + [_vi("Y", TensorProto.FLOAT, [1900, 2])], + [idx], + ) + feeds = {"X": _RNG.standard_normal((1900, 1, 2)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + + # Output shape must be [1900,2] — not [1,2] or [1900,1,2] + assert got["Y"].shape == (1900, 2), f"Expected (1900,2), got {got['Y'].shape}" + _assert_match(ref, got) + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze must restore output shape" + + +# ── GatherElements ────────────────────────────────────────────────────────────────────────────────── + + +def test_gatherelements_2d(): + """GatherElements data[4,8] axis=1 indices[4,8] → [4,8]; both expanded to 4D.""" + # indices: gather last 4 of 8 cols for each row + indices_arr = np.tile(np.arange(4, 8, dtype=np.int32), (4, 1)) # [4,4] + # shape must match data: [4,8] + indices_arr = np.concatenate([indices_arr, indices_arr], axis=1) # [4,8] + idx = _init("idx_ge", indices_arr) + model = _make_model( + [helper.make_node("GatherElements", ["X", "idx_ge"], ["Y"], name="ge", axis=1)], + "ge2d", + [_vi("X", TensorProto.FLOAT, [4, 8])], + [_vi("Y", TensorProto.FLOAT, [4, 8])], + [idx], + ) + feeds = {"X": _RNG.standard_normal((4, 8)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze must restore output shape" + + +# ── Expand ────────────────────────────────────────────────────────────────────────────────────────── + + +def test_expand_3d(): + """Expand [3,1,16] to [3,8,16]: shape init updated from [3,8,16] → [1,3,8,16].""" + target = _init("target_sh", np.array([3, 8, 16], dtype=np.int64)) + model = _make_model( + [helper.make_node("Expand", ["X", "target_sh"], ["Y"], name="ex")], + "ex3d", + [_vi("X", TensorProto.FLOAT, [3, 1, 16])], + [_vi("Y", TensorProto.FLOAT, [3, 8, 16])], + [target], + ) + feeds = {"X": _RNG.standard_normal((3, 1, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── Tile ──────────────────────────────────────────────────────────────────────────────────────────── + + +def test_tile_3d(): + """Tile [3,8,16] repeats=[1,2,1] → [3,16,16]: repeats updated from [1,2,1] → [1,1,2,1].""" + repeats = _init("repeats", np.array([1, 2, 1], dtype=np.int64)) + model = _make_model( + [helper.make_node("Tile", ["X", "repeats"], ["Y"], name="tl")], + "tl3d", + [_vi("X", TensorProto.FLOAT, [3, 8, 16])], + [_vi("Y", TensorProto.FLOAT, [3, 16, 16])], + [repeats], + ) + feeds = {"X": _RNG.standard_normal((3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── LpNormalization ───────────────────────────────────────────────────────────────────────────────── + + +def test_lpnorm_2d(): + """LpNormalization [8,16] axis set to -1; input unsqueezed, output squeezed.""" + model = _make_model( + [helper.make_node("LpNormalization", ["X"], ["Y"], name="lpn", axis=1, p=2)], + "lpn2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [8, 16])], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── Softmax ───────────────────────────────────────────────────────────────────────────────────────── + + +def test_softmax_2d(): + """Softmax [8,16] axis=1 → [8,16]; axis adjusted to 3.""" + model = _make_model( + [helper.make_node("Softmax", ["X"], ["Y"], name="sm", axis=1)], + "sm2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [8, 16])], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # Verify axis update + sm_node = next(n for n in tr.graph.node if n.op_type == "Softmax") + axis_val = next(a.i for a in sm_node.attribute if a.name == "axis") + assert axis_val == 3, f"Expected axis=3 after 2D→4D promotion, got {axis_val}" + + +# ── LogSoftmax ────────────────────────────────────────────────────────────────────────────────────── + + +def test_logsoftmax_3d(): + """LogSoftmax [3,8,16] axis=2 → [3,8,16]; axis adjusted to 3.""" + model = _make_model( + [helper.make_node("LogSoftmax", ["X"], ["Y"], name="lsm", axis=2)], + "lsm3d", + [_vi("X", TensorProto.FLOAT, [3, 8, 16])], + [_vi("Y", TensorProto.FLOAT, [3, 8, 16])], + ) + feeds = {"X": _RNG.standard_normal((3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── Concat ────────────────────────────────────────────────────────────────────────────────────────── + + +def test_concat_2d(): + """Concat two [8,16] tensors along axis=1 → [8,32]; axis adjusted to 3.""" + model = _make_model( + [helper.make_node("Concat", ["A", "B"], ["Y"], name="ct", axis=1)], + "ct2d", + [_vi("A", TensorProto.FLOAT, [8, 16]), _vi("B", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [8, 32])], + ) + feeds = { + "A": _RNG.standard_normal((8, 16)).astype(np.float32), + "B": _RNG.standard_normal((8, 16)).astype(np.float32), + } + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── Split ─────────────────────────────────────────────────────────────────────────────────────────── + + +def test_split_3d(): + """Split [3,8,16] along axis=2 into two [3,8,8] tensors; axis adjusted to 3.""" + split_sizes = _init("split_sizes", np.array([8, 8], dtype=np.int64)) + model = _make_model( + [helper.make_node("Split", ["X", "split_sizes"], ["Y1", "Y2"], name="sp", axis=2)], + "sp3d", + [_vi("X", TensorProto.FLOAT, [3, 8, 16])], + [_vi("Y1", TensorProto.FLOAT, [3, 8, 8]), _vi("Y2", TensorProto.FLOAT, [3, 8, 8])], + [split_sizes], + ) + feeds = {"X": _RNG.standard_normal((3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + +# ── InstanceNormalization ─────────────────────────────────────────────────────────────────────────── + + +def test_instancenorm_3d(): + """InstanceNorm [N=2,C=8,D=16]: trailing Unsqueeze(axis=3) before; Squeeze(axis=3) after.""" + n, c, d = 2, 8, 16 + scale = _init("in_scale", np.ones((c,), dtype=np.float32)) + bias = _init("in_bias", np.zeros((c,), dtype=np.float32)) + model = _make_model( + [helper.make_node("InstanceNormalization", ["X", "in_scale", "in_bias"], ["Y"], name="in")], + "in3d", + [_vi("X", TensorProto.FLOAT, [n, c, d])], + [_vi("Y", TensorProto.FLOAT, [n, c, d])], + [scale, bias], + ) + feeds = {"X": _RNG.standard_normal((n, c, d)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + # Verify Unsqueeze and Squeeze for trailing dim were added + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected" + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected" + + +# ── Generic handler (Relu) ────────────────────────────────────────────────────────────────────────── + + +def test_generic_relu_2d(): + """Relu [8,16] → [8,16]: generic handler unsqueezes input, squeezes output.""" + model = _make_model( + [helper.make_node("Relu", ["X"], ["Y"], name="rl")], + "rl2d", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [8, 16])], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected" + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected" + + +# ── 4D input — no transform ───────────────────────────────────────────────────────────────────────── + + +def test_4d_input_no_change(): + """Softmax with 4D input [1,3,8,16]: no Unsqueeze/Squeeze should be added.""" + model = _make_model( + [helper.make_node("Softmax", ["X"], ["Y"], name="sm4d", axis=3)], + "sm4d", + [_vi("X", TensorProto.FLOAT, [1, 3, 8, 16])], + [_vi("Y", TensorProto.FLOAT, [1, 3, 8, 16])], + ) + feeds = {"X": _RNG.standard_normal((1, 3, 8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + assert not any(n.op_type in ("Unsqueeze", "Squeeze") for n in tr.graph.node), ( + "No Unsqueeze/Squeeze should be added for 4D input" + ) + + +# ── MatMul — in _OP_SKIP ──────────────────────────────────────────────────────────────────────────── + + +def test_matmul_skipped(): + """MatMul in _OP_SKIP: Relu (2D) is transformed but MatMul is left unchanged.""" + w = _RNG.standard_normal((16, 12)).astype(np.float32) + model = _make_model( + [ + helper.make_node("Relu", ["X"], ["X_relu"], name="rl"), + helper.make_node("MatMul", ["X_relu", "W"], ["Y"], name="mm"), + ], + "mm_skip", + [_vi("X", TensorProto.FLOAT, [8, 16])], + [_vi("Y", TensorProto.FLOAT, [8, 12])], + [_init("W", w)], + ) + feeds = {"X": _RNG.standard_normal((8, 16)).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + mm_nodes = [n for n in tr.graph.node if n.op_type == "MatMul"] + assert len(mm_nodes) == 1, "MatMul must remain unchanged" + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Relu should have been wrapped" + + +# ── 5D Transpose collapse ─────────────────────────────────────────────────────────────────────────── + + +def test_transpose_5d_collapse(): + """Transpose [2,1,3,4,5] perm=[0,2,1,3,4] collapses to 4D with ORT parity. + + Algorithm: + current_shape = [orig_shape[p] for p in perm] = [2,3,1,4,5] + idx_to_remove = 2 (leftmost unary dim in the output) + removed_orig_dim = perm[2] = 1 + new_perm_raw = [0,2,3,4] → shift down values >=1 → [0,1,2,3] + + Transform: Squeeze(input, axis=1) → 4D → Transpose(perm=[0,1,2,3]) → + Unsqueeze(axis=2) → restores 5D output [2,3,1,4,5]. + ORT parity holds because the Unsqueeze restores the original output shape. + """ + perm = [0, 2, 1, 3, 4] + input_shape = [2, 1, 3, 4, 5] + output_shape_perm = [input_shape[p] for p in perm] # [2,3,1,4,5] + model = _make_model( + [helper.make_node("Transpose", ["X"], ["Y"], name="tr5d", perm=perm)], + "tr5d", + [_vi("X", TensorProto.FLOAT, input_shape)], + [_vi("Y", TensorProto.FLOAT, output_shape_perm)], + ) + feeds = {"X": _RNG.standard_normal(input_shape).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + _assert_match(ref, got) + + tr_nodes = list(tr.graph.node) + # Squeeze must be added to remove unary input dim, Unsqueeze to restore output shape + assert any(n.op_type == "Squeeze" for n in tr_nodes), ( + "Squeeze must be added to collapse unary input dim" + ) + assert any(n.op_type == "Unsqueeze" for n in tr_nodes), ( + "Unsqueeze must be added to restore 5D output shape" + ) + # Transpose must have the collapsed 4-element perm [0,1,2,3] + tr_node = next((n for n in tr_nodes if n.op_type == "Transpose"), None) + assert tr_node is not None, "Transpose must still exist after collapse" + perm_attr = next((a for a in tr_node.attribute if a.name == "perm"), None) + assert perm_attr is not None, "Transpose must have a perm attribute" + assert list(perm_attr.ints) == [0, 1, 2, 3], ( + f"Expected collapsed perm [0,1,2,3], got {list(perm_attr.ints)}" + ) + + +# ── Graph-output squeeze split ──────────────────────────────────────────────── +# +# When a node's output is BOTH a graph output AND an input to other nodes, +# _split_graph_output_squeezes creates two separate Squeeze nodes: +# • Squeeze → (preserves the graph-output binding) +# • Squeeze → _sq (for internal compute consumers) +# A value_info entry is added for _sq; consumers are rewired. +# ───────────────────────────────────────────────────────────────────────────── + + +def test_graph_output_split_single_internal_consumer(): + """Y is a graph output AND consumed by Relu → two Squeeze nodes, ORT parity. + + Model: X [1,3,4] → Add(X,X) → Y [1,3,4] ← graph output + Y → Relu(Y) → Z [1,3,4] ← graph output + """ + shape = [1, 3, 4] + model = _make_model( + [ + helper.make_node("Add", ["X", "X"], ["Y"], name="add"), + helper.make_node("Relu", ["Y"], ["Z"], name="relu"), + ], + "g", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape), _vi("Z", TensorProto.FLOAT, shape)], + ) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_convert_ops_to_4d(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + # Y must still be a graph output + go_names = {o.name for o in tr.graph.output} + assert "Y" in go_names, "Y must remain a graph output" + + # Exactly one Squeeze producing Y (graph output binding) + sq_y = [n for n in tr.graph.node if n.op_type == "Squeeze" and n.output[0] == "Y"] + assert len(sq_y) == 1, f"Expected 1 Squeeze → Y, got {len(sq_y)}" + + # Exactly one Squeeze producing Y_sq (internal consumers) + sq_y_sq = [n for n in tr.graph.node if n.op_type == "Squeeze" and n.output[0] == "Y_sq"] + assert len(sq_y_sq) == 1, f"Expected 1 Squeeze → Y_sq, got {len(sq_y_sq)}" + + # Both squeezes must share the same 4D input (Add's promoted output) + assert sq_y[0].input[0] == sq_y_sq[0].input[0], "Both squeezes must consume the same 4D tensor" + + # value_info must be registered for Y_sq + vi_names = {vi.name for vi in tr.graph.value_info} + assert "Y_sq" in vi_names, "value_info must be added for Y_sq" + + # No internal compute node may directly consume Y + for node in tr.graph.node: + for inp in node.input: + assert inp != "Y", f"Internal node {node.name!r} must not consume Y directly; use Y_sq" + + _assert_match(ref, got) + + +def test_graph_output_only_no_split(): + """Y is a graph output only (no internal consumer) → single Squeeze, no Y_sq created.""" + shape = [1, 3, 4] + model = _make_model( + [helper.make_node("Add", ["X", "X"], ["Y"], name="add")], + "g", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_convert_ops_to_4d(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + sq_y = [n for n in tr.graph.node if n.op_type == "Squeeze" and n.output[0] == "Y"] + assert len(sq_y) == 1, "Exactly one Squeeze → Y when Y has no internal consumers" + + sq_y_sq = [n for n in tr.graph.node if n.op_type == "Squeeze" and n.output[0] == "Y_sq"] + assert not sq_y_sq, "No Y_sq Squeeze must be created when Y is only a graph output" + + _assert_match(ref, got) + + +def test_graph_output_split_multiple_internal_consumers(): + """Y is a graph output AND consumed by two independent nodes → single Y_sq shared, ORT parity. + + Model: X [1,3,4] → Add(X,X) → Y (graph output) + Y → Relu(Y) → Z1 (graph output) + Y → Abs(Y) → Z2 (graph output) + """ + shape = [1, 3, 4] + model = _make_model( + [ + helper.make_node("Add", ["X", "X"], ["Y"], name="add"), + helper.make_node("Relu", ["Y"], ["Z1"], name="relu"), + helper.make_node("Abs", ["Y"], ["Z2"], name="absn"), + ], + "g", + [_vi("X", TensorProto.FLOAT, shape)], + [ + _vi("Y", TensorProto.FLOAT, shape), + _vi("Z1", TensorProto.FLOAT, shape), + _vi("Z2", TensorProto.FLOAT, shape), + ], + ) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_convert_ops_to_4d(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + # Y must remain a graph output; Y_sq for internal nodes + go_names = {o.name for o in tr.graph.output} + assert "Y" in go_names + + sq_y_sq = [n for n in tr.graph.node if n.op_type == "Squeeze" and n.output[0] == "Y_sq"] + assert sq_y_sq, "Y_sq Squeeze must exist for internal consumers" + + # No compute node may directly consume Y + for node in tr.graph.node: + for inp in node.input: + assert inp != "Y", f"Node {node.name!r} must not consume Y directly after split" + + _assert_match(ref, got) + + +# ── LayerNormalization ──────────────────────────────────────────────────────── + + +def test_layernorm_3d_data_input_promoted(): + """LayerNorm [2,4,8] → data unsqueezed to 4D; scale/bias kept 1D; output squeezed back; ORT parity.""" + n, c, d = 2, 4, 8 + scale = numpy_helper.from_array(np.ones(d, dtype=np.float32), name="scale") + bias = numpy_helper.from_array(np.zeros(d, dtype=np.float32), name="bias") + model = _make_model( + [helper.make_node("LayerNormalization", ["X", "scale", "bias"], ["Y"], name="ln", axis=-1)], + "ln3d", + [_vi("X", TensorProto.FLOAT, [n, c, d])], + [_vi("Y", TensorProto.FLOAT, [n, c, d])], + [scale, bias], + ) + feeds = {"X": _RNG.standard_normal((n, c, d)).astype(np.float32)} + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_convert_ops_to_4d(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + ln_node = next(n for n in tr.graph.node if n.op_type == "LayerNormalization") + + # Data input must be 4D (Unsqueeze inserted before LN) + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( + "Unsqueeze must be inserted to promote data input to 4D" + ) + unsq_out = next(n for n in tr.graph.node if n.op_type == "Unsqueeze").output[0] + assert ln_node.input[0] == unsq_out, "LN must read the 4D unsqueezed tensor" + + # Scale and bias must be unchanged (still the original 1D initializer names) + assert ln_node.input[1] == "scale", "scale must remain unchanged" + assert ln_node.input[2] == "bias", "bias must remain unchanged" + scale_init = next(i for i in tr.graph.initializer if i.name == "scale") + bias_init = next(i for i in tr.graph.initializer if i.name == "bias") + assert list(scale_init.dims) == [d], f"scale must stay 1D [{d}], got {list(scale_init.dims)}" + assert list(bias_init.dims) == [d], f"bias must stay 1D [{d}], got {list(bias_init.dims)}" + + # Squeeze inserted after LN to restore original rank + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be inserted to restore output rank" + ) + + _assert_match(ref, got) + + +def test_layernorm_2d_data_input_promoted(): + """LayerNorm [4,8] → data unsqueezed to 4D; scale/bias kept 1D; output squeezed back; ORT parity.""" + c, d = 4, 8 + scale = numpy_helper.from_array(np.ones(d, dtype=np.float32), name="scale2d") + bias = numpy_helper.from_array(np.zeros(d, dtype=np.float32), name="bias2d") + model = _make_model( + [ + helper.make_node( + "LayerNormalization", ["X", "scale2d", "bias2d"], ["Y"], name="ln2d", axis=-1 + ) + ], + "ln2d", + [_vi("X", TensorProto.FLOAT, [c, d])], + [_vi("Y", TensorProto.FLOAT, [c, d])], + [scale, bias], + ) + feeds = {"X": _RNG.standard_normal((c, d)).astype(np.float32)} + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_convert_ops_to_4d(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + ln_node = next(n for n in tr.graph.node if n.op_type == "LayerNormalization") + + # scale and bias must remain 1D + scale_init = next(i for i in tr.graph.initializer if i.name == "scale2d") + bias_init = next(i for i in tr.graph.initializer if i.name == "bias2d") + assert list(scale_init.dims) == [d] + assert list(bias_init.dims) == [d] + assert ln_node.input[1] == "scale2d" + assert ln_node.input[2] == "bias2d" + + _assert_match(ref, got) + + +# ── Flatten (handled inside convert_ops_to_4d via _handle_flatten) ─────────── +# +# Flatten(X, axis=a) → Unsqueeze(to 4D) → Reshape([1,1,outer,inner]) → Squeeze([outer,inner]) +# Input already 4D: Unsqueeze omitted. +# ───────────────────────────────────────────────────────────────────────────── + + +def _assert_flatten_chain(tr, *, expect_unsqueeze: bool) -> None: + """Check the expected structural Flatten conversion chain.""" + op_types = [n.op_type for n in tr.graph.node] + assert "Flatten" not in op_types, "Flatten must be removed" + assert "Reshape" in op_types, "Reshape must be present" + assert "Squeeze" in op_types, "Squeeze must be present" + if expect_unsqueeze: + assert "Unsqueeze" in op_types, "Unsqueeze must be present for non-4D input" + else: + assert "Unsqueeze" not in op_types, "Unsqueeze must NOT be inserted for 4D input" + + +@pytest.mark.parametrize( + ("shape", "axis", "output_shape"), + [ + ([1, 3, 4], 2, [3, 4]), + ([8, 16], 0, [1, 128]), + ([8, 16], 1, [8, 16]), + ([3, 8, 16], 1, [3, 128]), + ([3, 8, 16], -1, [24, 16]), + ([1, 3, 8, 16], 1, [1, 384]), + ], +) +def test_flatten_native_axes(shape, axis, output_shape): + """Execute Flatten across native ranks and verify the rank adapter chain.""" + model = _make_model( + [helper.make_node("Flatten", ["X"], ["Y"], name="fl", axis=axis)], + "flatten", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, output_shape)], + ) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + ref, got, tr = _apply_and_run(model, feeds) + assert got["Y"].shape == tuple(output_shape) + _assert_flatten_chain(tr, expect_unsqueeze=len(shape) < 4) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-6) + + +def test_five_dimensional_transpose_without_singleton_is_unchanged(): + """Skip a valid permutation whose native layout cannot collapse to four dimensions.""" + shape = [2, 3, 4, 5, 6] + model = _make_model( + [helper.make_node("Transpose", ["X"], ["Y"], perm=[4, 3, 2, 1, 0])], + "no_unary", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape[::-1])], + ) + before = model.SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.SerializeToString() == before + feeds = {"X": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + np.testing.assert_array_equal(_run_ort(converted, feeds)["Y"], np.transpose(feeds["X"])) + + +@pytest.mark.parametrize("target", [[0, -1], [0, 0, -1]]) +def test_unresolved_reshape_copied_dimensions_preserve_intermediate_layout(target): + """Leave copied-zero targets native when intermediate shape metadata is absent.""" + output_shape = [2, 12] if len(target) == 2 else [2, 3, 4] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Identity", ["X"], ["hidden_input"]), + helper.make_node("Reshape", ["hidden_input", "target"], ["hidden_output"]), + helper.make_node("Identity", ["hidden_output"], ["Y"]), + ], + "missing_metadata", + [_vi("X", TensorProto.FLOAT, [2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, output_shape)], + [_init("target", np.asarray(target, np.int64))], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.SerializeToString() == before + feeds = {"X": np.arange(24, dtype=np.float32).reshape(2, 3, 4)} + np.testing.assert_array_equal(_run_ort(converted, feeds)["Y"], feeds["X"].reshape(output_shape)) + + +@pytest.mark.parametrize("owner", ["private", "public", "peer", "capture"]) +def test_gather_cast_preserves_other_owners_and_updates_private_metadata(owner): + """Retarget only exclusive standard Cast outputs and preserve shared integer tensors.""" + nodes = [ + helper.make_node("Cast", ["source"], ["indices"], to=TensorProto.INT64), + helper.make_node("Gather", ["X", "indices"], ["Y"], axis=2), + ] + outputs = [_vi("Y", TensorProto.FLOAT, [1, 1, 2, 4])] + inits = [] + if owner == "public": + outputs.append(_vi("indices", TensorProto.INT64, [2])) + elif owner == "peer": + nodes.append(helper.make_node("Identity", ["indices"], ["other"])) + outputs.append(_vi("other", TensorProto.INT64, [2])) + elif owner == "capture": + branch = helper.make_graph( + [helper.make_node("Identity", ["indices"], ["out"])], + "branch", + [], + [_vi("out", TensorProto.INT64, [2])], + ) + inits.append(_init("condition", np.asarray(True))) + nodes.append( + helper.make_node("If", ["condition"], ["other"], then_branch=branch, else_branch=branch) + ) + outputs.append(_vi("other", TensorProto.INT64, [2])) + model = helper.make_model( + helper.make_graph( + nodes, + "cast_owner", + [_vi("X", TensorProto.FLOAT, [1, 1, 5, 4]), _vi("source", TensorProto.FLOAT, [None])], + outputs, + inits, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + model.graph.value_info.extend([_vi("indices", TensorProto.INT64, [2])]) + onnx.checker.check_model(model, full_check=True) + feeds = { + "X": np.arange(20, dtype=np.float32).reshape(1, 1, 5, 4), + "source": np.asarray([1, 3], np.float32), + } + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + cast = next(n for n in converted.graph.node if n.op_type == "Cast" and n.output[0] == "indices") + expected_type = TensorProto.INT32 if owner == "private" else TensorProto.INT64 + assert next(a.i for a in cast.attribute if a.name == "to") == expected_type + assert ( + next( + v for v in converted.graph.value_info if v.name == "indices" + ).type.tensor_type.elem_type + == expected_type + ) + + +@pytest.mark.parametrize("axis", [None, 0, -1]) +def test_fixed_dequantization_restores_public_and_convolution_bias_ranks(axis): + """Preserve native DQ outputs needed by public values and rank-one Conv biases.""" + scale = np.asarray(0.25, np.float32) if axis is None else np.asarray([0.25, 0.5], np.float32) + model = _make_model( + [ + helper.make_node( + "DequantizeLinear", + ["quantized", "scale", "zero"], + ["bias"], + **({} if axis is None else {"axis": axis}), + ), + helper.make_node("Conv", ["X", "weight", "bias"], ["Y"]), + ], + "dq_bias", + [_vi("X", TensorProto.FLOAT, [1, 1, 2, 2])], + [_vi("bias", TensorProto.FLOAT, [2]), _vi("Y", TensorProto.FLOAT, [1, 2, 2, 2])], + [ + _init("quantized", np.asarray([2, 4], np.int8)), + _init("scale", scale), + _init("zero", np.zeros(scale.shape, np.int8)), + _init("weight", np.ones((2, 1, 1, 1), np.float32)), + ], + ) + feeds = {"X": np.arange(4, dtype=np.float32).reshape(1, 1, 2, 2)} + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + assert ( + next(n for n in converted.graph.node if n.op_type == "DequantizeLinear").output[0] != "bias" + ) + + +def test_blocked_dequantization_layout_is_unchanged(): + """Leave blocked quantization native because its scale layout is rank-sensitive.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", ["data", "scale", "zero"], ["Y"], axis=1, block_size=2 + ) + ], + "blocked", + [], + [_vi("Y", TensorProto.FLOAT, [2, 4])], + [ + _init("data", np.arange(8, dtype=np.int8).reshape(2, 4)), + _init("scale", np.full((2, 2), 0.25, np.float32)), + _init("zero", np.zeros((2, 2), np.int8)), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + assert _apply_convert_ops_to_4d(model).SerializeToString() == before + + +@pytest.mark.parametrize("storage", ["raw", "typed"]) +def test_fixed_weight_promotion_preserves_payload_representation_and_metadata(storage): + """Change only dimension metadata when promoting a fixed elementwise weight.""" + data = np.arange(6, dtype=np.float32).reshape(2, 3) + tensor = ( + _init("weight", data) + if storage == "raw" + else helper.make_tensor("weight", TensorProto.FLOAT, [2, 3], data.reshape(-1)) + ) + tensor.doc_string = "Preserve caller metadata" + model = _make_model( + [helper.make_node("Add", ["X", "weight"], ["Y"])], + "weight_metadata", + [_vi("X", TensorProto.FLOAT, [2, 3])], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + [tensor], + ) + expected = copy.deepcopy(model.graph.initializer[0]) + expected.ClearField("dims") + converted = _apply_convert_ops_to_4d(model) + actual = copy.deepcopy(next(t for t in converted.graph.initializer if t.name == "weight")) + actual.ClearField("dims") + assert actual.SerializeToString() == expected.SerializeToString() + np.testing.assert_array_equal(_run_ort(converted, {"X": data})["Y"], data * 2) + + +@pytest.mark.parametrize( + ("operation", "mixed"), [("Relu", False), ("Add", False), ("Add", True), ("Mul", True)] +) +def test_scalar_inputs_preserve_native_values_with_optional_tensor_broadcast(operation, mixed): + """Choose a nonscalar data layout and leave entirely scalar operations native.""" + inputs = [_vi("scalar", TensorProto.FLOAT, [])] + node_inputs = ["scalar"] + feeds = {"scalar": np.asarray(2, np.float32)} + shape = [2, 3] if mixed else [] + if operation != "Relu": + inputs.append(_vi("other", TensorProto.FLOAT, shape)) + node_inputs.append("other") + feeds["other"] = ( + np.arange(6, dtype=np.float32).reshape(shape) if mixed else np.asarray(3, np.float32) + ) + model = _make_model( + [helper.make_node(operation, node_inputs, ["Y"])], + "scalar_layout", + inputs, + [_vi("Y", TensorProto.FLOAT, shape)], + ) + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize( + ("shape", "dtype", "opset"), [([], np.float32, 17), ([2, 3], np.float16, 19)] +) +def test_fixed_dequantization_preserves_scalar_and_half_precision_outputs(shape, dtype, opset): + """Restore native DQ ranks and keep the scale-selected output dtype.""" + data = np.asarray(4, np.int8) if not shape else np.arange(6, dtype=np.int8).reshape(shape) + model = helper.make_model( + helper.make_graph( + [helper.make_node("DequantizeLinear", ["data", "scale", "zero"], ["Y"])], + "dq_output_types", + [], + [_vi("Y", TensorProto.FLOAT if dtype == np.float32 else TensorProto.FLOAT16, shape)], + [ + _init("data", data), + _init("scale", np.asarray(0.25, dtype)), + _init("zero", np.asarray(0, np.int8)), + ], + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + _clamp_ir(model) + expected = _run_ort(model, {}) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, {})) + + +@pytest.mark.parametrize("storage", ["typed", "external"]) +def test_gather_elements_fixed_index_storage_is_preserved(tmp_path, storage): + """Use runtime index views without decoding unloaded sidecars or replacing tensor metadata.""" + indices = np.array([[0, 2, 1], [1, 0, 2]], np.int64) + tensor = ( + _init("indices", indices) + if storage == "external" + else helper.make_tensor("indices", TensorProto.INT64, [2, 3], indices.reshape(-1).tolist()) + ) + tensor.doc_string = "native index owner" + model = _make_model( + [helper.make_node("GatherElements", ["x", "indices"], ["y"], axis=1)], + "index_storage", + [_vi("x", TensorProto.FLOAT, [2, 3])], + [_vi("y", TensorProto.FLOAT, [2, 3])], + [tensor], + ) + native = tmp_path / "native.onnx" + if storage == "external": + onnx.save_model( + model, + str(native), + save_as_external_data=True, + all_tensors_to_one_file=True, + location="indices.bin", + size_threshold=0, + ) + else: + onnx.save_model(model, str(native)) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = ort.InferenceSession(str(native), providers=["CPUExecutionProvider"]).run( + None, feeds + )[0] + unloaded = onnx.load(str(native), load_external_data=False) + original = unloaded.graph.initializer[0].SerializeToString() + converted = _apply_convert_ops_to_4d(unloaded) + assert ( + next(t for t in converted.graph.initializer if t.name == "indices").SerializeToString() + == original + ) + target = tmp_path / "converted.onnx" + onnx.save_model(converted, str(target)) + onnx.checker.check_model(str(target), full_check=True) + actual = ort.InferenceSession(str(target), providers=["CPUExecutionProvider"]).run(None, feeds)[ + 0 + ] + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("operation", ["Gather", "Reshape", "Slice", "ReduceSum", "Expand", "Tile"]) +@pytest.mark.parametrize("decoy", [False, True]) +def test_unloaded_external_controls_preserve_model_directory( + tmp_path, monkeypatch, operation, decoy +): + """Never resolve unloaded control payloads against the process working directory.""" + specifications = { + "Gather": (["indices"], {"indices": [2, 0]}, [2, 2], {"axis": 1}), + "Reshape": (["shape"], {"shape": [3, 2]}, [3, 2], {}), + "Slice": ( + ["starts", "ends", "axes"], + {"starts": [0], "ends": [2], "axes": [1]}, + [2, 2], + {}, + ), + "ReduceSum": (["axes"], {"axes": [1]}, [2, 1], {"keepdims": 1}), + "Expand": (["shape"], {"shape": [2, 3]}, [2, 3], {}), + "Tile": (["repeats"], {"repeats": [1, 2]}, [2, 6], {}), + } + inputs, controls, output, attributes = specifications[operation] + model = _make_model( + [helper.make_node(operation, ["x", *inputs], ["y"], **attributes)], + "external_controls", + [_vi("x", TensorProto.FLOAT, [2, 3])], + [_vi("y", TensorProto.FLOAT, output)], + [_init(name, np.asarray(values, np.int64)) for name, values in controls.items()], + ) + native = tmp_path / "native.onnx" + onnx.save_model( + model, + str(native), + save_as_external_data=True, + all_tensors_to_one_file=True, + location="controls.bin", + size_threshold=0, + ) + payload = (tmp_path / "controls.bin").read_bytes() + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _run_ort(onnx.load(str(native)), feeds)["y"] + unloaded = onnx.load(str(native), load_external_data=False) + originals = {tensor.name: tensor.SerializeToString() for tensor in unloaded.graph.initializer} + working = tmp_path / "unrelated" + working.mkdir() + if decoy: + (working / "controls.bin").write_bytes(bytes(len(payload))) + monkeypatch.chdir(working) + converted = _apply_convert_ops_to_4d(unloaded) + assert { + tensor.name: tensor.SerializeToString() + for tensor in converted.graph.initializer + if tensor.name in originals + } == originals + if operation == "Gather": + assert any(node.op_type == "Cast" for node in converted.graph.node) + else: + assert len(converted.graph.node) == 1 + target = tmp_path / "converted.onnx" + onnx.save_model(converted, str(target)) + loaded = onnx.load(str(target)) + onnx.checker.check_model(loaded, full_check=True) + actual = _run_ort(loaded, feeds)["y"] + np.testing.assert_array_equal(actual, expected) + assert (tmp_path / "controls.bin").read_bytes() == payload + + +def test_external_weight_promotion_preserves_sidecar_bytes_and_relative_metadata(tmp_path): + """Promote unloaded external weights without reading or rewriting their payload sidecar.""" + weight = np.arange(6, dtype=np.float32).reshape(2, 3) + model = _make_model( + [helper.make_node("Add", ["X", "weight"], ["Y"])], + "external_weight", + [_vi("X", TensorProto.FLOAT, [2, 3])], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + [_init("weight", weight)], + ) + native = tmp_path / "native.onnx" + onnx.save_model( + model, + str(native), + save_as_external_data=True, + all_tensors_to_one_file=True, + location="weights.bin", + size_threshold=0, + ) + payload = (tmp_path / "weights.bin").read_bytes() + feeds = {"X": weight} + expected = ort.InferenceSession(str(native), providers=["CPUExecutionProvider"]).run( + None, feeds + )[0] + unloaded = onnx.load(str(native), load_external_data=False) + converted = _apply_convert_ops_to_4d(unloaded) + target = tmp_path / "converted.onnx" + onnx.save_model(converted, str(target)) + onnx.checker.check_model(str(target), full_check=True) + actual = ort.InferenceSession(str(target), providers=["CPUExecutionProvider"]).run(None, feeds)[ + 0 + ] + np.testing.assert_array_equal(actual, expected) + assert (tmp_path / "weights.bin").read_bytes() == payload + tensor = next(t for t in converted.graph.initializer if t.name == "weight") + assert tensor.data_location == TensorProto.EXTERNAL and not tensor.HasField("raw_data") + assert {entry.key: entry.value for entry in tensor.external_data}["location"] == "weights.bin" + + +@pytest.mark.parametrize("operation", ["Softmax", "Add", "DequantizeLinear"]) +@pytest.mark.parametrize("shared", [False, True]) +@pytest.mark.parametrize("legacy", [False, True]) +def test_promoted_initializers_refresh_value_info_without_changing_native_owners( + operation, shared, legacy +): + """Refresh promoted initializer annotations and keep shared native owners unchanged.""" + quantized = operation == "DequantizeLinear" + values = np.arange(6, dtype=np.int8 if quantized else np.float32).reshape(2, 3) + initializers = [_init("data", values)] + if quantized: + initializers.extend( + [_init("scale", np.asarray(0.25, np.float32)), _init("zero", np.asarray(0, np.int8))] + ) + nodes, outputs = [], [] + for index in range(2 if shared else 1): + output = f"y{index}" + nodes.append( + helper.make_node( + operation, + ["data", "scale", "zero"] + if quantized + else (["x", "data"] if operation == "Add" else ["data"]), + [output], + **({"axis": index} if operation == "Softmax" else {}), + ) + ) + outputs.append(_vi(output, TensorProto.FLOAT, [2, 3])) + inputs = [_vi("x", TensorProto.FLOAT, [2, 3])] if operation == "Add" else [] + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3) / 2} if inputs else {} + model = _make_model(nodes, "initializer_annotations", inputs, outputs, initializers) + model.graph.value_info.append( + _vi("data", TensorProto.INT8 if quantized else TensorProto.FLOAT, [2, 3]) + ) + if legacy: + model.ir_version = 3 + model.graph.input.extend( + _vi(tensor.name, tensor.data_type, list(tensor.dims)) + for tensor in model.graph.initializer + ) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + promoted = next(tensor for tensor in converted.graph.initializer if tensor.name == "data") + assert list(promoted.dims) == ([2, 3] if shared or legacy else [1, 1, 2, 3]) + annotation = next(value for value in converted.graph.value_info if value.name == "data") + assert [dim.dim_value for dim in annotation.type.tensor_type.shape.dim] == list(promoted.dims) + if legacy: + declaration = next(value for value in converted.graph.input if value.name == "data") + assert [dim.dim_value for dim in declaration.type.tensor_type.shape.dim] == list( + promoted.dims + ) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("keepdims", [0, 1]) +def test_scalar_reduction_preserves_native_rank(keepdims): + """Promote a scalar reduction internally and restore its zero-dimensional output.""" + model = _make_model( + [helper.make_node("ReduceSum", ["X"], ["Y"], keepdims=keepdims)], + "scalar_reduce", + [_vi("X", TensorProto.FLOAT, [])], + [_vi("Y", TensorProto.FLOAT, [])], + ) + feeds = {"X": np.asarray(3, np.float32)} + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("axes", [[0], [0, 1]]) +def test_runtime_reduction_axes_remain_native(axes): + """Leave unsupported runtime axes unchanged instead of reducing all promoted dimensions.""" + # The optional symbolic inferencer cannot handle runtime axes with keepdims=0. + model = helper.make_model( + helper.make_graph( + [helper.make_node("ReduceSum", ["X", "axes"], ["Y"], keepdims=0)], + "runtime_reduce", + [_vi("X", TensorProto.FLOAT, [1, 1, 1]), _vi("axes", TensorProto.INT64, [len(axes)])], + [_vi("Y", TensorProto.FLOAT, [1] * (3 - len(axes)))], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + feeds = {"X": np.asarray([[[3]]], np.float32), "axes": np.asarray(axes, np.int64)} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.SerializeToString() == before + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("owner", ["default", "capture"]) +def test_fixed_dequantization_preserves_runtime_and_captured_data_owners(owner): + """Keep native quantized layouts owned by callers or nested graphs.""" + data = np.arange(6, dtype=np.int8).reshape(2, 3) + inputs = [_vi("data", TensorProto.INT8, [2, 3])] if owner == "default" else [] + nodes = [helper.make_node("DequantizeLinear", ["data", "scale", "zero"], ["Y"])] + outputs = [_vi("Y", TensorProto.FLOAT, [2, 3])] + inits = [ + _init("data", data), + _init("scale", np.asarray(0.25, np.float32)), + _init("zero", np.asarray(0, np.int8)), + ] + if owner == "capture": + branch = helper.make_graph( + [helper.make_node("Identity", ["data"], ["out"])], + "branch", + [], + [_vi("out", TensorProto.INT8, [2, 3])], + ) + inits.append(_init("condition", np.asarray(True))) + nodes.append( + helper.make_node("If", ["condition"], ["view"], then_branch=branch, else_branch=branch) + ) + outputs.append(_vi("view", TensorProto.INT8, [2, 3])) + model = _make_model(nodes, "dq_owners", inputs, outputs, inits) + feeds = {"data": data + np.int8(10)} if owner == "default" else {} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.SerializeToString() == before + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("operation", ["Transpose", "Gather", "DequantizeLinear"]) +def test_standalone_converter_preserves_opaque_layout_functions(operation): + """Leave opaque functions intact even when their names match supported handlers.""" + attrs = {"perm": [1, 0]} if operation == "Transpose" else {} + inputs = [] if operation == "DequantizeLinear" else [_vi("X", TensorProto.FLOAT, [2, 3])] + inits = [_init("X", np.arange(6, dtype=np.float32).reshape(2, 3))] if not inputs else [] + function = helper.make_function( + "opaque", + operation, + ["x"], + ["y"], + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + attributes=list(attrs), + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["X"], ["Y"], domain="opaque", **attrs)], + "opaque_layout", + inputs, + [_vi("Y", TensorProto.FLOAT, [2, 3])], + inits, + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("opaque", 1)], + functions=[function], + ) + _clamp_ir(model) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"X": np.arange(6, dtype=np.float32).reshape(2, 3)} if inputs else {} + expected = ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.SerializeToString() == before + onnx.checker.check_model(converted, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(converted).run(None, feeds)[0], expected[0]) + + +def test_gather_does_not_retarget_an_opaque_cast_producer(): + """Prepare a standard private Cast while preserving the opaque producer's attributes.""" + function = helper.make_function( + "opaque", + "Cast", + ["x"], + ["y"], + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + attributes=["to"], + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Cast", ["source"], ["indices"], domain="opaque", to=TensorProto.INT64 + ), + helper.make_node("Gather", ["X", "indices"], ["Y"], axis=2), + ], + "opaque_cast", + [_vi("source", TensorProto.INT64, [2]), _vi("X", TensorProto.FLOAT, [1, 1, 5, 4])], + [_vi("Y", TensorProto.FLOAT, [1, 1, 2, 4]), _vi("indices", TensorProto.INT64, [2])], + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("opaque", 1)], + functions=[function], + ) + _clamp_ir(model) + model = onnx.shape_inference.infer_shapes(model) + feeds = { + "source": np.asarray([1, 3], np.int64), + "X": np.arange(20, dtype=np.float32).reshape(1, 1, 5, 4), + } + expected = ReferenceEvaluator(model).run(None, feeds) + producer = model.graph.node[0].SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.graph.node[0].SerializeToString() == producer + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip(ReferenceEvaluator(converted).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("omit_mean", [False, True]) +def test_layernorm_restores_optional_statistics_and_shared_public_views(omit_mean): + """Restore every optional normalization output and its internal consumers to native rank.""" + names = ["Y", "" if omit_mean else "Mean", "Inv"] + statistic = "Inv" if omit_mean else "Mean" + outputs = [ + _vi("Y", TensorProto.FLOAT, [2, 3, 4]), + _vi("Inv", TensorProto.FLOAT, [2, 1, 1]), + _vi("view", TensorProto.FLOAT, [2, 1, 1]), + ] + if not omit_mean: + outputs.insert(1, _vi("Mean", TensorProto.FLOAT, [2, 1, 1])) + model = _make_model( + [ + helper.make_node("LayerNormalization", ["X", "scale", "bias"], names, axis=1), + helper.make_node("Add", [statistic, "zero"], ["view"]), + ], + "layernorm_statistics", + [_vi("X", TensorProto.FLOAT, [2, 3, 4])], + outputs, + [ + _init("scale", np.ones((3, 4), np.float32)), + _init("bias", np.zeros((3, 4), np.float32)), + _init("zero", np.asarray(0, np.float32)), + ], + # The optional inferencer cannot handle an omitted middle output. + infer_symbolic_shapes=not omit_mean, + ) + feeds = {"X": np.random.default_rng(31).standard_normal((2, 3, 4)).astype(np.float32)} + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +def test_generated_rank_adapters_preserve_existing_control_value_and_node_names(): + """Allocate around simultaneous control, tensor and node-name collisions.""" + model = _make_model( + [ + helper.make_node("Relu", ["X"], ["Y"], name="relu"), + helper.make_node("Identity", ["other"], ["relu_unsq4d_0"], name="relu_sq_node_1"), + ], + "adapter_names", + [_vi("X", TensorProto.FLOAT, [2, 3]), _vi("other", TensorProto.FLOAT, [1])], + [ + _vi("Y", TensorProto.FLOAT, [2, 3]), + _vi("relu_unsq4d_0", TensorProto.FLOAT, [1]), + _vi("relu_unsq_axes_0", TensorProto.INT64, [1]), + ], + [_init("relu_unsq_axes_0", np.asarray([99], np.int64))], + ) + feeds = { + "X": np.arange(6, dtype=np.float32).reshape(2, 3) - 3, + "other": np.asarray([17], np.float32), + } + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("collision", [False, True]) +def test_shared_squeeze_preserves_partial_rank_metadata_and_existing_aliases(collision): + """Copy the native partial shape while allocating private consumer aliases.""" + nodes = [ + helper.make_node("Squeeze", ["X", "axes"], ["Y"], name="sq"), + helper.make_node("Identity", ["Y"], ["Z"]), + ] + inputs = [_vi("X", TensorProto.FLOAT, [1, None, 3, None])] + outputs = [ + _vi("Y", TensorProto.FLOAT, [None, 3, None]), + _vi("Z", TensorProto.FLOAT, [None, 3, None]), + ] + feeds = {"X": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + if collision: + nodes.append(helper.make_node("Identity", ["other"], ["Y_sq"], name="sq_internal")) + inputs.append(_vi("other", TensorProto.FLOAT, [1])) + outputs.append(_vi("Y_sq", TensorProto.FLOAT, [1])) + feeds["other"] = np.asarray([17], np.float32) + model = helper.make_model( + helper.make_graph( + nodes, "partial_squeeze", inputs, outputs, [_init("axes", np.asarray([0], np.int64))] + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + alias = next( + n.input[0] for n in converted.graph.node if n.op_type == "Identity" and n.output[0] == "Z" + ) + metadata = copy.deepcopy(next(v for v in converted.graph.value_info if v.name == alias)) + metadata.name = "Y" + assert metadata.SerializeToString() == model.graph.output[0].SerializeToString() + + +def test_opaque_squeeze_with_shared_public_output_is_not_split_as_a_standard_operator(): + """Keep opaque function semantics intact when its public result also feeds internal nodes.""" + function = helper.make_function( + "opaque", + "Squeeze", + ["x"], + ["y"], + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Squeeze", ["X"], ["Y"], domain="opaque"), + helper.make_node("Relu", ["Y"], ["Z"]), + ], + "opaque_shared_squeeze", + [_vi("X", TensorProto.FLOAT, [1, 2, 3])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 3]), _vi("Z", TensorProto.FLOAT, [1, 2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("opaque", 1)], + functions=[function], + ) + _clamp_ir(model) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"X": np.arange(6, dtype=np.float32).reshape(1, 2, 3) - 3} + expected = ReferenceEvaluator(model).run(None, feeds) + opaque = model.graph.node[0].SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.graph.node[0].SerializeToString() == opaque + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip(ReferenceEvaluator(converted).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("transpose", [False, True]) +def test_five_dimensional_reshape_restores_native_public_and_downstream_layouts(transpose): + """Restore the original five-dimensional contract after computing with a collapsed target.""" + shape = [1, 2, 3, 4, 5] + nodes = [helper.make_node("Reshape", ["X", "target"], ["middle" if transpose else "Y"])] + output_shape = shape + if transpose: + nodes.append(helper.make_node("Transpose", ["middle"], ["Y"], perm=[1, 2, 3, 4, 0])) + output_shape = [2, 3, 4, 5, 1] + model = _make_model( + nodes, + "reshape_five", + [_vi("X", TensorProto.FLOAT, [2, 3, 4, 5])], + [_vi("Y", TensorProto.FLOAT, output_shape)], + [_init("target", np.asarray(shape, np.int64))], + ) + feeds = {"X": np.arange(120, dtype=np.float32).reshape(2, 3, 4, 5)} + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize( + ("operation", "control"), [("Slice", [0]), ("Slice", [1]), ("Expand", [2, 3]), ("Tile", [2, 1])] +) +def test_runtime_structural_controls_preserve_native_intermediate_layouts(operation, control): + """Skip unsupported runtime controls even when their intermediate metadata is absent.""" + shape = [1, 3] if operation == "Expand" else [2, 3] + inits = [] + if operation == "Slice": + names = ["X", "starts", "ends", "control"] + inits = [ + _init("starts", np.asarray([1], np.int64)), + _init("ends", np.asarray([2], np.int64)), + ] + else: + names = ["X", "control"] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node(operation, names, ["hidden"]), + helper.make_node("Identity", ["hidden"], ["Y"]), + ], + "runtime_structural", + [_vi("X", TensorProto.FLOAT, shape), _vi("control", TensorProto.INT64, [len(control)])], + [_vi("Y", TensorProto.FLOAT, [None, None])], + inits, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + feeds = { + "X": np.arange(np.prod(shape), dtype=np.float32).reshape(shape), + "control": np.asarray(control, np.int64), + } + expected = _run_ort(model, feeds) + before = model.SerializeToString() + converted = _apply_convert_ops_to_4d(model) + assert converted.SerializeToString() == before + _assert_match(expected, _run_ort(converted, feeds)) + + +def test_fixed_expand_derives_output_rank_without_intermediate_metadata(): + """Restore the broadcast output rank from fixed control values and the native input shape.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Expand", ["X", "target"], ["hidden"]), + helper.make_node("Identity", ["hidden"], ["Y"]), + ], + "expand_without_metadata", + [_vi("X", TensorProto.FLOAT, [1, 3])], + [_vi("Y", TensorProto.FLOAT, [2, 2, 3])], + [_init("target", np.asarray([2, 2, 3], np.int64))], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + feeds = {"X": np.arange(3, dtype=np.float32).reshape(1, 3)} + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("operation", ["Relu", "Softmax"]) +def test_legacy_opset_without_input_axes_adapters_stays_native(operation): + """Preserve legacy schemas instead of inserting newer Squeeze/Unsqueeze inputs.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x"], ["y"])], + "legacy_adapter_schema", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 11)], + ) + model.ir_version = 9 + original = model.SerializeToString() + _apply_convert_ops_to_4d(model) + assert model.SerializeToString() == original + onnx.checker.check_model(model, full_check=True) + data = np.arange(6, dtype=np.float32).reshape(2, 3) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + output = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": data})[0] + assert output.shape == data.shape + + +@pytest.mark.parametrize("fixed", [False, True]) +@pytest.mark.parametrize( + ("shape", "index_shape", "axis", "supported"), + [ + ([8, 16], [2, 1], 0, False), + ([5, 4, 3], [1, 2], 1, False), + ([2, 3, 4], [1, 1, 2], 1, False), + ([8, 16], [1, 2], 0, True), + ([2, 3], [0], 0, True), + ], +) +def test_gather_requires_exact_native_output_restoration( + fixed, shape, index_shape, axis, supported +): + """Flatten only layouts whose public native output can be restored exactly.""" + indices = np.arange(np.prod(index_shape), dtype=np.int64).reshape(index_shape) % shape[axis] + inputs = [_vi("x", TensorProto.FLOAT, shape)] + initializers = [] + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + if fixed: + initializers.append(_init("indices", indices)) + else: + inputs.append(_vi("indices", TensorProto.INT64, index_shape)) + feeds["indices"] = indices + native = shape[:axis] + index_shape + shape[axis + 1 :] + model = _make_model( + [helper.make_node("Gather", ["x", "indices"], ["y"], axis=axis)], + "gather_layout", + inputs, + [_vi("y", TensorProto.FLOAT, native)], + initializers, + ) + original = model.SerializeToString() + expected = _run_ort(copy.deepcopy(model), feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + if not supported: + assert converted.SerializeToString() == original + else: + assert any(node.op_type == "Unsqueeze" for node in converted.graph.node) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("operation", ["Add", "Mul", "Where"]) +@pytest.mark.parametrize("fixed", [False, True]) +def test_generic_broadcast_output_without_intermediate_metadata(operation, fixed): + """Restore the maximum broadcast rank including fixed data branches.""" + inputs = [_vi("a", TensorProto.FLOAT, [3])] + initializers = [] + feeds = {"a": np.asarray([1, 2, 3], np.float32)} + matrix = np.arange(6, dtype=np.float32).reshape(2, 3) + if fixed: + initializers.append(_init("b", matrix)) + else: + inputs.append(_vi("b", TensorProto.FLOAT, [2, 3])) + feeds["b"] = matrix + operands = ["a", "b"] + if operation == "Where": + inputs.append(_vi("condition", TensorProto.BOOL, [])) + feeds["condition"] = np.asarray(False) + operands.insert(0, "condition") + model = helper.make_model( + helper.make_graph( + [ + helper.make_node(operation, operands, ["hidden"]), + helper.make_node("Relu", ["hidden"], ["y"]), + ], + "broadcast_without_metadata", + inputs, + [_vi("y", TensorProto.FLOAT, [2, 3])], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + assert not model.graph.value_info + expected = _run_ort(copy.deepcopy(model), feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +@pytest.mark.parametrize("owner", ["peer", "public", "capture", "two_gathers", "default", "legacy"]) +def test_gather_preserves_four_dimensional_initializer_owners(owner): + """Keep the original index type, shape, and runtime override semantics for every owner.""" + indices = np.asarray([0, 3], np.int64).reshape(1, 1, 1, 2) + nodes = [helper.make_node("Gather", ["x", "indices"], ["y"], axis=0)] + inputs = [_vi("x", TensorProto.FLOAT, [8])] + outputs = [_vi("y", TensorProto.FLOAT, [1, 1, 1, 2])] + if owner == "peer": + nodes.append(helper.make_node("Identity", ["indices"], ["peer"])) + outputs.append(_vi("peer", TensorProto.INT64, [1, 1, 1, 2])) + elif owner == "public": + outputs.append(_vi("indices", TensorProto.INT64, [1, 1, 1, 2])) + elif owner == "two_gathers": + nodes.append(helper.make_node("Gather", ["x", "indices"], ["z"], axis=0)) + outputs.append(_vi("z", TensorProto.FLOAT, [1, 1, 1, 2])) + elif owner == "capture": + branch = helper.make_graph( + [helper.make_node("Identity", ["indices"], ["branch_result"])], + "capture", + [], + [_vi("branch_result", TensorProto.INT64, [1, 1, 1, 2])], + ) + nodes.append( + helper.make_node( + "If", ["condition"], ["captured"], then_branch=branch, else_branch=branch + ) + ) + inputs.append(_vi("condition", TensorProto.BOOL, [])) + outputs.append(_vi("captured", TensorProto.INT64, [1, 1, 1, 2])) + if owner in {"default", "legacy"}: + inputs.append(_vi("indices", TensorProto.INT64, [1, 1, 1, 2])) + model = _make_model( + nodes, + "owned_gather_indices", + inputs, + outputs, + [_init("indices", indices)], + infer_symbolic_shapes=False, + ) + if owner == "legacy": + model.ir_version = 3 + original = model.graph.initializer[0].SerializeToString() + feeds = {"x": np.arange(8, dtype=np.float32)} + if owner == "capture": + feeds["condition"] = np.asarray(True) + cases = [feeds] + if owner == "default": + cases.append({**feeds, "indices": np.asarray([7, 5], np.int64).reshape(1, 1, 1, 2)}) + expected = [_run_ort(copy.deepcopy(model), case) for case in cases] + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + assert ( + next(t for t in converted.graph.initializer if t.name == "indices").SerializeToString() + == original + ) + onnx.checker.check_model(converted, full_check=True) + for case, reference in zip(cases, expected): + _assert_match(reference, _run_ort(converted, case)) + if owner == "two_gathers": + gathers = [node for node in converted.graph.node if node.op_type == "Gather"] + assert gathers[0].input[1] == gathers[1].input[1] + + +@pytest.mark.parametrize("missing", ["indices", "output", "both", "unequal_extents"]) +@pytest.mark.parametrize("fixed", [False, True]) +def test_gather_elements_restores_native_rank_without_intermediate_metadata(missing, fixed): + """Derive both adapter ranks from the schema, including unequal non-axis extents.""" + data_shape, index_shape, axis = [2, 3], [2, 3], 1 + if missing == "unequal_extents": + data_shape, index_shape, axis = [2, 1, 4], [1, 2, 4], 1 + indices = np.zeros(index_shape, np.int64) + nodes = [] + inputs = [_vi("x", TensorProto.FLOAT, data_shape)] + initializers = [] + if fixed: + initializers.append(_init("indices", indices)) + else: + inputs.append(_vi("indices", TensorProto.INT64, index_shape)) + index_name = "indices" + if missing in {"indices", "both"}: + nodes.append(helper.make_node("Identity", ["indices"], ["hidden_indices"])) + index_name = "hidden_indices" + result_name = "hidden_result" if missing in {"output", "both", "unequal_extents"} else "y" + nodes.append(helper.make_node("GatherElements", ["x", index_name], [result_name], axis=axis)) + if result_name != "y": + nodes.append(helper.make_node("Identity", [result_name], ["y"])) + model = _make_model( + nodes, + "missing_gather_elements_metadata", + inputs, + [_vi("y", TensorProto.FLOAT, index_shape)], + initializers, + infer_symbolic_shapes=False, + ) + feeds = {"x": np.arange(np.prod(data_shape), dtype=np.float32).reshape(data_shape)} + if not fixed: + feeds["indices"] = indices + expected = _run_ort(copy.deepcopy(model), feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +def test_expand_with_unknown_native_input_and_output_stays_unchanged(): + """Leave a fixed Expand native when neither operand rank nor result rank is known.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["hidden_input"]), + helper.make_node("Expand", ["hidden_input", "target"], ["hidden_output"]), + helper.make_node("Identity", ["hidden_output"], ["y"]), + ], + "unknown_expand_rank", + [_vi("x", TensorProto.FLOAT, [1, 3])], + [_vi("y", TensorProto.FLOAT, [2, 3])], + [_init("target", np.asarray([2, 3], np.int64))], + infer_symbolic_shapes=False, + ) + before = model.SerializeToString() + feeds = {"x": np.arange(3, dtype=np.float32).reshape(1, 3)} + expected = _run_ort(copy.deepcopy(model), feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + assert converted.SerializeToString() == before + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + + +def test_gather_does_not_retarget_a_cast_after_its_output_was_wrapped(): + """Keep the native Cast dtype when an earlier handler has renamed its output.""" + model = _make_model( + [ + helper.make_node( + "Cast", ["source"], ["indices"], name="source_cast", to=TensorProto.INT64 + ), + helper.make_node("Gather", ["X", "indices"], ["Y"], axis=2), + ], + "wrapped_cast_indices", + [_vi("X", TensorProto.FLOAT, [1, 1, 5, 4]), _vi("source", TensorProto.FLOAT, [2])], + [_vi("Y", TensorProto.FLOAT, [1, 1, 2, 4])], + ) + feeds = { + "X": np.arange(20, dtype=np.float32).reshape(1, 1, 5, 4), + "source": np.asarray([1, 3], np.float32), + } + expected = _run_ort(model, feeds) + converted = _apply_convert_ops_to_4d(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + _assert_match(expected, _run_ort(converted, feeds)) + cast = next(node for node in converted.graph.node if node.name == "source_cast") + assert next(attr.i for attr in cast.attribute if attr.name == "to") == TensorProto.INT64 + + +@pytest.mark.parametrize("public_control", [False, True]) +@pytest.mark.parametrize("operation", ["Add", "Split", "DequantizeLinear"]) +def test_shared_unchanged_controls_allow_data_rank_conversion(operation, public_control): + """Only initializer inputs whose dimensions change need private ownership.""" + shape = [1, 2, 3] + initializers = [] + nodes = [] + outputs = [] + if operation == "Add": + initializers = [_init("control", np.float32(2))] + controls = {"control"} + for index in range(2): + nodes.append(helper.make_node("Add", ["x", "control"], [f"y{index}"])) + outputs.append(_vi(f"y{index}", TensorProto.FLOAT, shape)) + elif operation == "Split": + initializers = [_init("control", np.array([1, 1], np.int64))] + controls = {"control"} + for index in range(2): + names = [f"y{index}a", f"y{index}b"] + nodes.append(helper.make_node("Split", ["x", "control"], names, axis=1)) + outputs.extend(_vi(name, TensorProto.FLOAT, [1, 1, 3]) for name in names) + else: + initializers = [ + _init("scale", np.array([0.5, 0.25], np.float32)), + _init("zero", np.array([0, 0], np.uint8)), + ] + controls = {"scale", "zero"} + for index in range(2): + initializers.append(_init(f"data{index}", np.arange(6, dtype=np.uint8).reshape(shape))) + nodes.append( + helper.make_node( + operation, [f"data{index}", "scale", "zero"], [f"y{index}"], axis=1 + ) + ) + outputs.append(_vi(f"y{index}", TensorProto.FLOAT, shape)) + inputs = [] if operation == "DequantizeLinear" else [_vi("x", TensorProto.FLOAT, shape)] + model = _make_model(nodes, "shared_unchanged_controls", inputs, outputs, initializers) + if public_control: + model.graph.input.extend( + _vi(tensor.name, tensor.data_type, list(tensor.dims)) + for tensor in model.graph.initializer + if tensor.name in controls + ) + feeds = {} if not inputs else {"x": np.arange(6, dtype=np.float32).reshape(shape)} + owners = {tensor.name: tensor for tensor in model.graph.initializer if tensor.name in controls} + before = {name: tensor.SerializeToString() for name, tensor in owners.items()} + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + _apply_convert_ops_to_4d(model) + inferred = onnx.shape_inference.infer_shapes(model, strict_mode=True) + onnx.checker.check_model(inferred, full_check=True) + ranks = { + info.name: len(info.type.tensor_type.shape.dim) + for info in [*inferred.graph.value_info, *inferred.graph.output] + } + converted = [node for node in model.graph.node if node.op_type == operation] + assert len(converted) == 2 + assert all( + ranks[name] == (3 if public_control else 4) for node in converted for name in node.output + ) + retained = {tensor.name: tensor for tensor in model.graph.initializer} + for name, owner in owners.items(): + assert retained[name] is owner + assert owner.SerializeToString() == before[name] + actual = _run_ort(model, feeds) + for name in expected: + np.testing.assert_array_equal(actual[name], expected[name])