From ebcbd33f15c66b13e46a3ce3caef974d66ff7190 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Wed, 28 Oct 2026 21:21:28 +0000 Subject: [PATCH] Extract static weight-layout folding into a separate draft Preserve the original implementation in an isolated draft so exporter-versus-runtime design choices can be discussed independently of #754. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/execution_providers.md | 10 +- src/mobius/_optimizations.py | 34 +- src/mobius/_passes/__init__.py | 5 + src/mobius/_passes/_fold_transpose.py | 4 +- src/mobius/_passes/_fold_transpose_test.py | 55 -- src/mobius/_passes/_fold_weight_layout.py | 268 ++++++ .../_passes/_fold_weight_layout_test.py | 872 ++++++++++++++++++ 7 files changed, 1181 insertions(+), 67 deletions(-) create mode 100644 src/mobius/_passes/_fold_weight_layout.py create mode 100644 src/mobius/_passes/_fold_weight_layout_test.py diff --git a/docs/execution_providers.md b/docs/execution_providers.md index dcc35e39d..8b0701748 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -142,11 +142,11 @@ absent or packed UINT8 zero points. Flat/row-major sidecars are supported; block size must be a power of two >= 16 and `accuracy_level` omitted or zero. Mixed bit widths, extra operands, and caller-owned function bodies are skipped. -Parameter Concats retain checkpoint names until binding, then initializer -Concats and matrix Transposes fold lazily without changing dtype. Missing leaves -are checked even with `fold_constants=False`. Generic Reshape/Transpose layout -folding is not part of this pipeline; those transformations remain explicit. -Overridable inputs remain untouched. +Post-weight folding is independent of experimental packing: parameter Concats +retain checkpoint names until binding, then Concat/Reshape/Transpose chains fold +lazily without changing dtype. Missing leaves are checked even with +`fold_constants=False`. Overridable inputs and MoE/QMoE/QDQ storage-layout paths +remain explicit to preserve loader names and byte layouts. **The INT4 numerical gate remains failed and unwaived:** large-dimension CUDA prefill/cached-decode logits exceeded `rtol=atol=1e-2`. Changing GEMM width can diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index f46efdbf3..e597fb217 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -21,8 +21,8 @@ Post-weight passes ------------------ :func:`fold_initializers_after_weights` should be called after weights are loaded. -It lazily folds initializer matrix Transpose and Concat nodes for projection -packing. Generic Reshape/Transpose layout chains remain explicit. +It lazily folds initializer Transpose, Concat, and static Reshape/Transpose +chains, independently of experimental projection packing. """ from __future__ import annotations @@ -47,6 +47,7 @@ import contextlib import dataclasses import logging +import math import warnings import onnx @@ -62,6 +63,7 @@ from mobius._passes import ( FoldConcatInitializersPass, FoldTransposedInitializerPass, + FoldWeightLayoutPass, Fp8KvCachePass, RemoveDeadGraphInputsPass, ) @@ -795,12 +797,14 @@ def _should_inline(func: ir.Function) -> bool: def fold_initializers_after_weights(model: ir.Model) -> None: - """Lazily fold initializer matrix Transpose and Concat nodes. + """Lazily fold weight Transpose, Concat, and static layout chains. - Preserve the shipped root Transpose naming convention and fold projection - packing to a fixed point. Generic weight layout chains remain explicit. + Fold small integer shape expressions first without exposing weight payloads + to the general constant folder. Preserve the shipped root Transpose naming + convention, then fold layouts and packing to a fixed point. """ common_passes.LiftConstantsToInitializersPass()(model) + hidden: dict[ir.Value, ir.TensorProtocol] = {} for graph in model.graphs(): for value in graph.initializers.values(): tensor = value.const_value @@ -814,9 +818,29 @@ def fold_initializers_after_weights(model: ir.Model) -> None: model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) ) + if tensor is not None and not ( + tensor.dtype == ir.DataType.INT64 + and all(isinstance(dim, int) for dim in tensor.shape) + and math.prod(tensor.shape.numpy()) <= 8192 + ): + hidden[value] = tensor + value.const_value = None + try: + # Partial evaluators also simplify Identity without reading its input. + # In particular, an output Identity transfers its public output name + # to the producer value before lazy layout folding. Size limits and + # should_fold do not gate those upstream symbolic simplifications. + onnxscript.optimizer._constant_folding.FoldConstantsPass( + shape_inference=False, input_size_limit=8192, output_size_limit=8192 + )(model) + finally: + for value, tensor in hidden.items(): + value.const_value = tensor + common_passes.LiftConstantsToInitializersPass()(model) common_passes.RemoveUnusedNodesPass()(model) lazy_passes = ( FoldTransposedInitializerPass(), + FoldWeightLayoutPass(), FoldConcatInitializersPass(), ) while True: diff --git a/src/mobius/_passes/__init__.py b/src/mobius/_passes/__init__.py index 36cfa31f7..499473ec4 100644 --- a/src/mobius/_passes/__init__.py +++ b/src/mobius/_passes/__init__.py @@ -36,6 +36,9 @@ graph initializer) into a single pre-concatenated initializer. Used to pre-pack QKV weight matrices after weight loading. +:class:`FoldWeightLayoutPass` + Folds static Reshape/Transpose chains into lazy initializers, including + higher-rank layout conversions of quantized weights, scales, and zero points. """ from __future__ import annotations @@ -43,11 +46,13 @@ __all__ = [ "FoldTransposedInitializerPass", "FoldConcatInitializersPass", + "FoldWeightLayoutPass", "Fp8KvCachePass", "RemoveDeadGraphInputsPass", ] from mobius._passes._fold_concat import FoldConcatInitializersPass from mobius._passes._fold_transpose import FoldTransposedInitializerPass +from mobius._passes._fold_weight_layout import FoldWeightLayoutPass from mobius._passes._fp8_kv_cache import Fp8KvCachePass from mobius._passes._remove_dead_inputs import RemoveDeadGraphInputsPass diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py index e14ddc033..52fb232dd 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -74,8 +74,8 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: # Register legacy ``weight_t`` names only for root-owned weights. # A Transpose inside If/Loop can capture a root weight and still use - # this pass. Skip child-local weights: registering them at the root - # would collide across sibling scopes. + # this pass. Child-local weights belong to FoldWeightLayoutPass: + # registering them at the root would collide across sibling scopes. if ( not inp.is_initializer() or inp.is_graph_input() diff --git a/src/mobius/_passes/_fold_transpose_test.py b/src/mobius/_passes/_fold_transpose_test.py index af7832895..dc64970ef 100644 --- a/src/mobius/_passes/_fold_transpose_test.py +++ b/src/mobius/_passes/_fold_transpose_test.py @@ -660,58 +660,3 @@ def test_transpose_inside_loop_body_folds_against_root_initializer(): w_t = model.graph.initializers["weight_t"] np.testing.assert_array_equal(w_t.const_value.numpy(), weight_data.T) - - -def test_post_weight_pipeline_preserves_generic_layout_chain(): - from mobius._optimizations import fold_initializers_after_weights - - data = np.arange(24, dtype=np.uint8).reshape(2, 12) - - def unexpected_read(): - raise AssertionError("generic layout weights must remain untouched") - - weight = ir.Value( - name="checkpoint.weight", - shape=ir.Shape(data.shape), - type=ir.TensorType(ir.DataType.UINT8), - const_value=ir.LazyTensor( - unexpected_read, dtype=ir.DataType.UINT8, shape=ir.Shape(data.shape) - ), - ) - shapes = [ - ir.Value( - name=name, - shape=ir.Shape([len(dimensions)]), - type=ir.TensorType(ir.DataType.INT64), - const_value=ir.tensor(np.array(dimensions, dtype=np.int64)), - ) - for name, dimensions in (("split_shape", (2, 2, 6)), ("final_shape", (2, 12))) - ] - reshape = ir.Node("", "Reshape", [weight, shapes[0]], num_outputs=1) - transpose = ir.Node( - "", - "Transpose", - [reshape.outputs[0]], - attributes=[ir.AttrInt64s("perm", [0, 2, 1])], - num_outputs=1, - ) - result = ir.Node("", "Reshape", [transpose.outputs[0], shapes[1]], num_outputs=1) - result.outputs[0].name = "reordered_weight" - nodes = [reshape, transpose, result] - model = ir.Model( - ir.Graph( - [], - [result.outputs[0]], - nodes=nodes, - initializers=[weight, *shapes], - opset_imports={"": 24}, - ), - ir_version=10, - ) - - fold_initializers_after_weights(model) - - assert list(model.graph) == nodes - assert model.graph.initializers["checkpoint.weight"] is weight - assert reshape.inputs[0] is weight - assert model.graph.outputs[0] is result.outputs[0] diff --git a/src/mobius/_passes/_fold_weight_layout.py b/src/mobius/_passes/_fold_weight_layout.py new file mode 100644 index 000000000..a208bd1a3 --- /dev/null +++ b/src/mobius/_passes/_fold_weight_layout.py @@ -0,0 +1,268 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lazily fold static weight Reshape/Transpose chains after checkpoint loading.""" + +from __future__ import annotations + +import logging +import math + +import onnx_ir as ir + +from mobius._constants import LAZY_SOURCE_FILES_META_KEY +from mobius._passes._dtype_utils import initializer_dtype + +logger = logging.getLogger(__name__) + + +def _reshape_shape(node: ir.Node, source_shape: tuple[int, ...]) -> tuple[int, ...] | None: + """Resolve a static shape, skip dynamic shapes, and reject invalid reshapes.""" + shape_value = node.inputs[1] + if shape_value is None or shape_value.is_graph_input(): + return None + if shape_value.is_initializer(): + shape_tensor = shape_value.const_value + else: + constant = shape_value.producer() + if ( + constant is None + or constant.domain != "" + or constant.op_type != "Constant" + or len(constant.attributes) != 1 + ): + return None + if "value" in constant.attributes: + shape_tensor = constant.attributes["value"].as_tensor() + elif "value_ints" in constant.attributes: + shape_tensor = ir.tensor( + constant.attributes["value_ints"].as_ints(), dtype=ir.DataType.INT64 + ) + else: + return None + if shape_tensor is None: + return None + if shape_tensor.dtype != ir.DataType.INT64 or len(shape_tensor.shape) != 1: + raise ValueError(f"Reshape {node.name!r} requires a one-dimensional INT64 shape") + requested = tuple(int(dim) for dim in shape_tensor.numpy()) + allowzero = node.attributes.get_int("allowzero", 0) + if allowzero not in (0, 1) or requested.count(-1) > 1 or any(d < -1 for d in requested): + raise ValueError(f"Invalid Reshape attributes or shape at {node.name!r}: {requested}") + resolved = list(requested) + if not allowzero: + for index, dim in enumerate(resolved): + if dim == 0: + if index >= len(source_shape): + raise ValueError(f"Reshape {node.name!r} copies a nonexistent dimension") + resolved[index] = source_shape[index] + size = math.prod(source_shape) + known_size = math.prod(dim for dim in resolved if dim != -1) + if -1 in resolved: + if known_size == 0 or size % known_size: + raise ValueError(f"Reshape {node.name!r} cannot infer dimension in {requested}") + resolved[resolved.index(-1)] = size // known_size + elif known_size != size: + raise ValueError( + f"Reshape {node.name!r} changes element count: {source_shape} to {requested}" + ) + return tuple(resolved) + + +def _storage_layout_values(model: ir.Model) -> set[ir.Value]: + """Find parameter views whose checkpoint storage ABI must remain explicit.""" + pending: list[ir.Value] = [] + for node in model.graph.all_nodes(): + if node.domain == "" and node.op_type == "DequantizeLinear": + parameters = node.inputs[:3] + elif node.domain == "com.microsoft" and node.op_type in ("MoE", "QMoE"): + # Skip activations/routing. Expert weights, biases, scales, and + # zero points retain their checkpoint-format names and row order; + # the explicit layout views adapt that storage to the runtime ABI. + parameters = node.inputs[2 : 14 if node.op_type == "QMoE" else 8] + else: + continue + pending.extend(value for value in parameters if value is not None) + protected: set[ir.Value] = set() + while pending: + value = pending.pop() + if value in protected: + continue + protected.add(value) + producer = value.producer() + if producer is None or producer.domain != "" or producer.overload: + continue + if producer.op_type in ( + "Identity", + "Reshape", + "Transpose", + "Cast", + "CastLike", + "Pad", + "Slice", + "Expand", + ): + pending.extend(item for item in producer.inputs[:1] if item is not None) + elif producer.op_type in ("Add", "Sub", "Mul", "Div", "Concat"): + # FP8 scales can carry an exact-zero execution dependency before + # DQ. Follow the scale path without evaluating its arithmetic. + pending.extend(item for item in producer.inputs if item is not None) + return protected + + +def _initializer_owners(value: ir.Value, graphs: tuple[ir.Graph, ...]) -> list[ir.Graph]: + """Find registries containing this exact value, not just its name. + + An If branch can capture a root initializer without owning it, and sibling + branches can own different values with the same name. Registry identity + determines where it is safe to remove a folded source. + """ + return [ + graph + for graph in graphs + if value.name is not None and graph.initializers.get(value.name) is value + ] + + +class FoldWeightLayoutPass(ir.passes.InPlacePass): + """Fold named static Reshape/Transpose outputs into lazy initializers. + + Graphs must be topologically ordered. A single traversal folds layout chains + without reading weight payloads: only constant Reshape shape vectors are read. + Materialization preserves the source dtype and uses no arithmetic or casts. + Overridable initializers, shared sources, and dynamic shapes remain untouched. + MoE/QMoE and QDQ parameter-layout paths retain checkpoint storage names and + byte layouts; their explicit runtime format conversions are not folded. + Dead folded inputs are removed in their owning graph; run dead-code elimination + afterward to remove unused Constant shape nodes. + Malformed static layout nodes raise ValueError rather than being skipped. + """ + + def call(self, model: ir.Model) -> ir.passes.PassResult: + folded = 0 + graphs = tuple(model.graphs()) + storage_layouts = _storage_layout_values(model) + + for graph in graphs: + for node in list(graph): + if ( + node.overload + or node.domain != "" + or node.op_type + not in ( + "Reshape", + "Transpose", + ) + ): + continue + expected_inputs = 2 if node.op_type == "Reshape" else 1 + if len(node.inputs) != expected_inputs or len(node.outputs) != 1: + raise ValueError(f"Malformed weight layout node {node.name!r}") + if node.outputs[0] in storage_layouts: + continue + source = node.inputs[0] + if source is None or not source.is_initializer() or source.is_graph_input(): + continue + # Keep tied tables shared instead of serializing a second full layout. + if ( + len(source.uses()) != 1 + or len(_initializer_owners(source, graphs)) != 1 + or any( + source in scope.inputs or source in scope.outputs for scope in graphs + ) + ): + continue + tensor = source.const_value + if tensor is None: + logger.warning( + "FoldWeightLayoutPass: skipping %r; initializer %r has no data", + node.name, + source.name, + ) + continue + output = node.outputs[0] + if not output.name or any(not isinstance(dim, int) for dim in tensor.shape): + logger.debug( + "FoldWeightLayoutPass: skipping unnamed or symbolic %r", node.name + ) + continue + source_shape = tuple(tensor.shape.numpy()) + dtype = initializer_dtype(source) + if dtype is None: + raise ValueError(f"Weight layout source {source.name!r} has no dtype") + is_transpose = node.op_type == "Transpose" + if is_transpose: + layout = tuple( + node.attributes.get_ints( + "perm", tuple(reversed(range(len(source_shape)))) + ) + ) + if sorted(layout) != list(range(len(source_shape))): + raise ValueError( + f"Invalid Transpose permutation at {node.name!r}: {layout}" + ) + shape = tuple(source_shape[index] for index in layout) + else: + resolved = _reshape_shape(node, source_shape) + if resolved is None: + logger.debug( + "FoldWeightLayoutPass: skipping dynamic shape at %r", node.name + ) + continue + layout = shape = resolved + + # Bind this iteration's values before the lazy callback runs later. + def materialize( + data: ir.TensorProtocol = tensor, + dimensions: tuple[int, ...] = layout, + transpose: bool = is_transpose, + data_type: ir.DataType = dtype, + ) -> ir.TensorProtocol: + array = data.numpy() + result = ( + array.transpose(dimensions) if transpose else array.reshape(dimensions) + ) + return ir.tensor(result, dtype=data_type) + + output_shape = ir.Shape(shape) + replacement = ir.Value( + name=output.name, + shape=output_shape, + type=ir.TensorType(dtype), + doc_string=output.doc_string, + metadata_props=dict(output.metadata_props), + ) + replacement.const_value = ir.LazyTensor( + materialize, dtype=dtype, shape=output_shape, name=output.name + ) + source_files = ( + (tensor.path,) + if isinstance(tensor, ir.ExternalTensor) + else tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) + if source_files: + replacement.const_value.meta[LAZY_SOURCE_FILES_META_KEY] = tuple( + source_files + ) + model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( + (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) + ) + output.replace_all_uses_with(replacement, replace_graph_outputs=True) + inputs = tuple(node.inputs) + graph.remove(node, safe=True) + # Downstream layout nodes can now capture this lazy initializer. + graph.register_initializer(replacement) + for value in inputs: + if ( + value is not None + and value.name is not None + and not value.uses() + and not any( + value in scope.inputs or value in scope.outputs for scope in graphs + ) + ): + registries = _initializer_owners(value, graphs) + if len(registries) == 1: + del registries[0].initializers[value.name] + folded += 1 + logger.debug("FoldWeightLayoutPass: folded %d weight layout nodes", folded) + return ir.passes.PassResult(model, modified=folded > 0) diff --git a/src/mobius/_passes/_fold_weight_layout_test.py b/src/mobius/_passes/_fold_weight_layout_test.py new file mode 100644 index 000000000..8bae303e0 --- /dev/null +++ b/src/mobius/_passes/_fold_weight_layout_test.py @@ -0,0 +1,872 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Value, laziness, and persistence coverage for static weight layout folding.""" + +from __future__ import annotations + +from pathlib import Path + +import ml_dtypes +import numpy as np +import onnx_ir as ir +import pytest +from onnx_ir.passes import common + +from mobius import ModelPackage +from mobius._constants import LAZY_SOURCE_FILES_META_KEY +from mobius._optimizations import fold_initializers_after_weights +from mobius._passes import FoldWeightLayoutPass +from mobius._testing.ort_inference import OnnxModelSession + + +def _initializer(name, data): + tensor = ir.tensor(data, name=name) + return ir.Value( + name=name, shape=tensor.shape, type=ir.TensorType(tensor.dtype), const_value=tensor + ) + + +def _model(data, *, first=(2, 2, 4, 3), perm=(0, 2, 1, 3), last=(2, 8, 3)): + weight = _initializer("weight", data) + shape1 = _initializer("shape1", np.array(first, dtype=np.int64)) + shape2 = _initializer("shape2", np.array(last, dtype=np.int64)) + reshape1 = ir.Node("", "Reshape", [weight, shape1], num_outputs=1) + transpose = ir.Node( + "", + "Transpose", + [reshape1.outputs[0]], + attributes=[ir.AttrInt64s("perm", perm)], + num_outputs=1, + ) + reshape2 = ir.Node("", "Reshape", [transpose.outputs[0], shape2], num_outputs=1) + identity = ir.Node("", "Identity", [reshape2.outputs[0]], num_outputs=1) + for node, name in zip( + (reshape1, transpose, reshape2, identity), + ("reshaped", "transposed", "layout", "output"), + strict=True, + ): + node.outputs[0].name = name + node.outputs[0].dtype = weight.dtype + graph = ir.Graph( + [], + [identity.outputs[0]], + nodes=[reshape1, transpose, reshape2, identity], + initializers=[weight, shape1, shape2], + opset_imports={"": 24}, + name="layout_test", + ) + return ir.Model(graph, ir_version=10) + + +@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, ml_dtypes.bfloat16]) +def test_quantized_layout_values_and_package_roundtrip(dtype, tmp_path): + data = np.arange(48).astype(dtype).reshape(2, 8, 3) + model = _model(data) + fold_initializers_after_weights(model) + # The shape-expression constant folder also resolves Identity symbolically, + # preserving the public output name on the eventual lazy initializer. + assert not list(model.graph) + assert set(model.graph.initializers) == {"output"} + folded = model.graph.initializers["output"] + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + expected = data.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3) + assert folded.dtype == ir.tensor(data).dtype + assert folded.const_value.numpy().tobytes() == expected.tobytes() + package = ModelPackage({"model": model}) + package.save(tmp_path, progress_bar=False) + loaded = ModelPackage.load(tmp_path) + restored = loaded["model"].graph.initializers["output"] + assert loaded["model"].graph.outputs[0] is restored + assert restored.dtype == folded.dtype + assert restored.shape == ir.Shape(expected.shape) + assert restored.const_value.numpy().tobytes() == expected.tobytes() + + +def test_bfloat16_special_value_bits_preserved(tmp_path): + bits = np.resize( + np.array([0, 0x8000, 0x7F80, 0xFF80, 0x7FC1, 0x7FFF, 1, 0x8001], dtype=np.uint16), + 48, + ).reshape(2, 8, 3) + model = _model(bits.view(ml_dtypes.bfloat16)) + fold_initializers_after_weights(model) + folded = model.graph.initializers["output"] + assert set(model.graph.initializers) == {"output"} + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + actual = folded.const_value.numpy() + expected = bits.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3) + assert actual.tobytes() == expected.tobytes() + ModelPackage({"model": model}).save(tmp_path, progress_bar=False) + restored_graph = ModelPackage.load(tmp_path)["model"].graph + restored = restored_graph.initializers["output"] + assert restored_graph.outputs[0] is restored + assert restored.dtype == ir.DataType.BFLOAT16 + assert restored.const_value.numpy().tobytes() == expected.tobytes() + + +def test_weight_payload_remains_lazy_and_pass_is_idempotent(): + data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) + model = _model(data) + source = model.graph.initializers["weight"] + calls = [] + + def read(): + calls.append(True) + return ir.tensor(data) + + source.const_value = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) + assert FoldWeightLayoutPass()(model).modified + assert not calls + common.RemoveUnusedNodesPass()(model) + assert set(model.graph.initializers) == {"layout"} + assert not FoldWeightLayoutPass()(model).modified + assert not calls + actual = model.graph.initializers["layout"].const_value.numpy() + assert len(calls) == 1 + assert ( + actual.tobytes() + == data.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3).tobytes() + ) + + +@pytest.mark.parametrize("dynamic", ["weight", "shape1"]) +@pytest.mark.parametrize("overridable", [False, True]) +def test_dynamic_or_overridable_inputs_are_not_folded(dynamic, overridable): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + value = model.graph.initializers[dynamic] + model.graph.inputs.append(value) + if not overridable: + del model.graph.initializers[dynamic] + value.const_value = None + assert not FoldWeightLayoutPass()(model).modified + assert len(model.graph) == 4 + + +def test_shared_sources_are_not_duplicated(): + data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) + model = _model(data) + source = model.graph.initializers["weight"] + extra = ir.Node("", "Identity", [source], num_outputs=1) + extra.outputs[0].name = "original_weight" + extra.outputs[0].dtype = source.dtype + model.graph.append(extra) + model.graph.outputs.append(extra.outputs[0]) + layout = next(node for node in model.graph if node.outputs[0].name == "layout") + model.graph.outputs.append(layout.outputs[0]) + assert not FoldWeightLayoutPass()(model).modified + common.RemoveUnusedNodesPass()(model) + assert [output.name for output in model.graph.outputs] == [ + "output", + "original_weight", + "layout", + ] + assert extra.inputs[0] is source + assert source.name in model.graph.initializers + assert model.graph.outputs[2] is layout.outputs[0] + assert set(model.graph.initializers) == {"weight", "shape1", "shape2"} + assert source.const_value.numpy().tobytes() == data.tobytes() + + +@pytest.mark.parametrize( + ("domain", "op_type", "parameter_slot", "protected"), + [ + ("com.microsoft", "MoE", 2, True), + ("com.microsoft", "QMoE", 2, True), + ("com.microsoft", "QMoE", 3, True), + ("", "DequantizeLinear", 0, True), + ("", "DequantizeLinear", 1, True), + ("com.microsoft", "MatMulNBits", 1, False), + ("", "MatMul", 1, False), + ], +) +def test_post_weight_storage_layout_paths_preserve_named_sources( + domain, op_type, parameter_slot, protected +): + model = _model(np.arange(48, dtype=np.float32).reshape(2, 8, 3)) + source = model.graph.initializers["weight"] + + def unexpected_read(): + raise AssertionError("eligibility and layout folding must not read weight payloads") + + tensor = ir.LazyTensor(unexpected_read, dtype=source.dtype, shape=source.shape) + source.const_value = tensor + # Only the consumer's parameter position matters to this structural test; + # complete model/schema and public save/load contracts are covered by the + # existing QMoE/FP8/streaming model tests, unchanged by this regression fix. + inputs = [] + for index in range(parameter_slot): + argument = ir.Value( + name=f"argument_{index}", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([1]), + ) + model.graph.inputs.append(argument) + inputs.append(argument) + inputs.append(model.graph.outputs[0]) + consumer = ir.Node(domain, op_type, inputs, num_outputs=1) + consumer.outputs[0].name = "consumer_output" + model.graph.append(consumer) + model.graph.outputs.append(consumer.outputs[0]) + fold_initializers_after_weights(model) + if protected: + assert model.graph.initializers["weight"] is source + assert source.const_value is tensor + assert sum(node.op_type == "Reshape" for node in model.graph) == 2 + assert sum(node.op_type == "Transpose" for node in model.graph) == 1 + else: + assert "weight" not in model.graph.initializers + assert not any(node.op_type in ("Reshape", "Transpose") for node in model.graph) + + +@pytest.mark.parametrize("with_dependency", [False, True]) +def test_qdq_scalar_scale_storage_survives_pipeline_and_package_roundtrip( + tmp_path, with_dependency +): + data = np.array([[0.5]], dtype=np.float32) + scale = _initializer("checkpoint_scale", data) + reads = [] + + def read_scale(): + reads.append(True) + return ir.tensor(data) + + tensor = ir.LazyTensor(read_scale, dtype=scale.dtype, shape=scale.shape) + scale.const_value = tensor + scalar_shape = _initializer("scalar_shape", np.array([], dtype=np.int64)) + reshape = ir.Node("", "Reshape", [scale, scalar_shape], num_outputs=1) + reshape.outputs[0].name = "runtime_scalar_scale" + reshape.outputs[0].dtype = ir.DataType.FLOAT + reshape.outputs[0].shape = ir.Shape([]) + codes = ir.Value(name="codes", type=ir.TensorType(ir.DataType.UINT8), shape=ir.Shape([4])) + inputs = [codes] + nodes = [reshape] + runtime_scale = reshape.outputs[0] + if with_dependency: + dependency = ir.Value( + name="dependency", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([]) + ) + inputs.append(dependency) + add = ir.Node("", "Add", [runtime_scale, dependency], num_outputs=1) + add.outputs[0].name = "dependent_scale" + add.outputs[0].dtype = ir.DataType.FLOAT + add.outputs[0].shape = ir.Shape([]) + nodes.append(add) + runtime_scale = add.outputs[0] + dequantize = ir.Node("", "DequantizeLinear", [codes, runtime_scale], num_outputs=1) + dequantize.outputs[0].name = "output" + dequantize.outputs[0].dtype = ir.DataType.FLOAT + dequantize.outputs[0].shape = ir.Shape([4]) + nodes.append(dequantize) + model = ir.Model( + ir.Graph( + inputs, + [dequantize.outputs[0]], + nodes=nodes, + initializers=[scale, scalar_shape], + opset_imports={"": 21}, + ), + ir_version=10, + ) + fold_initializers_after_weights(model) + assert model.graph.initializers["checkpoint_scale"] is scale + assert scale.const_value is tensor + assert not reads + assert reshape.graph is model.graph + assert reshape.inputs[0] is scale + # The protected source stays in checkpoint orientation, not the scalar + # runtime format. This is a storage contract, not just numerical parity. + assert scale.shape == ir.Shape([1, 1]) + assert scale.dtype == ir.DataType.FLOAT + fold_initializers_after_weights(model) + assert model.graph.initializers["checkpoint_scale"] is scale + assert not reads + + ModelPackage({"model": model}).save(tmp_path, progress_bar=False) + restored = ModelPackage.load(tmp_path)["model"] + restored_scale = restored.graph.initializers["checkpoint_scale"] + assert restored_scale.shape == scale.shape + assert restored_scale.dtype == scale.dtype + assert restored_scale.const_value.numpy().tobytes() == data.tobytes() + for candidate in (model, restored): + session = OnnxModelSession(candidate, device="cpu") + for dependency_value in (0.0, 0.25) if with_dependency else (0.0,): + code_data = np.array([0, 1, 127, 255], dtype=np.uint8) + feeds: dict[str, np.ndarray] = {"codes": code_data} + if with_dependency: + feeds["dependency"] = np.asarray(dependency_value, dtype=np.float32) + actual = session.run(feeds)["output"] + np.testing.assert_array_equal( + actual, code_data.astype(np.float32) * (0.5 + dependency_value) + ) + + +def test_unshared_graph_output_layout_name_survives(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + layout = next(node for node in model.graph if node.outputs[0].name == "layout") + model.graph.outputs.append(layout.outputs[0]) + assert FoldWeightLayoutPass()(model).modified + assert model.graph.outputs[1] is model.graph.initializers["layout"] + + +@pytest.mark.parametrize("op_type", ["Reshape", "Transpose"]) +def test_graph_output_source_is_not_duplicated(op_type): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + if op_type == "Transpose": + first = next(iter(model.graph)) + first.op_type = "Transpose" + first.replace_input_with(1, None) + first.resize_inputs(1) + first.attributes["perm"] = ir.AttrInt64s("perm", [2, 1, 0]) + model.graph.outputs.append(model.graph.initializers["weight"]) + assert not FoldWeightLayoutPass()(model).modified + assert "layout" not in model.graph.initializers + + +@pytest.mark.parametrize("style", ["value", "value_ints"]) +@pytest.mark.parametrize("shared", [False, True]) +def test_post_weight_pipeline_folds_constant_shape_nodes(style, shared): + data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) + model = _model(data) + if shared: + extra = ir.Node("", "Identity", [model.graph.initializers["weight"]], num_outputs=1) + extra.outputs[0].name = "shared" + model.graph.append(extra) + model.graph.outputs.append(extra.outputs[0]) + for name in ("shape1", "shape2"): + shape = model.graph.initializers[name] + attribute = ( + ir.AttrTensor("value", shape.const_value) + if style == "value" + else ir.AttrInt64s("value_ints", shape.const_value.numpy().tolist()) + ) + constant = ir.Node("", "Constant", [], attributes=[attribute], num_outputs=1) + constant.outputs[0].name = name + model.graph.insert_before(next(iter(model.graph)), constant) + shape.replace_all_uses_with(constant.outputs[0]) + del model.graph.initializers[name] + fold_initializers_after_weights(model) + if shared: + assert set(model.graph.initializers) == {"weight"} + assert sum(node.op_type == "Reshape" for node in model.graph) == 2 + assert sum(node.op_type == "Transpose" for node in model.graph) == 1 + return + assert not list(model.graph) + assert set(model.graph.initializers) == {"output"} + folded = model.graph.initializers["output"] + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + expected = data.reshape(2, 2, 4, 3).swapaxes(1, 2).reshape(2, 8, 3) + assert folded.const_value.numpy().tobytes() == expected.tobytes() + + +def test_tied_embedding_pipeline_keeps_one_table(tmp_path): + weight = _initializer("weight", np.arange(24, dtype=np.float16).reshape(8, 3)) + indices = ir.Value( + name="indices", shape=ir.Shape([2]), type=ir.TensorType(ir.DataType.INT64) + ) + gather = ir.Node("", "Gather", [weight, indices], num_outputs=1) + transpose = ir.Node( + "", "Transpose", [weight], attributes=[ir.AttrInt64s("perm", [1, 0])], num_outputs=1 + ) + matmul = ir.Node("", "MatMul", [gather.outputs[0], transpose.outputs[0]], num_outputs=1) + for node, name in ((gather, "embedding"), (transpose, "head"), (matmul, "output")): + node.outputs[0].name = name + node.outputs[0].dtype = ir.DataType.FLOAT16 + graph = ir.Graph( + [indices], + [matmul.outputs[0]], + nodes=[gather, transpose, matmul], + initializers=[weight], + opset_imports={"": 24}, + name="tied_table", + ) + model = ir.Model(graph, ir_version=10) + fold_initializers_after_weights(model) + assert set(graph.initializers) == {"weight"} + assert transpose.inputs[0] is gather.inputs[0] + ModelPackage({"model": model}).save(tmp_path, progress_bar=False) + loaded = ModelPackage.load(tmp_path)["model"].graph + assert set(loaded.initializers) == {"weight"} + assert [node.op_type for node in loaded] == ["Gather", "Transpose", "MatMul"] + + +@pytest.mark.parametrize( + ("source", "requested", "allowzero", "expected"), + [ + ((2, 3), (0, -1), 0, (2, 3)), + ((2, 3), (-1, 2), 0, (3, 2)), + ((0, 3), (0, 3), 1, (0, 3)), + ((0, 3), (-1, 3), 0, (0, 3)), + ((1,), (), 0, ()), + ], +) +def test_reshape_zero_inferred_and_scalar_shapes(source, requested, allowzero, expected): + weight = _initializer("weight", np.zeros(source, dtype=np.float16)) + shape = _initializer("shape", np.array(requested, dtype=np.int64)) + node = ir.Node( + "", + "Reshape", + [weight, shape], + attributes=[ir.AttrInt64("allowzero", allowzero)], + num_outputs=1, + ) + node.outputs[0].name = "output" + graph = ir.Graph( + [], + [node.outputs[0]], + nodes=[node], + initializers=[weight, shape], + opset_imports={"": 24}, + ) + model = ir.Model(graph, ir_version=10) + assert FoldWeightLayoutPass()(model).modified + assert model.graph.outputs[0].shape == ir.Shape(expected) + assert model.graph.outputs[0].const_value.numpy().shape == expected + + +@pytest.mark.parametrize("requested", [(0, -1), (-1, -1), (4,), (-2,)]) +def test_invalid_static_reshape_fails_explicitly(requested): + model = _model(np.zeros((0, 3), dtype=np.float16), first=requested) + with pytest.raises(ValueError, match="Reshape"): + FoldWeightLayoutPass()(model) + + +def test_default_transpose_and_missing_declared_dtype(): + data = np.arange(24, dtype=np.float16).reshape(2, 3, 4) + weight = _initializer("weight", data) + weight.type = None + node = ir.Node("", "Transpose", [weight], num_outputs=1) + node.outputs[0].name = "output" + graph = ir.Graph([], [node.outputs[0]], nodes=[node], initializers=[weight]) + model = ir.Model(graph, ir_version=10) + assert FoldWeightLayoutPass()(model).modified + output = graph.outputs[0] + assert output.dtype == ir.DataType.FLOAT16 + assert output.const_value.numpy().tobytes() == data.transpose().tobytes() + + +def test_invalid_permutation_and_conflicting_dtype_fail(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3), perm=(0, 0, 1, 2)) + with pytest.raises(ValueError, match="permutation"): + FoldWeightLayoutPass()(model) + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + model.graph.initializers["weight"].dtype = ir.DataType.FLOAT + with pytest.raises(ValueError, match="declares dtype"): + FoldWeightLayoutPass()(model) + + +def test_custom_domain_layout_op_is_not_folded(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + transpose = next(node for node in model.graph if node.op_type == "Transpose") + transpose.domain = "custom" + FoldWeightLayoutPass()(model) + assert transpose in list(model.graph) + assert transpose.outputs[0].producer() is transpose + + +def test_missing_weights_are_warned_and_left_unmodified(caplog): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + model.graph.initializers["weight"].const_value = None + assert not FoldWeightLayoutPass()(model).modified + assert "has no data" in caplog.text + + +def test_invalid_shape_tensor_dtype_fails_explicitly(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + shape = model.graph.initializers["shape1"] + shape.dtype = ir.DataType.INT32 + shape.const_value = ir.tensor(np.array([2, 2, 4, 3], dtype=np.int32)) + with pytest.raises(ValueError, match="INT64 shape"): + FoldWeightLayoutPass()(model) + + +def test_nested_graph_initializers_are_registered_in_their_scope(): + first = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)).graph + second = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)).graph + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + node = ir.Node( + "", + "If", + [condition], + attributes=[ir.AttrGraph("then_branch", first), ir.AttrGraph("else_branch", second)], + num_outputs=1, + ) + node.outputs[0].name = "result" + graph = ir.Graph([condition], [node.outputs[0]], nodes=[node], opset_imports={"": 24}) + model = ir.Model(graph, ir_version=10) + assert FoldWeightLayoutPass()(model).modified + assert not graph.initializers + for branch in (first, second): + assert [item.op_type for item in branch] == ["Identity"] + assert set(branch.initializers) == {"layout"} + + +@pytest.mark.parametrize("parent_capture", [False, True]) +def test_post_weight_pipeline_preserves_overridable_concat_input(tmp_path, parent_capture): + import onnxruntime as ort + + data = np.arange(8, dtype=np.float32) + weight = _initializer("W", data) + shape = _initializer("shape", np.array([2, 4], dtype=np.int64)) + bias = _initializer("B", np.zeros((1, 4), dtype=np.float32)) + reshape = ir.Node("", "Reshape", [weight, shape], num_outputs=1) + reshape.outputs[0].name = "layout" + reshape.outputs[0].shape = ir.Shape([2, 4]) + reshape.outputs[0].dtype = ir.DataType.FLOAT + concat = ir.Node( + "", + "Concat", + [reshape.outputs[0], bias], + [ir.AttrInt64("axis", 0)], + num_outputs=1, + ) + concat.outputs[0].name = "output" + concat.outputs[0].shape = ir.Shape([3, 4]) + concat.outputs[0].dtype = ir.DataType.FLOAT + graph = ir.Graph( + [bias], + concat.outputs, + nodes=[reshape, concat], + initializers=[weight, shape, bias], + name="overridable_concat", + opset_imports={"": 24}, + ) + if parent_capture: + graph.inputs.clear() + del graph.initializers["B"] + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + fallback = ir.Node("", "Identity", [bias], num_outputs=1) + fallback.outputs[0].name = "fallback" + fallback.outputs[0].shape = bias.shape + fallback.outputs[0].dtype = bias.dtype + other = ir.Graph([], fallback.outputs, nodes=[fallback], name="other_branch") + select = ir.Node( + "", + "If", + [condition], + [ir.AttrGraph("then_branch", graph), ir.AttrGraph("else_branch", other)], + num_outputs=1, + ) + select.outputs[0].name = "result" + select.outputs[0].shape = ir.Shape([None, 4]) + select.outputs[0].dtype = ir.DataType.FLOAT + graph = ir.Graph( + [condition, bias], + select.outputs, + nodes=[select], + initializers=[bias], + name="captured_overridable_concat", + opset_imports={"": 24}, + ) + model = ir.Model(graph, ir_version=10) + fold_initializers_after_weights(model) + assert all(node.op_type != "Reshape" for node in model.graph.all_nodes()) + assert sum(node.op_type == "Concat" for node in model.graph.all_nodes()) == 1 + assert model.graph.initializers["B"] is bias + assert bias in model.graph.inputs + + ModelPackage({"model": model}).save(tmp_path, progress_bar=False, max_workers=1) + restored = ModelPackage.load(tmp_path)["model"] + for index, candidate in enumerate((model, restored)): + model_path = tmp_path / f"runtime_{index}.onnx" + ir.save(candidate, model_path) + # The test wrapper filters out overridable initializers via get_inputs(); + # use ORT directly so B is actually supplied to the runtime. + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + assert [value.name for value in session.get_overridable_initializers()] == ["B"] + for override in (None, np.full((1, 4), 17, dtype=np.float32)): + feeds = {"condition": np.asarray(True)} if parent_capture else {} + if override is not None: + feeds["B"] = override + actual = session.run(None, feeds)[0] + expected = np.concatenate( + [ + data.reshape(2, 4), + np.zeros((1, 4), dtype=np.float32) if override is None else override, + ] + ) + np.testing.assert_array_equal(actual, expected) + + +def test_post_weight_pipeline_concat_layouts_are_isolated_in_if_branches(tmp_path): + branches = [] + expected_outputs = [] + for offset in (0, 100): + data = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + offset + zeros = np.zeros((1, 4, 3), dtype=np.float32) + # Each lexical scope deliberately uses the same names for different + # payloads. The rank-3 Transpose is folded by FoldWeightLayoutPass, + # which exposes an all-initializer Concat in each branch. + weight = _initializer("W", data) + padding = _initializer("Z", zeros) + transpose = ir.Node( + "", "Transpose", [weight], [ir.AttrInt64s("perm", [0, 2, 1])], num_outputs=1 + ) + transpose.outputs[0].name = "layout" + transpose.outputs[0].shape = ir.Shape([2, 4, 3]) + transpose.outputs[0].dtype = ir.DataType.FLOAT + concat = ir.Node( + "", + "Concat", + [transpose.outputs[0], padding], + [ir.AttrInt64("axis", 0)], + num_outputs=1, + ) + concat.outputs[0].name = "branch_output" + concat.outputs[0].shape = ir.Shape([3, 4, 3]) + concat.outputs[0].dtype = ir.DataType.FLOAT + branches.append( + ir.Graph( + [], + [concat.outputs[0]], + nodes=[transpose, concat], + initializers=[weight, padding], + name=f"branch_{offset}", + ) + ) + expected_outputs.append(np.concatenate([data.transpose(0, 2, 1), zeros], axis=0)) + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + select = ir.Node( + "", + "If", + [condition], + [ + ir.AttrGraph("then_branch", branches[0]), + ir.AttrGraph("else_branch", branches[1]), + ], + num_outputs=1, + ) + select.outputs[0].name = "result" + select.outputs[0].shape = ir.Shape([3, 4, 3]) + select.outputs[0].dtype = ir.DataType.FLOAT + model = ir.Model( + ir.Graph( + [condition], + [select.outputs[0]], + nodes=[select], + opset_imports={"": 24}, + name="layout_concat_if", + ), + ir_version=10, + ) + + fold_initializers_after_weights(model) + assert not model.graph.initializers + packed_name = "layout__Z__axis_0__concat" + packed_values = [] + for branch, expected in zip(branches, expected_outputs, strict=True): + assert not list(branch) + assert set(branch.initializers) == {packed_name} + packed = branch.initializers[packed_name] + assert branch.outputs[0] is packed + assert isinstance(packed.const_value, ir.LazyTensor) + np.testing.assert_array_equal(packed.const_value.numpy(), expected) + packed_values.append(packed) + assert packed_values[0] is not packed_values[1] + fold_initializers_after_weights(model) + assert not model.graph.initializers + assert all( + branch.initializers[packed_name] is packed + for branch, packed in zip(branches, packed_values, strict=True) + ) + + output_directory = tmp_path / "package" + ModelPackage({"model": model}).save( + str(output_directory), max_workers=1, progress_bar=False + ) + restored = ModelPackage.load(str(output_directory))["model"] + assert not restored.graph.initializers + restored_select = next(iter(restored.graph)) + for attribute, expected in zip( + ("then_branch", "else_branch"), expected_outputs, strict=True + ): + branch = restored_select.attributes[attribute].as_graph() + assert not list(branch) + assert set(branch.initializers) == {packed_name} + packed = branch.initializers[packed_name] + assert branch.outputs[0] is packed + np.testing.assert_array_equal(packed.const_value.numpy(), expected) + + # Both branches must execute their own payload, before and after the public + # package roundtrip, rather than reusing a same-named sibling initializer. + for candidate in (model, restored): + session = OnnxModelSession(candidate, device="cpu") + for take_then, expected in zip((True, False), expected_outputs, strict=True): + result = session.run({"condition": np.asarray(take_then, dtype=np.bool_)}) + np.testing.assert_array_equal(result["result"], expected) + + +@pytest.mark.parametrize("retained", [None, "shared", "graph_output"]) +def test_parent_capture_initializer_ownership_and_persisted_payload(tmp_path, retained): + data = np.arange(6, dtype=np.float32).reshape(2, 3) + source = _initializer("parent_weight", data) + transpose = ir.Node( + "", "Transpose", [source], [ir.AttrInt64s("perm", [1, 0])], num_outputs=1 + ) + transpose.outputs[0].name = "captured_layout" + transpose.outputs[0].shape = ir.Shape([3, 2]) + transpose.outputs[0].dtype = ir.DataType.FLOAT + identity = ir.Node("", "Identity", [transpose.outputs[0]], num_outputs=1) + identity.outputs[0].name = "branch_output" + identity.outputs[0].shape = ir.Shape([3, 2]) + identity.outputs[0].dtype = ir.DataType.FLOAT + branch = ir.Graph([], [identity.outputs[0]], nodes=[transpose, identity]) + other = _initializer("other_weight", np.zeros((3, 2), dtype=np.float32)) + other_branch = ir.Graph([], [other], nodes=[], initializers=[other]) + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + select = ir.Node( + "", + "If", + [condition], + [ir.AttrGraph("then_branch", branch), ir.AttrGraph("else_branch", other_branch)], + num_outputs=1, + ) + select.outputs[0].name = "output" + select.outputs[0].shape = ir.Shape([3, 2]) + select.outputs[0].dtype = ir.DataType.FLOAT + graph = ir.Graph( + [condition], + [select.outputs[0]], + nodes=[select], + initializers=[source], + opset_imports={"": 24}, + ) + if retained == "shared": + extra = ir.Node("", "Identity", [source], num_outputs=1) + extra.outputs[0].name = "shared_source" + extra.outputs[0].shape = source.shape + extra.outputs[0].dtype = source.dtype + graph.append(extra) + graph.outputs.append(extra.outputs[0]) + elif retained == "graph_output": + graph.outputs.append(source) + model = ir.Model(graph, ir_version=10) + result = FoldWeightLayoutPass()(model) + assert result.modified == (retained is None) + if retained is None: + assert "parent_weight" not in graph.initializers + assert set(branch.initializers) == {"captured_layout"} + assert not FoldWeightLayoutPass()(model).modified + np.testing.assert_array_equal( + branch.initializers["captured_layout"].const_value.numpy(), data.T + ) + else: + assert graph.initializers["parent_weight"] is source + assert not branch.initializers + assert transpose.inputs[0] is source + # Exercise the public serializer, not a protobuf encoding. The retained + # parent case must stay shared; the unshared case must serialize only the + # folded child weight, not both it and the dead parent payload. + fold_initializers_after_weights(model) + ModelPackage({"model": model}).save(str(tmp_path), max_workers=1, progress_bar=False) + restored = ModelPackage.load(str(tmp_path))["model"] + persisted = [value for scope in restored.graphs() for value in scope.initializers.values()] + assert sum(value.const_value.numpy().nbytes for value in persisted) == data.nbytes * 2 + if retained is None: + assert "parent_weight" not in restored.graph.initializers + restored_branch = next(iter(restored.graph)).attributes["then_branch"].as_graph() + assert set(restored_branch.initializers) == {"captured_layout"} + np.testing.assert_array_equal( + restored_branch.initializers["captured_layout"].const_value.numpy(), data.T + ) + else: + np.testing.assert_array_equal( + restored.graph.initializers["parent_weight"].const_value.numpy(), data + ) + + +@pytest.mark.parametrize("width", [3, 3 * 8192]) +def test_shape_constant_folding_exposes_layout_without_reading_large_weights(width): + data = np.arange(16 * width, dtype=np.float32).reshape(2, 8, width) + model = _model(data, first=(2, 2, 4, width), last=(2, 8, width)) + source = model.graph.initializers["weight"] + reads = [] + + def read(): + reads.append(True) + return ir.tensor(data) + + source.const_value = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) + shape = model.graph.initializers["shape1"] + offset = _initializer("shape_offset", np.zeros(4, dtype=np.int64)) + model.graph.register_initializer(offset) + expression = ir.Node("", "Add", [shape, offset], num_outputs=1) + expression.outputs[0].name = "computed_shape" + expression.outputs[0].shape = shape.shape + expression.outputs[0].dtype = ir.DataType.INT64 + first = next(iter(model.graph)) + model.graph.insert_before(first, expression) + first.replace_input_with(1, expression.outputs[0]) + fold_initializers_after_weights(model) + assert not reads + assert set(model.graph.initializers) == {"output"} + assert not list(model.graph) + folded = model.graph.initializers["output"] + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + fold_initializers_after_weights(model) + assert not reads + assert model.graph.initializers["output"] is folded + assert model.graph.outputs[0] is folded + expected = data.reshape(2, 2, 4, width).transpose(0, 2, 1, 3).reshape(2, 8, width) + np.testing.assert_array_equal(folded.const_value.numpy(), expected) + assert len(reads) == 1 + + +def test_folded_layout_keeps_pruned_external_source_ownership(tmp_path, monkeypatch): + # The 48-element fixture used elsewhere is only 192 bytes: ir.save keeps + # tensors below its 256-byte default threshold inline, so it has no file + # backing to protect. This fixture must exercise a real ExternalTensor. + width = 24 + data = np.arange(16 * width, dtype=np.float32).reshape(2, 8, width) + base = tmp_path / "base" + model = _model(data, first=(2, 2, 4, width), last=(2, 8, width)) + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + before = {path.name: path.read_bytes() for path in base.iterdir()} + package = ModelPackage.load(str(base)) + model = package["model"] + source = model.graph.initializers["weight"].const_value + assert isinstance(source, ir.ExternalTensor) + source_path = Path(source.path).resolve() + assert source_path == (base / "model.onnx.data").resolve() + + def unexpected_read(*args, **kwargs): + raise AssertionError("source-alias refusal must precede lazy payload reads") + + with monkeypatch.context() as patch: + patch.setattr(ir.ExternalTensor, "numpy", unexpected_read) + fold_initializers_after_weights(model) + folded_graph = model.graph + assert set(folded_graph.initializers) == {"output"} + folded = folded_graph.initializers["output"] + assert folded_graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + assert source_path in { + Path(path).resolve() for path in model.meta[LAZY_SOURCE_FILES_META_KEY] + } + assert source_path in { + Path(path).resolve() + for path in folded.const_value.meta[LAZY_SOURCE_FILES_META_KEY] + } + with pytest.raises(ValueError, match="aliases lazy source file"): + package.save(str(base), max_workers=1, progress_bar=False) + assert source_path in package._native_streaming_source_files + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + safe = tmp_path / "safe" + package.save(str(safe), max_workers=1, progress_bar=False) + restored_graph = ModelPackage.load(str(safe))["model"].graph + assert set(restored_graph.initializers) == {"output"} + restored = restored_graph.initializers["output"] + assert restored_graph.outputs[0] is restored + expected = data.reshape(2, 2, 4, width).transpose(0, 2, 1, 3).reshape(2, 8, width) + np.testing.assert_array_equal(restored.const_value.numpy(), expected)