#!/usr/bin/env python3 """Add tiny block-boundary outputs to a DiT graph and record CPU ORT references.""" from __future__ import annotations import argparse import hashlib import json import os from pathlib import Path from typing import Any INPUT_SHAPES = { "latent": (1, 8192, 16), "camera": (1, 1, 5), "t": (1,), "feature1": (1, 4101, 1280), "feature2": (1, 4101, 128), } FIXTURE_FILES = (*INPUT_SHAPES, "pred_latent", "pred_camera") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--graph", type=Path, required=True) parser.add_argument("--sidecar", type=Path, required=True) parser.add_argument("--input-fixture-dir", type=Path, required=True) parser.add_argument("--output-graph", type=Path, required=True) parser.add_argument("--output-fixture-dir", type=Path, required=True) parser.add_argument( "--probe-set", choices=("blocks", "noise0", "context0"), default="blocks", help=( "Block progression, detailed first noise-refiner probes, or bounded " "first context-refiner attention probes (default: %(default)s)." ), ) parser.add_argument( "--inline-node-prefix", action="append", default=[], help=( "Inline external initializers consumed by nodes whose names start with this prefix. " "May be repeated; intended for diagnosing browser external-data boundaries." ), ) parser.add_argument("--session-threads", type=int, default=8) return parser.parse_args() def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: while block := stream.read(8 * 1024 * 1024): digest.update(block) return digest.hexdigest() def replace_with_hard_link(source: Path, destination: Path) -> None: destination.parent.mkdir(parents=True, exist_ok=True) if destination.exists(): if os.path.samefile(source, destination): return destination.unlink() os.link(source, destination) def probe_sources( probe_set: str, available: set[str], shapes: dict[str, list[int]], ) -> list[tuple[str, str, list[int] | None]]: vector = [1, 4, 16] if probe_set == "context0": context_vector = [1, 4, 1024] attention_rows = [1, 16, 4, 4101] return [ ( "context0_block_input", "/flow_model/context_refiner.0/norm1/Cast_output_0", context_vector, ), ( "context0_attention_input", "/flow_model/context_refiner.0/norm1/Cast_1_output_0", context_vector, ), ( "context0_qkv", "/flow_model/context_refiner.0/attn/Reshape_output_0", None, ), ( "context0_q_after_rope", "/flow_model/context_refiner.0/attn/Reshape_2_output_0", None, ), ( "context0_k_after_rope", "/flow_model/context_refiner.0/attn/Reshape_4_output_0", None, ), ( "context0_v", "/flow_model/context_refiner.0/attn/Squeeze_2_output_0", None, ), ( "context0_q_normalized", "/flow_model/context_refiner.0/attn/q_norm/Slice_1_output_0", None, ), ( "context0_k_normalized", "/flow_model/context_refiner.0/attn/k_norm/Slice_1_output_0", None, ), ( "context0_q_transposed", "/flow_model/context_refiner.0/attn/Transpose_output_0", attention_rows[:-1] + [64], ), ( "context0_k_transposed", "/flow_model/context_refiner.0/attn/Transpose_1_output_0", [1, 16, 4101, 64], ), ( "context0_v_transposed", "/flow_model/context_refiner.0/attn/Transpose_2_output_0", [1, 16, 4101, 64], ), ( "context0_scaled_logits_rows", "/flow_model/context_refiner.0/attn/MatMul_output_0", attention_rows, ), ( "context0_probabilities_rows", "/flow_model/context_refiner.0/attn/Softmax_output_0", attention_rows, ), ( "context0_weighted_value_rows", "/flow_model/context_refiner.0/attn/MatMul_1_output_0", [1, 16, 4, 64], ), ( "context0_pre_projection", "/flow_model/context_refiner.0/attn/Reshape_5_output_0", context_vector, ), ( "context0_post_projection", "/flow_model/context_refiner.0/attn/out/Reshape_output_0", context_vector, ), ( "context0_residual", "/flow_model/context_refiner.0/Add_output_0", context_vector, ), ] if probe_set == "noise0": attention = [1, 2, 4, 8] full_attention_candidates = [ name for name, shape in shapes.items() if name.startswith("/flow_model/noise_refiner.0/attn/Concat") and len(shape) == 4 and shape[1:] == [16, 8192, 64] ] if len(full_attention_candidates) != 1: raise ValueError( "Expected exactly one full 16-head attention concatenation, found " f"{full_attention_candidates}" ) full_attention = full_attention_candidates[0] attention_head_groups = [ name for name, shape in shapes.items() if name.startswith("/flow_model/noise_refiner.0/attn/Concat") and len(shape) == 4 and shape[1:] == [8, 8192, 64] ] projected = next( candidate for candidate in ( "/flow_model/noise_refiner.0/attn/out/Reshape_1_output_0", "/flow_model/noise_refiner.0/attn/out/Reshape_output_0", "/flow_model/noise_refiner.0/attn/out/Add_output_0", ) if candidate in available ) result = [ ("input_positioned", "/flow_model/Add_1_output_0", vector), ("noise0_modulated", "/flow_model/noise_refiner.0/Add_3_output_0", vector), ("noise0_qkv", "/flow_model/noise_refiner.0/attn/qkv/Add_output_0", vector), ("noise0_qkv_full", "/flow_model/noise_refiner.0/attn/qkv/Add_output_0", None), ("noise0_q_norm", "/flow_model/noise_refiner.0/attn/q_norm/Cast_1_output_0", attention), ("noise0_q_norm_full", "/flow_model/noise_refiner.0/attn/q_norm/Cast_1_output_0", None), ("noise0_k_norm", "/flow_model/noise_refiner.0/attn/k_norm/Cast_1_output_0", attention), ("noise0_k_norm_full", "/flow_model/noise_refiner.0/attn/k_norm/Cast_1_output_0", None), ("noise0_q_transposed", "/flow_model/noise_refiner.0/attn/Transpose_output_0", attention), ("noise0_q_transposed_full", "/flow_model/noise_refiner.0/attn/Transpose_output_0", None), ("noise0_k_transposed_full", "/flow_model/noise_refiner.0/attn/Transpose_1_output_0", None), ("noise0_v_transposed_full", "/flow_model/noise_refiner.0/attn/Transpose_2_output_0", None), ("noise0_q_chunk0", "/flow_model/noise_refiner.0/attn/Cast_6_output_0", attention), ("noise0_score0", "/flow_model/noise_refiner.0/attn/MatMul_output_0", attention), ("noise0_softmax0", "/flow_model/noise_refiner.0/attn/Softmax_output_0", attention), ("noise0_value0", "/flow_model/noise_refiner.0/attn/MatMul_1_output_0", attention), ("noise0_attention_all", full_attention, attention), ("noise0_attention_reshaped", "/flow_model/noise_refiner.0/attn/Reshape_5_output_0", vector), ("noise0_attention_projected", projected, vector), ("noise0_attention_gated", "/flow_model/noise_refiner.0/Mul_7_output_0", vector), ("noise0_attention_residual", "/flow_model/noise_refiner.0/Add_4_output_0", vector), ("noise0_mlp_normalized", "/flow_model/noise_refiner.0/norm2/Cast_1_output_0", vector), ("noise0_mlp_hidden", "/flow_model/noise_refiner.0/mlp/mlp/mlp.1/Gelu_output_0", vector), ("noise0_output", "/flow_model/noise_refiner.0/Add_7_output_0", vector), ] if ( "/flow_model/noise_refiner.0/attn/out/MatMul_output_0" in available and "/flow_model/noise_refiner.0/attn/out/Reshape_1_output_0" in available ): projection_steps = [ ( f"noise0_projection_input_tile_{index}", ( "/flow_model/noise_refiner.0/attn/out/Slice_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Slice_{index}_output_0" ), [4, 8], ) for index in range(4) ] projection_steps.extend( ( f"noise0_projection_partial_{index}", ( "/flow_model/noise_refiner.0/attn/out/MatMul_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/MatMul_{index}_output_0" ), [4, 8], ) for index in range(4) ) projection_steps.extend( ( name, source, [4, 8], ) for name, source in ( ("noise0_projection_sum_01", "/flow_model/noise_refiner.0/attn/out/Add_output_0"), ("noise0_projection_sum_23", "/flow_model/noise_refiner.0/attn/out/Add_1_output_0"), ("noise0_projection_sum", "/flow_model/noise_refiner.0/attn/out/Add_2_output_0"), ("noise0_projection_bias", "/flow_model/noise_refiner.0/attn/out/Add_3_output_0"), ("noise0_projection_cast", "/flow_model/noise_refiner.0/attn/out/Cast_4_output_0"), ) ) result[12:12] = projection_steps elif "/flow_model/noise_refiner.0/attn/out/MatMul_output_0" in available: result[12:12] = [ ( "noise0_projection_channels_first", "/flow_model/noise_refiner.0/attn/out/Slice_output_0", [1, 8, 4], ), ( "noise0_projection_weight_left", "/flow_model/noise_refiner.0/attn/out/MatMul_output_0", [1, 8, 4], ), ( "noise0_projection_weight_left_transposed", "/flow_model/noise_refiner.0/attn/out/Transpose_1_output_0", [1, 4, 8], ), ] elif shapes.get("/flow_model/noise_refiner.0/attn/out/Add_9_output_0") == [1, 1, 256, 256]: fixed_256_tree = [ *[ ( f"noise0_attention_head_group_{index}_full", source, None, ) for index, source in enumerate(attention_head_groups) ], ( "noise0_attention_all_full", full_attention, None, ), ( "noise0_attention_reshaped_full", "/flow_model/noise_refiner.0/attn/Reshape_5_output_0", None, ), ( "noise0_projection_first_token_full", "/flow_model/noise_refiner.0/attn/out/Slice_output_0", None, ), ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), ( "noise0_projection_special_products_full", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", None, ), ( "noise0_projection_products_transposed", "/flow_model/noise_refiner.0/attn/out/Transpose_2_output_0", [1, 1, 4, 16], ), ( "noise0_projection_products_transposed_full", "/flow_model/noise_refiner.0/attn/out/Transpose_2_output_0", None, ), *[ ( f"noise0_projection_tree_{index}", ( "/flow_model/noise_refiner.0/attn/out/Add_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Add_{index}_output_0" ), [1, 1, 4, 16], ) for index in range(10) ], ( "noise0_projection_tree_0_full", "/flow_model/noise_refiner.0/attn/out/Add_output_0", None, ), ( "noise0_projection_tree_1_full", "/flow_model/noise_refiner.0/attn/out/Add_1_output_0", None, ), ( "noise0_projection_fixed_gather_left", "/flow_model/noise_refiner.0/attn/out/Gather_1_output_0", [1, 1, 4, 16], ), ( "noise0_projection_fixed_gather_right", "/flow_model/noise_refiner.0/attn/out/Gather_2_output_0", [1, 1, 4, 16], ), ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/Gather_17_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_10_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 4, 16], ), ] result[12:12] = fixed_256_tree elif shapes.get("/flow_model/noise_refiner.0/attn/out/Add_9_output_0") == [1, 1, 128, 256]: fixed_tree = [ ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), ( "noise0_projection_products_transposed", "/flow_model/noise_refiner.0/attn/out/Transpose_2_output_0", [1, 1, 4, 16], ), *[ ( f"noise0_projection_tree_{index}", ( "/flow_model/noise_refiner.0/attn/out/Add_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Add_{index}_output_0" ), [1, 1, 4, 16], ) for index in range(10) ], ( "noise0_projection_tree_2_full", "/flow_model/noise_refiner.0/attn/out/Add_2_output_0", None, ), ( "noise0_projection_fixed_gather_left", "/flow_model/noise_refiner.0/attn/out/Gather_1_output_0", [1, 1, 4, 16], ), ( "noise0_projection_fixed_gather_right", "/flow_model/noise_refiner.0/attn/out/Gather_2_output_0", [1, 1, 4, 16], ), ( "noise0_projection_fixed_masked_left", "/flow_model/noise_refiner.0/attn/out/Mul_1_output_0", [1, 1, 4, 16], ), ( "noise0_projection_fixed_masked_right", "/flow_model/noise_refiner.0/attn/out/Mul_2_output_0", [1, 1, 4, 16], ), ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/Gather_15_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_10_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 4, 16], ), ] result[12:12] = fixed_tree elif "/flow_model/noise_refiner.0/attn/out/Gather_15_output_0" in available: gathered_tree = [ ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), *[ ( f"noise0_projection_tree_{index}", ( "/flow_model/noise_refiner.0/attn/out/Add_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Add_{index}_output_0" ), [1, 1, 4, min(8, 512 >> index)] if index < 3 else [1, 1, 4, 128], ) for index in range(10) ], ( "noise0_projection_gathered_left", "/flow_model/noise_refiner.0/attn/out/Gather_1_output_0", [1, 1, 4, 8], ), ( "noise0_projection_gathered_right", "/flow_model/noise_refiner.0/attn/out/Gather_2_output_0", [1, 1, 4, 8], ), ( "noise0_projection_padded_left", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 1, 4, 128], ), ( "noise0_projection_padded_right", "/flow_model/noise_refiner.0/attn/out/Concat_1_output_0", [1, 1, 4, 128], ), ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/Gather_15_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_10_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_14_output_0", [1, 4, 16], ), ] result[12:12] = gathered_tree elif "/flow_model/noise_refiner.0/attn/out/Concat_14_output_0" in available: padded_tree = [ ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), *[ ( f"noise0_projection_tree_{index}", ( "/flow_model/noise_refiner.0/attn/out/Add_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Add_{index}_output_0" ), [1, 1, 4, min(8, 512 >> index)] if index < 3 else [1, 1, 4, 128], ) for index in range(10) ], ( "noise0_projection_padded_left", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 1, 4, 128], ), ( "noise0_projection_padded_right", "/flow_model/noise_refiner.0/attn/out/Concat_1_output_0", [1, 1, 4, 128], ), ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/Gather_1_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_10_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_14_output_0", [1, 4, 16], ), ] result[12:12] = padded_tree elif shapes.get("/flow_model/noise_refiner.0/attn/out/Add_9_output_0") == [1, 1, 1, 256]: transposed_tree = [ ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), ( "noise0_projection_products_transposed", "/flow_model/noise_refiner.0/attn/out/Transpose_2_output_0", [1, 1, 4, 16], ), ( "noise0_projection_first_left", "/flow_model/noise_refiner.0/attn/out/Slice_1_output_0", [1, 1, 4, 16], ), ( "noise0_projection_first_right", "/flow_model/noise_refiner.0/attn/out/Slice_2_output_0", [1, 1, 4, 16], ), *[ ( f"noise0_projection_tree_{index}", ( "/flow_model/noise_refiner.0/attn/out/Add_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Add_{index}_output_0" ), [1, 1, min(4, 512 >> index), 16], ) for index in range(10) ], ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/Squeeze_1_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_10_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 4, 16], ), ] result[12:12] = transposed_tree elif "/flow_model/noise_refiner.0/attn/out/Add_9_output_0" in available: tree_levels = [ ( f"noise0_projection_tree_{index}", ( "/flow_model/noise_refiner.0/attn/out/Add_output_0" if index == 0 else f"/flow_model/noise_refiner.0/attn/out/Add_{index}_output_0" ), [1, 1, 4, min(8, 512 >> index)], ) for index in range(10) ] result[12:12] = [ ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), *tree_levels, ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/Squeeze_1_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_10_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 4, 16], ), ] elif "/flow_model/noise_refiner.0/attn/out/ReduceSum_output_0" in available: result[12:12] = [ ( "noise0_projection_special_products", "/flow_model/noise_refiner.0/attn/out/Mul_output_0", [1, 1, 4, 8], ), ( "noise0_projection_special_sum", "/flow_model/noise_refiner.0/attn/out/ReduceSum_output_0", [1, 1, 16], ), ( "noise0_projection_special_bias", "/flow_model/noise_refiner.0/attn/out/Add_output_0", [1, 1, 16], ), ( "noise0_projection_bulk_tail", "/flow_model/noise_refiner.0/attn/out/Gather_output_0", [1, 4, 16], ), ( "noise0_projection_corrected_chunk", "/flow_model/noise_refiner.0/attn/out/Concat_output_0", [1, 4, 16], ), ] return result result = [ ("input_linear", "/flow_model/input_layer/Add_output_0", vector), ("input_positioned", "/flow_model/Add_1_output_0", vector), ("condition_dino", "/flow_model/cond_embedder/Add_output_0", vector), ("condition_vae", "/flow_model/cond_embedder2/Add_output_0", vector), ("condition_combined", "/flow_model/Add_output_0", vector), ( "noise_refiner_00_attention_residual", "/flow_model/noise_refiner.0/Add_4_output_0", vector, ), ("noise_refiner_00", "/flow_model/noise_refiner.0/Add_7_output_0", vector), ( "noise_refiner_01_attention_residual", "/flow_model/noise_refiner.1/Add_4_output_0", vector, ), ("noise_refiner_01", "/flow_model/noise_refiner.1/Add_7_output_0", vector), ( "context_refiner_00_attention_residual", "/flow_model/context_refiner.0/Add_output_0", vector, ), ("context_refiner_00", "/flow_model/context_refiner.0/Add_1_output_0", vector), ( "context_refiner_01_attention_residual", "/flow_model/context_refiner.1/Add_output_0", vector, ), ("context_refiner_01", "/flow_model/context_refiner.1/Add_1_output_0", vector), ("combined_tokens", "/flow_model/Concat_output_0", vector), ( "block_00_attention_residual", "/flow_model/blocks.0/Add_4_output_0", vector, ), ] result.extend( (f"block_{index:02d}", f"/flow_model/blocks.{index}/Add_7_output_0", vector) for index in range(24) ) return result def external_data_fields(tensor: Any) -> dict[str, str]: return {entry.key: entry.value for entry in tensor.external_data} def inline_selected_initializers(model: Any, sidecar: Path, node_prefixes: list[str]) -> list[str]: if not node_prefixes: return [] selected_inputs = { input_name for node in model.graph.node if any(node.name.startswith(prefix) for prefix in node_prefixes) for input_name in node.input } inlined: list[str] = [] with sidecar.open("rb") as stream: for tensor in model.graph.initializer: if tensor.name not in selected_inputs or not tensor.external_data: continue fields = external_data_fields(tensor) location = fields.get("location") if location and Path(location).name != sidecar.name: raise ValueError( f"Initializer {tensor.name!r} points to {location!r}, not {sidecar.name!r}." ) offset = int(fields.get("offset", "0")) length = int(fields["length"]) stream.seek(offset) raw_data = stream.read(length) if len(raw_data) != length: raise EOFError( f"Initializer {tensor.name!r} read {len(raw_data)} bytes at {offset}; expected {length}." ) tensor.raw_data = raw_data tensor.ClearField("external_data") tensor.data_location = 0 inlined.append(tensor.name) if not inlined: raise ValueError(f"No external initializers matched node prefixes: {node_prefixes}") return inlined def main() -> None: try: import numpy as np import onnx import onnxruntime as ort from onnx import helper except ImportError as exc: raise SystemExit(f"ONNX, ONNX Runtime, and NumPy are required: {exc}") from exc args = parse_args() graph = args.graph.expanduser().resolve() sidecar = args.sidecar.expanduser().resolve() fixture_dir = args.input_fixture_dir.expanduser().resolve() output_graph = args.output_graph.expanduser().resolve() output_fixture = args.output_fixture_dir.expanduser().resolve() for path in (graph, sidecar): if not path.is_file(): raise FileNotFoundError(path) model = onnx.load_model(str(graph), load_external_data=False) available = {value for node in model.graph.node for value in node.output} inferred_model = onnx.shape_inference.infer_shapes(model) shapes = { value.name: [dimension.dim_value for dimension in value.type.tensor_type.shape.dim] for value in (*inferred_model.graph.value_info, *inferred_model.graph.output) if value.type.HasField("tensor_type") } # ONNX shape inference loses the fixed batch value at attention # concatenations created by the statically unrolled chunk loops. All public # inputs are fixed to batch one, so make that contract explicit on probes. for shape in shapes.values(): if shape and shape[0] == 0: shape[0] = 1 sources = probe_sources(args.probe_set, available, shapes) missing = [source for _, source, _ in sources if source not in available] if missing: raise KeyError(f"Graph is missing probe values: {missing}") inlined_initializers = inline_selected_initializers(model, sidecar, args.inline_node_prefix) include_final_outputs = args.probe_set == "blocks" if not include_final_outputs: del model.graph.output[:] for index, (name, source, ends) in enumerate(sources): prefix = f"triposplat_probe_{index:02d}" output_shape = shapes[source] if ends is None else ends cast_input = source if ends is not None: constants = { "starts": [0] * len(ends), "ends": ends, "axes": list(range(len(ends))), "steps": [1] * len(ends), } inputs = [] for suffix, values in constants.items(): initializer_name = f"{prefix}_{suffix}" model.graph.initializer.append( helper.make_tensor( initializer_name, onnx.TensorProto.INT64, [len(values)], values, ) ) inputs.append(initializer_name) slice_output = f"{prefix}_slice" model.graph.node.append( helper.make_node( "Slice", [source, *inputs], [slice_output], name=f"/{prefix}/Slice", ) ) cast_input = slice_output output_name = f"probe_{name}" model.graph.node.append( helper.make_node( "Cast", [cast_input], [output_name], name=f"/{prefix}/Cast", to=onnx.TensorProto.FLOAT, ) ) model.graph.output.append( helper.make_tensor_value_info( output_name, onnx.TensorProto.FLOAT, output_shape, ) ) sidecar_name = output_graph.name + ".data" for tensor in model.graph.initializer: for entry in tensor.external_data: if entry.key == "location": entry.value = sidecar_name metadata = {entry.key: entry.value for entry in model.metadata_props} metadata["triposplat.diagnostic_probes"] = str(len(sources)) if inlined_initializers: metadata["triposplat.diagnostic_inlined_initializers"] = ",".join(inlined_initializers) del model.metadata_props[:] for key, value in sorted(metadata.items()): entry = model.metadata_props.add() entry.key = key entry.value = value output_graph.parent.mkdir(parents=True, exist_ok=True) onnx.save_model(model, str(output_graph), save_as_external_data=False) output_sidecar = output_graph.parent / sidecar_name replace_with_hard_link(sidecar, output_sidecar) output_fixture.mkdir(parents=True, exist_ok=True) for name in FIXTURE_FILES: replace_with_hard_link(fixture_dir / f"{name}.f32", output_fixture / f"{name}.f32") feeds = { name: np.fromfile(fixture_dir / f"{name}.f32", dtype="