diff --git a/src/winml/modelkit/quant/fp16.py b/src/winml/modelkit/quant/fp16.py index 0294b8076..ff207b414 100644 --- a/src/winml/modelkit/quant/fp16.py +++ b/src/winml/modelkit/quant/fp16.py @@ -4,24 +4,2886 @@ # -------------------------------------------------------------------------- """FP16 conversion utility for ONNX models. -Provides a single entry point for FP32→FP16 model conversion, used by +Provides a single entry point for FP32->FP16 model conversion, used by the quantizer's ``mode="fp16"`` path. """ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from copy import deepcopy +from dataclasses import dataclass +from math import prod +from typing import TYPE_CHECKING, cast from google.protobuf.message import EncodeError if TYPE_CHECKING: - from onnx import ModelProto + from collections.abc import Sequence + + from onnx import ( + AttributeProto, + FunctionProto, + GraphProto, + ModelProto, + NodeProto, + TensorProto, + TypeProto, + ValueInfoProto, + ) + from onnx.defs import OpSchema logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class _InitializerOutput: + """A graph output supplied directly by an initializer in an ORT-traversed graph.""" + + graph_index: int + name: str + output_index: int + has_consumers: bool + + +def _tensor_data_is_loaded(initializer: TensorProto) -> bool: + """Whether a FLOAT tensor carries resident data rather than only a sidecar ref.""" + if initializer.raw_data or initializer.float_data: + return True + return prod(initializer.dims) == 0 + + +def _effective_blocked_ops(op_block_list: list[str] | None) -> set[str]: + """Return the op types ORT skips for this wrapper's exposed block-list option.""" + from onnxruntime.transformers.float16 import DEFAULT_OP_BLOCK_LIST + + return set(DEFAULT_OP_BLOCK_LIST if op_block_list is None else op_block_list) + + +def _ort_inference_preflight_model(model: ModelProto) -> ModelProto: + """Run ORT's normal shape-inference preflight on an isolated clone.""" + from onnx import ModelProto as ONNXModelProto + from onnx import shape_inference + + candidate: object = model + if not isinstance(candidate, ONNXModelProto): + return model + try: + return shape_inference.infer_shapes(deepcopy(model)) + except EncodeError: + return model + + +def _graph_tensor_names(graph: GraphProto) -> set[str]: + """Collect tensor names in one ONNX lexical scope.""" + names = { + value.name + for values in ( + getattr(graph, "input", []), + getattr(graph, "output", []), + getattr(graph, "value_info", []), + ) + for value in values + } + names.update(initializer.name for initializer in getattr(graph, "initializer", [])) + names.update(sparse.values.name for sparse in getattr(graph, "sparse_initializer", [])) + names.update( + name for node in getattr(graph, "node", []) for name in (*node.input, *node.output) if name + ) + return names + + +def _all_node_names(model: ModelProto, op_block_list: list[str] | None) -> set[str]: + """Collect generated-name collisions whose skipped nodes need conversion.""" + return { + node.name + for graph in _ort_traversed_graphs(model, op_block_list) + for node in getattr(graph, "node", []) + if node.name and not _node_is_conversion_neutral(model, graph, node) + } + + +def _node_is_conversion_neutral(model: ModelProto, graph: GraphProto, node: NodeProto) -> bool: + """Whether ORT can skip a colliding node without changing its precision.""" + if node.attribute or any( + input_name and not _node_input_is_proven_non_float(model, graph, node, input_index) + for input_index, input_name in enumerate(node.input) + ): + return False + for output_name in node.output: + if not output_name: + continue + output_types = _graph_declared_types(graph, output_name) + if not output_types or any( + not _type_proto_is_concrete(value_type) or _type_proto_contains_float_tensor(value_type) + for value_type in output_types + ): + return False + return True + + +def _all_graphs(model: ModelProto) -> list[GraphProto]: + """Return the top-level graph and all nested attribute graphs.""" + return [model.graph, *_iter_nested_graphs(model)] + + +def _iter_nested_graphs(model: ModelProto) -> list[GraphProto]: + """Return nested attribute graphs without relying on tensor-name scope.""" + from onnx import AttributeProto + + nested: list[GraphProto] = [] + pending = [model.graph] + while pending: + graph = pending.pop() + for node in graph.node: + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + nested.append(attribute.g) + pending.append(attribute.g) + elif attribute.type == AttributeProto.GRAPHS: + nested.extend(attribute.graphs) + pending.extend(attribute.graphs) + return nested + + +def _ort_traversed_graphs(model: ModelProto, op_block_list: list[str] | None) -> list[GraphProto]: + """Return graphs ORT's FP16 converter visits for the given op block list.""" + blocked_ops = _effective_blocked_ops(op_block_list) + traversed: list[GraphProto] = [] + pending = [model.graph] + while pending: + graph = pending.pop() + traversed.append(graph) + pending.extend(_iter_ort_child_graphs(graph, blocked_ops)) + return traversed + + +def _direct_initializer_outputs_in_graph( + graph: GraphProto, + *, + data_types: set[int] | None = None, +) -> list[TensorProto]: + """Return direct graph-output initializers, optionally filtered by data type.""" + if not hasattr(graph, "output") or not hasattr(graph, "initializer"): + return [] + + produced = {name for node in getattr(graph, "node", []) for name in node.output if name} + output_names = {output.name for output in graph.output} + return [ + initializer + for initializer in graph.initializer + if initializer.name in output_names + and initializer.name not in produced + and (data_types is None or initializer.data_type in data_types) + ] + + +def _direct_initializer_outputs( + model: ModelProto, + *, + data_types: set[int] | None = None, + graphs: list[GraphProto] | None = None, +) -> list[TensorProto]: + """Return direct initializer-backed outputs across the requested graph scopes.""" + return [ + initializer + for graph in (graphs if graphs is not None else _all_graphs(model)) + for initializer in _direct_initializer_outputs_in_graph(graph, data_types=data_types) + ] + + +def _has_nested_initializer_outputs(model: ModelProto, op_block_list: list[str] | None) -> bool: + """Whether any traversed nested graph output is supplied directly by a FLOAT initializer.""" + from onnx import TensorProto + + return any( + _direct_initializer_outputs_in_graph(graph, data_types={TensorProto.FLOAT}) + for graph in _ort_traversed_graphs(model, op_block_list)[1:] + ) + + +def _has_float_sparse_initializers(model: ModelProto, op_block_list: list[str] | None) -> bool: + """Whether any ORT-traversed graph has sparse FLOAT initializer values.""" + from onnx import TensorProto + + return any( + sparse.values.data_type == TensorProto.FLOAT + for graph in _ort_traversed_graphs(model, op_block_list) + for sparse in getattr(graph, "sparse_initializer", []) + ) + + +def _reject_sparse_initializer_tensor_metadata( + model: ModelProto, op_block_list: list[str] | None +) -> None: + """Reject sparse initializer metadata that conflicts with ONNX sparse typing.""" + for graph in _ort_traversed_graphs(model, op_block_list): + sparse_names = {sparse.values.name for sparse in getattr(graph, "sparse_initializer", [])} + if not sparse_names: + continue + for value_info in ( + *getattr(graph, "input", []), + *getattr(graph, "output", []), + *getattr(graph, "value_info", []), + ): + if value_info.name in sparse_names and value_info.type.HasField("tensor_type"): + graph_name = graph.name or "" + msg = ( + f"Sparse initializer '{graph_name}.{value_info.name}' has " + "tensor_type metadata; sparse initializer metadata must use " + "sparse_tensor_type." + ) + raise RuntimeError(msg) + + +def _reject_duplicate_float_initializer_names( + model: ModelProto, op_block_list: list[str] | None +) -> None: + """Reject FLOAT initializer names that ORT's global conversion map cannot scope.""" + from onnx import TensorProto + + seen: set[str] = set() + duplicates: set[str] = set() + for graph in _ort_traversed_graphs(model, op_block_list): + if not hasattr(graph, "initializer"): + continue + for initializer in graph.initializer: + if initializer.data_type != TensorProto.FLOAT: + continue + if initializer.name in seen: + duplicates.add(initializer.name) + else: + seen.add(initializer.name) + if duplicates: + names = ", ".join(sorted(duplicates)) + msg = ( + "FP16 conversion cannot safely process duplicate FLOAT initializer " + f"names across graph scopes: {names}." + ) + raise RuntimeError(msg) + + +def _ort_graphs_with_parents( + model: ModelProto, blocked_ops: set[str] +) -> tuple[list[GraphProto], dict[int, GraphProto | None]]: + """Return ORT's graph BFS registration order and lexical parents.""" + graphs = [model.graph] + parents: dict[int, GraphProto | None] = {id(model.graph): None} + for graph in graphs: + for child in _iter_ort_child_graphs(graph, blocked_ops): + parents[id(child)] = graph + graphs.append(child) + return graphs, parents + + +def _input_resolves_to_initializer( + graph: GraphProto, + name: str, + owner: GraphProto, + parents: dict[int, GraphProto | None], +) -> bool: + """Whether a graph-local input name resolves to the requested initializer.""" + current: GraphProto | None = graph + while current is not None: + if any(initializer.name == name for initializer in getattr(current, "initializer", [])): + return current is owner + if ( + any(value.name == name for value in getattr(current, "input", [])) + or any( + sparse.values.name == name for sparse in getattr(current, "sparse_initializer", []) + ) + or any( + output_name == name + for node in getattr(current, "node", []) + for output_name in node.output + if output_name + ) + ): + return False + current = parents.get(id(current)) + return False + + +def _initializer_tracking_analysis( + model: ModelProto, + *, + keep_io_types: bool, + blocked_ops: set[str], + graphs: list[GraphProto], + parents: dict[int, GraphProto | None], +) -> tuple[dict[str, GraphProto], set[str], set[str], set[str], set[str]]: + """Model ORT and lexical FLOAT initializer conversion decisions.""" + from onnx import TensorProto + + name_mapping, io_casts = _ort_keep_io_name_mapping(model, keep_io_types=keep_io_types) + owners: dict[str, GraphProto] = {} + ort_fp16_initializers: set[str] = set() + lexical_fp16_initializers: set[str] = set() + mismatched_names: set[str] = set() + converted_initializer_inputs: set[str] = set() + for graph in graphs: + float_initializers = { + initializer.name: initializer + for initializer in getattr(graph, "initializer", []) + if initializer.data_type == TensorProto.FLOAT + } + owners.update(dict.fromkeys(float_initializers, graph)) + converted_initializer_inputs.update( + value.name + for value in getattr(graph, "input", []) + if value.name in float_initializers + and value.name not in name_mapping + and value.type.HasField("tensor_type") + and value.type.tensor_type.elem_type == TensorProto.FLOAT + ) + if not keep_io_types or graph is not model.graph: + lexical_fp16_initializers.update( + initializer.name + for initializer in _direct_initializer_outputs_in_graph( + graph, data_types={TensorProto.FLOAT} + ) + if not _initializer_output_has_consumers(graph, initializer.name, blocked_ops) + ) + for node in getattr(graph, "node", []): + if node.name in io_casts: + continue + for input_index, original_input_name in enumerate(node.input): + input_name = name_mapping.get(original_input_name, original_input_name) + owner = owners.get(input_name) + if owner is None: + continue + uses_fp16 = not _node_expects_fp32_input(node, input_index, blocked_ops) + if uses_fp16: + ort_fp16_initializers.add(input_name) + if _input_resolves_to_initializer(graph, input_name, owner, parents): + if uses_fp16: + lexical_fp16_initializers.add(input_name) + else: + mismatched_names.add(input_name) + return ( + owners, + ort_fp16_initializers, + lexical_fp16_initializers, + mismatched_names, + converted_initializer_inputs, + ) + + +def _reject_scope_unsafe_initializer_tracking( + model: ModelProto, + *, + keep_io_types: bool, + op_block_list: list[str] | None, +) -> None: + """Reject lexical conflation only when it changes initializer conversion.""" + from onnx import TensorProto + + blocked_ops = _effective_blocked_ops(op_block_list) + graphs, parents = _ort_graphs_with_parents(model, blocked_ops) + ( + owners, + ort_fp16_initializers, + lexical_fp16_initializers, + mismatched_names, + converted_initializer_inputs, + ) = _initializer_tracking_analysis( + model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + graphs=graphs, + parents=parents, + ) + + unsafe_names = sorted( + name + for name in mismatched_names + if (name in ort_fp16_initializers) != (name in lexical_fp16_initializers) + and ( + _initializer_output_has_consumers(owners[name], name, blocked_ops) + or any( + value.name == name + for value in ( + *getattr(owners[name], "input", []), + *getattr(owners[name], "output", []), + ) + ) + ) + ) + if unsafe_names: + names = ", ".join(unsafe_names) + msg = ( + f"FLOAT initializer names are consumed through different lexical " + f"bindings and change ORT's conversion decision: {names}; ORT's " + "initializer tracking is not scope-aware." + ) + raise RuntimeError(msg) + divergent_inputs = sorted( + converted_initializer_inputs - ort_fp16_initializers - lexical_fp16_initializers + ) + if divergent_inputs: + names = ", ".join(divergent_inputs) + msg = ( + "FLOAT initializer-backed graph input declarations convert to FP16 " + f"while their default initializers remain FLOAT: {names}." + ) + raise RuntimeError(msg) + + name_mapping, _ = _ort_keep_io_name_mapping(model, keep_io_types=keep_io_types) + repaired_outputs = { + initializer.name + for graph in graphs + if not keep_io_types or graph is not model.graph + for initializer in _direct_initializer_outputs_in_graph( + graph, data_types={TensorProto.FLOAT} + ) + if not _initializer_output_has_consumers(graph, initializer.name, blocked_ops) + } + incompatible_metadata = sorted( + name + for name, owner in owners.items() + if name not in ort_fp16_initializers + and name not in repaired_outputs + and name not in name_mapping + and any( + value.name == name and _value_info_enters_ort_global_list(value) + for values in ( + getattr(owner, "input", []), + getattr(owner, "output", []), + getattr(owner, "value_info", []), + ) + for value in values + ) + ) + if incompatible_metadata: + names = ", ".join(incompatible_metadata) + msg = ( + "ORT converts FLOAT initializer metadata to FP16 while retaining " + f"the initializer in FP32: {names}." + ) + raise RuntimeError(msg) + + +def _reject_unloaded_external_initializer_outputs( + model: ModelProto, op_block_list: list[str] | None +) -> None: + """Reject direct FLOAT output initializers whose external data is not resident.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + for initializer in _direct_initializer_outputs( + model, + data_types={TensorProto.FLOAT}, + graphs=_ort_traversed_graphs(model, op_block_list), + ): + if uses_external_data(initializer) and not _tensor_data_is_loaded(initializer): + msg = ( + f"Initializer-backed output '{initializer.name}' uses unloaded external data; " + "load external weights before FP16 conversion." + ) + raise RuntimeError(msg) + + +def _internalize_external_initializer_outputs( + model: ModelProto, op_block_list: list[str] | None +) -> None: + """Drop stale external metadata for resident direct output initializer data.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + for initializer in _direct_initializer_outputs( + model, + data_types={TensorProto.FLOAT}, + graphs=_ort_traversed_graphs(model, op_block_list), + ): + if uses_external_data(initializer): + del initializer.external_data[:] + initializer.data_location = TensorProto.DEFAULT + + +def _ort_converted_initializer_names( + model: ModelProto, + *, + keep_io_types: bool, + blocked_ops: set[str], +) -> set[str]: + """Return FLOAT initializer names ORT selects for FP16 conversion.""" + graphs, parents = _ort_graphs_with_parents(model, blocked_ops) + _, converted, _, _, _ = _initializer_tracking_analysis( + model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + graphs=graphs, + parents=parents, + ) + return converted + + +def _internalize_selected_external_initializers( + model: ModelProto, + names: set[str], + op_block_list: list[str] | None, +) -> None: + """Drop external metadata after selected weights are loaded in memory.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + for graph in _ort_traversed_graphs(model, op_block_list): + for initializer in graph.initializer: + if ( + initializer.name in names + and initializer.data_type == TensorProto.FLOAT + and uses_external_data(initializer) + ): + del initializer.external_data[:] + initializer.data_location = TensorProto.DEFAULT + + +def _graph_node_consumes_name(graph: GraphProto, name: str) -> bool: + """Whether a node in this graph directly consumes the requested name.""" + return any( + input_name == name + for node in getattr(graph, "node", []) + for input_name in node.input + if input_name + ) + + +def _graph_defines_name(graph: GraphProto, name: str) -> bool: + """Whether a nested graph shadows an outer-scope name.""" + return ( + any(value.name == name for value in getattr(graph, "input", [])) + or any(initializer.name == name for initializer in getattr(graph, "initializer", [])) + or any(sparse.values.name == name for sparse in getattr(graph, "sparse_initializer", [])) + or any( + output_name == name + for node in getattr(graph, "node", []) + for output_name in node.output + if output_name + ) + ) + + +def _graph_node_references_name(graph: GraphProto, name: str) -> bool: + """Whether any node input or output in this graph mentions the requested name.""" + return any( + value_name == name + for node in getattr(graph, "node", []) + for value_name in (*node.input, *node.output) + ) + + +def _graph_keep_io_mapping_declares_name(graph: GraphProto, name: str) -> bool: + """Whether a graph formal can be hit by ORT's global keep-I/O map.""" + return any(value.name == name for value in getattr(graph, "input", [])) + + +def _graph_keep_io_mapping_references_name(graph: GraphProto, name: str) -> bool: + """Whether ORT's global keep-I/O map may rewrite this graph's local name.""" + return _graph_node_references_name(graph, name) or _graph_keep_io_mapping_declares_name( + graph, name + ) + + +def _graph_processed_node_references_name( + model: ModelProto, + graph: GraphProto, + name: str, + io_casts: set[str], +) -> bool: + """Whether ORT processes a node reference subject to global mapping.""" + return any( + value_name == name + for node in graph.node + if node.name not in io_casts or not _node_is_conversion_neutral(model, graph, node) + for value_name in (*node.input, *node.output) + ) + + +def _ort_skips_node_attributes(node: NodeProto, blocked_ops: set[str]) -> bool: + """Whether ORT skips a node's graph-valued attributes.""" + from onnxruntime.transformers.float16 import ALWAYS_FLOAT_INPUTS + + return node.op_type in blocked_ops or node.op_type in ALWAYS_FLOAT_INPUTS + + +def _iter_ort_child_graphs(graph: GraphProto, blocked_ops: set[str]) -> list[GraphProto]: + """Return child graphs ORT traverses from this graph.""" + from onnx import AttributeProto + + children: list[GraphProto] = [] + for node in getattr(graph, "node", []): + if _ort_skips_node_attributes(node, blocked_ops): + continue + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + children.append(attribute.g) + elif attribute.type == AttributeProto.GRAPHS: + children.extend(attribute.graphs) + return children + + +def _iter_all_child_graphs(graph: GraphProto) -> list[GraphProto]: + """Return all direct child graphs, including those ORT skips under blocked nodes.""" + from onnx import AttributeProto + + children: list[GraphProto] = [] + for node in getattr(graph, "node", []): + for attribute in node.attribute: + if attribute.type == AttributeProto.GRAPH: + children.append(attribute.g) + elif attribute.type == AttributeProto.GRAPHS: + children.extend(attribute.graphs) + return children + + +def _descendant_has_free_consumer(graph: GraphProto, name: str, blocked_ops: set[str]) -> bool: + """Whether a traversed descendant consumes an outer name without shadowing it.""" + if _graph_defines_name(graph, name): + return False + if _graph_node_consumes_name(graph, name): + return True + return any( + _descendant_has_free_consumer(child, name, blocked_ops) + for child in _iter_ort_child_graphs(graph, blocked_ops) + ) + + +def _descendant_has_node_reference(graph: GraphProto, name: str, blocked_ops: set[str]) -> bool: + """Whether any ORT-traversed descendant mentions a name globally mapped by ORT.""" + if _graph_keep_io_mapping_references_name(graph, name): + return True + return any( + _descendant_has_node_reference(child, name, blocked_ops) + for child in _iter_ort_child_graphs(graph, blocked_ops) + ) + + +def _descendant_has_shadowed_node_reference( + model: ModelProto, + graph: GraphProto, + name: str, + blocked_ops: set[str], + io_casts: set[str], + *, + shadowed: bool = False, +) -> bool: + """Whether a traversed descendant uses a local name ORT would globally rewrite.""" + shadowed = ( + shadowed + or _graph_defines_name(graph, name) + or _graph_keep_io_mapping_declares_name(graph, name) + ) + if shadowed and ( + any( + value.name == name and _type_proto_enters_ort_global_list(value.type) + for value in graph.input + ) + or _graph_processed_node_references_name(model, graph, name, io_casts) + ): + return True + return any( + _descendant_has_shadowed_node_reference( + model, + child, + name, + blocked_ops, + io_casts, + shadowed=shadowed, + ) + for child in _iter_ort_child_graphs(graph, blocked_ops) + ) + + +def _descendant_has_shadowed_mapped_alias( + model: ModelProto, + graph: GraphProto, + source_name: str, + mapped_name: str, + blocked_ops: set[str], + io_casts: set[str], + *, + mapped_name_shadowed: bool = False, +) -> bool: + """Whether mapping a free capture would bind it to a nested target alias.""" + mapped_name_shadowed = mapped_name_shadowed or _graph_defines_name(graph, mapped_name) + if mapped_name_shadowed and _graph_processed_node_references_name( + model, graph, source_name, io_casts + ): + return True + return any( + _descendant_has_shadowed_mapped_alias( + model, + child, + source_name, + mapped_name, + blocked_ops, + io_casts, + mapped_name_shadowed=mapped_name_shadowed, + ) + for child in _iter_ort_child_graphs(graph, blocked_ops) + ) + + +def _descendant_has_free_consumer_in_any_graph(graph: GraphProto, name: str) -> bool: + """Whether any descendant consumes an outer name without shadowing it.""" + if _graph_defines_name(graph, name): + return False + if _graph_node_consumes_name(graph, name): + return True + return any( + _descendant_has_free_consumer_in_any_graph(child, name) + for child in _iter_all_child_graphs(graph) + ) + + +def _has_blocked_free_consumer(graph: GraphProto, name: str, blocked_ops: set[str]) -> bool: + """Whether an ORT-skipped child graph can still capture an outer initializer.""" + for node in getattr(graph, "node", []): + children = _iter_all_child_graphs_from_node(node) + if _ort_skips_node_attributes(node, blocked_ops): + if any(_descendant_has_free_consumer_in_any_graph(child, name) for child in children): + return True + continue + if any( + not _graph_defines_name(child, name) + and _has_blocked_free_consumer(child, name, blocked_ops) + for child in children + ): + return True + return False + + +def _node_expects_fp32_input(node: NodeProto, input_index: int, blocked_ops: set[str]) -> bool: + """Whether ORT leaves this node input in FP32 during FP16 conversion.""" + from onnxruntime.transformers.float16 import ALWAYS_FLOAT_INPUTS + + return node.op_type in blocked_ops or input_index in ALWAYS_FLOAT_INPUTS.get(node.op_type, []) + + +def _schema_parameter_at_index( + parameters: Sequence[OpSchema.FormalParameter], index: int +) -> OpSchema.FormalParameter | None: + """Resolve a fixed or trailing variadic ONNX schema parameter.""" + from onnx import defs + + if index < len(parameters): + return parameters[index] + if parameters and parameters[-1].option == defs.OpSchema.FormalParameterOption.Variadic: + return parameters[-1] + return None + + +def _node_schema(model: ModelProto, node: NodeProto) -> OpSchema | None: + """Resolve a node's schema at the model's imported opset version.""" + from onnx import defs + + version = next( + ( + opset.version + for opset in getattr(model, "opset_import", []) + if opset.domain == node.domain + ), + None, + ) + if version is None: + return None + try: + return defs.get_schema(node.op_type, version, node.domain) + except defs.SchemaError: + return None + + +def _schema_parameters_share_concrete_type( + first: OpSchema.FormalParameter, + second: OpSchema.FormalParameter, +) -> bool: + """Whether schema parameters require one concrete runtime type.""" + from onnx import defs + + if first.type_str != second.type_str: + return False + return all( + parameter.option != defs.OpSchema.FormalParameterOption.Variadic or parameter.is_homogeneous + for parameter in (first, second) + ) + + +def _schema_parameter_allowed_types( + schema: OpSchema, + parameter: OpSchema.FormalParameter, +) -> set[str] | None: + """Resolve a schema parameter's allowed type strings.""" + constraint = next( + ( + constraint + for constraint in schema.type_constraints + if constraint.type_param_str == parameter.type_str + ), + None, + ) + if constraint is not None: + return set(constraint.allowed_type_strs) + return {parameter.type_str} if "(" in parameter.type_str else None + + +def _schema_payload_type(type_str: str) -> str: + """Strip ONNX container wrappers to their payload type.""" + value = type_str + while True: + if value.startswith(("seq(", "optional(")) and value.endswith(")"): + value = value[value.index("(") + 1 : -1] + continue + if value.startswith("map(") and value.endswith(")"): + depth = 0 + for index, character in enumerate(value[4:-1], start=4): + if character == "(": + depth += 1 + elif character == ")": + depth -= 1 + elif character == "," and depth == 0: + value = value[index + 1 : -1] + break + else: + return value + continue + return value + + +def _schema_parameters_share_payload_domain( + schema: OpSchema, + first: OpSchema.FormalParameter, + second: OpSchema.FormalParameter, +) -> bool: + """Whether two parameters allow the same normalized payload types.""" + first_types = _schema_parameter_allowed_types(schema, first) + second_types = _schema_parameter_allowed_types(schema, second) + if first_types is None or second_types is None: + return False + return {_schema_payload_type(type_str) for type_str in first_types} == { + _schema_payload_type(type_str) for type_str in second_types + } + + +def _graph_declared_types(graph: GraphProto, name: str) -> list[TypeProto]: + """Return graph metadata types declared for a binding.""" + return [ + value.type + for values in ( + getattr(graph, "input", []), + getattr(graph, "output", []), + getattr(graph, "value_info", []), + ) + for value in values + if value.name == name + ] + + +def _child_graph_input_index( + schema: OpSchema | None, + node: NodeProto, + child: GraphProto, + input_index: int, +) -> int | None: + """Map a node input to an aligned child formal input.""" + from onnx import defs + + if ( + schema is None + or not schema.inputs + or schema.inputs[-1].option != defs.OpSchema.FormalParameterOption.Variadic + ): + return None + if len(child.input) == len(node.input): + return input_index + fixed_prefix = len(schema.inputs) - 1 + if len(child.input) == len(node.input) - fixed_prefix and input_index >= fixed_prefix: + return input_index - fixed_prefix + return None + + +def _child_graph_feedback_output_index( + model: ModelProto, + node: NodeProto, + child: GraphProto, + input_index: int, +) -> int | None: + """Map a variadic node input to its positional child feedback output.""" + from onnx import defs + + schema = _node_schema(model, node) + if ( + schema is None + or not schema.inputs + or schema.inputs[-1].option != defs.OpSchema.FormalParameterOption.Variadic + ): + return None + input_prefix = len(schema.inputs) - 1 + output_prefix = len(child.output) - len(node.output) + if output_prefix <= 0 or input_prefix != output_prefix + 1 or input_index < input_prefix: + return None + state_index = input_index - input_prefix + candidate = output_prefix + state_index + input_parameter = _schema_parameter_at_index(schema.inputs, input_index) + output_parameter = _schema_parameter_at_index(schema.outputs, state_index) + if ( + state_index >= len(node.output) + or candidate >= len(child.output) + or input_parameter is None + or output_parameter is None + or input_parameter.type_str != output_parameter.type_str + ): + return None + return candidate + + +def _node_input_actual_types( + model: ModelProto, + graph: GraphProto, + node: NodeProto, + input_index: int, +) -> list[TypeProto]: + """Return concrete type sources aligned to a node input.""" + schema = _node_schema(model, node) + value_types = _graph_declared_types(graph, node.input[input_index]) + value_types.extend( + child.input[child_index].type + for child in _iter_all_child_graphs_from_node(node) + if (child_index := _child_graph_input_index(schema, node, child, input_index)) is not None + ) + return value_types + + +def _node_input_is_proven_non_float( + model: ModelProto, + graph: GraphProto, + node: NodeProto, + input_index: int, +) -> bool: + """Whether concrete schema-aligned evidence proves a non-FLOAT input.""" + schema = _node_schema(model, node) + evidence = _node_input_actual_types(model, graph, node, input_index) + parameter = ( + _schema_parameter_at_index(schema.inputs, input_index) if schema is not None else None + ) + if schema is not None and parameter is not None and parameter.type_str: + for other_index, input_name in enumerate(node.input): + if other_index == input_index or not input_name: + continue + other = _schema_parameter_at_index(schema.inputs, other_index) + if other is not None and _schema_parameters_share_concrete_type(parameter, other): + evidence.extend(_node_input_actual_types(model, graph, node, other_index)) + for output_index, output_name in enumerate(node.output): + if not output_name: + continue + other = _schema_parameter_at_index(schema.outputs, output_index) + if other is not None and _schema_parameters_share_concrete_type(parameter, other): + evidence.extend(_graph_declared_types(graph, output_name)) + concrete_types = [value_type for value_type in evidence if _type_proto_is_concrete(value_type)] + return bool(concrete_types) and all( + not _type_proto_contains_float_tensor(value_type) for value_type in concrete_types + ) + + +def _node_input_is_precision_coupled( + model: ModelProto, + graph: GraphProto, + node: NodeProto, + input_index: int, + blocked_ops: set[str], +) -> bool: + """Whether schema relationships can propagate an input's precision.""" + if _node_input_is_proven_non_float(model, graph, node, input_index): + return False + if any( + _type_proto_enters_ort_global_list(value_type) + for value_type in _node_input_actual_types(model, graph, node, input_index) + ): + return True + schema = _node_schema(model, node) + if any( + child.input[child_index].name + and _binding_has_precision_coupled_consumer( + model, + child, + child.input[child_index].name, + blocked_ops, + ) + for child in _iter_all_child_graphs_from_node(node) + if (child_index := _child_graph_input_index(schema, node, child, input_index)) is not None + ): + return True + if schema is None: + return True + parameter = _schema_parameter_at_index(schema.inputs, input_index) + if parameter is None or not parameter.type_str: + return True + related_parameters: list[tuple[OpSchema.FormalParameter, list[TypeProto]]] = [] + for other_index, input_name in enumerate(node.input): + if other_index == input_index or not input_name: + continue + other = _schema_parameter_at_index(schema.inputs, other_index) + if other is not None: + related_parameters.append( + ( + other, + _node_input_actual_types(model, graph, node, other_index), + ) + ) + for output_index, output_name in enumerate(node.output): + if not output_name: + continue + other = _schema_parameter_at_index(schema.outputs, output_index) + if other is not None: + related_parameters.append( + ( + other, + _graph_declared_types(graph, output_name), + ) + ) + if any( + _schema_parameters_share_concrete_type(parameter, other) for other, _ in related_parameters + ): + return True + + input_types = _schema_parameter_allowed_types(schema, parameter) + if input_types is None: + return True + for other, actual_types in related_parameters: + other_types = _schema_parameter_allowed_types(schema, other) + if other_types is None: + return True + if any( + input_type != other_type + and _schema_payload_type(input_type) == _schema_payload_type(other_type) + for input_type in input_types + for other_type in other_types + ) and any(_type_proto_enters_ort_global_list(value_type) for value_type in actual_types): + return True + return False + + +def _binding_has_precision_coupled_consumer( + model: ModelProto, + graph: GraphProto, + name: str, + blocked_ops: set[str], +) -> bool: + """Whether a traversed consumer propagates the binding's precision.""" + for node in getattr(graph, "node", []): + if any( + input_name == name + and not _node_expects_fp32_input(node, input_index, blocked_ops) + and _node_input_is_precision_coupled( + model, + graph, + node, + input_index, + blocked_ops, + ) + for input_index, input_name in enumerate(node.input) + ): + return True + children = _iter_all_child_graphs_from_node(node) + if _ort_skips_node_attributes(node, blocked_ops): + continue + if any( + not _graph_defines_name(child, name) + and _binding_has_precision_coupled_consumer(model, child, name, blocked_ops) + for child in children + ): + return True + return False + + +def _binding_has_unconverted_payload_consumer( + model: ModelProto, + graph: GraphProto, + name: str, + blocked_ops: set[str], +) -> bool: + """Whether an FP16 boundary would conflict with container metadata.""" + for node in getattr(graph, "node", []): + schema = _node_schema(model, node) + for input_index, input_name in enumerate(node.input): + if ( + input_name != name + or _node_expects_fp32_input(node, input_index, blocked_ops) + or schema is None + ): + continue + parameter = _schema_parameter_at_index(schema.inputs, input_index) + if parameter is None: + continue + for output_index, output_name in enumerate(node.output): + if not output_name: + continue + output_parameter = _schema_parameter_at_index(schema.outputs, output_index) + if output_parameter is None or not _schema_parameters_share_payload_domain( + schema, parameter, output_parameter + ): + continue + if any( + _type_proto_contains_float_tensor(value_type) + and not _type_proto_enters_ort_global_list(value_type) + for value_type in _graph_declared_types(graph, output_name) + ): + return True + if _ort_skips_node_attributes(node, blocked_ops): + continue + if any( + not _graph_defines_name(child, name) + and _binding_has_unconverted_payload_consumer(model, child, name, blocked_ops) + for child in _iter_all_child_graphs_from_node(node) + ): + return True + return False + + +def _binding_is_proven_non_float(model: ModelProto, graph: GraphProto, name: str) -> bool: + """Whether a binding's consumers prove that its type is non-FLOAT.""" + declared = [ + value_type + for value_type in _graph_declared_types(graph, name) + if _type_proto_is_concrete(value_type) + ] + if declared: + return all(not _type_proto_contains_float_tensor(value_type) for value_type in declared) + for node in getattr(graph, "node", []): + if any( + input_name == name and _node_input_is_proven_non_float(model, graph, node, input_index) + for input_index, input_name in enumerate(node.input) + ): + return True + return False + + +def _sparse_initializer_consumer_types( + graph: GraphProto, name: str, blocked_ops: set[str] +) -> tuple[bool, bool]: + """Return whether a sparse initializer has FP16 and/or FP32 consumers.""" + has_fp16_consumer = False + has_fp32_consumer = False + for node in getattr(graph, "node", []): + for input_index, input_name in enumerate(node.input): + if input_name != name: + continue + if _node_expects_fp32_input(node, input_index, blocked_ops): + has_fp32_consumer = True + else: + has_fp16_consumer = True + + children = _iter_all_child_graphs_from_node(node) + if _ort_skips_node_attributes(node, blocked_ops): + if any(_descendant_has_free_consumer_in_any_graph(child, name) for child in children): + has_fp32_consumer = True + continue + for child in children: + if _graph_defines_name(child, name): + continue + child_has_fp16, child_has_fp32 = _sparse_initializer_consumer_types( + child, name, blocked_ops + ) + has_fp16_consumer = has_fp16_consumer or child_has_fp16 + has_fp32_consumer = has_fp32_consumer or child_has_fp32 + return has_fp16_consumer, has_fp32_consumer + + +def _iter_all_child_graphs_from_node(node: NodeProto) -> list[GraphProto]: + """Return all child graphs attached to a node.""" + return [ + child for attribute in node.attribute for child in _iter_graphs_from_attribute(attribute) + ] + + +def _iter_graphs_from_attribute(attribute: AttributeProto) -> list[GraphProto]: + """Return graphs carried by one node attribute.""" + from onnx import AttributeProto as ONNXAttributeProto + + if attribute.type == ONNXAttributeProto.GRAPH: + return [attribute.g] + if attribute.type == ONNXAttributeProto.GRAPHS: + return list(attribute.graphs) + return [] + + +def _ort_traversed_attributes(model: ModelProto, blocked_ops: set[str]) -> list[AttributeProto]: + """Return attributes ORT's FP16 BFS processes.""" + return [ + attribute + for graph in _ort_traversed_graphs(model, list(blocked_ops)) + for node in graph.node + if not _ort_skips_node_attributes(node, blocked_ops) + for attribute in node.attribute + ] + + +def _attribute_requires_type_validation( + attribute: AttributeProto, +) -> bool: + """Whether ORT leaves potentially tensor-defining FLOAT data unchanged.""" + from onnx import AttributeProto as ONNXAttributeProto + from onnx import TensorProto + + if attribute.type in { + ONNXAttributeProto.FLOAT, + ONNXAttributeProto.FLOATS, + }: + return True + if attribute.type == ONNXAttributeProto.SPARSE_TENSOR: + return attribute.sparse_tensor.values.data_type == TensorProto.FLOAT + if attribute.type == ONNXAttributeProto.SPARSE_TENSORS: + return any( + sparse.values.data_type == TensorProto.FLOAT for sparse in attribute.sparse_tensors + ) + if attribute.type == ONNXAttributeProto.TYPE_PROTO: + return _type_proto_contains_float_tensor(attribute.tp) + if attribute.type == ONNXAttributeProto.TYPE_PROTOS: + return any( + _type_proto_contains_float_tensor(value_type) for value_type in attribute.type_protos + ) + return False + + +def _external_float_attribute_tensors( + model: ModelProto, blocked_ops: set[str] +) -> list[TensorProto]: + """Return traversed FLOAT tensor attributes backed by external data.""" + from onnx import AttributeProto as ONNXAttributeProto + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + tensors = [] + for attribute in _ort_traversed_attributes(model, blocked_ops): + if attribute.type == ONNXAttributeProto.TENSOR: + tensors.append(attribute.t) + elif attribute.type == ONNXAttributeProto.TENSORS: + tensors.extend(attribute.tensors) + return [ + tensor + for tensor in tensors + if tensor.data_type == TensorProto.FLOAT and uses_external_data(tensor) + ] + + +def _internalize_external_float_attribute_tensors(model: ModelProto, blocked_ops: set[str]) -> None: + """Drop stale external metadata before ORT converts tensor attributes.""" + from onnx import TensorProto + + for tensor in _external_float_attribute_tensors(model, blocked_ops): + del tensor.external_data[:] + tensor.data_location = TensorProto.DEFAULT + + +def _initializer_output_has_consumers(graph: GraphProto, name: str, blocked_ops: set[str]) -> bool: + """Resolve direct-output initializer consumers using ONNX lexical scopes.""" + if _graph_node_consumes_name(graph, name): + return True + return any( + _descendant_has_free_consumer(child, name, blocked_ops) + for child in _iter_ort_child_graphs(graph, blocked_ops) + ) + + +def _ort_keep_io_name_mapping( + model: ModelProto, *, keep_io_types: bool +) -> tuple[dict[str, str], set[str]]: + """Return ORT's global top-level I/O tensor mapping and generated Cast names.""" + from onnx import TensorProto + + if not keep_io_types: + return {}, set() + + name_mapping: dict[str, str] = {} + io_casts: set[str] = set() + for io_kind, values in ( + ("input", getattr(model.graph, "input", [])), + ("output", getattr(model.graph, "output", [])), + ): + for value_index, value in enumerate(values): + if ( + not value.type.HasField("tensor_type") + or value.type.tensor_type.elem_type != TensorProto.FLOAT + ): + continue + name_mapping[value.name] = f"graph_{io_kind}_cast_{value_index}" + io_casts.add(f"graph_{io_kind}_cast{value_index}") + return name_mapping, io_casts + + +def _reject_unpreserved_float_container_io(model: ModelProto, *, keep_io_types: bool) -> None: + """Reject FLOAT container I/O that ORT cannot preserve.""" + if not keep_io_types: + return + names = sorted( + value.name + for values in ( + getattr(model.graph, "input", []), + getattr(model.graph, "output", []), + ) + for value in values + if value.type.WhichOneof("value") != "tensor_type" + and _type_proto_enters_ort_global_list(value.type) + ) + if names: + joined_names = ", ".join(names) + msg = ( + f"keep_io_types cannot preserve FLOAT container graph I/O: {joined_names}; " + "convert with keep_io_types=False or expose tensor I/O." + ) + raise RuntimeError(msg) + + +def _reject_shared_keep_io_names(model: ModelProto, *, keep_io_types: bool) -> None: + """Reject FLOAT input/output aliases that overwrite ORT's global I/O map.""" + from onnx import TensorProto + + if not keep_io_types: + return + float_inputs = { + value.name + for value in getattr(model.graph, "input", []) + if value.type.HasField("tensor_type") + and value.type.tensor_type.elem_type == TensorProto.FLOAT + } + float_output_names = [ + value.name + for value in getattr(model.graph, "output", []) + if value.type.HasField("tensor_type") + and value.type.tensor_type.elem_type == TensorProto.FLOAT + ] + repeated_outputs = sorted( + name for name in set(float_output_names) if float_output_names.count(name) > 1 + ) + if repeated_outputs: + names = ", ".join(repeated_outputs) + msg = ( + "Top-level repeated FLOAT output names cannot be safely converted with " + f"keep_io_types=True; ORT overwrites their generated aliases: {names}." + ) + raise RuntimeError(msg) + float_outputs = set(float_output_names) + shared_names = float_inputs & float_outputs + if shared_names: + names = ", ".join(sorted(shared_names)) + msg = ( + "Top-level FLOAT names used as both input and output cannot be " + "safely converted with keep_io_types=True; ORT's global I/O " + f"mapping overwrites entries: {names}." + ) + raise RuntimeError(msg) + + +def _reject_scope_unsafe_keep_io_mappings( + model: ModelProto, + *, + keep_io_types: bool, + blocked_ops: set[str], +) -> None: + """Reject nested shadowing that ORT's global keep-I/O name mapping corrupts.""" + name_mapping, io_casts = _ort_keep_io_name_mapping(model, keep_io_types=keep_io_types) + for name, mapped_name in name_mapping.items(): + if any( + _descendant_has_shadowed_node_reference(model, child, name, blocked_ops, io_casts) + for child in _iter_ort_child_graphs(model.graph, blocked_ops) + ): + msg = ( + f"Top-level keep_io_types name '{name}' is referenced by a " + "traversed nested graph with the same local name; ORT's " + "I/O name mapping is not scope-aware." + ) + raise RuntimeError(msg) + if any( + _descendant_has_shadowed_mapped_alias( + model, + child, + name, + mapped_name, + blocked_ops, + io_casts, + ) + for child in _iter_ort_child_graphs(model.graph, blocked_ops) + ): + msg = ( + f"Top-level keep_io_types name '{name}' maps to generated alias " + f"'{mapped_name}', but a traversed nested graph binds that alias " + "locally; ORT's I/O name mapping would change the captured value." + ) + raise RuntimeError(msg) + + +def _lexical_binding_owner_id( + graph: GraphProto, + name: str, + parents: dict[int, GraphProto | None], + generated_top_names: set[str], +) -> int | None: + """Resolve a tensor name to its defining graph, including virtual ORT I/O aliases.""" + current: GraphProto | None = graph + while current is not None: + parent = parents.get(id(current)) + if _graph_defines_name(current, name) or (parent is None and name in generated_top_names): + return id(current) + current = parent + return None + + +def _lexical_binding_owner( + graph: GraphProto, + name: str, + parents: dict[int, GraphProto | None], +) -> GraphProto | None: + """Resolve a name to its defining graph.""" + current: GraphProto | None = graph + while current is not None: + if _graph_defines_name(current, name): + return current + current = parents.get(id(current)) + return None + + +def _value_info_enters_ort_global_list(value_info: ValueInfoProto) -> bool: + """Whether ORT records this value metadata for late blocked-node processing.""" + return _type_proto_enters_ort_global_list(value_info.type) + + +def _type_proto_enters_ort_global_list(value_type: TypeProto) -> bool: + """Whether ORT converts and records this declared FLOAT type.""" + from onnx import TensorProto + + if value_type.HasField("tensor_type") and value_type.tensor_type.elem_type == TensorProto.FLOAT: + return True + return ( + value_type.HasField("sequence_type") + and value_type.sequence_type.elem_type.HasField("tensor_type") + and value_type.sequence_type.elem_type.tensor_type.elem_type == TensorProto.FLOAT + ) + + +def _type_proto_is_concrete(value_type: TypeProto) -> bool: + """Whether ONNX metadata declares a concrete value type.""" + from onnx import TensorProto + + type_kind = value_type.WhichOneof("value") + if type_kind == "tensor_type": + return value_type.tensor_type.elem_type != TensorProto.UNDEFINED + if type_kind == "sparse_tensor_type": + return value_type.sparse_tensor_type.elem_type != TensorProto.UNDEFINED + if type_kind == "sequence_type": + return _type_proto_is_concrete(value_type.sequence_type.elem_type) + if type_kind == "optional_type": + return _type_proto_is_concrete(value_type.optional_type.elem_type) + if type_kind == "map_type": + return value_type.map_type.key_type != TensorProto.UNDEFINED and _type_proto_is_concrete( + value_type.map_type.value_type + ) + if type_kind == "opaque_type": + return bool(value_type.opaque_type.domain or value_type.opaque_type.name) + return False + + +def _type_proto_contains_float_tensor(value_type: TypeProto) -> bool: + """Whether a declared type contains a FLOAT tensor payload.""" + from onnx import TensorProto + + type_kind = value_type.WhichOneof("value") + if type_kind == "tensor_type": + return value_type.tensor_type.elem_type == TensorProto.FLOAT + if type_kind == "sparse_tensor_type": + return value_type.sparse_tensor_type.elem_type == TensorProto.FLOAT + if type_kind == "sequence_type": + return _type_proto_contains_float_tensor(value_type.sequence_type.elem_type) + if type_kind == "optional_type": + return _type_proto_contains_float_tensor(value_type.optional_type.elem_type) + if type_kind == "map_type": + return _type_proto_contains_float_tensor(value_type.map_type.value_type) + return False + + +def _type_proto_has_unconverted_float_container( + value_type: TypeProto, +) -> bool: + """Whether ORT leaves a FLOAT-bearing container declaration unchanged.""" + type_kind = value_type.WhichOneof("value") + if type_kind in {"optional_type", "map_type"}: + return _type_proto_contains_float_tensor(value_type) + if type_kind == "sequence_type": + return not _type_proto_enters_ort_global_list( + value_type + ) and _type_proto_contains_float_tensor(value_type) + return False + + +def _has_unconverted_float_container_declarations( + model: ModelProto, op_block_list: list[str] | None +) -> bool: + """Whether a traversed graph declares a FLOAT-bearing container.""" + return any( + _type_proto_has_unconverted_float_container(value.type) + for graph in _ort_traversed_graphs(model, op_block_list) + for values in ( + getattr(graph, "input", []), + getattr(graph, "output", []), + getattr(graph, "value_info", []), + ) + for value in values + ) + + +def _ort_global_value_info_bindings( + model: ModelProto, + *, + keep_io_types: bool, + blocked_ops: set[str], + graphs: list[GraphProto], + parents: dict[int, GraphProto | None], +) -> tuple[dict[str, int | None], dict[str, str], set[str], dict[str, str]]: + """Simulate the first lexical owner ORT registers for global value metadata.""" + from onnx import TensorProto + + first_owner: dict[str, int | None] = {} + first_type_kind: dict[str, str] = {} + generated_top_names: set[str] = set() + name_mapping, io_casts = _ort_keep_io_name_mapping(model, keep_io_types=keep_io_types) + if keep_io_types: + for io_kind, values in ( + ("input", getattr(model.graph, "input", [])), + ("output", getattr(model.graph, "output", [])), + ): + for value_index, value in enumerate(values): + if ( + value.type.HasField("tensor_type") + and value.type.tensor_type.elem_type == TensorProto.FLOAT + ): + generated_name = f"graph_{io_kind}_cast_{value_index}" + generated_top_names.add(generated_name) + first_owner.setdefault(generated_name, id(model.graph)) + first_type_kind.setdefault(generated_name, "tensor_type") + + graph_io_to_skip = set(name_mapping) + for graph in graphs: + for values in ( + getattr(graph, "input", []), + getattr(graph, "output", []), + getattr(graph, "value_info", []), + ): + for value_info in values: + if ( + value_info.name not in graph_io_to_skip + and _value_info_enters_ort_global_list(value_info) + and value_info.name not in first_owner + ): + first_owner[value_info.name] = _lexical_binding_owner_id( + graph, + value_info.name, + parents, + generated_top_names, + ) + first_type_kind[value_info.name] = value_info.type.WhichOneof("value") or "" + + initializer_owners: dict[str, int] = {} + fp16_initializers: set[str] = set() + for graph in graphs: + for initializer in getattr(graph, "initializer", []): + if initializer.data_type == TensorProto.FLOAT: + initializer_owners[initializer.name] = id(graph) + for node in getattr(graph, "node", []): + if node.name in io_casts: + continue + for input_index, original_name in enumerate(node.input): + input_name = name_mapping.get(original_name, original_name) + if input_name in initializer_owners and not _node_expects_fp32_input( + node, input_index, blocked_ops + ): + fp16_initializers.add(input_name) + + for name, owner in initializer_owners.items(): + if name in fp16_initializers: + first_owner.setdefault(name, owner) + first_type_kind.setdefault(name, "tensor_type") + return first_owner, name_mapping, generated_top_names, first_type_kind + + +def _reject_scope_unsafe_value_info_lookups( + model: ModelProto, + *, + keep_io_types: bool, + op_block_list: list[str] | None, +) -> None: + """Reject late Casts that ORT resolves through the wrong global metadata.""" + blocked_ops = _effective_blocked_ops(op_block_list) + graphs, parents = _ort_graphs_with_parents(model, blocked_ops) + first_owner, name_mapping, generated_top_names, first_type_kind = ( + _ort_global_value_info_bindings( + model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + graphs=graphs, + parents=parents, + ) + ) + _, _, fp16_initializers, _, _ = _initializer_tracking_analysis( + model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + graphs=graphs, + parents=parents, + ) + owners_by_id = {id(graph): graph for graph in graphs} + + def lookup_preserves_top_binding(graph: GraphProto, name: str) -> bool: + intended_owner = _lexical_binding_owner_id(graph, name, parents, generated_top_names) + top_owner = _lexical_binding_owner_id(model.graph, name, parents, generated_top_names) + return ( + intended_owner is not None + and intended_owner == top_owner + and first_owner[name] == intended_owner + ) + + def binding_converts_to_fp16(graph: GraphProto, name: str) -> bool: + owner_id = _lexical_binding_owner_id(graph, name, parents, generated_top_names) + owner = owners_by_id.get(owner_id) if owner_id is not None else None + if owner is None or any( + sparse.values.name == name for sparse in getattr(owner, "sparse_initializer", []) + ): + return False + return _binding_converts_to_fp16( + model, + owner, + name, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + name_mapping=name_mapping, + fp16_initializers=fp16_initializers, + parents=parents, + ) + + reserved_late_tensors: set[str] = set() + reserved_late_nodes: set[str] = set() + existing_top_tensors = _graph_tensor_names(model.graph) | generated_top_names + _, generated_io_nodes = _ort_keep_io_name_mapping(model, keep_io_types=keep_io_types) + existing_top_nodes = { + node.name for node in getattr(model.graph, "node", []) if node.name + } | generated_io_nodes + + def reserve_late_tensor( + graph: GraphProto, + node: NodeProto, + io_kind: str, + value_index: int, + ) -> None: + tensor_name = f"{node.name}_{io_kind}_cast_{value_index}" + node_name = f"{node.name}_{io_kind}_cast{value_index}" + shadowed = False + current = graph + while current is not model.graph: + if _graph_defines_name(current, tensor_name): + shadowed = True + break + parent = parents.get(id(current)) + if parent is None: + break + current = parent + if tensor_name in existing_top_tensors or tensor_name in reserved_late_tensors or shadowed: + msg = ( + f"ORT late Cast tensor '{tensor_name}' for node '{node.name}' " + "collides with an existing or generated lexical binding." + ) + raise RuntimeError(msg) + if node_name in existing_top_nodes or node_name in reserved_late_nodes: + msg = ( + f"ORT late Cast node '{node_name}' for node '{node.name}' " + "collides with an existing or generated top-level node." + ) + raise RuntimeError(msg) + reserved_late_tensors.add(tensor_name) + reserved_late_nodes.add(node_name) + + for graph in graphs: + for node in getattr(graph, "node", []): + for input_index, original_name in enumerate(node.input): + if not _node_expects_fp32_input(node, input_index, blocked_ops): + continue + input_name = name_mapping.get(original_name, original_name) + if input_name not in first_owner: + if binding_converts_to_fp16(graph, input_name): + msg = ( + "ORT cannot add a required precision-boundary Cast " + f"for blocked or mixed-type node '{node.name}' input " + f"'{input_name}' because of missing FLOAT metadata." + ) + raise RuntimeError(msg) + continue + if not lookup_preserves_top_binding(graph, input_name): + scope = "nested " if graph is not model.graph else "" + msg = ( + "ORT's global value-info lookup for blocked or mixed-type " + f"{scope}node '{node.name}' cannot preserve the lexical input " + f"binding for '{input_name}'." + ) + raise RuntimeError(msg) + if first_type_kind[input_name] != "tensor_type": + msg = ( + "ORT's late Cast path only supports tensor values; " + f"blocked or mixed-type node '{node.name}' resolves " + f"non-tensor input '{input_name}'." + ) + raise RuntimeError(msg) + reserve_late_tensor(graph, node, "input", input_index) + + if node.op_type not in blocked_ops: + continue + for output_index, original_name in enumerate(node.output): + output_name = name_mapping.get(original_name, original_name) + if not output_name: + continue + if output_name not in first_owner: + if _binding_has_precision_coupled_consumer( + model, graph, output_name, blocked_ops + ): + msg = ( + "ORT cannot add a required precision-boundary Cast " + f"for blocked node '{node.name}' output '{output_name}' " + "because of missing FLOAT metadata." + ) + raise RuntimeError(msg) + continue + if not lookup_preserves_top_binding(graph, output_name): + scope = "nested " if graph is not model.graph else "" + msg = ( + "ORT's global value-info lookup for blocked or mixed-type " + f"{scope}node '{node.name}' cannot preserve the lexical output " + f"binding for '{output_name}'." + ) + raise RuntimeError(msg) + if first_type_kind[output_name] != "tensor_type": + msg = ( + "ORT's late Cast path only supports tensor values; " + f"blocked node '{node.name}' resolves non-tensor output " + f"'{output_name}'." + ) + raise RuntimeError(msg) + if _binding_has_unconverted_payload_consumer( + model, graph, output_name, blocked_ops + ): + msg = ( + f"ORT's FP16 boundary for blocked node '{node.name}' " + f"output '{output_name}' conflicts with an unconverted " + "container payload declaration." + ) + raise RuntimeError(msg) + reserve_late_tensor(graph, node, "output", output_index) + + +def _reject_generated_io_cast_name_collisions( + model: ModelProto, + *, + keep_io_types: bool, + op_block_list: list[str] | None, +) -> None: + """Reject names that collide with ORT's deterministic top-level I/O Casts.""" + from onnx import TensorProto + + if not keep_io_types: + return + + tensor_names = _graph_tensor_names(model.graph) + node_names = _all_node_names(model, op_block_list) + for io_kind, values in ( + ("input", getattr(model.graph, "input", [])), + ("output", getattr(model.graph, "output", [])), + ): + for value_index, value in enumerate(values): + if ( + not value.type.HasField("tensor_type") + or value.type.tensor_type.elem_type != TensorProto.FLOAT + ): + continue + generated_tensor = f"graph_{io_kind}_cast_{value_index}" + generated_node = f"graph_{io_kind}_cast{value_index}" + collisions = [ + name + for name, existing_names in ( + (generated_tensor, tensor_names), + (generated_node, node_names), + ) + if name in existing_names + ] + if collisions: + msg = ( + "FP16 conversion cannot safely allocate ORT " + f"{io_kind} Cast names for '{value.name}'; existing names " + f"collide: {', '.join(collisions)}." + ) + raise RuntimeError(msg) + + +def _capture_safe_initializer_outputs( + model: ModelProto, + *, + keep_io_types: bool, + op_block_list: list[str] | None, +) -> list[_InitializerOutput]: + """Capture safe direct initializer outputs or fail before ORT mutates the model. + + Top-level shared outputs are allowed for pure-FP16 conversion when ORT + converts their consumers consistently; keep-I/O conversion still rejects + them because ORT rewrites consumer inputs to the generated output-cast alias. + """ + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + if not hasattr(model.graph, "input") or not hasattr(model.graph, "output"): + return [] + + blocked_ops = _effective_blocked_ops(op_block_list) + traversed_graphs = _ort_traversed_graphs(model, op_block_list) + + captured: list[_InitializerOutput] = [] + for graph_index, graph in enumerate(traversed_graphs): + produced = {name for node in getattr(graph, "node", []) for name in node.output if name} + graph_inputs = {value.name for value in getattr(graph, "input", [])} + initializers = { + initializer.name: initializer for initializer in getattr(graph, "initializer", []) + } + for output_index, output in enumerate(getattr(graph, "output", [])): + initializer = initializers.get(output.name) + if ( + initializer is None + or output.name in produced + or initializer.data_type != TensorProto.FLOAT + ): + continue + + if output.name in graph_inputs: + msg = ( + f"Initializer-backed output '{output.name}' is also a graph input; " + "FP16 conversion cannot preserve overridable-initializer semantics." + ) + raise RuntimeError(msg) + has_consumers = _initializer_output_has_consumers(graph, output.name, blocked_ops) + if keep_io_types and graph_index == 0 and has_consumers: + msg = ( + f"Initializer-backed output '{output.name}' has internal consumers; " + "FP16 conversion cannot safely preserve keep_io_types semantics." + ) + raise RuntimeError(msg) + if ( + keep_io_types + and graph_index == 0 + and not has_consumers + and any( + _descendant_has_node_reference(child, output.name, blocked_ops) + for child in _iter_ort_child_graphs(graph, blocked_ops) + ) + ): + msg = ( + f"Initializer-backed output '{output.name}' is referenced by a " + "traversed nested graph with the same local name; ORT's " + "keep_io_types output mapping is not scope-aware." + ) + raise RuntimeError(msg) + if (not keep_io_types or graph_index != 0) and _has_blocked_free_consumer( + graph, output.name, blocked_ops + ): + msg = ( + f"Initializer-backed output '{output.name}' is captured by a " + "blocked subgraph; FP16 conversion cannot safely change its " + "initializer type while that subgraph remains FP32." + ) + raise RuntimeError(msg) + if uses_external_data(initializer) and not _tensor_data_is_loaded(initializer): + msg = ( + f"Initializer-backed output '{output.name}' uses unloaded external data; " + "load external weights before FP16 conversion." + ) + raise RuntimeError(msg) + + captured.append( + _InitializerOutput(graph_index, output.name, output_index, has_consumers) + ) + return captured + + +def _initializer_data_type(graph: GraphProto, name: str) -> int: + """Return the initializer data type for a captured output in its graph.""" + return next(value.data_type for value in graph.initializer if value.name == name) + + +def _set_direct_output_elem_type(graph: GraphProto, name: str, data_type: int) -> None: + """Set a direct graph output's tensor element type.""" + for output in graph.output: + if output.name == name and output.type.HasField("tensor_type"): + output.type.tensor_type.elem_type = data_type + return + + +def _convert_output_initializer_to_fp16(graph: GraphProto, name: str) -> None: + """Convert a captured output's resident FLOAT initializer in place.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + from onnxruntime.transformers.float16 import convert_tensor_float_to_float16 + + initializer = next(value for value in graph.initializer if value.name == name) + if uses_external_data(initializer): + del initializer.external_data[:] + initializer.data_location = TensorProto.DEFAULT + initializer.CopyFrom(cast("TensorProto", convert_tensor_float_to_float16(initializer))) + + +def _convert_sparse_initializer_to_fp16(graph: GraphProto, name: str) -> None: + """Convert a sparse FLOAT initializer's values tensor to FLOAT16.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + from onnxruntime.transformers.float16 import convert_tensor_float_to_float16 + + sparse_initializer = next( + sparse for sparse in graph.sparse_initializer if sparse.values.name == name + ) + if uses_external_data(sparse_initializer.values): + del sparse_initializer.values.external_data[:] + sparse_initializer.values.data_location = TensorProto.DEFAULT + sparse_initializer.values.CopyFrom( + cast("TensorProto", convert_tensor_float_to_float16(sparse_initializer.values)) + ) + for value_info in ( + *getattr(graph, "input", []), + *getattr(graph, "output", []), + *getattr(graph, "value_info", []), + ): + if ( + value_info.name == name + and value_info.type.HasField("sparse_tensor_type") + and value_info.type.sparse_tensor_type.elem_type == TensorProto.FLOAT + ): + value_info.type.sparse_tensor_type.elem_type = TensorProto.FLOAT16 + + +def _has_sparse_graph_io(graph: GraphProto, name: str) -> bool: + """Whether a sparse initializer name is part of graph input/output metadata.""" + return any( + value_info.name == name and value_info.type.HasField("sparse_tensor_type") + for value_info in (*getattr(graph, "input", []), *getattr(graph, "output", [])) + ) + + +def _has_sparse_graph_output(graph: GraphProto, name: str) -> bool: + """Whether a sparse initializer name is a graph output.""" + return any( + value_info.name == name and value_info.type.HasField("sparse_tensor_type") + for value_info in getattr(graph, "output", []) + ) + + +def _has_sparse_graph_input(graph: GraphProto, name: str) -> bool: + """Whether a sparse initializer name is a graph input.""" + return any( + value_info.name == name and value_info.type.HasField("sparse_tensor_type") + for value_info in getattr(graph, "input", []) + ) + + +def _is_kept_top_level_sparse_edge( + model: ModelProto, graph: GraphProto, name: str, *, keep_io_types: bool +) -> bool: + """Whether a nested output feeds an unchanged public sparse output.""" + from onnx import TensorProto + + if not keep_io_types: + return False + current_graph = graph + current_name = name + visited: set[int] = set() + while current_graph is not model.graph: + if id(current_graph) in visited: + return False + visited.add(id(current_graph)) + output_index = next( + ( + index + for index, output in enumerate(getattr(current_graph, "output", [])) + if output.name == current_name + ), + None, + ) + if output_index is None: + return False + parents = [ + (parent_graph, node) + for parent_graph in _all_graphs(model) + for node in getattr(parent_graph, "node", []) + if any(child is current_graph for child in _iter_all_child_graphs_from_node(node)) + ] + if len(parents) != 1: + return False + parent_graph, node = parents[0] + if len(node.output) != len(current_graph.output) or output_index >= len(node.output): + return False + parent_name = node.output[output_index] + if not parent_name or _graph_node_consumes_name(parent_graph, parent_name): + return False + current_graph = parent_graph + current_name = parent_name + return any( + output.name == current_name + and output.type.HasField("sparse_tensor_type") + and output.type.sparse_tensor_type.elem_type == TensorProto.FLOAT + for output in getattr(model.graph, "output", []) + ) + + +def _repair_sparse_float_initializers( + model: ModelProto, op_block_list: list[str] | None, *, keep_io_types: bool +) -> None: + """Convert sparse FLOAT initializers when ORT converted every consumer to FP16.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + blocked_ops = _effective_blocked_ops(op_block_list) + for graph_index, graph in enumerate(_ort_traversed_graphs(model, op_block_list)): + for sparse_initializer in getattr(graph, "sparse_initializer", []): + values = sparse_initializer.values + if values.data_type != TensorProto.FLOAT: + continue + has_fp16_consumer, has_fp32_consumer = _sparse_initializer_consumer_types( + graph, values.name, blocked_ops + ) + kept_sparse_edge = graph_index != 0 and _is_kept_top_level_sparse_edge( + model, + graph, + values.name, + keep_io_types=keep_io_types, + ) + has_fp16_output = ( + (not keep_io_types or graph_index != 0) + and _has_sparse_graph_output(graph, values.name) + and not kept_sparse_edge + ) + if ( + graph_index != 0 + and (has_fp16_consumer or has_fp16_output) + and _has_sparse_graph_output(graph, values.name) + ): + msg = ( + f"Sparse FLOAT initializer '{values.name}' is a nested sparse " + "graph output; FP16 conversion cannot safely propagate its type " + "through the parent graph edge." + ) + raise RuntimeError(msg) + if (has_fp16_consumer or has_fp16_output) and has_fp32_consumer: + msg = ( + f"Sparse FLOAT initializer '{values.name}' has both FP16 and " + "FP32 consumers after conversion; FP16 conversion cannot safely " + "choose one initializer type." + ) + raise RuntimeError(msg) + if not has_fp16_consumer and not has_fp16_output: + continue + if has_fp16_output and _has_sparse_graph_input(graph, values.name): + msg = ( + f"Sparse FLOAT initializer '{values.name}' is also sparse graph input " + "and output; FP16 conversion cannot preserve overridable-initializer " + "semantics." + ) + raise RuntimeError(msg) + if keep_io_types and _has_sparse_graph_io(graph, values.name): + msg = ( + f"Sparse FLOAT initializer '{values.name}' is also sparse graph I/O; " + "FP16 conversion cannot preserve keep_io_types semantics." + ) + raise RuntimeError(msg) + if uses_external_data(values) and not _tensor_data_is_loaded(values): + msg = ( + f"Sparse FLOAT initializer '{values.name}' uses unloaded external data; " + "load external weights before FP16 conversion." + ) + raise RuntimeError(msg) + _convert_sparse_initializer_to_fp16(graph, values.name) + + +def _internalize_output_initializer(graph: GraphProto, name: str) -> None: + """Drop stale external metadata after resident bytes were loaded.""" + from onnx import TensorProto + from onnx.external_data_helper import uses_external_data + + initializer = next(value for value in graph.initializer if value.name == name) + if uses_external_data(initializer): + del initializer.external_data[:] + initializer.data_location = TensorProto.DEFAULT + + +def _format_data_type(data_type: int) -> str: + """Format ONNX tensor data type values for diagnostics.""" + from onnx import TensorProto + + try: + return TensorProto.DataType.Name(data_type) + except ValueError: + return str(data_type) + + +def _validate_initializer_output_types(model: ModelProto) -> None: + """Reject direct initializer-backed outputs whose declared type diverged.""" + mismatches: list[str] = [] + for graph in _all_graphs(model): + if not hasattr(graph, "output") or not hasattr(graph, "initializer"): + continue + produced = {name for node in graph.node for name in node.output if name} + initializers = {initializer.name: initializer for initializer in graph.initializer} + for output in graph.output: + initializer = initializers.get(output.name) + if initializer is None or output.name in produced: + continue + if not output.type.HasField("tensor_type"): + continue + elem_type = output.type.tensor_type.elem_type + if elem_type != initializer.data_type: + graph_name = graph.name or "" + mismatches.append( + f"{graph_name}.{output.name} declares {_format_data_type(elem_type)} " + f"but initializer is {_format_data_type(initializer.data_type)}" + ) + if mismatches: + msg = ( + "FP16 conversion produced initializer-backed outputs with mismatched " + f"types: {'; '.join(mismatches)}." + ) + raise RuntimeError(msg) + + +def _remove_orphan_output_casts( + model: ModelProto, + captured: list[_InitializerOutput], +) -> None: + """Remove ORT output Casts whose inputs cannot have producers by construction.""" + from onnx import TensorProto + + if not captured: + return + + remove_indices: list[int] = [] + orphan_inputs: set[str] = set() + for item in captured: + generated_tensor = f"graph_output_cast_{item.output_index}" + generated_node = f"graph_output_cast{item.output_index}" + matches = [ + (index, node) + for index, node in enumerate(model.graph.node) + if node.name == generated_node + and node.op_type == "Cast" + and list(node.input) == [generated_tensor] + and list(node.output) == [item.name] + and any( + attribute.name == "to" and attribute.i == TensorProto.FLOAT + for attribute in node.attribute + ) + ] + if len(matches) != 1: + msg = f"Expected one ORT graph-output Cast for initializer-backed output '{item.name}'." + raise RuntimeError(msg) + remove_indices.append(matches[0][0]) + orphan_inputs.add(generated_tensor) + + for index in sorted(remove_indices, reverse=True): + del model.graph.node[index] + for item in captured: + _internalize_output_initializer(model.graph, item.name) + retained = [value for value in model.graph.value_info if value.name not in orphan_inputs] + del model.graph.value_info[:] + model.graph.value_info.extend(retained) + + +def _graph_free_references(graph: GraphProto) -> set[str]: + """Return values referenced by a nested graph but defined in an outer scope.""" + local_names = {value.name for value in getattr(graph, "input", []) if value.name} + local_names.update( + initializer.name for initializer in getattr(graph, "initializer", []) if initializer.name + ) + local_names.update( + sparse.values.name + for sparse in getattr(graph, "sparse_initializer", []) + if sparse.values.name + ) + local_names.update( + output_name + for node in getattr(graph, "node", []) + for output_name in node.output + if output_name + ) + + references = { + input_name for node in getattr(graph, "node", []) for input_name in node.input if input_name + } + references.update(output.name for output in getattr(graph, "output", []) if output.name) + for node in getattr(graph, "node", []): + for child in _iter_all_child_graphs_from_node(node): + references.update(_graph_free_references(child)) + return references - local_names + + +def _nested_graph_input_source( + model: ModelProto, graph: GraphProto, name: str +) -> tuple[GraphProto, str, int | None] | None: + """Resolve a nested formal input through a structurally aligned parent node.""" + input_index = next( + (index for index, value in enumerate(getattr(graph, "input", [])) if value.name == name), + None, + ) + if input_index is None: + return None + for parent_graph in _all_graphs(model): + for node in getattr(parent_graph, "node", []): + if not any(child is graph for child in _iter_all_child_graphs_from_node(node)): + continue + if len(node.input) != len(graph.input) or input_index >= len(node.input): + return None + source_name = node.input[input_index] + if not source_name: + return None + feedback_index = _child_graph_feedback_output_index(model, node, graph, input_index) + return parent_graph, source_name, feedback_index + return None + + +def _node_output_precision_sources( + model: ModelProto, + node: NodeProto, + output_index: int, + blocked_ops: set[str], +) -> list[str] | None: + """Return direct inputs proven to carry an output's precision payload.""" + schema = _node_schema(model, node) + if schema is None: + return None + children = _iter_all_child_graphs_from_node(node) + if children: + sources = [] + for input_index, input_name in enumerate(node.input): + if not input_name: + continue + input_prefix = len(schema.inputs) - 1 + if input_index - input_prefix != output_index: + continue + aligned = [ + (child, child_input_index, feedback_index) + for child in children + if (child_input_index := _child_graph_input_index(schema, node, child, input_index)) + is not None + and ( + feedback_index := _child_graph_feedback_output_index( + model, node, child, input_index + ) + ) + is not None + ] + if len(aligned) == len(children) and all( + _binding_preserves_precision_from( + model, + child, + child.output[feedback_index].name, + child.input[child_input_index].name, + blocked_ops, + {}, + ) + for child, child_input_index, feedback_index in aligned + ): + sources.append(input_name) + return sources or None + output_parameter = _schema_parameter_at_index(schema.outputs, output_index) + if output_parameter is None or not output_parameter.type_str: + return None + return [ + input_name + for input_index, input_name in enumerate(node.input) + if input_name + and (input_parameter := _schema_parameter_at_index(schema.inputs, input_index)) is not None + and ( + _schema_parameters_share_concrete_type(input_parameter, output_parameter) + or _schema_parameters_share_payload_domain(schema, input_parameter, output_parameter) + ) + ] + + +def _binding_preserves_precision_from( + model: ModelProto, + graph: GraphProto, + name: str, + source_name: str, + blocked_ops: set[str], + name_mapping: dict[str, str], +) -> bool: + """Whether a single-source producer chain preserves a binding's precision.""" + current_name = name + visited: set[str] = set() + while current_name != source_name: + if current_name in name_mapping or any( + _type_proto_enters_ort_global_list(value_type) + for value_type in _graph_declared_types(graph, current_name) + ): + return False + if current_name in visited: + return False + visited.add(current_name) + producers = [ + (node, output_index) + for node in getattr(graph, "node", []) + if node.op_type not in blocked_ops + for output_index, output_name in enumerate(node.output) + if output_name == current_name + ] + if len(producers) != 1: + return False + node, output_index = producers[0] + sources = _node_output_precision_sources(model, node, output_index, blocked_ops) + if sources is None or len(sources) != 1: + return False + current_name = sources[0] + return True + + +def _producer_chain_converts_to_fp16( + model: ModelProto, + graph: GraphProto, + name: str, + *, + keep_io_types: bool, + blocked_ops: set[str], + name_mapping: dict[str, str], + fp16_initializers: set[str], + parents: dict[int, GraphProto | None], +) -> bool: + """Trace single-source producer chains without consuming Python stack.""" + from onnx import TensorProto + + current_graph = graph + current_name = name + through_producer = False + visited: set[tuple[int, str]] = set() + while True: + owner = _lexical_binding_owner(current_graph, current_name, parents) + if owner is None: + return True + current_graph = owner + binding = (id(current_graph), current_name) + if binding in visited: + return True + visited.add(binding) + if _binding_is_proven_non_float(model, current_graph, current_name): + return False + if any( + initializer.name == current_name + for initializer in getattr(current_graph, "initializer", []) + if initializer.data_type == TensorProto.FLOAT + ): + return current_name in fp16_initializers + if any( + sparse.values.name == current_name + for sparse in getattr(current_graph, "sparse_initializer", []) + if sparse.values.data_type == TensorProto.FLOAT + ): + has_fp16_consumer, _ = _sparse_initializer_consumer_types( + current_graph, + current_name, + blocked_ops, + ) + has_fp16_output = ( + not keep_io_types or current_graph is not model.graph + ) and _has_sparse_graph_output(current_graph, current_name) + return has_fp16_consumer or has_fp16_output + if current_name in name_mapping: + return through_producer + value_types = _graph_declared_types(current_graph, current_name) + if any(_type_proto_enters_ort_global_list(value_type) for value_type in value_types): + return True + if any(_type_proto_contains_float_tensor(value_type) for value_type in value_types): + if current_graph is not model.graph and any( + value.name == current_name for value in getattr(current_graph, "input", []) + ): + return _binding_converts_to_fp16( + model, + current_graph, + current_name, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + name_mapping=name_mapping, + fp16_initializers=fp16_initializers, + parents=parents, + ) + producers = [ + (node, output_index) + for node in getattr(current_graph, "node", []) + if node.op_type not in blocked_ops + for output_index, output_name in enumerate(node.output) + if output_name == current_name + ] + if not producers: + return False + if len(producers) != 1: + return True + node, output_index = producers[0] + sources = _node_output_precision_sources(model, node, output_index, blocked_ops) + if sources is None: + return True + if not sources: + return False + if len(sources) != 1: + return True + current_name = sources[0] + through_producer = True + continue + if any(_type_proto_is_concrete(value_type) for value_type in value_types): + return False + return any( + current_name in node.output and node.op_type not in blocked_ops + for node in getattr(current_graph, "node", []) + ) + + +def _binding_converts_to_fp16( + model: ModelProto, + graph: GraphProto, + name: str, + *, + keep_io_types: bool, + blocked_ops: set[str], + name_mapping: dict[str, str], + fp16_initializers: set[str], + parents: dict[int, GraphProto | None], +) -> bool: + """Whether ORT or wrapper repair changes an outer FLOAT binding to FP16.""" + from onnx import TensorProto + + owner = _lexical_binding_owner(graph, name, parents) + if owner is None: + return True + graph = owner + if _binding_is_proven_non_float(model, graph, name): + return False + if any( + initializer.name == name + for initializer in getattr(graph, "initializer", []) + if initializer.data_type == TensorProto.FLOAT + ): + return name in fp16_initializers + if any( + sparse.values.name == name + for sparse in getattr(graph, "sparse_initializer", []) + if sparse.values.data_type == TensorProto.FLOAT + ): + has_fp16_consumer, _ = _sparse_initializer_consumer_types(graph, name, blocked_ops) + has_fp16_output = ( + not keep_io_types or graph is not model.graph + ) and _has_sparse_graph_output(graph, name) + return has_fp16_consumer or has_fp16_output + if name in name_mapping: + return False + metadata = [ + value_info + for values in ( + getattr(graph, "input", []), + getattr(graph, "output", []), + getattr(graph, "value_info", []), + ) + for value_info in values + if value_info.name == name + ] + if any(_value_info_enters_ort_global_list(value_info) for value_info in metadata): + return True + if any(_type_proto_contains_float_tensor(value_info.type) for value_info in metadata): + if graph is not model.graph and any( + value.name == name for value in getattr(graph, "input", []) + ): + source = _nested_graph_input_source(model, graph, name) + if source is None: + return True + parent_graph, source_name, feedback_index = source + if _binding_converts_to_fp16( + model, + parent_graph, + source_name, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + name_mapping=name_mapping, + fp16_initializers=fp16_initializers, + parents=parents, + ): + return True + if feedback_index is None: + return True + feedback = graph.output[feedback_index] + return not _type_proto_contains_float_tensor( + feedback.type + ) or not _binding_preserves_precision_from( + model, + graph, + feedback.name, + name, + blocked_ops, + name_mapping, + ) + return _producer_chain_converts_to_fp16( + model, + graph, + name, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + name_mapping=name_mapping, + fp16_initializers=fp16_initializers, + parents=parents, + ) + if any(_type_proto_is_concrete(value_info.type) for value_info in metadata): + return False + return any( + name in node.output and node.op_type not in blocked_ops + for node in getattr(graph, "node", []) + ) + + +def _local_function_executed_attributes( + model: ModelProto, node: NodeProto +) -> list[AttributeProto] | None: + """Resolve supplied or default attributes referenced by a local function.""" + function = next( + ( + candidate + for candidate in getattr(model, "functions", []) + if candidate.domain == node.domain + and candidate.name == node.op_type + and getattr(candidate, "overload", "") == getattr(node, "overload", "") + ), + None, + ) + if function is None: + return None + referenced_attributes: set[str] = set() + pending_nodes = list(function.node) + while pending_nodes: + function_node = pending_nodes.pop() + for attribute in function_node.attribute: + if attribute.ref_attr_name: + referenced_attributes.add(attribute.ref_attr_name) + for child in _iter_graphs_from_attribute(attribute): + pending_nodes.extend(child.node) + supplied = {attribute.name: attribute for attribute in node.attribute} + defaults = {attribute.name: attribute for attribute in function.attribute_proto} + return [ + attribute + for name in referenced_attributes + if (attribute := supplied.get(name, defaults.get(name))) is not None + ] + + +def _function_contains_concrete_float(function: FunctionProto) -> bool: + """Whether an unvisited function body stores a concrete FLOAT value.""" + from onnx import AttributeProto as ONNXAttributeProto + from onnx import TensorProto + + if any( + not any(node.input) + and any( + attribute.type + in { + ONNXAttributeProto.FLOAT, + ONNXAttributeProto.FLOATS, + } + for attribute in node.attribute + ) + for node in function.node + ): + return True + if any(_type_proto_contains_float_tensor(value.type) for value in function.value_info): + return True + + pending_attributes = [ + attribute + for node in function.node + for attribute in node.attribute + if not attribute.ref_attr_name + ] + pending_graphs: list[GraphProto] = [] + while pending_attributes or pending_graphs: + while pending_attributes: + attribute = pending_attributes.pop() + if ( + attribute.type == ONNXAttributeProto.TENSOR + and attribute.t.data_type == TensorProto.FLOAT + ): + return True + if attribute.type == ONNXAttributeProto.TENSORS and any( + tensor.data_type == TensorProto.FLOAT for tensor in attribute.tensors + ): + return True + if ( + attribute.type == ONNXAttributeProto.SPARSE_TENSOR + and attribute.sparse_tensor.values.data_type == TensorProto.FLOAT + ): + return True + if attribute.type == ONNXAttributeProto.SPARSE_TENSORS and any( + sparse.values.data_type == TensorProto.FLOAT for sparse in attribute.sparse_tensors + ): + return True + if ( + attribute.type == ONNXAttributeProto.TYPE_PROTO + and _type_proto_contains_float_tensor(attribute.tp) + ): + return True + if attribute.type == ONNXAttributeProto.TYPE_PROTOS and any( + _type_proto_contains_float_tensor(value_type) + for value_type in attribute.type_protos + ): + return True + pending_graphs.extend(_iter_graphs_from_attribute(attribute)) + while pending_graphs: + graph = pending_graphs.pop() + if any( + not any(node.input) + and any( + attribute.type + in { + ONNXAttributeProto.FLOAT, + ONNXAttributeProto.FLOATS, + } + for attribute in node.attribute + ) + for node in graph.node + ): + return True + if any( + initializer.data_type == TensorProto.FLOAT for initializer in graph.initializer + ) or any( + sparse.values.data_type == TensorProto.FLOAT for sparse in graph.sparse_initializer + ): + return True + if any( + _type_proto_contains_float_tensor(value.type) + for values in ( + graph.input, + graph.output, + graph.value_info, + ) + for value in values + ): + return True + pending_attributes.extend( + attribute + for node in graph.node + for attribute in node.attribute + if not attribute.ref_attr_name + ) + return False + + +def _reject_unconverted_local_function_float_data(model: ModelProto, blocked_ops: set[str]) -> None: + """Reject concrete FLOAT data in function bodies ORT never traverses.""" + functions = { + ( + function.domain, + function.name, + getattr(function, "overload", ""), + ): function + for function in getattr(model, "functions", []) + } + pending = [ + function + for graph in _ort_traversed_graphs(model, list(blocked_ops)) + for node in graph.node + if node.op_type not in blocked_ops + and ( + function := functions.get( + ( + node.domain, + node.op_type, + getattr(node, "overload", ""), + ) + ) + ) + is not None + and ( + any( + _type_proto_contains_float_tensor(value_type) + for input_index in range(len(node.input)) + for value_type in _node_input_actual_types(model, graph, node, input_index) + ) + or any( + _type_proto_contains_float_tensor(value_type) + for output_name in node.output + if output_name + for value_type in _graph_declared_types(graph, output_name) + ) + ) + ] + visited: set[tuple[str, str, str]] = set() + while pending: + function = pending.pop() + key = ( + function.domain, + function.name, + getattr(function, "overload", ""), + ) + if key in visited: + continue + visited.add(key) + if _function_contains_concrete_float(function): + msg = ( + f"ORT does not convert concrete FLOAT data in local function " + f"'{function.domain}::{function.name}'." + ) + raise RuntimeError(msg) + pending.extend( + nested + for node in function.node + if ( + nested := functions.get( + ( + node.domain, + node.op_type, + getattr(node, "overload", ""), + ) + ) + ) + is not None + ) + + +def _reject_blocked_subgraph_converted_captures( + model: ModelProto, + *, + keep_io_types: bool, + op_block_list: list[str] | None, +) -> None: + """Reject free captures whose FLOAT binding changes while ORT skips the child.""" + blocked_ops = _effective_blocked_ops(op_block_list) + graphs, parents = _ort_graphs_with_parents(model, blocked_ops) + _, _, fp16_initializers, _, _ = _initializer_tracking_analysis( + model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + graphs=graphs, + parents=parents, + ) + name_mapping, _ = _ort_keep_io_name_mapping(model, keep_io_types=keep_io_types) + owners_by_id = {id(graph): graph for graph in graphs} + for graph in graphs: + for node in getattr(graph, "node", []): + if not _ort_skips_node_attributes(node, blocked_ops): + continue + executed_attributes = _local_function_executed_attributes(model, node) + children = ( + child + for attribute in ( + node.attribute if executed_attributes is None else executed_attributes + ) + for child in _iter_graphs_from_attribute(attribute) + ) + for child in children: + for name in _graph_free_references(child): + owner_id = _lexical_binding_owner_id(graph, name, parents, set()) + owner = owners_by_id.get(owner_id) if owner_id is not None else None + if owner is None or not _binding_converts_to_fp16( + model, + owner, + name, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + name_mapping=name_mapping, + fp16_initializers=fp16_initializers, + parents=parents, + ): + continue + msg = ( + f"A blocked subgraph under node '{node.name}' captures FLOAT " + f"value '{name}' that converts to FP16 while ORT skips the " + "subgraph." + ) + raise RuntimeError(msg) + + +def _validate_local_function_conversion(model: ModelProto) -> None: + """Validate converted local functions through their expanded call sites.""" + if not getattr(model, "functions", []): + return + + from onnx import checker, shape_inference + from onnx.inliner import inline_local_functions + + inlined = inline_local_functions(model) + try: + checker.check_model(inlined) + inferred = shape_inference.infer_shapes(inlined, check_type=True, strict_mode=True) + checker.check_model(inferred) + except ( + checker.ValidationError, + shape_inference.InferenceError, + ) as error: + msg = ( + "ORT leaves local function bodies unconverted, producing " + "incompatible FP16 call-site types." + ) + raise RuntimeError(msg) from error + + +def _validate_converted_types(model: ModelProto) -> None: + """Reject converted graphs whose concrete types no longer agree.""" + from onnx import checker, shape_inference + + try: + checker.check_model(model) + inferred = shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + checker.check_model(inferred) + except ( + checker.ValidationError, + shape_inference.InferenceError, + ) as error: + msg = "FP16 conversion produced incompatible FP16 types." + raise RuntimeError(msg) from error + + +def _node_dependencies(node: NodeProto) -> set[str]: + """Return explicit inputs and outer values captured by the node's subgraphs.""" + dependencies = {input_name for input_name in node.input if input_name} + for child in _iter_all_child_graphs_from_node(node): + dependencies.update(_graph_free_references(child)) + return dependencies + + +def _graph_topological_sort(graph: GraphProto) -> None: + """Topologically sort nodes while treating dense and sparse initializers as inputs.""" + deps = {initializer.name for initializer in getattr(graph, "initializer", [])} + deps.update(sparse.values.name for sparse in getattr(graph, "sparse_initializer", [])) + deps.update(value.name for value in getattr(graph, "input", [])) + node_dependencies = [_node_dependencies(node) for node in graph.node] + + sorted_indices: set[int] = set() + sorted_nodes = [] + last_blocked_node = None + previous_count = -1 + while len(sorted_indices) != len(graph.node): + if len(sorted_indices) == previous_count: + break + previous_count = len(sorted_indices) + for node_index, node in enumerate(graph.node): + if node_index in sorted_indices: + continue + if node_dependencies[node_index] <= deps: + sorted_nodes.append(node) + sorted_indices.add(node_index) + deps.update(output for output in node.output if output) + else: + last_blocked_node = node.name + + if len(sorted_indices) != len(graph.node): + msg = ( + "Graph is not a DAG: " + f"len(sorted_node_set)={len(sorted_indices)}, " + f"len(graph.node)={len(graph.node)}, " + f"failed at node {last_blocked_node}" + ) + raise RuntimeError(msg) + + del graph.node[:] + graph.node.extend(sorted_nodes) + + def convert_to_fp16( model: ModelProto, *, @@ -31,34 +2893,130 @@ def convert_to_fp16( """Convert an ONNX model from FP32 to FP16 precision. Uses onnxruntime.transformers.float16.convert_float_to_float16 internally. - No new dependencies — ORT is already a project dependency. + The successful conversion mutates and returns ``model`` as before. - Note: ORT's converter mutates the model in-place and returns the same object. - - Args: - model: Input ONNX ModelProto (will be mutated in-place by ORT). - keep_io_types: If True, preserve FP32 model inputs/outputs by inserting - Cast nodes at boundaries. Recommended for CPU-safe inference. - op_block_list: Op types to keep in FP32 (e.g., ["LayerNorm", "Softmax"]). - When None, ORT uses its DEFAULT_OP_BLOCK_LIST which includes ops - known to be numerically unsafe in FP16 (e.g., TopK, CumSum, etc.). - - Returns: - The converted model (same object as input due to ORT in-place mutation). + ORT assumes each graph output has a node producer. For a safe top-level + output supplied only by a dense FLOAT initializer, keep-I/O conversion adds + a Cast with no producer; remove that exact Cast. Pure-FP16 conversion changes + the output declaration but not its initializer, so convert that initializer + explicitly. """ from onnx import TensorProto from onnxruntime.transformers.float16 import convert_float_to_float16 - # Skip if model is already FP16 (check floating-point initializer dtypes) - fp32_types = {TensorProto.FLOAT, TensorProto.DOUBLE, TensorProto.BFLOAT16} - initializers = model.graph.initializer - if initializers: - float_inits = [t for t in initializers if t.data_type in fp32_types | {TensorProto.FLOAT16}] - if float_inits and all(t.data_type == TensorProto.FLOAT16 for t in float_inits): - logger.info("Model is already FP16 — skipping conversion.") - return model + _reject_sparse_initializer_tensor_metadata(model, op_block_list) + _reject_duplicate_float_initializer_names(model, op_block_list) + io_preflight_model = _ort_inference_preflight_model(model) + blocked_ops = _effective_blocked_ops(op_block_list) + _reject_unpreserved_float_container_io( + io_preflight_model, + keep_io_types=keep_io_types, + ) + _reject_scope_unsafe_keep_io_mappings( + io_preflight_model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + ) + _reject_generated_io_cast_name_collisions( + io_preflight_model, + keep_io_types=keep_io_types, + op_block_list=op_block_list, + ) + _reject_scope_unsafe_initializer_tracking( + io_preflight_model, + keep_io_types=keep_io_types, + op_block_list=op_block_list, + ) + _reject_unconverted_local_function_float_data(io_preflight_model, blocked_ops) + _reject_blocked_subgraph_converted_captures( + io_preflight_model, + keep_io_types=keep_io_types, + op_block_list=op_block_list, + ) + _reject_scope_unsafe_value_info_lookups( + io_preflight_model, + keep_io_types=keep_io_types, + op_block_list=op_block_list, + ) + captured = _capture_safe_initializer_outputs( + model, keep_io_types=keep_io_types, op_block_list=op_block_list + ) + selected_external_names = _ort_converted_initializer_names( + io_preflight_model, + keep_io_types=keep_io_types, + blocked_ops=blocked_ops, + ) + selected_external_names.update( + item.name for item in captured if not keep_io_types or item.graph_index != 0 + ) + from onnx.external_data_helper import uses_external_data + selected_external_initializers = [ + initializer + for graph in _ort_traversed_graphs(model, op_block_list) + for initializer in graph.initializer + if initializer.name in selected_external_names + and initializer.data_type == TensorProto.FLOAT + and uses_external_data(initializer) + ] + unloaded_external = sorted( + initializer.name + for initializer in selected_external_initializers + if not _tensor_data_is_loaded(initializer) + ) + if unloaded_external: + names = ", ".join(unloaded_external) + msg = ( + f"FLOAT initializers use unloaded external data: {names}; " + "load external weights before FP16 conversion." + ) + raise RuntimeError(msg) + external_attribute_tensors = _external_float_attribute_tensors(model, blocked_ops) + unloaded_attribute_tensors = sorted( + tensor.name or "" + for tensor in external_attribute_tensors + if not _tensor_data_is_loaded(tensor) + ) + if unloaded_attribute_tensors: + names = ", ".join(unloaded_attribute_tensors) + msg = ( + f"FLOAT tensor attributes use unloaded external data: {names}; " + "load external weights before FP16 conversion." + ) + raise RuntimeError(msg) + requires_attribute_validation = any( + _attribute_requires_type_validation(attribute) + for attribute in _ort_traversed_attributes(io_preflight_model, blocked_ops) + ) + requires_container_validation = _has_unconverted_float_container_declarations( + io_preflight_model, op_block_list + ) + _reject_shared_keep_io_names( + io_preflight_model, + keep_io_types=keep_io_types, + ) + needs_safe_conversion = ( + bool(captured) + or _has_nested_initializer_outputs(model, op_block_list) + or _has_float_sparse_initializers(model, op_block_list) + or bool(getattr(model, "functions", [])) + or bool(selected_external_initializers) + or bool(external_attribute_tensors) + or requires_attribute_validation + or requires_container_validation + ) + if needs_safe_conversion: + _reject_unloaded_external_initializer_outputs(model, op_block_list) original_nodes = len(model.graph.node) + conversion_model = deepcopy(model) if needs_safe_conversion else model + if needs_safe_conversion: + _internalize_external_initializer_outputs(conversion_model, op_block_list) + _internalize_selected_external_initializers( + conversion_model, + selected_external_names, + op_block_list, + ) + _internalize_external_float_attribute_tensors(conversion_model, blocked_ops) logger.info("Converting model to FP16...") if keep_io_types: @@ -68,7 +3026,7 @@ def convert_to_fp16( try: converted: ModelProto = convert_float_to_float16( - model, + conversion_model, keep_io_types=keep_io_types, op_block_list=op_block_list, ) @@ -79,19 +3037,36 @@ def convert_to_fp16( "large ONNX models that use external data." ) converted = convert_float_to_float16( - model, + conversion_model, keep_io_types=keep_io_types, disable_shape_infer=True, op_block_list=op_block_list, ) - # ORT's converter appends Cast nodes at the end of the node list (for - # keep_io_types), which breaks topological ordering. Re-sort the graph - # using ORT's own topological sort utility. + converted_graphs = _ort_traversed_graphs(converted, op_block_list) if keep_io_types: - from onnxruntime.transformers.onnx_model import OnnxModel + _remove_orphan_output_casts(converted, [item for item in captured if item.graph_index == 0]) + + for item in captured: + if keep_io_types and item.graph_index == 0: + continue + graph = converted_graphs[item.graph_index] + if item.has_consumers and _initializer_data_type(graph, item.name) == TensorProto.FLOAT16: + _internalize_output_initializer(graph, item.name) + elif not item.has_consumers: + _set_direct_output_elem_type(graph, item.name, TensorProto.FLOAT16) + _convert_output_initializer_to_fp16(graph, item.name) + + _repair_sparse_float_initializers(converted, op_block_list, keep_io_types=keep_io_types) + _graph_topological_sort(converted.graph) + _validate_initializer_output_types(converted) + if requires_attribute_validation or external_attribute_tensors or requires_container_validation: + _validate_converted_types(converted) + _validate_local_function_conversion(converted) - OnnxModel.graph_topological_sort(converted.graph) + if converted is not model: + model.CopyFrom(converted) + converted = model converted_nodes = len(converted.graph.node) if converted_nodes != original_nodes: diff --git a/tests/unit/optim/test_fp16.py b/tests/unit/optim/test_fp16.py index 2e63ab193..caa2f4486 100644 --- a/tests/unit/optim/test_fp16.py +++ b/tests/unit/optim/test_fp16.py @@ -15,73 +15,5719 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np -from onnx import ModelProto, TensorProto, helper, numpy_helper +import onnxruntime as ort +from google.protobuf.message import EncodeError +from onnx import ( + AttributeProto, + GraphProto, + ModelProto, + SparseTensorProto, + TensorProto, + checker, + helper, + numpy_helper, + shape_inference, +) + +from winml.modelkit.quant.fp16 import convert_to_fp16 + + +if TYPE_CHECKING: + import pytest + + +# ============================================================================= +# HELPERS +# ============================================================================= + + +def _build_simple_fp32_model() -> ModelProto: + """Build a simple FP32 model: out = x + weight.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]) + out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4]) + weight = numpy_helper.from_array(np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight") + add = helper.make_node("Add", ["x", "weight"], ["out"], name="add") + graph = helper.make_graph([add], "simple", [x], [out], [weight]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_multi_op_fp32_model() -> ModelProto: + """Build a model with multiple ops: out = Relu(x + weight).""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]) + out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4]) + weight = numpy_helper.from_array(np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight") + add = helper.make_node("Add", ["x", "weight"], ["add_out"], name="add") + relu = helper.make_node("Relu", ["add_out"], ["out"], name="relu") + graph = helper.make_graph([add, relu], "multi_op", [x], [out], [weight]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_scalar_float_attribute_model() -> ModelProto: + """Build a FLOAT-producing scalar attribute ORT leaves unchanged.""" + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, []) + constant = helper.make_node( + "Constant", + [], + ["output"], + name="constant", + value_float=3.5, + ) + graph = helper.make_graph([constant], "scalar_float_attribute", [], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_sparse_float_attribute_model() -> ModelProto: + """Build a sparse FLOAT tensor attribute ORT leaves unchanged.""" + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [2]) + values = numpy_helper.from_array(np.array([3.5], dtype=np.float32), "values") + indices = numpy_helper.from_array(np.array([[1]], dtype=np.int64), "indices") + sparse = helper.make_sparse_tensor(values, indices, [2]) + constant = helper.make_node( + "Constant", + [], + ["output"], + name="constant", + sparse_value=sparse, + ) + graph = helper.make_graph([constant], "sparse_float_attribute", [], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_external_float_attribute_model(*, clear_data: bool) -> ModelProto: + """Build a FLOAT tensor attribute carrying external-data metadata.""" + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + value = numpy_helper.from_array(np.array([3.5], dtype=np.float32), "value") + if clear_data: + value.ClearField("raw_data") + del value.float_data[:] + value.data_location = TensorProto.EXTERNAL + location = value.external_data.add() + location.key = "location" + location.value = "value.bin" + length = value.external_data.add() + length.key = "length" + length.value = "4" + constant = helper.make_node("Constant", [], ["output"], name="constant", value=value) + graph = helper.make_graph([constant], "external_float_attribute", [], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_external_int_name_collision_model() -> ModelProto: + """Reuse a selected FLOAT weight name for nested external INT data.""" + model = _build_simple_fp32_model() + weight = model.graph.initializer[0] + weight.data_location = TensorProto.EXTERNAL + weight_location = weight.external_data.add() + weight_location.key = "location" + weight_location.value = "weight.bin" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + integer_output = helper.make_tensor_value_info("integer_output", TensorProto.INT64, [1]) + + def _branch(name: str, initializer_name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.INT64, [1]) + initializer = numpy_helper.from_array(np.array([7], dtype=np.int64), initializer_name) + if name == "then": + initializer.ClearField("raw_data") + del initializer.int32_data[:] + del initializer.int64_data[:] + initializer.data_location = TensorProto.EXTERNAL + location = initializer.external_data.add() + location.key = "location" + location.value = "integer.bin" + identity = helper.make_node( + "Identity", + [initializer_name], + [branch_output.name], + name=f"{name}_identity", + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + conditional = helper.make_node( + "If", + ["condition"], + ["integer_output"], + name="if", + then_branch=_branch("then", "weight"), + else_branch=_branch("else", "other"), + ) + model.graph.input.append(condition) + model.graph.output.append(integer_output) + model.graph.node.append(conditional) + model.graph.name = "nested_external_int_name_collision" + return model + + +def _build_omitted_optional_output_model() -> ModelProto: + """Build optional inputs and outputs that use the empty sentinel.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2]) + dropped = helper.make_tensor_value_info("dropped", TensorProto.FLOAT, [2]) + clipped = helper.make_tensor_value_info("clipped", TensorProto.FLOAT, [2]) + drop = helper.make_node("Dropout", ["x"], ["dropped", ""], name="drop") + clip = helper.make_node("Clip", ["x", "", ""], ["clipped"], name="clip") + graph = helper.make_graph( + [drop, clip], + "omitted_optional_output", + [x], + [dropped, clipped], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_initializer_backed_output_model() -> ModelProto: + """Build a graph whose output is supplied directly by an initializer.""" + out = helper.make_tensor_value_info("constant_output", TensorProto.FLOAT, [1, 2]) + value = numpy_helper.from_array( + np.array([[1.0001, 2.0003]], dtype=np.float32), "constant_output" + ) + graph = helper.make_graph([], "initializer_output", [], [out], [value]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_shared_initializer_output_model() -> ModelProto: + """Build a graph where an initializer is both an output and a node input.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2]) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1, 2]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2]) + value = numpy_helper.from_array(np.array([[1.0, 2.0]], dtype=np.float32), "shared") + add = helper.make_node("Add", ["x", "shared"], ["y"], name="add") + graph = helper.make_graph([add], "shared_initializer", [x], [shared, y], [value]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_overridable_shared_initializer_output_model() -> ModelProto: + """Build a graph input/output initializer that can be overridden by callers.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2]) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1, 2]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2]) + value = numpy_helper.from_array(np.array([[1.0, 2.0]], dtype=np.float32), "shared") + add = helper.make_node("Add", ["x", "shared"], ["y"], name="add") + graph = helper.make_graph([add], "overridable_shared", [x, shared], [shared, y], [value]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_initializer_output_model() -> ModelProto: + """Build an If whose branch outputs are supplied by initializers.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + + def _branch(name: str, output_name: str, value: float): + branch_output = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), output_name) + return helper.make_graph([], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", "then_output", 1.0), + else_branch=_branch("else", "else_output", 2.0), + ) + graph = helper.make_graph([node], "nested_initializer", [condition], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_consumed_initializer_output_model() -> ModelProto: + """Build an If whose branch initializer outputs are also consumed locally.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + + def _branch(name: str, value: float) -> GraphProto: + initializer_name = f"{name}_value" + branch_output = helper.make_tensor_value_info(initializer_name, TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), initializer_name) + identity = helper.make_node( + "Identity", [initializer_name], [f"{name}_used"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", 1.0), + else_branch=_branch("else", 2.0), + ) + graph = helper.make_graph([node], "nested_consumed_initializer", [condition], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_consumed_initializer_output_with_top_level_input_collision_model() -> ModelProto: + """Build nested local initializer consumers shadowing a kept top-level input.""" + same = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + + def _branch(name: str, value: float, *, collides: bool) -> GraphProto: + output_name = "same" if collides else f"{name}_value" + branch_output = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), output_name) + identity = helper.make_node( + "Identity", [output_name], [f"{name}_used"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", 1.0, collides=True), + else_branch=_branch("else", 2.0, collides=False), + ) + graph = helper.make_graph([node], "nested_input_collision", [same, condition], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_fp32_identity_with_fp16_nested_initializer_model() -> ModelProto: + """Build FP32 top-level I/O with only nested FP16 floating initializers.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + branch_output = helper.make_tensor_value_info("branch_value", TensorProto.FLOAT16, [1]) + nested_output = helper.make_tensor_value_info("nested_output", TensorProto.FLOAT16, [1]) + initializer = numpy_helper.from_array(np.array([1.0], dtype=np.float16), "branch_value") + branch = helper.make_graph([], "branch", [], [branch_output], [initializer]) + identity = helper.make_node("Identity", ["x"], ["y"], name="identity") + if_node = helper.make_node( + "If", + ["condition"], + ["nested_output"], + name="if", + then_branch=branch, + else_branch=branch, + ) + graph = helper.make_graph( + [identity, if_node], + "fp32_top_level_with_fp16_nested_initializer", + [x, condition], + [y, nested_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_sparse_shape_reshape_model() -> ModelProto: + """Build a Reshape graph that consumes a sparse INT64 shape initializer.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [4]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [2, 2]) + sparse_shape = SparseTensorProto() + sparse_shape.values.CopyFrom(numpy_helper.from_array(np.array([2, 2], dtype=np.int64))) + sparse_shape.indices.CopyFrom(numpy_helper.from_array(np.array([[0], [1]], dtype=np.int64))) + sparse_shape.dims.extend([2]) + sparse_shape.values.name = "shape" + reshape = helper.make_node("Reshape", ["x", "shape"], ["output"], name="reshape") + graph = helper.make_graph([reshape], "sparse_shape", [x], [output]) + graph.sparse_initializer.append(sparse_shape) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_sparse_float_add_model() -> ModelProto: + """Build an Add graph that consumes a sparse FLOAT initializer.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [2]) + sparse_weight = SparseTensorProto() + sparse_weight.values.CopyFrom(numpy_helper.from_array(np.array([1.0, 2.0], dtype=np.float32))) + sparse_weight.indices.CopyFrom(numpy_helper.from_array(np.array([[0], [1]], dtype=np.int64))) + sparse_weight.dims.extend([2]) + sparse_weight.values.name = "weight" + add = helper.make_node("Add", ["x", "weight"], ["output"], name="add") + graph = helper.make_graph([add], "sparse_float_add", [x], [output]) + graph.sparse_initializer.append(sparse_weight) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_sparse_float_add_model_with_value_info() -> ModelProto: + """Build an Add graph with sparse FLOAT initializer metadata.""" + model = _build_sparse_float_add_model() + model.graph.value_info.append( + helper.make_sparse_tensor_value_info("weight", TensorProto.FLOAT, [2]) + ) + return model + + +def _build_sparse_float_add_model_with_tensor_value_info() -> ModelProto: + """Build an Add graph with tensor metadata for a sparse FLOAT initializer.""" + model = _build_sparse_float_add_model() + model.graph.value_info.append(helper.make_tensor_value_info("weight", TensorProto.FLOAT, [2])) + return model + + +def _build_retained_float_initializer_metadata_model() -> ModelProto: + """Build an always-FLOAT initializer with convertible tensor metadata.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 2]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1, 4, 4]) + scales = numpy_helper.from_array( + np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), + "scales", + ) + resize = helper.make_node( + "Resize", + ["x", "", "scales"], + ["output"], + name="resize", + mode="nearest", + ) + graph = helper.make_graph( + [resize], + "retained_float_initializer_metadata", + [x], + [output], + [scales], + value_info=[helper.make_tensor_value_info("scales", TensorProto.FLOAT, [4])], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_sparse_float_add_model_with_io_metadata() -> ModelProto: + """Build an Add graph whose sparse FLOAT initializer is also graph sparse I/O.""" + model = _build_sparse_float_add_model() + model.graph.input.append(helper.make_sparse_tensor_value_info("weight", TensorProto.FLOAT, [2])) + model.graph.output.append( + helper.make_sparse_tensor_value_info("weight", TensorProto.FLOAT, [2]) + ) + return model + + +def _build_sparse_float_add_model_with_tensor_io_metadata() -> ModelProto: + """Build an Add graph whose sparse FLOAT initializer is also graph tensor I/O.""" + model = _build_sparse_float_add_model() + model.graph.input.append(helper.make_tensor_value_info("weight", TensorProto.FLOAT, [2])) + model.graph.output.append(helper.make_tensor_value_info("weight", TensorProto.FLOAT, [2])) + return model + + +def _build_direct_sparse_float_output_model() -> ModelProto: + """Build a graph with a direct sparse FLOAT initializer output.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2]) + dense_output = helper.make_tensor_value_info("dense_output", TensorProto.FLOAT, [2]) + sparse_output = helper.make_sparse_tensor_value_info("sparse_output", TensorProto.FLOAT, [2]) + sparse_value = SparseTensorProto() + sparse_value.values.CopyFrom(numpy_helper.from_array(np.array([1.0, 2.0], dtype=np.float32))) + sparse_value.indices.CopyFrom(numpy_helper.from_array(np.array([[0], [1]], dtype=np.int64))) + sparse_value.dims.extend([2]) + sparse_value.values.name = "sparse_output" + identity = helper.make_node("Identity", ["x"], ["dense_output"], name="identity") + graph = helper.make_graph( + [identity], "direct_sparse_output", [x], [dense_output, sparse_output] + ) + graph.sparse_initializer.append(sparse_value) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_direct_sparse_float_tensor_output_model() -> ModelProto: + """Build a graph with a direct sparse FLOAT initializer and tensor output metadata.""" + model = _build_direct_sparse_float_output_model() + sparse_output = next(value for value in model.graph.output if value.name == "sparse_output") + sparse_output.CopyFrom(helper.make_tensor_value_info("sparse_output", TensorProto.FLOAT, [2])) + return model + + +def _build_if_with_direct_sparse_float_outputs_model() -> ModelProto: + """Build an If whose branches return direct sparse FLOAT initializers.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_sparse_tensor_value_info("output", TensorProto.FLOAT, [2]) + + def _branch(name: str, values: list[float]) -> GraphProto: + output_name = f"{name}_output" + branch_output = helper.make_sparse_tensor_value_info(output_name, TensorProto.FLOAT, [2]) + sparse_value = SparseTensorProto() + sparse_value.values.CopyFrom(numpy_helper.from_array(np.array(values, dtype=np.float32))) + sparse_value.indices.CopyFrom(numpy_helper.from_array(np.array([[0], [1]], dtype=np.int64))) + sparse_value.dims.extend([2]) + sparse_value.values.name = output_name + graph = helper.make_graph([], name, [], [branch_output]) + graph.sparse_initializer.append(sparse_value) + return graph + + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", [1.0, 2.0]), + else_branch=_branch("else", [3.0, 4.0]), + ) + graph = helper.make_graph([conditional], "nested_sparse_outputs", [condition], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_deep_if_with_direct_sparse_float_outputs_model() -> ModelProto: + """Build two nested If levels ending in direct sparse FLOAT outputs.""" + outer_condition = helper.make_tensor_value_info("outer_condition", TensorProto.BOOL, []) + inner_condition = helper.make_tensor_value_info("inner_condition", TensorProto.BOOL, []) + output = helper.make_sparse_tensor_value_info("output", TensorProto.FLOAT, [2]) + + def _leaf(name: str, values: list[float]) -> GraphProto: + output_name = f"{name}_output" + branch_output = helper.make_sparse_tensor_value_info(output_name, TensorProto.FLOAT, [2]) + sparse_value = SparseTensorProto() + sparse_value.values.CopyFrom( + numpy_helper.from_array(np.array(values, dtype=np.float32), output_name) + ) + sparse_value.indices.CopyFrom(numpy_helper.from_array(np.array([[0], [1]], dtype=np.int64))) + sparse_value.dims.extend([2]) + graph = helper.make_graph([], name, [], [branch_output]) + graph.sparse_initializer.append(sparse_value) + return graph + + def _outer_branch(name: str, offset: float) -> GraphProto: + branch_output = helper.make_sparse_tensor_value_info( + f"{name}_output", TensorProto.FLOAT, [2] + ) + inner = helper.make_node( + "If", + ["inner_condition"], + [branch_output.name], + name=f"{name}_if", + then_branch=_leaf(f"{name}_then", [offset + 1.0, offset + 2.0]), + else_branch=_leaf(f"{name}_else", [offset + 3.0, offset + 4.0]), + ) + return helper.make_graph([inner], name, [], [branch_output]) + + outer = helper.make_node( + "If", + ["outer_condition"], + ["output"], + name="outer_if", + then_branch=_outer_branch("outer_then", 0.0), + else_branch=_outer_branch("outer_else", 4.0), + ) + graph = helper.make_graph( + [outer], + "deep_nested_sparse_outputs", + [outer_condition, inner_condition], + [output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _iter_attribute_graphs(model: ModelProto) -> list[GraphProto]: + """Return nested graphs stored on node attributes.""" + graphs: list[GraphProto] = [] + for node in model.graph.node: + for attribute in node.attribute: + if attribute.g.name: + graphs.append(attribute.g) + graphs.extend(attribute.graphs) + return graphs + + +def _mark_initializers_as_external(graph: GraphProto, *, clear_data: bool) -> None: + """Mark all graph initializers as external, optionally without resident bytes.""" + for initializer in graph.initializer: + if clear_data: + initializer.ClearField("raw_data") + del initializer.float_data[:] + initializer.data_location = TensorProto.EXTERNAL + location = initializer.external_data.add() + location.key = "location" + location.value = f"{initializer.name}.bin" + + +def _build_lexically_captured_initializer_output_model() -> ModelProto: + """Build an output initializer captured only by nested If branches.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([1.0], dtype=np.float32), "shared") + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", ["shared"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + node = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [node], "lexical_capture", [condition], [shared, output], [initializer] + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_initializer_output_name_collision_model() -> ModelProto: + """Build a legal graph with a name that collides with ORT's generated alias.""" + existing = helper.make_tensor_value_info("graph_output_cast_0", TensorProto.FLOAT, [1]) + constant_output = helper.make_tensor_value_info("constant_output", TensorProto.FLOAT, [1]) + result = helper.make_tensor_value_info("result", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([1.0], dtype=np.float32), "constant_output") + identity = helper.make_node("Identity", ["graph_output_cast_0"], ["result"], name="identity") + graph = helper.make_graph( + [identity], "name_collision", [existing], [constant_output, result], [initializer] + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_initializer_output_node_name_collision_model() -> ModelProto: + """Build a graph with a user node named like ORT's output Cast node.""" + model = _build_initializer_backed_output_model() + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1]) + model.graph.input.append(x) + model.graph.output.append(y) + model.graph.node.append(helper.make_node("Identity", ["x"], ["y"], name="graph_output_cast0")) + return model + + +def _build_nested_node_name_collision_model() -> ModelProto: + """Build a nested user node named like ORT's top-level output Cast.""" + model = _build_initializer_backed_output_model() + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + nested_output = helper.make_tensor_value_info("nested_output", TensorProto.FLOAT, [1]) + + def _branch(name: str, value: float) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), f"{name}_value") + identity = helper.make_node( + "Identity", + [f"{name}_value"], + [f"{name}_output"], + name="graph_output_cast0" if name == "then" else f"{name}_identity", + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["nested_output"], + name="if", + then_branch=_branch("then", 1.0), + else_branch=_branch("else", 2.0), + ) + model.graph.input.append(condition) + model.graph.output.append(nested_output) + model.graph.node.append(node) + return model + + +def _build_regular_output_nested_node_name_collision_model() -> ModelProto: + """Build a normal output whose generated Cast name collides in a nested graph.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + nested_output = helper.make_tensor_value_info("nested_output", TensorProto.FLOAT, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", + ["x"], + [f"{name}_output"], + name="graph_output_cast0" if name == "then" else f"{name}_identity", + ) + return helper.make_graph([identity], name, [], [branch_output]) + + identity = helper.make_node("Identity", ["x"], ["output"], name="identity") + conditional = helper.make_node( + "If", + ["condition"], + ["nested_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [identity, conditional], + "regular_output_nested_node_collision", + [x, condition], + [output, nested_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_neutral_nested_node_name_collision_model() -> ModelProto: + """Build a generated-name collision on a conversion-neutral INT node.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + integer = helper.make_tensor_value_info("integer", TensorProto.INT64, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + nested_output = helper.make_tensor_value_info("nested_output", TensorProto.INT64, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.INT64, [1]) + identity = helper.make_node( + "Identity", + ["integer"], + [branch_output.name], + name=("graph_output_cast0" if name == "then" else f"{name}_identity"), + ) + return helper.make_graph([identity], name, [], [branch_output]) + + identity = helper.make_node("Identity", ["x"], ["output"], name="identity") + conditional = helper.make_node( + "If", + ["condition"], + ["nested_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [identity, conditional], + "neutral_nested_node_name_collision", + [x, integer, condition], + [output, nested_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_neutral_shadowed_keep_io_collision_model() -> ModelProto: + """Shadow kept FLOAT I/O only inside a skipped neutral INT node.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + integer = helper.make_tensor_value_info("integer", TensorProto.INT64, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.INT64, [1]) + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_tensor_value_info("x", TensorProto.INT64, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.INT64, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="condition_feedback", + ), + helper.make_node( + "Identity", + ["x"], + ["body_state_out"], + name="graph_input_cast0", + ), + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_state_out], + ) + identity = helper.make_node("Identity", ["x"], ["output"], name="identity") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "integer"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [identity, loop], + "neutral_shadowed_keep_io_collision", + [x, trip_count, condition, integer], + [output, loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_inferred_output_nested_node_name_collision_model() -> ModelProto: + """Build an inferred FLOAT output whose generated Cast name collides.""" + model = _build_regular_output_nested_node_name_collision_model() + model.graph.output[0].CopyFrom(helper.make_empty_tensor_value_info("output")) + return model + + +def _build_nested_local_generated_tensor_alias_model() -> ModelProto: + """Build a Loop with a local input matching a top-level generated tensor alias.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_state = helper.make_tensor_value_info("graph_output_cast_0", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ), + helper.make_node( + "Identity", + ["graph_output_cast_0"], + ["body_state_out"], + name="body_state", + ), + ], + "body", + [iteration, body_condition, local_state], + [body_condition_out, body_state_out], + ) + identity = helper.make_node("Identity", ["x"], ["output"], name="identity") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [identity, loop], + "nested_local_generated_alias", + [trip_count, condition, loop_state, x], + [output, loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_kept_input_free_capture_value_info_model() -> ModelProto: + """Annotate a free kept-input capture with nested value_info.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", + ["x"], + [branch_output.name], + name=f"{name}_identity", + ) + return helper.make_graph( + [identity], + name, + [], + [branch_output], + value_info=[helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])], + ) + + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [conditional], + "kept_input_free_capture_value_info", + [x, condition], + [output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_blocked_generated_tensor_alias_model() -> ModelProto: + """Build a blocked nested node that consumes a generated tensor alias locally.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.INT64, [1]) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.INT64, [1]) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_state = helper.make_tensor_value_info("graph_output_cast_0", TensorProto.INT64, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.INT64, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="body_condition", + ), + helper.make_node( + "Identity", + ["graph_output_cast_0"], + ["body_state_out"], + name="body_state", + ), + ], + "body", + [iteration, body_condition, local_state], + [body_condition_out, body_state_out], + ) + relu = helper.make_node("Relu", ["x"], ["output"], name="relu") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [relu, loop], + "nested_blocked_generated_alias", + [trip_count, condition, loop_state, x], + [output, loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_mixed_generated_tensor_alias_model() -> ModelProto: + """Build a mixed-type nested node that consumes a generated tensor alias locally.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_scales = helper.make_tensor_value_info("graph_output_cast_0", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + resize_input = helper.make_tensor("resize_input", TensorProto.FLOAT, [1], [2.0]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="body_condition", + ), + helper.make_node( + "Resize", + ["resize_input", "", "graph_output_cast_0"], + ["body_state_out"], + name="body_resize", + ), + ], + "body", + [iteration, body_condition, local_scales], + [body_condition_out, body_state_out], + initializer=[resize_input], + ) + relu = helper.make_node("Relu", ["x"], ["output"], name="relu") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [relu, loop], + "nested_mixed_generated_alias", + [trip_count, condition, loop_state, x], + [output, loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_inferred_nested_mixed_name_collision_model() -> ModelProto: + """Build a mixed nested input collision visible only after shape inference.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_scale = helper.make_tensor_value_info("loop_scale", TensorProto.FLOAT, [1]) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [None]) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_scale = helper.make_empty_tensor_value_info("same") + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [None]) + resize_input = helper.make_tensor("resize_input", TensorProto.FLOAT, [1], [2.0]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="body_condition", + ), + helper.make_node( + "Resize", + ["resize_input", "", "same"], + ["body_state_out"], + name="body_resize", + ), + ], + "body", + [iteration, body_condition, local_scale], + [body_condition_out, body_state_out], + initializer=[resize_input], + ) + identity = helper.make_node("Identity", ["x"], ["same"], name="top_identity") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_scale"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [identity, loop], + "inferred_nested_mixed_collision", + [trip_count, condition, loop_scale, x], + [loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_top_blocked_global_value_info_collision_model() -> ModelProto: + """Build a top blocked input colliding with nested FLOAT metadata.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + same = helper.make_tensor_value_info("same", TensorProto.INT64, [1]) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + integer_output = helper.make_tensor_value_info("integer_output", TensorProto.INT64, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_same = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [ + helper.make_node( + "And", + ["body_condition", "body_condition"], + ["body_condition_out"], + name="body_condition", + ), + helper.make_node("Relu", ["same"], ["body_state_out"], name="body_relu"), + ], + "body", + [iteration, body_condition, local_same], + [body_condition_out, body_state_out], + ) + blocked = helper.make_node("Identity", ["same"], ["integer_output"]) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [blocked, loop], + "top_blocked_global_collision", + [trip_count, condition, same, loop_state], + [integer_output, loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_duplicate_late_cast_alias_model() -> ModelProto: + """Build two unnamed blocked nodes that generate identical late Cast aliases.""" + a = helper.make_tensor_value_info("a", TensorProto.FLOAT, [1]) + b = helper.make_tensor_value_info("b", TensorProto.FLOAT, [1]) + c = helper.make_tensor_value_info("c", TensorProto.FLOAT, [1]) + d = helper.make_tensor_value_info("d", TensorProto.FLOAT, [1]) + graph = helper.make_graph( + [ + helper.make_node("Identity", ["a"], ["b"]), + helper.make_node("Identity", ["c"], ["d"]), + ], + "duplicate_late_cast_alias", + [a, c], + [b, d], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_late_cast_node_name_collision_model() -> ModelProto: + """Build a blocked node whose generated Cast node name already exists.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + other = helper.make_tensor_value_info("other", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + other_output = helper.make_tensor_value_info("other_output", TensorProto.FLOAT, [1]) + blocked = helper.make_node("Abs", ["x"], ["output"], name="blocked") + existing = helper.make_node( + "Identity", + ["other"], + ["other_output"], + name="blocked_input_cast0", + ) + graph = helper.make_graph( + [blocked, existing], + "late_cast_node_name_collision", + [x, other], + [output, other_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_subgraph_float_capture_model() -> ModelProto: + """Build blocked branches that capture a converted outer FLOAT value.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + relu_output = helper.make_tensor_value_info("relu_output", TensorProto.FLOAT, [1]) + if_output = helper.make_tensor_value_info("if_output", TensorProto.FLOAT, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", ["x"], [branch_output.name], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + relu = helper.make_node("Relu", ["x"], ["relu_output"], name="relu") + conditional = helper.make_node( + "If", + ["condition"], + ["if_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [relu, conditional], + "blocked_subgraph_float_capture", + [x, condition], + [relu_output, if_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_untyped_blocked_subgraph_capture_model() -> ModelProto: + """Build a blocked capture whose FLOAT producer has no declared metadata.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + relu_output = helper.make_tensor_value_info("relu_output", TensorProto.FLOAT, [1]) + if_output = helper.make_tensor_value_info("if_output", TensorProto.FLOAT, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", ["x"], [branch_output.name], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + constant = helper.make_node( + "Constant", + [], + ["x"], + name="constant", + value=helper.make_tensor("value", TensorProto.FLOAT, [1], [2.0]), + ) + relu = helper.make_node("Relu", ["x"], ["relu_output"], name="relu") + conditional = helper.make_node( + "If", + ["condition"], + ["if_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [constant, relu, conditional], + "untyped_blocked_subgraph_capture", + [condition], + [relu_output, if_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_uninferred_custom_op_blocked_capture_model() -> ModelProto: + """Build a custom type-preserving output omitted by successful inference.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + if_output = helper.make_tensor_value_info("if_output", TensorProto.FLOAT, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", ["y"], [branch_output.name], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + gelu = helper.make_node("Gelu", ["x"], ["y"], name="gelu", domain="com.microsoft") + conditional = helper.make_node( + "If", + ["condition"], + ["if_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [gelu, conditional], + "uninferred_custom_op_blocked_capture", + [x, condition], + [if_output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_missing_metadata_blocked_edge_model() -> ModelProto: + """Build an uninferred custom edge that can cross a blocked-node boundary.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + bias = numpy_helper.from_array(np.array([1.0], dtype=np.float32), "bias") + gelu = helper.make_node("Gelu", ["x"], ["hidden"], name="gelu", domain="com.microsoft") + add = helper.make_node("Add", ["hidden", "bias"], ["output"], name="add") + graph = helper.make_graph( + [gelu, add], + "missing_metadata_blocked_edge", + [x], + [output], + [bias], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_existing_fp16_boundary_model() -> ModelProto: + """Build an uninferred blocked output with an explicit FP16 boundary.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1]) + gelu = helper.make_node("Gelu", ["x"], ["hidden"], name="gelu", domain="com.microsoft") + boundary = helper.make_node( + "Cast", ["hidden"], ["output"], name="boundary", to=TensorProto.FLOAT16 + ) + graph = helper.make_graph( + [gelu, boundary], + "existing_fp16_boundary", + [x], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_missing_metadata_equal_consumer_model() -> ModelProto: + """Build a blocked uninferred output coupled to a converted sibling input.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + other = helper.make_tensor_value_info("other", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, [1]) + gelu = helper.make_node("Gelu", ["x"], ["hidden"], name="gelu", domain="com.microsoft") + equal = helper.make_node("Equal", ["hidden", "other"], ["output"], name="equal") + graph = helper.make_graph( + [gelu, equal], + "missing_metadata_equal_consumer", + [x, other], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_missing_metadata_sequence_consumer_model() -> ModelProto: + """Build a blocked uninferred tensor feeding a FLOAT sequence output.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, [1]) + gelu = helper.make_node("Gelu", ["x"], ["hidden"], name="gelu", domain="com.microsoft") + sequence = helper.make_node("SequenceConstruct", ["hidden"], ["output"], name="sequence") + graph = helper.make_graph( + [gelu, sequence], + "missing_metadata_sequence_consumer", + [x], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_missing_metadata_sequence_map_consumer_model() -> ModelProto: + """Build a blocked tensor feeding a SequenceMap additional input.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + sequence = helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, [1]) + output = helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, [1]) + element = helper.make_tensor_value_info("element", TensorProto.FLOAT, [1]) + additional = helper.make_tensor_value_info("additional", TensorProto.FLOAT, [1]) + mapped = helper.make_tensor_value_info("mapped", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [helper.make_node("Add", ["element", "additional"], ["mapped"])], + "body", + [element, additional], + [mapped], + ) + gelu = helper.make_node("Gelu", ["x"], ["hidden"], name="gelu", domain="com.microsoft") + sequence_map = helper.make_node( + "SequenceMap", + ["sequence", "hidden"], + ["output"], + name="sequence_map", + body=body, + ) + graph = helper.make_graph( + [gelu, sequence_map], + "missing_metadata_sequence_map_consumer", + [x, sequence], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_int_sequence_map_consumer_model() -> ModelProto: + """Build a blocked uninferred INT output used by SequenceMap.""" + text = helper.make_tensor_value_info("text", TensorProto.STRING, [1]) + sequence = helper.make_tensor_sequence_value_info("sequence", TensorProto.INT32, [1]) + output = helper.make_tensor_sequence_value_info("output", TensorProto.INT32, [1]) + element = helper.make_tensor_value_info("element", TensorProto.INT32, [1]) + additional = helper.make_tensor_value_info("additional", TensorProto.INT32, [1]) + mapped = helper.make_tensor_value_info("mapped", TensorProto.INT32, [1]) + body = helper.make_graph( + [helper.make_node("Add", ["element", "additional"], ["mapped"])], + "body", + [element, additional], + [mapped], + ) + murmur = helper.make_node( + "MurmurHash3", + ["text"], + ["hidden"], + name="murmur", + domain="com.microsoft", + positive=0, + ) + sequence_map = helper.make_node( + "SequenceMap", + ["sequence", "hidden"], + ["output"], + name="sequence_map", + body=body, + ) + graph = helper.make_graph( + [murmur, sequence_map], + "int_sequence_map_consumer", + [text, sequence], + [output], + value_info=[helper.make_tensor_value_info("hidden", TensorProto.INT32, [1])], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_int_identity_consumer_model() -> ModelProto: + """Build a blocked INT output consumed by a same-type schema edge.""" + boxes = helper.make_tensor_value_info("boxes", TensorProto.FLOAT, [1, 2, 4]) + scores = helper.make_tensor_value_info("scores", TensorProto.FLOAT, [1, 1, 2]) + selected = helper.make_tensor_value_info("selected", TensorProto.INT64, [None, 3]) + output = helper.make_tensor_value_info("output", TensorProto.INT64, [None, 3]) + max_output = numpy_helper.from_array(np.array(2, dtype=np.int64), "max_output") + nms = helper.make_node( + "NonMaxSuppression", + ["boxes", "scores", "max_output"], + ["selected"], + name="nms", + ) + identity = helper.make_node("Identity", ["selected"], ["output"], name="identity") + graph = helper.make_graph( + [nms, identity], + "int_identity_consumer", + [boxes, scores], + [output], + [max_output], + value_info=[selected], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_uninferred_int_identity_consumer_model() -> ModelProto: + """Build an uninferred INT output with a concrete same-type consumer output.""" + text = helper.make_tensor_value_info("text", TensorProto.STRING, [1]) + output = helper.make_tensor_value_info("output", TensorProto.INT32, [1]) + murmur = helper.make_node( + "MurmurHash3", + ["text"], + ["hidden"], + name="murmur", + domain="com.microsoft", + positive=0, + ) + identity = helper.make_node("Identity", ["hidden"], ["output"], name="identity") + graph = helper.make_graph( + [murmur, identity], + "uninferred_int_identity_consumer", + [text], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_function_int_consumer_model() -> ModelProto: + """Build a concrete INT edge consumed by a local function without a schema.""" + text = helper.make_tensor_value_info("text", TensorProto.STRING, [1]) + hidden = helper.make_tensor_value_info("hidden", TensorProto.INT32, [1]) + output = helper.make_tensor_value_info("output", TensorProto.INT32, [1]) + murmur = helper.make_node( + "MurmurHash3", + ["text"], + ["hidden"], + name="murmur", + domain="com.microsoft", + positive=0, + ) + function_node = helper.make_node( + "IntIdentity", + ["hidden"], + ["output"], + name="function", + domain="local.test", + ) + graph = helper.make_graph( + [murmur, function_node], + "function_int_consumer", + [text], + [output], + value_info=[hidden], + ) + function = helper.make_function( + "local.test", + "IntIdentity", + ["input"], + ["output"], + [helper.make_node("Identity", ["input"], ["output"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + helper.make_opsetid("local.test", 1), + ], + ) + model.functions.append(function) + return model + + +def _build_function_with_unrelated_graph_attribute_model() -> ModelProto: + """Build a local function carrying an unrelated graph-valued attribute.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + attribute_input = helper.make_tensor_value_info("attribute_input", TensorProto.INT32, [1]) + attribute_output = helper.make_tensor_value_info("attribute_input", TensorProto.INT32, [1]) + attribute_graph = helper.make_graph( + [], + "unused_attribute", + [attribute_input], + [attribute_output], + ) + gelu = helper.make_node( + "Gelu", + ["x"], + ["hidden"], + name="gelu", + domain="com.microsoft", + ) + function_node = helper.make_node( + "FloatIdentity", + ["hidden"], + ["output"], + name="function", + domain="local.test", + unused_graph=attribute_graph, + ) + graph = helper.make_graph( + [gelu, function_node], + "function_with_unrelated_graph_attribute", + [x], + [output], + ) + function = helper.make_function( + "local.test", + "FloatIdentity", + ["input"], + ["output"], + [helper.make_node("Identity", ["input"], ["output"])], + [helper.make_opsetid("", 17)], + attributes=["unused_graph"], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + helper.make_opsetid("local.test", 1), + ], + ) + model.functions.append(function) + return model + + +def _build_function_with_float_constant_model() -> ModelProto: + """Build a local function whose concrete FLOAT body is not converted by ORT.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + one = numpy_helper.from_array(np.array([1.0], dtype=np.float32)) + function = helper.make_function( + "local.test", + "AddOne", + ["input"], + ["function_output"], + [ + helper.make_node("Constant", [], ["one"], name="one", value=one), + helper.make_node( + "Add", + ["input", "one"], + ["function_output"], + name="add", + ), + ], + [helper.make_opsetid("", 17)], + value_info=[helper.make_tensor_value_info("one", TensorProto.FLOAT, [1])], + ) + invocation = helper.make_node( + "AddOne", + ["x"], + ["output"], + name="function", + domain="local.test", + ) + graph = helper.make_graph( + [invocation], + "function_with_float_constant", + [x], + [output], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("local.test", 1), + ], + ) + model.functions.append(function) + return model + + +def _build_function_with_contrib_float_constant_model() -> ModelProto: + """Build a concrete FLOAT function body ONNX inference cannot validate.""" + model = _build_function_with_float_constant_model() + function = model.functions[0] + function.node[1].CopyFrom( + helper.make_node( + "Gelu", + ["one"], + ["function_output"], + name="gelu", + domain="com.microsoft", + ) + ) + function.opset_import.append(helper.make_opsetid("com.microsoft", 1)) + model.opset_import.append(helper.make_opsetid("com.microsoft", 1)) + model.graph.name = "function_with_contrib_float_constant" + return model + + +def _build_function_with_scalar_contrib_float_model() -> ModelProto: + """Build scalar FLOAT storage before an uninferred function-body op.""" + model = _build_function_with_contrib_float_constant_model() + function = model.functions[0] + function.node[0].CopyFrom( + helper.make_node( + "Constant", + [], + ["one"], + name="one", + value_float=1.0, + ) + ) + del function.value_info[:] + model.graph.name = "function_with_scalar_contrib_float" + return model + + +def _build_always_float_function_graph_attribute_model() -> ModelProto: + """Build an always-float op type with a graph attribute ORT skips.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + local_input = helper.make_tensor_value_info("local_input", TensorProto.FLOAT, [1]) + local_output = helper.make_tensor_value_info("local_output", TensorProto.FLOAT, [1]) + attribute_graph = helper.make_graph( + [ + helper.make_node( + "Identity", + ["local_input"], + ["local_output"], + name="local_identity", + ) + ], + "unused_attribute", + [local_input], + [local_output], + ) + function_node = helper.make_node( + "GroupNorm", + ["x"], + ["output"], + name="function", + domain="local.test", + unused_graph=attribute_graph, + ) + graph = helper.make_graph( + [function_node], + "always_float_function_graph_attribute", + [x], + [output], + ) + function = helper.make_function( + "local.test", + "GroupNorm", + ["input"], + ["output"], + [helper.make_node("Identity", ["input"], ["output"])], + [helper.make_opsetid("", 17)], + attributes=["unused_graph"], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("local.test", 1), + ], + ) + model.functions.append(function) + return model + + +def _build_always_float_function_capture_model() -> ModelProto: + """Build executed skipped attributes that capture a converted initializer.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + converted = helper.make_tensor_value_info("converted", TensorProto.FLOAT, [1]) + selected = helper.make_tensor_value_info("selected", TensorProto.FLOAT, [1]) + shared = numpy_helper.from_array(np.array([2.0], dtype=np.float32), "shared") + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", + ["shared"], + [branch_output.name], + name=f"{name}_identity", + ) + return helper.make_graph( + [identity], + name, + [], + [branch_output], + value_info=[helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1])], + ) + + function_if = helper.make_node("If", ["condition"], ["output"], name="if") + then_branch = helper.make_attribute_ref("then_branch", AttributeProto.GRAPH) + then_branch.ref_attr_name = "then_branch" + else_branch = helper.make_attribute_ref("else_branch", AttributeProto.GRAPH) + else_branch.ref_attr_name = "else_branch" + function_if.attribute.extend([then_branch, else_branch]) + function = helper.make_function( + "local.test", + "GroupNorm", + ["condition"], + ["output"], + [function_if], + [helper.make_opsetid("", 17)], + attributes=["then_branch", "else_branch"], + ) + function_node = helper.make_node( + "GroupNorm", + ["condition"], + ["selected"], + name="function", + domain="local.test", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + identity = helper.make_node("Identity", ["shared"], ["converted"], name="identity") + graph = helper.make_graph( + [identity, function_node], + "always_float_function_capture", + [condition], + [converted, selected], + [shared], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("local.test", 1), + ], + ) + model.functions.append(function) + return model + + +def _build_always_float_nested_function_capture_model() -> ModelProto: + """Nest referenced invocation attributes inside a function body graph.""" + model = _build_always_float_function_capture_model() + function = model.functions[0] + invocation = next(node for node in model.graph.node if node.name == "function") + for attribute in invocation.attribute: + if attribute.type != AttributeProto.GRAPH: + continue + attribute.g.node[0].CopyFrom( + helper.make_node( + "IsNaN", + ["shared"], + [attribute.g.output[0].name], + name=f"{attribute.g.name}_is_nan", + ) + ) + attribute.g.output[0].type.tensor_type.elem_type = TensorProto.BOOL + model.graph.output[1].type.tensor_type.elem_type = TensorProto.BOOL + + def _outer_branch(name: str) -> GraphProto: + branch_output = helper.make_empty_tensor_value_info(f"{name}_output") + inner_if = helper.make_node( + "If", + ["condition"], + [branch_output.name], + name=f"{name}_inner_if", + ) + then_branch = helper.make_attribute_ref("then_branch", AttributeProto.GRAPH) + then_branch.ref_attr_name = "then_branch" + else_branch = helper.make_attribute_ref("else_branch", AttributeProto.GRAPH) + else_branch.ref_attr_name = "else_branch" + inner_if.attribute.extend([then_branch, else_branch]) + return helper.make_graph([inner_if], name, [], [branch_output]) + + function.node[0].CopyFrom( + helper.make_node( + "If", + ["condition"], + ["output"], + name="outer_if", + then_branch=_outer_branch("outer_then"), + else_branch=_outer_branch("outer_else"), + ) + ) + model.graph.name = "always_float_nested_function_capture" + return model + + +def _build_always_float_default_function_capture_model() -> ModelProto: + """Supply executed function graph attributes through default values.""" + model = _build_always_float_function_capture_model() + function = model.functions[0] + function_node = next(node for node in model.graph.node if node.name == "function") + function.attribute_proto.extend(function_node.attribute) + del function_node.attribute[:] + model.graph.name = "always_float_default_function_capture" + return model + + +def _build_always_float_unused_function_capture_model() -> ModelProto: + """Build an unused function attribute that captures converted FLOAT.""" + model = _build_always_float_function_graph_attribute_model() + attribute_graph = next( + attribute.g + for attribute in model.graph.node[0].attribute + if attribute.type == AttributeProto.GRAPH + ) + del attribute_graph.input[:] + attribute_graph.node[0].input[0] = "shared" + attribute_graph.value_info.append( + helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + ) + shared = numpy_helper.from_array(np.array([3.0], dtype=np.float32), "shared") + converted = helper.make_tensor_value_info("converted", TensorProto.FLOAT, [1]) + model.graph.initializer.append(shared) + model.graph.node.append(helper.make_node("Relu", ["shared"], ["converted"], name="relu")) + model.graph.output.append(converted) + model.graph.name = "always_float_unused_function_capture" + return model + + +def _build_skipped_duplicate_initializer_model() -> ModelProto: + """Build duplicate FLOAT names where ORT skips the nested scope.""" + model = _build_always_float_function_graph_attribute_model() + model.graph.initializer.append( + numpy_helper.from_array(np.array([1.0], dtype=np.float32), "duplicate") + ) + attribute_graph = next( + attribute.g + for attribute in model.graph.node[0].attribute + if attribute.type == AttributeProto.GRAPH + ) + attribute_graph.initializer.append( + numpy_helper.from_array(np.array([2.0], dtype=np.float32), "duplicate") + ) + model.graph.name = "skipped_duplicate_initializer" + return model + + +def _build_scan8_int_state_model() -> ModelProto: + """Build Scan-8 with a fixed input before heterogeneous variadic inputs.""" + text = helper.make_tensor_value_info("text", TensorProto.STRING, [1]) + scan_input = helper.make_tensor_value_info("scan_input", TensorProto.FLOAT, [1, 2]) + final_state = helper.make_tensor_value_info("final_state", TensorProto.INT32, [1]) + scan_output = helper.make_tensor_value_info("scan_output", TensorProto.FLOAT, [1, 2]) + body_state = helper.make_tensor_value_info("body_state", TensorProto.INT32, []) + body_input = helper.make_tensor_value_info("body_input", TensorProto.FLOAT, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.INT32, []) + body_output = helper.make_tensor_value_info("body_output", TensorProto.FLOAT, []) + body = helper.make_graph( + [ + helper.make_node("Identity", ["body_state"], ["body_state_out"]), + helper.make_node("Identity", ["body_input"], ["body_output"]), + ], + "body", + [body_state, body_input], + [body_state_out, body_output], + ) + murmur = helper.make_node( + "MurmurHash3", + ["text"], + ["hidden"], + name="murmur", + domain="com.microsoft", + positive=0, + ) + scan = helper.make_node( + "Scan", + ["", "hidden", "scan_input"], + ["final_state", "scan_output"], + name="scan", + body=body, + num_scan_inputs=1, + ) + graph = helper.make_graph( + [murmur, scan], + "scan8_int_state", + [text, scan_input], + [final_state, scan_output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 8), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_scan8_float_state_model() -> ModelProto: + """Build Scan-8 with an uninferred FLOAT state requiring a boundary.""" + model = _build_scan8_int_state_model() + model.graph.node[0].CopyFrom( + helper.make_node( + "Gelu", + ["x"], + ["hidden"], + name="gelu", + domain="com.microsoft", + ) + ) + model.graph.input[0].CopyFrom(helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])) + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("final_state", TensorProto.FLOAT, [1]) + ) + body = next( + attribute.g for attribute in model.graph.node[1].attribute if attribute.name == "body" + ) + body.input[0].type.tensor_type.elem_type = TensorProto.FLOAT + body.output[0].type.tensor_type.elem_type = TensorProto.FLOAT + model.graph.name = "scan8_float_state" + return model + + +def _build_scan8_without_state_model() -> ModelProto: + """Build Scan-8 whose only variadic input is a scan input.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2]) + output = helper.make_tensor_value_info("output", TensorProto.INT32, [1, 2]) + body_input = helper.make_empty_tensor_value_info("body_input") + body_output = helper.make_tensor_value_info("body_output", TensorProto.INT32, []) + bias = numpy_helper.from_array(np.array(1.0, dtype=np.float32), "bias") + body = helper.make_graph( + [ + helper.make_node("Add", ["body_input", "bias"], ["added"]), + helper.make_node( + "Cast", + ["added"], + ["body_output"], + to=TensorProto.INT32, + ), + ], + "body", + [body_input], + [body_output], + [bias], + value_info=[helper.make_tensor_value_info("bias", TensorProto.FLOAT, [])], + ) + gelu = helper.make_node( + "Gelu", + ["x"], + ["hidden"], + name="gelu", + domain="com.microsoft", + ) + scan = helper.make_node( + "Scan", + ["", "hidden"], + ["output"], + name="scan", + body=body, + num_scan_inputs=1, + ) + graph = helper.make_graph( + [gelu, scan], + "scan8_without_state", + [x], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 8), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_loop_float_state_model() -> ModelProto: + """Build a Loop with an untyped formal and declared FLOAT feedback.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_empty_tensor_value_info("body_state") + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + ), + helper.make_node("Identity", ["body_state"], ["body_state_out"]), + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_state_out], + ) + gelu = helper.make_node( + "Gelu", + ["x"], + ["hidden"], + name="gelu", + domain="com.microsoft", + ) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "hidden"], + ["output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [gelu, loop], + "loop_float_state", + [trip_count, condition, x], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 11), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_optional_float_blocked_capture_model() -> ModelProto: + """Build a blocked branch that captures an Optional wrapping converted FLOAT.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + if_output = helper.make_tensor_value_info("if_output", TensorProto.FLOAT, [1]) + optional_info = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + opt = helper.make_value_info("opt", optional_info) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + get_element = helper.make_node( + "OptionalGetElement", + ["opt"], + [branch_output.name], + name=f"{name}_get_element", + ) + return helper.make_graph([get_element], name, [], [branch_output]) + + optional = helper.make_node("Optional", ["x"], ["opt"], name="optional") + conditional = helper.make_node( + "If", + ["condition"], + ["if_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [optional, conditional], + "optional_float_blocked_capture", + [x, condition], + [if_output], + value_info=[opt], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_float_optional_output_model() -> ModelProto: + """Build a blocked FLOAT output wrapped by an unchanged Optional.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + output = helper.make_value_info("output", optional_type) + gelu = helper.make_node( + "Gelu", + ["x"], + ["hidden"], + name="gelu", + domain="com.microsoft", + ) + optional = helper.make_node("Optional", ["hidden"], ["output"], name="optional") + graph = helper.make_graph( + [gelu, optional], + "blocked_float_optional_output", + [x], + [output], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_float_optional_output_model() -> ModelProto: + """Build converted FLOAT wrapped by unchanged Optional metadata.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + output = helper.make_value_info("output", optional_type) + optional = helper.make_node("Optional", ["x"], ["output"], name="optional") + graph = helper.make_graph([optional], "float_optional_output", [x], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_blocked_tensor_optional_consumer_model() -> ModelProto: + """Build an inferred blocked tensor output wrapped by an Optional.""" + model = _build_blocked_float_optional_output_model() + model.graph.node[0].CopyFrom(helper.make_node("Abs", ["x"], ["hidden"], name="abs")) + model.graph.name = "blocked_tensor_optional_consumer" + return model + + +def _build_blocked_float_optional_capture_model() -> ModelProto: + """Build a blocked child capturing an Optional with an FP32 payload.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + optional_info = helper.make_value_info("optional", optional_type) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + get_element = helper.make_node( + "OptionalGetElement", + ["optional"], + [branch_output.name], + name=f"{name}_get_element", + ) + return helper.make_graph([get_element], name, [], [branch_output]) + + gelu = helper.make_node( + "Gelu", + ["x"], + ["hidden"], + name="gelu", + domain="com.microsoft", + ) + optional = helper.make_node("Optional", ["hidden"], ["optional"], name="optional") + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [gelu, optional, conditional], + "blocked_float_optional_capture", + [x, condition], + [output], + value_info=[optional_info], + ) + return helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + + +def _build_empty_optional_blocked_capture_model() -> ModelProto: + """Build a blocked branch capturing a source-less empty Optional.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + converted = helper.make_tensor_value_info("converted", TensorProto.FLOAT, [1]) + tensor_type = helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + optional_type = helper.make_optional_type_proto(tensor_type) + optional_info = helper.make_value_info("optional", optional_type) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + ["optional"], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + make_optional = helper.make_node( + "Optional", [], ["optional"], name="optional", type=tensor_type + ) + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + relu = helper.make_node("Relu", ["x"], ["converted"], name="relu") + graph = helper.make_graph( + [make_optional, conditional, relu], + "empty_optional_blocked_capture", + [condition, x], + [output, converted], + value_info=[optional_info], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_nested_optional_input_blocked_capture_model() -> ModelProto: + """Build a blocked capture of a nested Optional formal input.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_optional = helper.make_value_info("body_optional", optional_type) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_optional_out = helper.make_value_info("body_optional", optional_type) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + ["body_optional"], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + body_if = helper.make_node( + "If", + ["body_condition"], + ["body_if_output"], + name="body_if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="body_condition", + ), + body_if, + ], + "body", + [iteration, body_condition, body_optional], + [body_condition_out, body_optional_out], + value_info=[helper.make_tensor_value_info("body_if_output", TensorProto.BOOL, [])], + ) + make_optional = helper.make_node("Optional", ["x"], ["optional"], name="make_optional") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "optional"], + ["final_optional"], + name="loop", + body=body, + ) + get_output = helper.make_node( + "OptionalHasElement", ["final_optional"], ["output"], name="get_output" + ) + graph = helper.make_graph( + [make_optional, loop, get_output], + "nested_optional_input_blocked_capture", + [trip_count, condition, x], + [output], + value_info=[ + helper.make_value_info("optional", optional_type), + helper.make_value_info("final_optional", optional_type), + ], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_top_optional_input_blocked_capture_model() -> ModelProto: + """Build a nested blocked capture sourced from an unchanged top Optional input.""" + model = _build_nested_optional_input_blocked_capture_model() + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + del model.graph.node[0] + del model.graph.input[2] + model.graph.input.append(helper.make_value_info("optional", optional_type)) + retained = [value for value in model.graph.value_info if value.name != "optional"] + del model.graph.value_info[:] + model.graph.value_info.extend(retained) + model.graph.name = "top_optional_input_blocked_capture" + return model + + +def _build_identity_optional_feedback_blocked_capture_model() -> ModelProto: + """Return an unchanged Optional state through a differently named Identity.""" + model = _build_top_optional_input_blocked_capture_model() + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + loop = model.graph.node[0] + body = next(attribute.g for attribute in loop.attribute if attribute.g.name) + body.output[1].CopyFrom(helper.make_value_info("body_optional_out", optional_type)) + body.node.append( + helper.make_node( + "Identity", + ["body_optional"], + ["body_optional_out"], + name="optional_feedback", + ) + ) + model.graph.name = "identity_optional_feedback_blocked_capture" + return model + + +def _build_rewrapped_optional_feedback_blocked_capture_model() -> ModelProto: + """Rewrap a converted FLOAT element as Optional loop feedback.""" + model = _build_top_optional_input_blocked_capture_model() + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + loop = model.graph.node[0] + body = next(attribute.g for attribute in loop.attribute if attribute.g.name) + body.output[1].CopyFrom(helper.make_value_info("body_optional_out", optional_type)) + body.node.extend( + [ + helper.make_node( + "OptionalGetElement", + ["body_optional"], + ["feedback_element"], + name="unwrap_feedback", + ), + helper.make_node( + "Optional", + ["feedback_element"], + ["body_optional_out"], + name="rewrap_feedback", + ), + ] + ) + body.value_info.append( + helper.make_tensor_value_info("feedback_element", TensorProto.FLOAT, [1]) + ) + model.graph.name = "rewrapped_optional_feedback_blocked_capture" + return model + + +def _build_loop_optional_output_blocked_capture_model() -> ModelProto: + """Pass an Optional through Loop before a blocked capture.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + converted = helper.make_tensor_value_info("converted", TensorProto.FLOAT, [1]) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + optional = helper.make_value_info("optional", optional_type) + final_optional = helper.make_value_info("final_optional", optional_type) + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_optional = helper.make_value_info("body_optional", optional_type) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_optional_out = helper.make_value_info("body_optional_out", optional_type) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="condition_feedback", + ), + helper.make_node( + "Identity", + ["body_optional"], + ["body_optional_out"], + name="optional_feedback", + ), + ], + "body", + [iteration, body_condition, body_optional], + [body_condition_out, body_optional_out], + ) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + ["final_optional"], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "optional"], + ["final_optional"], + name="loop", + body=body, + ) + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + relu = helper.make_node("Relu", ["x"], ["converted"], name="relu") + graph = helper.make_graph( + [loop, conditional, relu], + "loop_optional_output_blocked_capture", + [trip_count, condition, optional, x], + [output, converted], + value_info=[final_optional], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_optional_identity_blocked_capture_model() -> ModelProto: + """Build an unchanged Optional pass-through captured by a blocked child.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + optional = helper.make_value_info("optional", optional_type) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + converted = helper.make_tensor_value_info("converted", TensorProto.FLOAT, [1]) + passed_optional = helper.make_value_info("passed_optional", optional_type) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + ["passed_optional"], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + identity = helper.make_node( + "Identity", + ["optional"], + ["passed_optional"], + name="identity", + ) + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + relu = helper.make_node("Relu", ["x"], ["converted"], name="relu") + graph = helper.make_graph( + [identity, conditional, relu], + "optional_identity_blocked_capture", + [condition, optional, x], + [output, converted], + value_info=[passed_optional], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_long_optional_identity_chain_model( + length: int = 1010, +) -> ModelProto: + """Build a long precision-preserving Optional producer chain.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + optional = helper.make_value_info("optional", optional_type) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + nodes = [] + value_info = [] + source = "optional" + for index in range(length): + target = f"optional_{index}" + nodes.append( + helper.make_node( + "Identity", + [source], + [target], + name=f"identity_{index}", + ) + ) + value_info.append(helper.make_value_info(target, optional_type)) + source = target + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + [source], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + nodes.append( + helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + ) + graph = helper.make_graph( + nodes, + "long_optional_identity_chain", + [condition, optional], + [output], + value_info=value_info, + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_free_optional_identity_blocked_capture_model() -> ModelProto: + """Build a nested Identity sourced from a converted outer Optional.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + state = helper.make_tensor_value_info("state", TensorProto.BOOL, []) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + optional = helper.make_value_info("optional", optional_type) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_tensor_value_info("body_state", TensorProto.BOOL, []) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.BOOL, []) + passed_optional = helper.make_value_info("passed_optional", optional_type) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + ["passed_optional"], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + body_if = helper.make_node( + "If", + ["body_state"], + ["body_state_out"], + name="body_if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["optional"], + ["passed_optional"], + name="optional_identity", + ), + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="condition_identity", + ), + body_if, + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_state_out], + value_info=[passed_optional], + ) + make_optional = helper.make_node("Optional", ["x"], ["optional"], name="make_optional") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "state"], + ["output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [make_optional, loop], + "free_optional_identity_blocked_capture", + [trip_count, condition, state, x], + [output], + value_info=[optional], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_optional_feedback_blocked_capture_model() -> ModelProto: + """Build a nested Optional input changed by a later loop-carried value.""" + model = _build_top_optional_input_blocked_capture_model() + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + model.graph.input.append(helper.make_tensor_value_info("feedback", TensorProto.FLOAT, [1])) + loop = model.graph.node[0] + body = next(attribute.g for attribute in loop.attribute if attribute.g.name) + body.output[1].CopyFrom(helper.make_value_info("body_optional_out", optional_type)) + body.node.append( + helper.make_node( + "Optional", + ["feedback"], + ["body_optional_out"], + name="make_feedback_optional", + ) + ) + model.graph.name = "optional_feedback_blocked_capture" + return model + + +def _build_mispositioned_optional_feedback_model() -> ModelProto: + """Build two loop states with unchanged binding returned in the wrong slot.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + optional_type = helper.make_optional_type_proto( + helper.make_tensor_type_proto(TensorProto.FLOAT, [1]) + ) + optional_a = helper.make_value_info("optional_a", optional_type) + optional_b = helper.make_value_info("optional_b", optional_type) + feedback = helper.make_tensor_value_info("feedback", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.BOOL, []) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_a = helper.make_value_info("body_a", optional_type) + body_b = helper.make_value_info("body_b", optional_type) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + converted_a = helper.make_value_info("converted_a", optional_type) + returned_b = helper.make_value_info("body_a", optional_type) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.BOOL, []) + has_element = helper.make_node( + "OptionalHasElement", + ["body_a"], + [branch_output.name], + name=f"{name}_has_element", + ) + return helper.make_graph([has_element], name, [], [branch_output]) + + body_if = helper.make_node( + "If", + ["body_condition"], + ["body_if_output"], + name="body_if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + body = helper.make_graph( + [ + helper.make_node( + "Identity", + ["body_condition"], + ["body_condition_out"], + name="body_condition", + ), + body_if, + helper.make_node( + "Optional", + ["feedback"], + ["converted_a"], + name="make_converted_a", + ), + ], + "body", + [iteration, body_condition, body_a, body_b], + [body_condition_out, converted_a, returned_b], + value_info=[helper.make_tensor_value_info("body_if_output", TensorProto.BOOL, [])], + ) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "optional_a", "optional_b"], + ["final_a", "final_b"], + name="loop", + body=body, + ) + has_output = helper.make_node("OptionalHasElement", ["final_a"], ["output"], name="has_output") + graph = helper.make_graph( + [loop, has_output], + "mispositioned_optional_feedback", + [trip_count, condition, optional_a, optional_b, feedback], + [output], + value_info=[ + helper.make_value_info("final_a", optional_type), + helper.make_value_info("final_b", optional_type), + ], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + + +def _build_scan_mispositioned_optional_feedback_model() -> ModelProto: + """Add a scan output to the wrong-slot optional feedback model.""" + model = _build_mispositioned_optional_feedback_model() + loop = model.graph.node[0] + loop.output.append("scan_output") + body = next(attribute.g for attribute in loop.attribute if attribute.g.name) + body.output.append(helper.make_tensor_value_info("body_if_output", TensorProto.BOOL, [])) + model.graph.name = "scan_mispositioned_optional_feedback" + return model + + +def _build_scan_top_optional_input_blocked_capture_model() -> ModelProto: + """Add a scan output without changing the pass-through loop state.""" + model = _build_top_optional_input_blocked_capture_model() + loop = model.graph.node[0] + loop.output.append("scan_output") + body = next(attribute.g for attribute in loop.attribute if attribute.g.name) + body.output.append(helper.make_tensor_value_info("body_if_output", TensorProto.BOOL, [])) + model.graph.name = "scan_top_optional_input_blocked_capture" + return model + + +def _build_blocked_float_sequence_input_model() -> ModelProto: + """Build a blocked node consuming a sequence of FLOAT tensors.""" + sequence = helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, [1]) + length = helper.make_tensor_value_info("length", TensorProto.INT64, []) + node = helper.make_node("SequenceLength", ["sequence"], ["length"], name="sequence_length") + graph = helper.make_graph([node], "blocked_float_sequence_input", [sequence], [length]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_float_sequence_io_model() -> ModelProto: + """Build a FLOAT tensor sequence exposed directly as graph I/O.""" + sequence_input = helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, [1]) + sequence_output = helper.make_tensor_sequence_value_info("output", TensorProto.FLOAT, [1]) + identity = helper.make_node("Identity", ["sequence"], ["output"], name="identity") + graph = helper.make_graph( + [identity], + "float_sequence_io", + [sequence_input], + [sequence_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_overridable_blocked_initializer_model() -> ModelProto: + """Build a graph-input initializer consumed only by a blocked node.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([2.0], dtype=np.float32), "x") + blocked = helper.make_node("Identity", ["x"], ["output"], name="blocked") + graph = helper.make_graph( + [blocked], "overridable_blocked_initializer", [x], [output], [initializer] + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_duplicate_float_output_name_model() -> ModelProto: + """Build a legal graph that exposes one FLOAT value through two outputs.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + first = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1]) + second = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1]) + identity = helper.make_node("Identity", ["x"], ["y"], name="identity") + graph = helper.make_graph([identity], "duplicate_float_output", [x], [first, second]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_shared_float_input_output_model() -> ModelProto: + """Build a direct FLOAT pass-through with one public input/output name.""" + graph_input = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + graph_output = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + graph = helper.make_graph([], "shared_float_io", [graph_input], [graph_output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_shadowed_initializer_output_model() -> ModelProto: + """Build legal nested outputs that shadow an outer initializer name.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + outer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "same") + + def _branch(name: str, value: float) -> GraphProto: + branch_output = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), "same") + return helper.make_graph([], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", 1.0), + else_branch=_branch("else", 2.0), + ) + graph = helper.make_graph([node], "shadow", [condition], [output], [outer]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_if_duplicate_local_initializer_model() -> ModelProto: + """Build blocked If branches with duplicate local FLOAT initializer names.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + + def _branch(name: str, value: float) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), "same") + identity = helper.make_node( + "Identity", ["same"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", 1.0), + else_branch=_branch("else", 2.0), + ) + graph = helper.make_graph([node], "blocked_if_duplicate_local", [condition], [output]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_if_shadowed_output_initializer_model() -> ModelProto: + """Build a top-level output initializer shadowed by blocked branch-local values.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "shared") + + def _branch(name: str, value: float) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + local = numpy_helper.from_array(np.array([value], dtype=np.float32), "shared") + identity = helper.make_node( + "Identity", ["shared"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output], [local]) + + node = helper.make_node( + "If", + ["condition"], + ["if_output"], + name="if", + then_branch=_branch("then", 1.0), + else_branch=_branch("else", 2.0), + ) + graph = helper.make_graph([node], "blocked_if_shadowed", [condition], [shared], [initializer]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_traversed_shadowed_output_initializer_model() -> ModelProto: + """Build traversed Loop body names that shadow a top-level output initializer.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_carried = helper.make_tensor_value_info("loop_carried", TensorProto.FLOAT, [1]) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "shared") + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_value_out = helper.make_tensor_value_info("body_value_out", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ), + helper.make_node("Identity", ["shared"], ["body_value_out"], name="body_value"), + ], + "body", + [iteration, body_condition, local_shared], + [body_condition_out, body_value_out], + ) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_carried"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [loop], + "traversed_shadowed", + [trip_count, condition, loop_carried], + [shared, loop_output], + [initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_loop_state_input_shadowing_output_initializer_model() -> ModelProto: + """Build a Loop whose local state input shadows a top-level output initializer.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + same = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "same") + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + body_initializer = numpy_helper.from_array(np.array([1.0], dtype=np.float32), "body_state_out") + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ) + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_state_out], + [body_initializer], + ) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [loop], + "loop_state_input_shadowing", + [trip_count, condition, loop_state], + [same, loop_output], + [initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_loop_value_info_shadowing_output_initializer_model() -> ModelProto: + """Build a Loop whose local value_info shadows a top-level output initializer.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + same = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "same") + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_tensor_value_info("body_state", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ), + helper.make_node("Identity", ["body_state"], ["body_state_out"], name="body_state"), + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_state_out], + ) + body.value_info.append(helper.make_tensor_value_info("same", TensorProto.FLOAT, [1])) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [loop], + "loop_value_info_shadowing", + [trip_count, condition, loop_state], + [same, loop_output], + [initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_output_initializer_with_top_level_io_name_collision_model() -> ModelProto: + """Build a nested direct initializer output matching an unrelated top-level input.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + same = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_tensor_value_info("body_state", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_output = helper.make_tensor_value_info("same", TensorProto.FLOAT, [1]) + body_initializer = numpy_helper.from_array(np.array([7.0], dtype=np.float32), "same") + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ) + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_output], + [body_initializer], + ) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [loop], + "nested_io_name_collision", + [trip_count, condition, loop_state, same], + [loop_output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_initializer_then_mapped_free_capture_model() -> ModelProto: + """Build a nested initializer followed by a mapped capture of kept top-level I/O.""" + model = _build_nested_output_initializer_with_top_level_io_name_collision_model() + captured_output = helper.make_tensor_value_info("captured_output", TensorProto.FLOAT, [1]) + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", ["same"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + conditional = helper.make_node( + "If", + ["condition"], + ["captured_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + model.graph.node.append(conditional) + model.graph.output.append(captured_output) + return model + + +def _build_mapped_capture_shadowed_by_generated_alias_model() -> ModelProto: + """Build a mapped free capture whose target alias has a nested local binding.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + + def _branch(name: str, *, shadow_alias: bool) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node("Identity", ["x"], [f"{name}_output"], name=f"{name}_identity") + initializers = ( + [numpy_helper.from_array(np.array([9.0], dtype=np.float32), "graph_input_cast_0")] + if shadow_alias + else [] + ) + return helper.make_graph([identity], name, [], [branch_output], initializer=initializers) + + conditional = helper.make_node( + "If", + ["condition"], + ["output"], + name="if", + then_branch=_branch("then", shadow_alias=True), + else_branch=_branch("else", shadow_alias=False), + ) + graph = helper.make_graph( + [conditional], + "mapped_capture_shadowed_by_alias", + [x, condition], + [output], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_nested_blocked_ordinary_name_collision_model() -> ModelProto: + """Build nested blocked input shadowing an earlier global value-info name.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + same = helper.make_tensor_value_info("same", TensorProto.FLOAT, []) + selected = helper.make_tensor_value_info("selected", TensorProto.INT64, [None, 3]) + + def _branch(name: str, threshold_name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info( + f"{name}_selected", TensorProto.INT64, [None, 3] + ) + boxes = numpy_helper.from_array( + np.array([[[0.0, 0.0, 1.0, 1.0], [2.0, 2.0, 3.0, 3.0]]], dtype=np.float32), + f"{name}_boxes", + ) + scores = numpy_helper.from_array( + np.array([[[0.9, 0.8]]], dtype=np.float32), f"{name}_scores" + ) + max_output = numpy_helper.from_array(np.array(2, dtype=np.int64), f"{name}_max_output") + iou_threshold = numpy_helper.from_array( + np.array(0.5, dtype=np.float32), f"{name}_iou_threshold" + ) + score_threshold = numpy_helper.from_array(np.array(0.85, dtype=np.float32), threshold_name) + nms = helper.make_node( + "NonMaxSuppression", + [ + boxes.name, + scores.name, + max_output.name, + iou_threshold.name, + score_threshold.name, + ], + [branch_output.name], + name=f"{name}_nms", + ) + return helper.make_graph( + [nms], + name, + [], + [branch_output], + initializer=[ + boxes, + scores, + max_output, + iou_threshold, + score_threshold, + ], + ) + + conditional = helper.make_node( + "If", + ["condition"], + ["selected"], + name="if", + then_branch=_branch("then", "same"), + else_branch=_branch("else", "else_score_threshold"), + ) + graph = helper.make_graph( + [conditional], + "nested_blocked_ordinary_collision", + [condition, same], + [selected], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_if_free_capture_output_initializer_model() -> ModelProto: + """Build a blocked If branch that free-captures a top-level output initializer.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + if_output = helper.make_tensor_value_info("if_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "shared") + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + identity = helper.make_node( + "Identity", ["shared"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + node = helper.make_node( + "If", + ["condition"], + ["if_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + graph = helper.make_graph( + [node], + "blocked_if_free_capture", + [condition], + [shared, if_output], + [initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_if_shadowed_free_capture_model() -> ModelProto: + """Build blocked branches that capture a local shadow, not the top-level initializer.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + top_initializer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "shared") + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + body_state = helper.make_tensor_value_info("body_state", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + local_shared = numpy_helper.from_array(np.array([7.0], dtype=np.float16), "shared") + + def _branch(name: str) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT16, [1]) + identity = helper.make_node( + "Identity", ["shared"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output]) + + blocked_if = helper.make_node( + "If", + ["body_condition"], + ["if_output"], + name="if", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ), + helper.make_node("Identity", ["body_state"], ["body_state_out"], name="body_state"), + blocked_if, + ], + "body", + [iteration, body_condition, body_state], + [body_condition_out, body_state_out], + [local_shared], + ) + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [loop], + "blocked_shadowed_capture", + [trip_count, condition, loop_state], + [shared, loop_output], + [top_initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_duplicate_non_output_initializer_name_model() -> ModelProto: + """Build duplicate initializer names across scopes without initializer-backed outputs.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + top_output = helper.make_tensor_value_info("top_output", TensorProto.FLOAT, [1]) + nested_output = helper.make_tensor_value_info("nested_output", TensorProto.FLOAT, [1]) + outer = numpy_helper.from_array(np.array([9.0], dtype=np.float32), "same") + top_identity = helper.make_node("Identity", ["same"], ["top_output"], name="top_identity") + + def _branch(name: str, value: float) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), "same") + identity = helper.make_node( + "Identity", ["same"], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + node = helper.make_node( + "If", + ["condition"], + ["nested_output"], + name="if", + then_branch=_branch("then", 1.0), + else_branch=_branch("else", 2.0), + ) + graph = helper.make_graph( + [top_identity, node], + "duplicate_non_output", + [condition], + [top_output, nested_output], + [outer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_initializer_consumer_model() -> ModelProto: + """Build an output initializer consumed only by an FP32-blocked node.""" + shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + copied = helper.make_tensor_value_info("copied", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([1.0001], dtype=np.float32), "shared") + identity = helper.make_node("Identity", ["shared"], ["copied"], name="identity") + graph = helper.make_graph([identity], "blocked_consumer", [], [shared, copied], [initializer]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_blocked_initializer_with_nested_input_shadow_model() -> ModelProto: + """Build a blocked initializer consumer plus an unrelated nested local shadow.""" + trip_count = helper.make_tensor_value_info("trip_count", TensorProto.INT64, []) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + loop_state = helper.make_tensor_value_info("loop_state", TensorProto.FLOAT, [1]) + blocked_output = helper.make_tensor_value_info("blocked_output", TensorProto.FLOAT, [1]) + loop_output = helper.make_tensor_value_info("loop_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([1.0001], dtype=np.float32), "shared") + + iteration = helper.make_tensor_value_info("iteration", TensorProto.INT64, []) + body_condition = helper.make_tensor_value_info("body_condition", TensorProto.BOOL, []) + local_shared = helper.make_tensor_value_info("shared", TensorProto.FLOAT, [1]) + body_condition_out = helper.make_tensor_value_info("body_condition_out", TensorProto.BOOL, []) + body_state_out = helper.make_tensor_value_info("body_state_out", TensorProto.FLOAT, [1]) + body = helper.make_graph( + [ + helper.make_node( + "Identity", ["body_condition"], ["body_condition_out"], name="body_condition" + ), + helper.make_node("Relu", ["shared"], ["body_state_out"], name="body_relu"), + ], + "body", + [iteration, body_condition, local_shared], + [body_condition_out, body_state_out], + ) + blocked = helper.make_node("Abs", ["shared"], ["blocked_output"], name="blocked_abs") + loop = helper.make_node( + "Loop", + ["trip_count", "condition", "loop_state"], + ["loop_output"], + name="loop", + body=body, + ) + graph = helper.make_graph( + [blocked, loop], + "blocked_initializer_nested_shadow", + [trip_count, condition, loop_state], + [blocked_output, loop_output], + [initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _build_unused_initializer_with_nested_input_shadow_model() -> ModelProto: + """Build an unused initializer shadowed by a consumed nested formal.""" + model = _build_blocked_initializer_with_nested_input_shadow_model() + del model.graph.node[0] + model.graph.output[0].CopyFrom(model.graph.output[1]) + del model.graph.output[1:] + model.graph.name = "unused_initializer_nested_shadow" + return model + + +def _build_non_float_initializer_before_nested_float_initializer_model() -> ModelProto: + """Build an earlier non-FLOAT binding sharing a later nested FLOAT initializer name.""" + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + integer_output = helper.make_tensor_value_info("integer_output", TensorProto.INT64, [1]) + nested_output = helper.make_tensor_value_info("nested_output", TensorProto.FLOAT, [1]) + integer_initializer = numpy_helper.from_array(np.array([7], dtype=np.int64), "same") + + def _branch(name: str, initializer_name: str, value: float) -> GraphProto: + branch_output = helper.make_tensor_value_info(f"{name}_output", TensorProto.FLOAT, [1]) + initializer = numpy_helper.from_array(np.array([value], dtype=np.float32), initializer_name) + identity = helper.make_node( + "Identity", [initializer_name], [f"{name}_output"], name=f"{name}_identity" + ) + return helper.make_graph([identity], name, [], [branch_output], [initializer]) + + integer_identity = helper.make_node( + "Identity", ["same"], ["integer_output"], name="integer_identity" + ) + conditional = helper.make_node( + "If", + ["condition"], + ["nested_output"], + name="if", + then_branch=_branch("then", "same", 1.5), + else_branch=_branch("else", "other", 2.5), + ) + graph = helper.make_graph( + [integer_identity, conditional], + "initializer_registration_order", + [condition], + [integer_output, nested_output], + [integer_initializer], + ) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +# ============================================================================= +# CONVERT_TO_FP16 TESTS +# ============================================================================= + + +class TestConvertToFP16: + """Test convert_to_fp16 utility function.""" + + def test_converts_weights_to_fp16(self) -> None: + """FP16 conversion converts float32 initializers to float16.""" + model = _build_simple_fp32_model() + result = convert_to_fp16(model) + + has_fp16 = any(init.data_type == TensorProto.FLOAT16 for init in result.graph.initializer) + assert has_fp16, "Expected at least one FP16 initializer after conversion" + + def test_success_mutates_and_returns_input_model(self) -> None: + """A successful conversion preserves the wrapper's in-place API.""" + model = _build_simple_fp32_model() + + result = convert_to_fp16(model) + + assert result is model + assert any( + initializer.data_type == TensorProto.FLOAT16 for initializer in model.graph.initializer + ) + + def test_scalar_float_tensor_attribute_is_rejected(self) -> None: + """ORT cannot convert scalar FLOAT storage in tensor attributes.""" + model = _build_scalar_float_attribute_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {}) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "incompatible FP16 types"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_sparse_float_tensor_attribute_is_rejected(self) -> None: + """ORT cannot convert sparse FLOAT storage in tensor attributes.""" + model = _build_sparse_float_attribute_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "incompatible FP16 types"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_loaded_external_tensor_attribute_is_internalized( + self, + ) -> None: + """Resident tensor attributes drop stale external metadata.""" + model = _build_external_float_attribute_model(clear_data=False) + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + value = result.graph.node[0].attribute[0].t + assert value.data_type == TensorProto.FLOAT16 + assert value.data_location == TensorProto.DEFAULT + assert not value.external_data + assert len(value.raw_data) == 2 + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {}) + assert output.dtype == np.float16 + + def test_unloaded_external_tensor_attribute_is_rejected( + self, + ) -> None: + """Selected tensor attributes require resident external data.""" + model = _build_external_float_attribute_model(clear_data=True) + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "unloaded external data"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_omitted_optional_output_is_not_a_binding(self) -> None: + """Empty optional output and input sentinels are never connected.""" + model = _build_omitted_optional_output_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Dropout"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + dropped, clipped = session.run(None, {"x": np.array([1.0, 2.0], dtype=np.float16)}) + assert dropped.dtype == np.float16 + assert clipped.dtype == np.float16 + + def test_default_keeps_io_types(self) -> None: + """Default keep_io_types=True preserves FP32 model I/O.""" + model = _build_simple_fp32_model() + result = convert_to_fp16(model, keep_io_types=True) + + for inp in result.graph.input: + assert inp.type.tensor_type.elem_type == TensorProto.FLOAT + for outp in result.graph.output: + assert outp.type.tensor_type.elem_type == TensorProto.FLOAT + + def test_float_sequence_io_is_rejected_when_io_types_are_kept(self) -> None: + """Unsupported FLOAT container I/O is not silently converted to FP16.""" + model = _build_float_sequence_io_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + {"sequence": [np.array([2.0], dtype=np.float32)]}, + ) + assert output[0].dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "FLOAT container graph I/O"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_float_optional_input_is_preserved_when_io_types_are_kept(self) -> None: + """Container I/O that ORT leaves unchanged remains supported.""" + model = _build_top_optional_input_blocked_capture_model() + + result = convert_to_fp16( + model, + keep_io_types=True, + op_block_list=["If"], + ) + + optional_input = next(value for value in result.graph.input if value.name == "optional") + assert ( + optional_input.type.optional_type.elem_type.tensor_type.elem_type == TensorProto.FLOAT + ) + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + def test_keep_io_types_false_converts_io(self) -> None: + """With keep_io_types=False, model I/O becomes FP16.""" + model = _build_simple_fp32_model() + result = convert_to_fp16(model, keep_io_types=False) + + for inp in result.graph.input: + assert inp.type.tensor_type.elem_type == TensorProto.FLOAT16 + for outp in result.graph.output: + assert outp.type.tensor_type.elem_type == TensorProto.FLOAT16 + + def test_initializer_backed_output_stays_fp32_when_io_types_are_kept(self) -> None: + """An initializer graph output remains valid FP32 when preserving I/O.""" + model = _build_initializer_backed_output_model() + + result = convert_to_fp16(model, keep_io_types=True) + + output = result.graph.output[0] + assert output.type.tensor_type.elem_type == TensorProto.FLOAT + assert any( + initializer.data_type == TensorProto.FLOAT for initializer in result.graph.initializer + ) + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + np.testing.assert_array_equal( + session.run(None, {})[0], + np.array([[1.0001, 2.0003]], dtype=np.float32), + ) + + def test_initializer_backed_output_converts_data_when_io_types_are_not_kept(self) -> None: + """A pure-FP16 output converts its backing initializer as well as its type.""" + model = _build_initializer_backed_output_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + output = result.graph.output[0] + initializer = result.graph.initializer[0] + assert output.type.tensor_type.elem_type == TensorProto.FLOAT16 + assert initializer.data_type == TensorProto.FLOAT16 + shape_inference.infer_shapes(result, strict_mode=True) + + def test_initializer_backed_output_with_fp16_consumer_converts_to_fp16(self) -> None: + """Pure-FP16 conversion allows an output initializer consumed by FP16-capable nodes.""" + model = _build_shared_initializer_output_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert all( + output.type.tensor_type.elem_type == TensorProto.FLOAT16 + for output in result.graph.output + ) + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + shared, y = session.run(None, {"x": np.array([[3.0, 4.0]], dtype=np.float16)}) + np.testing.assert_array_equal(shared, np.array([[1.0, 2.0]], dtype=np.float16)) + np.testing.assert_array_equal(y, np.array([[4.0, 6.0]], dtype=np.float16)) + + def test_overridable_initializer_output_is_rejected_when_io_types_are_not_kept( + self, + ) -> None: + """Graph input/output initializer aliases are rejected before ORT optimizer crashes.""" + model = _build_overridable_shared_initializer_output_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "also a graph input"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_scan8_scan_output_is_not_feedback_type_evidence(self) -> None: + """A scan output cannot type an unrelated untyped scan input.""" + model = _build_scan8_without_state_model() + original = model.SerializeToString() + + checker.check_model(model) + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu"], + ) + assert model.SerializeToString() == original + + def test_nested_initializer_output_with_fp16_consumer_converts_to_fp16(self) -> None: + """Nested initializer outputs are allowed when ORT converts them consistently.""" + model = _build_nested_consumed_initializer_output_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + np.testing.assert_array_equal( + session.run(None, {"condition": np.array(True)})[0], + np.array([1.0], dtype=np.float16), + ) + + def test_fp16_initializers_do_not_skip_fp32_graph_conversion(self) -> None: + """FP16 initializers alone do not prove graph I/O and nodes are already FP16.""" + model = _build_fp32_identity_with_fp16_nested_initializer_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.input[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_sparse_initializer_inputs_are_available_to_topological_sort(self) -> None: + """Sorting after conversion treats sparse initializer names as available values.""" + model = _build_sparse_shape_reshape_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_sparse_float_initializer_converts_with_fp16_consumers(self) -> None: + """Sparse FLOAT initializers are converted when their consumers become FP16.""" + model = _build_sparse_float_add_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + assert result.graph.sparse_initializer[0].values.data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + np.testing.assert_array_equal( + session.run(None, {"x": np.array([3.0, 4.0], dtype=np.float16)})[0], + np.array([4.0, 6.0], dtype=np.float16), + ) + + def test_sparse_float_initializer_value_info_converts_with_values(self) -> None: + """Sparse FLOAT value_info metadata is kept consistent with converted values.""" + model = _build_sparse_float_add_model_with_value_info() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + sparse_info = next(value for value in result.graph.value_info if value.name == "weight") + assert sparse_info.type.sparse_tensor_type.elem_type == TensorProto.FLOAT16 + assert result.graph.sparse_initializer[0].values.data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_sparse_float_initializer_tensor_value_info_is_rejected_before_mutation( + self, + ) -> None: + """Tensor metadata for sparse FLOAT initializers is rejected before ORT fails.""" + model = _build_sparse_float_add_model_with_tensor_value_info() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "sparse initializer metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_sparse_float_initializer_io_alias_is_rejected_when_io_types_are_not_kept( + self, + ) -> None: + """Sparse graph input/output initializer aliases are rejected before mutation.""" + model = _build_sparse_float_add_model_with_io_metadata() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "also sparse graph input"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_sparse_float_initializer_tensor_io_alias_is_rejected_before_mutation( + self, + ) -> None: + """Tensor graph input/output sparse initializer metadata is rejected pre-ORT.""" + model = _build_sparse_float_add_model_with_tensor_io_metadata() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "sparse initializer metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_retained_float_initializer_metadata_is_rejected(self) -> None: + """Converted metadata cannot describe an initializer retained in FP32.""" + model = _build_retained_float_initializer_metadata_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array( + [[[[1.0, 2.0], [3.0, 4.0]]]], + dtype=np.float32, + ) + }, + ) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "initializer metadata"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_sparse_float_initializer_io_metadata_is_rejected_when_io_is_kept(self) -> None: + """Sparse graph I/O dtypes are not silently changed when keep_io_types=True.""" + model = _build_sparse_float_add_model_with_io_metadata() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "sparse graph I/O"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_sparse_float_initializer_tensor_io_metadata_is_rejected_before_keep_io( + self, + ) -> None: + """Tensor graph I/O sparse initializer metadata is rejected before ORT fails.""" + model = _build_sparse_float_add_model_with_tensor_io_metadata() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "sparse initializer metadata"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_direct_sparse_float_output_converts_when_io_types_are_not_kept(self) -> None: + """Direct sparse FLOAT outputs convert without requiring node consumers.""" + model = _build_direct_sparse_float_output_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + sparse_output = next( + value for value in result.graph.output if value.name == "sparse_output" + ) + assert sparse_output.type.sparse_tensor_type.elem_type == TensorProto.FLOAT16 + assert result.graph.sparse_initializer[0].values.data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_direct_sparse_float_tensor_output_is_rejected_before_mutation( + self, + ) -> None: + """Direct sparse FLOAT outputs with tensor metadata are rejected pre-ORT.""" + model = _build_direct_sparse_float_tensor_output_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "sparse initializer metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_direct_sparse_float_outputs_are_rejected_before_mutation( + self, + ) -> None: + """Nested sparse outputs are rejected when parent edge types cannot be repaired.""" + model = _build_if_with_direct_sparse_float_outputs_model() + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "nested sparse graph output"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_direct_sparse_float_outputs_are_preserved_when_io_is_kept( + self, + ) -> None: + """A kept top-level sparse output leaves its nested sparse edges in FLOAT.""" + model = _build_if_with_direct_sparse_float_outputs_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert result.graph.output[0].type.sparse_tensor_type.elem_type == TensorProto.FLOAT + for graph in _iter_attribute_graphs(result): + assert graph.output[0].type.sparse_tensor_type.elem_type == TensorProto.FLOAT + assert graph.sparse_initializer[0].values.data_type == TensorProto.FLOAT + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_deep_nested_sparse_outputs_are_preserved_when_io_is_kept( + self, + ) -> None: + """A kept sparse edge remains FLOAT through every enclosing graph.""" + model = _build_deep_if_with_direct_sparse_float_outputs_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert result.graph.output[0].type.sparse_tensor_type.elem_type == TensorProto.FLOAT + for graph in _iter_attribute_graphs(result): + assert graph.output[0].type.sparse_tensor_type.elem_type == TensorProto.FLOAT + for sparse in graph.sparse_initializer: + assert sparse.values.data_type == TensorProto.FLOAT + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_direct_sparse_float_output_with_fp32_consumer_is_rejected(self) -> None: + """Direct sparse FP16 outputs are mixed uses when blocked consumers stay FP32.""" + model = _build_sparse_float_add_model_with_io_metadata() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "both FP16 and FP32"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Add"]) + assert model.SerializeToString() == original + + def test_unloaded_nested_external_initializer_output_is_rejected_before_mutation( + self, + ) -> None: + """Unloaded nested external backing data is rejected before dtype metadata changes.""" + model = _build_nested_consumed_initializer_output_model() + for graph in _iter_attribute_graphs(model): + _mark_initializers_as_external(graph, clear_data=True) + original = model.SerializeToString() + + with np.testing.assert_raises_regex( + RuntimeError, + "load external weights before FP16 conversion", + ): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_loaded_nested_external_initializer_output_is_internalized(self) -> None: + """Loaded nested external output data no longer points to stale sidecars.""" + model = _build_nested_consumed_initializer_output_model() + for graph in _iter_attribute_graphs(model): + _mark_initializers_as_external(graph, clear_data=False) + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + for graph in _iter_attribute_graphs(result): + for initializer in graph.initializer: + assert initializer.data_type == TensorProto.FLOAT16 + assert initializer.data_location == TensorProto.DEFAULT + assert not initializer.external_data + + def test_non_float_external_initializer_output_metadata_is_preserved(self) -> None: + """Non-FLOAT direct output initializers are outside the FP16 repair path.""" + model = _build_initializer_backed_output_model() + int_output = helper.make_tensor_value_info("int_output", TensorProto.INT64, [1]) + int_value = numpy_helper.from_array(np.array([7], dtype=np.int64), "int_output") + int_value.ClearField("raw_data") + int_value.data_location = TensorProto.EXTERNAL + location = int_value.external_data.add() + location.key = "location" + location.value = "int_output.bin" + model.graph.output.append(int_output) + model.graph.initializer.append(int_value) + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + int_initializer = next( + initializer + for initializer in result.graph.initializer + if initializer.name == "int_output" + ) + assert int_initializer.data_type == TensorProto.INT64 + assert int_initializer.data_location == TensorProto.EXTERNAL + assert [(entry.key, entry.value) for entry in int_initializer.external_data] == [ + ("location", "int_output.bin") + ] + + def test_multiple_initializer_backed_outputs_are_all_converted(self) -> None: + """Every initializer-backed output is repaired independently.""" + model = _build_initializer_backed_output_model() + second_output = helper.make_tensor_value_info("second_output", TensorProto.FLOAT, [1, 2]) + second_value = numpy_helper.from_array( + np.array([[3.0, 4.0]], dtype=np.float32), "second_output" + ) + model.graph.output.append(second_output) + model.graph.initializer.append(second_value) + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert all( + output.type.tensor_type.elem_type == TensorProto.FLOAT16 + for output in result.graph.output + ) + assert all( + initializer.data_type == TensorProto.FLOAT16 for initializer in result.graph.initializer + ) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_multiple_initializer_outputs_keep_exact_fp32_values(self) -> None: + """Removing several orphan Casts preserves every FP32 model output.""" + model = _build_initializer_backed_output_model() + second_output = helper.make_tensor_value_info("second_output", TensorProto.FLOAT, [1, 2]) + second_value = numpy_helper.from_array( + np.array([[3.0005, 4.0007]], dtype=np.float32), "second_output" + ) + model.graph.output.append(second_output) + model.graph.initializer.append(second_value) + + result = convert_to_fp16(model, keep_io_types=True) + + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + first, second = session.run(None, {}) + np.testing.assert_array_equal(first, np.array([[1.0001, 2.0003]], dtype=np.float32)) + np.testing.assert_array_equal(second, np.array([[3.0005, 4.0007]], dtype=np.float32)) + + def test_overridable_initializer_output_is_rejected_before_mutation(self) -> None: + """A graph-input initializer keeps its caller override semantics.""" + model = _build_initializer_backed_output_model() + model.graph.input.append( + helper.make_tensor_value_info("constant_output", TensorProto.FLOAT, [1, 2]) + ) + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "also a graph input"): + convert_to_fp16(model, keep_io_types=True) + + assert model.SerializeToString() == original + + def test_collision_uses_original_mixed_output_index(self) -> None: + """ORT Cast collision checks count preceding non-FLOAT outputs.""" + model = _build_initializer_backed_output_model() + int_output = helper.make_tensor_value_info("int_output", TensorProto.INT64, [1]) + int_value = numpy_helper.from_array(np.array([1], dtype=np.int64), "int_output") + model.graph.output.insert(0, int_output) + model.graph.initializer.append(int_value) + model.graph.input.append( + helper.make_tensor_value_info("graph_output_cast_1", TensorProto.FLOAT, [1]) + ) + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "graph_output_cast_1"): + convert_to_fp16(model, keep_io_types=True) + + assert model.SerializeToString() == original + + def test_unloaded_external_initializer_output_is_rejected(self) -> None: + """Conversion refuses external backing data that was not loaded.""" + model = _build_initializer_backed_output_model() + initializer = model.graph.initializer[0] + initializer.ClearField("raw_data") + initializer.data_location = TensorProto.EXTERNAL + location = initializer.external_data.add() + location.key = "location" + location.value = "weights.data" + + original_output_type = model.graph.output[0].type.tensor_type.elem_type + original_initializer_type = initializer.data_type + + with np.testing.assert_raises_regex( + RuntimeError, + "load external weights before FP16 conversion", + ): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert model.graph.output[0].type.tensor_type.elem_type == original_output_type + assert model.graph.initializer[0].data_type == original_initializer_type + + def test_loaded_external_initializer_output_is_converted(self) -> None: + """Resident tensor bytes are valid even if external metadata remains.""" + model = _build_initializer_backed_output_model() + initializer = model.graph.initializer[0] + initializer.data_location = TensorProto.EXTERNAL + location = initializer.external_data.add() + location.key = "location" + location.value = "weights.data" + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + + def test_loaded_external_initializer_output_is_internalized_when_io_is_kept(self) -> None: + """Resident output data no longer points to a stale sidecar after repair.""" + model = _build_initializer_backed_output_model() + initializer = model.graph.initializer[0] + initializer.data_location = TensorProto.EXTERNAL + location = initializer.external_data.add() + location.key = "location" + location.value = "weights.data" + + result = convert_to_fp16(model, keep_io_types=True) + + repaired = result.graph.initializer[0] + assert repaired.data_location == TensorProto.DEFAULT + assert not repaired.external_data + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + np.testing.assert_array_equal( + session.run(None, {})[0], + np.array([[1.0001, 2.0003]], dtype=np.float32), + ) + + def test_loaded_external_weight_is_internalized_before_conversion( + self, + ) -> None: + """Resident external FLOAT weights convert without stale sidecars.""" + model = _build_simple_fp32_model() + _mark_initializers_as_external(model.graph, clear_data=False) + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + + initializer = result.graph.initializer[0] + assert initializer.data_type == TensorProto.FLOAT16 + assert initializer.data_location == TensorProto.DEFAULT + assert not initializer.external_data + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + {"x": np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float16)}, + ) + assert output.dtype == np.float16 + + def test_unloaded_external_weight_is_rejected_before_mutation( + self, + ) -> None: + """Selected external FLOAT weights require resident tensor data.""" + model = _build_simple_fp32_model() + _mark_initializers_as_external(model.graph, clear_data=True) + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "unloaded external data"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_external_non_float_name_collision_is_not_selected( + self, + ) -> None: + """Global FLOAT selection does not internalize same-named INT data.""" + model = _build_nested_external_int_name_collision_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + then_graph = next( + attribute.g + for attribute in result.graph.node[-1].attribute + if attribute.g.name == "then" + ) + integer = then_graph.initializer[0] + assert integer.data_type == TensorProto.INT64 + assert integer.data_location == TensorProto.EXTERNAL + assert integer.external_data + + def test_shared_initializer_output_is_rejected_before_mutation(self) -> None: + """Shared initializer-output semantics are rejected before mutation.""" + model = _build_shared_initializer_output_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "has internal consumers"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_initializer_outputs_are_converted_when_io_types_are_kept(self) -> None: + """Nested direct output initializers are repaired in traversed graphs.""" + model = _build_nested_initializer_output_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT + for graph in _iter_attribute_graphs(result): + assert all( + output.type.tensor_type.elem_type == TensorProto.FLOAT16 for output in graph.output + ) + assert all( + initializer.data_type == TensorProto.FLOAT16 for initializer in graph.initializer + ) + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + np.testing.assert_array_equal( + session.run(None, {"condition": np.array(True)})[0], + np.array([1.0], dtype=np.float32), + ) + + def test_lexical_nested_consumers_are_rejected_before_mutation(self) -> None: + """Lexically shared output initializers are rejected before mutation.""" + model = _build_lexically_captured_initializer_output_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "has internal consumers"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_generated_tensor_name_collision_is_rejected_before_mutation(self) -> None: + """Repair allocates a fresh alias instead of duplicating an existing name.""" + model = _build_initializer_output_name_collision_model() + + original = model.SerializeToString() + with np.testing.assert_raises_regex(RuntimeError, "existing names collide"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_generated_cast_node_name_collision_is_rejected_before_mutation(self) -> None: + """A user node occupying ORT's deterministic Cast name fails safely.""" + model = _build_initializer_output_node_name_collision_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "existing names collide"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert model.SerializeToString() == original + + def test_nested_generated_node_name_collision_is_rejected_before_mutation(self) -> None: + """Nested nodes also participate in ORT's global generated-name set.""" + model = _build_nested_node_name_collision_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "existing names collide"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert model.SerializeToString() == original + + def test_regular_output_nested_node_name_collision_is_rejected_before_mutation( + self, + ) -> None: + """Generated I/O Cast names are reserved even without initializer outputs.""" + model = _build_regular_output_nested_node_name_collision_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "graph_output_cast0"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_neutral_nested_node_name_collision_is_allowed(self) -> None: + """Skipping an already-concrete INT node cannot alter precision.""" + model = _build_neutral_nested_node_name_collision_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, nested_output = session.run( + None, + { + "x": np.array([2.0], dtype=np.float32), + "integer": np.array([7], dtype=np.int64), + "condition": np.array(True), + }, + ) + np.testing.assert_array_equal(output, np.array([2.0], dtype=np.float32)) + np.testing.assert_array_equal(nested_output, np.array([7], dtype=np.int64)) + + def test_neutral_collision_can_use_shadowed_int_formal( + self, + ) -> None: + """A skipped neutral node leaves its local INT formal untouched.""" + model = _build_neutral_shadowed_keep_io_collision_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, loop_output = session.run( + None, + { + "x": np.array([2.0], dtype=np.float32), + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "integer": np.array([7], dtype=np.int64), + }, + ) + np.testing.assert_array_equal(output, np.array([2.0], dtype=np.float32)) + np.testing.assert_array_equal(loop_output, np.array([7], dtype=np.int64)) + + def test_inferred_output_nested_node_name_collision_is_rejected_before_mutation( + self, + ) -> None: + """Generated Cast reservations use the same inferred I/O types as ORT.""" + model = _build_inferred_output_nested_node_name_collision_model() + inferred = shape_inference.infer_shapes(model, strict_mode=True) + assert inferred.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT + checker.check_model(inferred) + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "graph_output_cast0"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_local_generated_tensor_alias_does_not_collide(self) -> None: + """A nested local binding may legally shadow a generated top-level alias.""" + model = _build_nested_local_generated_tensor_alias_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, loop_output = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "loop_state": np.array([2.0], dtype=np.float32), + "x": np.array([3.0], dtype=np.float32), + }, + ) + np.testing.assert_array_equal(output, np.array([3.0], dtype=np.float32)) + np.testing.assert_array_equal(loop_output, np.array([2.0], dtype=np.float32)) + + def test_nested_blocked_generated_tensor_alias_is_rejected_before_mutation( + self, + ) -> None: + """ORT's blocked-node lookup is global even for a nested local binding.""" + model = _build_nested_blocked_generated_tensor_alias_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "blocked or mixed-type nested node"): + convert_to_fp16(model, keep_io_types=True, op_block_list=["Identity"]) + assert model.SerializeToString() == original + + def test_nested_mixed_generated_tensor_alias_is_rejected_before_mutation( + self, + ) -> None: + """ORT's mixed-input lookup is global even for a nested local binding.""" + model = _build_nested_mixed_generated_tensor_alias_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "blocked or mixed-type nested node"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_inferred_nested_mixed_name_collision_is_rejected_before_mutation( + self, + ) -> None: + """Late lookup safety uses the same shape inference metadata as ORT.""" + model = _build_inferred_nested_mixed_name_collision_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (loop_output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "loop_scale": np.array([2.0], dtype=np.float32), + "x": np.array([3.0], dtype=np.float32), + }, + ) + assert loop_output.shape == (2,) + + with np.testing.assert_raises_regex(RuntimeError, "global value-info"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_top_blocked_global_value_info_collision_is_rejected(self) -> None: + """Top-level blocked nodes also use ORT's global metadata lookup.""" + model = _build_top_blocked_global_value_info_collision_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "global value-info"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Identity"]) + assert model.SerializeToString() == original + + def test_duplicate_late_cast_alias_is_rejected_before_mutation(self) -> None: + """Blocked nodes must not allocate the same deterministic Cast tensor name.""" + model = _build_duplicate_late_cast_alias_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "late Cast tensor"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Identity"]) + assert model.SerializeToString() == original + + def test_late_cast_node_name_collision_is_rejected_before_mutation(self) -> None: + """Generated late Cast node names must be unique in the top-level graph.""" + model = _build_late_cast_node_name_collision_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "late Cast node"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Abs"]) + assert model.SerializeToString() == original + + def test_blocked_subgraph_float_capture_is_rejected_before_mutation(self) -> None: + """Skipped branch graphs cannot keep FLOAT captures that become FP16.""" + model = _build_blocked_subgraph_float_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_untyped_blocked_capture_is_rejected_when_inference_falls_back( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Inference fallback conservatively rejects unresolved node captures.""" + model = _build_untyped_blocked_subgraph_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + def fail_inference(*args: object, **kwargs: object) -> None: + raise EncodeError("simulated serialization failure") + + monkeypatch.setattr(shape_inference, "infer_shapes", fail_inference) + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_uninferred_custom_op_blocked_capture_is_rejected(self) -> None: + """Successful but incomplete inference still requires conservative safety.""" + model = _build_uninferred_custom_op_blocked_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + inferred = shape_inference.infer_shapes(model, strict_mode=True) + assert all(value.name != "y" for value in inferred.graph.value_info) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([1.0], dtype=np.float32), + "condition": np.array(True), + }, + ) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_nested_value_info_can_annotate_kept_input_capture( + self, + ) -> None: + """value_info alone does not create a nested lexical binding.""" + model = _build_kept_input_free_capture_value_info_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=True, + op_block_list=[], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([2.0], dtype=np.float32), + "condition": np.array(True), + }, + ) + np.testing.assert_array_equal(output, np.array([2.0], dtype=np.float32)) + + def test_empty_metadata_blocked_capture_is_rejected(self) -> None: + """Name-only metadata does not prove a captured producer stays non-FLOAT.""" + model = _build_uninferred_custom_op_blocked_capture_model() + model.graph.value_info.append(helper.make_empty_tensor_value_info("y")) + original = model.SerializeToString() + + checker.check_model(model) + inferred = shape_inference.infer_shapes(model, strict_mode=True) + y = next(value for value in inferred.graph.value_info if value.name == "y") + assert y.type.tensor_type.elem_type == TensorProto.UNDEFINED + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_missing_metadata_blocked_input_is_rejected(self) -> None: + """A blocked input needs metadata to receive an FP16-to-FLOAT boundary Cast.""" + model = _build_missing_metadata_blocked_edge_model() + original = model.SerializeToString() + + checker.check_model(model) + inferred = shape_inference.infer_shapes(model, strict_mode=True) + assert all(value.name != "hidden" for value in inferred.graph.value_info) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([1.0], dtype=np.float32)}) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Add"]) + assert model.SerializeToString() == original + + def test_missing_metadata_blocked_output_is_rejected(self) -> None: + """A blocked output needs metadata to receive a FLOAT-to-FP16 boundary Cast.""" + model = _build_missing_metadata_blocked_edge_model() + original = model.SerializeToString() + + checker.check_model(model) + inferred = shape_inference.infer_shapes(model, strict_mode=True) + assert all(value.name != "hidden" for value in inferred.graph.value_info) + + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Gelu"]) + assert model.SerializeToString() == original + + def test_existing_fp16_boundary_does_not_require_metadata(self) -> None: + """An explicit type boundary can safely consume a blocked FLOAT output.""" + model = _build_existing_fp16_boundary_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([1.0], dtype=np.float16)}) + assert output.dtype == np.float16 + + def test_missing_metadata_equal_sibling_coupling_is_rejected(self) -> None: + """A same-type sibling input still requires an FP16 output boundary.""" + model = _build_missing_metadata_equal_consumer_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([1.0], dtype=np.float32), + "other": np.array([1.0], dtype=np.float32), + }, + ) + assert output.dtype == np.bool_ + + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Gelu"]) + assert model.SerializeToString() == original + + def test_missing_metadata_sequence_coupling_is_rejected(self) -> None: + """A tensor input and sequence output can share payload precision.""" + model = _build_missing_metadata_sequence_consumer_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([1.0], dtype=np.float32)}) + assert output[0].dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Gelu"]) + assert model.SerializeToString() == original + + def test_missing_metadata_sequence_map_coupling_is_rejected(self) -> None: + """Partially overlapping schema constraints can still couple payload precision.""" + model = _build_missing_metadata_sequence_map_consumer_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([1.0], dtype=np.float32), + "sequence": [np.array([2.0], dtype=np.float32)], + }, + ) + assert output[0].dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Gelu"]) + assert model.SerializeToString() == original + + def test_int_sequence_map_consumer_does_not_require_float_boundary(self) -> None: + """Concrete child formal types prove an uninferred edge is non-FLOAT.""" + model = _build_int_sequence_map_consumer_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["MurmurHash3"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "text": np.array(["hello"]), + "sequence": [np.array([1], dtype=np.int32)], + }, + ) + assert output[0].dtype == np.int32 + + def test_int_same_type_consumer_does_not_require_float_boundary(self) -> None: + """Concrete non-FLOAT edges override same-type schema relationships.""" + model = _build_int_identity_consumer_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["NonMaxSuppression"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "boxes": np.array( + [ + [ + [0.0, 0.0, 1.0, 1.0], + [2.0, 2.0, 3.0, 3.0], + ] + ], + dtype=np.float16, + ), + "scores": np.array([[[0.9, 0.8]]], dtype=np.float16), + }, + ) + assert output.dtype == np.int64 + + def test_concrete_consumer_output_proves_uninferred_input_is_int(self) -> None: + """A concrete same-type output can prove an uninferred input is non-FLOAT.""" + model = _build_uninferred_int_identity_consumer_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["MurmurHash3"], + ) + + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"text": np.array(["hello"])}) + assert output.dtype == np.int32 + + def test_concrete_function_input_is_int_without_registered_schema(self) -> None: + """Concrete non-FLOAT evidence is sufficient for a local function input.""" + model = _build_function_int_consumer_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["MurmurHash3"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"text": np.array(["hello"])}) + assert output.dtype == np.int32 + + def test_unconverted_float_constant_in_local_function_is_rejected( + self, + ) -> None: + """ORT does not visit concrete FLOAT values in FunctionProto bodies.""" + model = _build_function_with_float_constant_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([3.0], dtype=np.float32)}) + np.testing.assert_array_equal(output, np.array([4.0], dtype=np.float32)) + + with np.testing.assert_raises_regex(RuntimeError, "local function"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_unconverted_contrib_float_function_is_rejected(self) -> None: + """Concrete FLOAT function bodies cannot rely on ONNX-only inference.""" + model = _build_function_with_contrib_float_constant_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([3.0], dtype=np.float32)}) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "local function"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_unconverted_scalar_contrib_float_function_is_rejected( + self, + ) -> None: + """Scalar FLOAT storage in an uninferred function stays FP32.""" + model = _build_function_with_scalar_contrib_float_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([3.0], dtype=np.float32)}) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "local function"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_unrelated_graph_attribute_does_not_supply_input_type(self) -> None: + """Graph attributes require schema-proven input alignment.""" + model = _build_function_with_unrelated_graph_attribute_model() + original = model.SerializeToString() + + checker.check_model(model) + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu"], + ) + assert model.SerializeToString() == original + + def test_always_float_op_graph_attribute_is_not_traversed(self) -> None: + """Preflight mirrors ORT when an always-float op owns a graph attribute.""" + model = _build_always_float_function_graph_attribute_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Identity"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([2.0], dtype=np.float16)}) + assert output.dtype == np.float16 + + def test_scan8_child_formal_proves_variadic_state_is_int(self) -> None: + """A fixed node-only input does not shift heterogeneous child formals.""" + model = _build_scan8_int_state_model() -from winml.modelkit.quant.fp16 import convert_to_fp16 + checker.check_model(model) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["MurmurHash3"], + ) + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + final_state, scan_output = session.run( + None, + { + "text": np.array(["hello"]), + "scan_input": np.array([[1.0, 2.0]], dtype=np.float16), + }, + ) + assert final_state.dtype == np.int32 + assert scan_output.dtype == np.float16 + def test_scan8_float_state_requires_precision_boundary_metadata(self) -> None: + """A FLOAT child formal proves positional state coupling.""" + model = _build_scan8_float_state_model() + original = model.SerializeToString() -# ============================================================================= -# HELPERS -# ============================================================================= + checker.check_model(model) + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu"], + ) + assert model.SerializeToString() == original + def test_loop_feedback_output_proves_float_state_coupling(self) -> None: + """A positionally aligned FLOAT feedback output requires a boundary.""" + model = _build_loop_float_state_model() + original = model.SerializeToString() -def _build_simple_fp32_model() -> ModelProto: - """Build a simple FP32 model: out = x + weight.""" - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]) - out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4]) - weight = numpy_helper.from_array(np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight") - add = helper.make_node("Add", ["x", "weight"], ["out"], name="add") - graph = helper.make_graph([add], "simple", [x], [out], [weight]) - return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + checker.check_model(model) + with np.testing.assert_raises_regex(RuntimeError, "missing FLOAT metadata"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu"], + ) + assert model.SerializeToString() == original + def test_metadata_free_int_input_to_blocked_node_is_preserved(self) -> None: + """Concrete consumer output evidence prevents a spurious FLOAT Cast.""" + model = _build_uninferred_int_identity_consumer_model() -def _build_multi_op_fp32_model() -> ModelProto: - """Build a model with multiple ops: out = Relu(x + weight).""" - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4]) - out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [1, 4]) - weight = numpy_helper.from_array(np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32), "weight") - add = helper.make_node("Add", ["x", "weight"], ["add_out"], name="add") - relu = helper.make_node("Relu", ["add_out"], ["out"], name="relu") - graph = helper.make_graph([add, relu], "multi_op", [x], [out], [weight]) - return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Identity"], + ) + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"text": np.array(["hello"])}) + assert output.dtype == np.int32 -# ============================================================================= -# CONVERT_TO_FP16 TESTS -# ============================================================================= + def test_optional_float_blocked_capture_is_rejected(self) -> None: + """A blocked child cannot retain FLOAT for an Optional that becomes FP16.""" + model = _build_optional_float_blocked_capture_model() + original = model.SerializeToString() + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([2.0], dtype=np.float32), + "condition": np.array(True), + }, + ) + np.testing.assert_array_equal(output, np.array([2.0], dtype=np.float32)) -class TestConvertToFP16: - """Test convert_to_fp16 utility function.""" + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original - def test_converts_weights_to_fp16(self) -> None: - """FP16 conversion converts float32 initializers to float16.""" - model = _build_simple_fp32_model() - result = convert_to_fp16(model) + def test_kept_input_optional_producer_uses_generated_fp16_alias(self) -> None: + """Producer tracing follows the FP16 side of a kept input Cast.""" + model = _build_optional_float_blocked_capture_model() + original = model.SerializeToString() - has_fp16 = any(init.data_type == TensorProto.FLOAT16 for init in result.graph.initializer) - assert has_fp16, "Expected at least one FP16 initializer after conversion" + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16( + model, + keep_io_types=True, + op_block_list=["If"], + ) + assert model.SerializeToString() == original - def test_default_keeps_io_types(self) -> None: - """Default keep_io_types=True preserves FP32 model I/O.""" - model = _build_simple_fp32_model() - result = convert_to_fp16(model, keep_io_types=True) + def test_blocked_float_can_feed_unconverted_optional_output(self) -> None: + """Optional metadata does not force its FLOAT payload to FP16.""" + model = _build_blocked_float_optional_output_model() - for inp in result.graph.input: - assert inp.type.tensor_type.elem_type == TensorProto.FLOAT - for outp in result.graph.output: - assert outp.type.tensor_type.elem_type == TensorProto.FLOAT + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu"], + ) - def test_keep_io_types_false_converts_io(self) -> None: - """With keep_io_types=False, model I/O becomes FP16.""" - model = _build_simple_fp32_model() - result = convert_to_fp16(model, keep_io_types=False) + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([2.0], dtype=np.float16)}) + assert output.dtype == np.float32 - for inp in result.graph.input: - assert inp.type.tensor_type.elem_type == TensorProto.FLOAT16 - for outp in result.graph.output: - assert outp.type.tensor_type.elem_type == TensorProto.FLOAT16 + def test_converted_float_cannot_feed_unconverted_optional_output( + self, + ) -> None: + """Optional payload declarations must match converted tensor inputs.""" + model = _build_float_optional_output_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, check_type=True, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([2.0], dtype=np.float32)}) + assert output.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "incompatible FP16 types"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_blocked_tensor_cast_cannot_feed_unconverted_optional(self) -> None: + """A late output Cast cannot change an Optional payload silently.""" + model = _build_blocked_tensor_optional_consumer_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + with np.testing.assert_raises_regex(RuntimeError, "container payload declaration"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Abs"], + ) + assert model.SerializeToString() == original + + def test_blocked_float_optional_capture_remains_fp32(self) -> None: + """A blocked child may capture an Optional whose payload stays FP32.""" + model = _build_blocked_float_optional_capture_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["Gelu", "If"], + ) + + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([2.0], dtype=np.float16), + "condition": np.array(True), + }, + ) + assert output.dtype == np.float16 + + def test_empty_optional_blocked_capture_remains_source_less(self) -> None: + """An input-less Optional has no FLOAT producer to convert.""" + model = _build_empty_optional_blocked_capture_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, converted = session.run( + None, + { + "condition": np.array(True), + "x": np.array([2.0], dtype=np.float16), + }, + ) + assert output == np.array(False) + assert converted.dtype == np.float16 + + def test_loop_optional_output_preserves_top_input_precision(self) -> None: + """A graph-bearing producer can pass through unchanged Optional state.""" + model = _build_loop_optional_output_blocked_capture_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, converted = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + "x": np.array([3.0], dtype=np.float16), + }, + ) + assert output == np.array(True) + assert converted.dtype == np.float16 + + def test_always_float_executed_attribute_capture_is_rejected(self) -> None: + """Skipped executed attributes cannot retain a converted capture.""" + model = _build_always_float_function_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + converted, selected = session.run(None, {"condition": np.array(True)}) + assert converted.dtype == np.float32 + assert selected.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_always_float_nested_attribute_capture_is_rejected(self) -> None: + """Function attribute references inside body graphs still execute.""" + model = _build_always_float_nested_function_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + converted, selected = session.run(None, {"condition": np.array(True)}) + assert converted.dtype == np.float32 + assert selected.dtype == np.bool_ + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_always_float_default_attribute_capture_is_rejected(self) -> None: + """Default graph attributes execute when invocation values are absent.""" + model = _build_always_float_default_function_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + converted, selected = session.run(None, {"condition": np.array(True)}) + assert converted.dtype == np.float32 + assert selected.dtype == np.float32 + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + assert model.SerializeToString() == original + + def test_always_float_unused_attribute_capture_is_ignored(self) -> None: + """Only local-function attributes referenced by its body execute.""" + model = _build_always_float_unused_function_capture_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, converted = session.run(None, {"x": np.array([2.0], dtype=np.float16)}) + assert output.dtype == np.float16 + assert converted.dtype == np.float16 + + def test_skipped_scope_duplicate_initializer_is_allowed(self) -> None: + """Initializer name checks visit only scopes ORT traverses.""" + model = _build_skipped_duplicate_initializer_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([2.0], dtype=np.float16)}) + assert output.dtype == np.float16 + + def test_nested_optional_input_blocked_capture_is_rejected(self) -> None: + """A nested FLOAT container input may receive a converted parent value.""" + model = _build_nested_optional_input_blocked_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "x": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_top_optional_input_nested_blocked_capture_is_preserved(self) -> None: + """A top Optional input remains FLOAT when carried into a nested graph.""" + model = _build_top_optional_input_blocked_capture_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + def test_identity_optional_feedback_is_preserved(self) -> None: + """A type-preserving producer may rename unchanged feedback.""" + model = _build_identity_optional_feedback_blocked_capture_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + def test_rewrapped_optional_feedback_conversion_is_rejected(self) -> None: + """Converted tensor intermediates change rewrapped feedback payloads.""" + model = _build_rewrapped_optional_feedback_blocked_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + assert model.SerializeToString() == original + + def test_optional_identity_blocked_capture_is_preserved(self) -> None: + """An unchanged Optional pass-through stays independent of FP16 tensors.""" + model = _build_optional_identity_blocked_capture_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + output, converted = session.run( + None, + { + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + "x": np.array([3.0], dtype=np.float16), + }, + ) + assert output == np.array(True) + assert converted.dtype == np.float16 + + def test_long_optional_identity_chain_is_iterative(self) -> None: + """Long precision-preserving chains do not consume Python stack.""" + model = _build_long_optional_identity_chain_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + def test_free_optional_identity_uses_outer_precision_source(self) -> None: + """Producer tracing resolves free inputs to their lexical owner.""" + model = _build_free_optional_identity_blocked_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "state": np.array(True), + "x": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + assert model.SerializeToString() == original + + def test_optional_loop_feedback_conversion_is_rejected(self) -> None: + """Later loop-carried values must be included in nested input analysis.""" + model = _build_optional_feedback_blocked_capture_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(2, dtype=np.int64), + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + "feedback": np.array([3.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_mispositioned_optional_feedback_is_rejected(self) -> None: + """A same-named output in another state slot is not unchanged feedback.""" + model = _build_mispositioned_optional_feedback_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(2, dtype=np.int64), + "condition": np.array(True), + "optional_a": np.array([1.0], dtype=np.float32), + "optional_b": np.array([2.0], dtype=np.float32), + "feedback": np.array([3.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_scan_output_does_not_shift_optional_feedback_slot(self) -> None: + """Scan outputs cannot change loop-state feedback alignment.""" + model = _build_scan_mispositioned_optional_feedback_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_scan_output_preserves_unchanged_optional_feedback_slot(self) -> None: + """Scan outputs do not hide an unchanged loop-carried binding.""" + model = _build_scan_top_optional_input_blocked_capture_model() + + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["If"], + ) + + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "optional": np.array([2.0], dtype=np.float32), + }, + ) + assert output == np.array(True) + + def test_blocked_float_sequence_input_is_rejected_before_mutation(self) -> None: + """ORT's late tensor Cast cannot consume a sequence value.""" + model = _build_blocked_float_sequence_input_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (length,) = session.run(None, {"sequence": [np.array([1.0], dtype=np.float32)]}) + assert length == 1 + + with np.testing.assert_raises_regex(RuntimeError, "non-tensor"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["SequenceLength"], + ) + assert model.SerializeToString() == original + + def test_overridable_blocked_initializer_is_rejected_before_mutation(self) -> None: + """A converted graph input cannot retain a FLOAT default initializer.""" + model = _build_overridable_blocked_initializer_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + + with np.testing.assert_raises_regex(RuntimeError, "initializer-backed graph input"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Identity"]) + assert model.SerializeToString() == original + + def test_duplicate_float_output_name_is_rejected_before_mutation(self) -> None: + """ORT cannot map repeated kept outputs to two generated Cast aliases.""" + model = _build_duplicate_float_output_name_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + first, second = session.run(None, {"x": np.array([3.0], dtype=np.float32)}) + np.testing.assert_array_equal(first, second) + + with np.testing.assert_raises_regex(RuntimeError, "repeated FLOAT output"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_shared_float_input_output_name_is_rejected_before_mutation(self) -> None: + """ORT overwrites keep-I/O mappings when one name is both input and output.""" + model = _build_shared_float_input_output_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run(None, {"x": np.array([3.0], dtype=np.float32)}) + np.testing.assert_array_equal(output, np.array([3.0], dtype=np.float32)) + + with np.testing.assert_raises_regex(RuntimeError, "both input and output"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_initializer_output_shadowing_is_rejected_before_mutation(self) -> None: + """Nested shadowing is rejected before ORT's global initializer map.""" + model = _build_nested_shadowed_initializer_output_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "duplicate FLOAT initializer names"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_duplicate_non_output_initializers_are_rejected_before_mutation(self) -> None: + """Duplicate FLOAT initializer names are rejected before ORT mutates the graph.""" + model = _build_duplicate_non_output_initializer_name_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "duplicate FLOAT initializer names"): + convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + assert model.SerializeToString() == original + + def test_duplicate_initializers_under_blocked_if_are_not_rejected(self) -> None: + """Duplicate local initializers under blocked nodes are not traversed by ORT.""" + model = _build_blocked_if_duplicate_local_initializer_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_blocked_shadowed_consumers_do_not_make_output_initializer_shared(self) -> None: + """Blocked branch-local values do not consume a top-level output initializer.""" + model = _build_blocked_if_shadowed_output_initializer_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=["If"]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT + assert result.graph.initializer[0].data_type == TensorProto.FLOAT + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + np.testing.assert_array_equal( + session.run(None, {"condition": np.array(True)})[0], + np.array([9.0], dtype=np.float32), + ) + + def test_traversed_shadowed_output_initializer_references_are_rejected(self) -> None: + """Traversed nested refs with the same name hit ORT's global I/O name mapping.""" + model = _build_traversed_shadowed_output_initializer_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "scope-aware"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_input_shadowing_initializer_output_is_rejected_before_mutation( + self, + ) -> None: + """Traversed nested input declarations can hit ORT's global I/O name mapping.""" + model = _build_loop_state_input_shadowing_output_initializer_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "scope-aware"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_value_info_does_not_shadow_initializer_output( + self, + ) -> None: + """A value_info annotation alone does not create a local binding.""" + model = _build_loop_value_info_shadowing_output_initializer_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + checker.check_model(result) + shape_inference.infer_shapes(result, check_type=True, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + same, loop_output = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "loop_state": np.array([2.0], dtype=np.float32), + }, + ) + np.testing.assert_array_equal(same, np.array([9.0], dtype=np.float32)) + np.testing.assert_array_equal(loop_output, np.array([2.0], dtype=np.float32)) + + def test_nested_shadowed_top_level_input_references_are_rejected(self) -> None: + """Top-level input casts are also globally mapped by ORT keep_io_types.""" + model = _build_nested_consumed_initializer_output_with_top_level_input_collision_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "scope-aware"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_output_initializer_matching_kept_top_level_io_is_repaired(self) -> None: + """Nested direct output repair overrides ORT's global top-level I/O skip name.""" + model = _build_nested_output_initializer_with_top_level_io_name_collision_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + [body] = _iter_attribute_graphs(result) + assert body.output[1].type.tensor_type.elem_type == TensorProto.FLOAT16 + assert body.initializer[0].data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_mapped_free_capture_does_not_bind_to_earlier_nested_initializer(self) -> None: + """ORT tracks the mapped top-level capture separately from a nested initializer.""" + model = _build_nested_initializer_then_mapped_free_capture_model() + + result = convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + + assert result.graph.output[-1].type.tensor_type.elem_type == TensorProto.FLOAT + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_mapped_capture_shadowed_by_generated_alias_is_rejected(self) -> None: + """A nested binding must not intercept a generated top-level I/O alias.""" + model = _build_mapped_capture_shadowed_by_generated_alias_model() + original = model.SerializeToString() + + checker.check_model(model) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (output,) = session.run( + None, + { + "x": np.array([3.0], dtype=np.float32), + "condition": np.array(True), + }, + ) + np.testing.assert_array_equal(output, np.array([3.0], dtype=np.float32)) + + with np.testing.assert_raises_regex(RuntimeError, "generated alias"): + convert_to_fp16(model, keep_io_types=True, op_block_list=[]) + assert model.SerializeToString() == original + + def test_nested_blocked_ordinary_name_collision_is_rejected(self) -> None: + """A blocked nested input must retain its local lexical binding.""" + model = _build_nested_blocked_ordinary_name_collision_model() + original = model.SerializeToString() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + session = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (selected,) = session.run( + None, + { + "condition": np.array(True), + "same": np.array(0.0, dtype=np.float32), + }, + ) + assert selected.shape == (1, 3) + + with np.testing.assert_raises_regex(RuntimeError, "global value-info"): + convert_to_fp16( + model, + keep_io_types=False, + op_block_list=["NonMaxSuppression"], + ) + assert model.SerializeToString() == original + + def test_blocked_free_capture_prevents_pure_fp16_initializer_repair(self) -> None: + """Blocked subgraphs can still consume outer initializers at runtime.""" + model = _build_blocked_if_free_capture_output_initializer_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "blocked subgraph"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + assert model.SerializeToString() == original + + def test_blocked_free_capture_of_shadowed_value_does_not_reject_outer_initializer(self) -> None: + """Blocked descendants that capture local shadows do not consume outer initializers.""" + model = _build_blocked_if_shadowed_free_capture_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=["If"]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + [body] = _iter_attribute_graphs(result) + assert body.initializer[0].data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_unconsumed_nested_initializer_outputs_are_converted_to_fp16(self) -> None: + """Unconsumed nested direct output initializers are repaired after conversion.""" + model = _build_nested_initializer_output_model() + model.graph.initializer.append( + numpy_helper.from_array(np.array([1.0], dtype=np.float16), "top_level_fp16") + ) + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT16 + for graph in _iter_attribute_graphs(result): + assert all( + output.type.tensor_type.elem_type == TensorProto.FLOAT16 for output in graph.output + ) + assert all( + initializer.data_type == TensorProto.FLOAT16 for initializer in graph.initializer + ) + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + + def test_blocked_fp32_consumer_does_not_round_trip_through_fp16(self) -> None: + """Blocked consumers are rejected instead of silently losing precision.""" + model = _build_blocked_initializer_consumer_model() + original = model.SerializeToString() + + with np.testing.assert_raises_regex(RuntimeError, "initializer metadata"): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Identity"]) + assert model.SerializeToString() == original + + def test_nested_local_input_cannot_change_outer_initializer_precision(self) -> None: + """ORT's global initializer tracker must not treat local shadows as consumers.""" + model = _build_blocked_initializer_with_nested_input_shadow_model() + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + original = model.SerializeToString() + + with np.testing.assert_raises_regex( + RuntimeError, "initializer tracking is not scope-aware" + ): + convert_to_fp16(model, keep_io_types=False, op_block_list=["Abs"]) + assert model.SerializeToString() == original + + def test_harmless_initializer_consumer_misattribution_is_allowed(self) -> None: + """A false FP16 consumer is harmless when the initializer already converts.""" + model = _build_blocked_initializer_with_nested_input_shadow_model() + model.graph.node[0].op_type = "Relu" + model.graph.node[0].name = "top_relu" + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + blocked_output, loop_output = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "loop_state": np.array([2.0], dtype=np.float16), + }, + ) + np.testing.assert_allclose(blocked_output, np.array([1.0], dtype=np.float16), atol=1e-3) + np.testing.assert_array_equal(loop_output, np.array([2.0], dtype=np.float16)) + + def test_unused_initializer_shadow_misattribution_is_allowed(self) -> None: + """An unobservable initializer may be converted by a local name collision.""" + model = _build_unused_initializer_with_nested_input_shadow_model() + + checker.check_model(model) + shape_inference.infer_shapes(model, strict_mode=True) + result = convert_to_fp16( + model, + keep_io_types=False, + op_block_list=[], + ) + + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + (loop_output,) = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "loop_state": np.array([2.0], dtype=np.float16), + }, + ) + np.testing.assert_array_equal(loop_output, np.array([2.0], dtype=np.float16)) + + def test_output_repair_makes_initializer_misattribution_harmless(self) -> None: + """A direct output repair can independently require initializer conversion.""" + model = _build_loop_state_input_shadowing_output_initializer_model() + [body] = _iter_attribute_graphs(model) + del body.initializer[:] + body.node.append(helper.make_node("Relu", ["same"], ["body_state_out"], name="body_state")) + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.initializer[0].data_type == TensorProto.FLOAT16 + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + same, loop_output = session.run( + None, + { + "trip_count": np.array(1, dtype=np.int64), + "condition": np.array(True), + "loop_state": np.array([2.0], dtype=np.float16), + }, + ) + np.testing.assert_array_equal(same, np.array([9.0], dtype=np.float16)) + np.testing.assert_array_equal(loop_output, np.array([2.0], dtype=np.float16)) + + def test_later_nested_float_initializer_does_not_rebind_earlier_non_float_input( + self, + ) -> None: + """Initializer tracking follows ORT registration order across graph scopes.""" + model = _build_non_float_initializer_before_nested_float_initializer_model() + + result = convert_to_fp16(model, keep_io_types=False, op_block_list=[]) + + assert result.graph.initializer[0].data_type == TensorProto.INT64 + same = next( + initializer + for graph in _iter_attribute_graphs(result) + for initializer in graph.initializer + if initializer.name == "same" + ) + assert same.data_type == TensorProto.FLOAT16 + checker.check_model(result) + session = ort.InferenceSession( + result.SerializeToString(), providers=["CPUExecutionProvider"] + ) + integer_output, nested_output = session.run(None, {"condition": np.array(True)}) + np.testing.assert_array_equal(integer_output, np.array([7], dtype=np.int64)) + np.testing.assert_array_equal(nested_output, np.array([1.5], dtype=np.float16)) + + def test_initializer_output_repair_preserves_unrelated_casts(self) -> None: + """Repair removes only ORT's orphan output Cast, not user graph Casts.""" + model = _build_initializer_backed_output_model() + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + cast_output = helper.make_tensor_value_info("cast_output", TensorProto.INT32, [1]) + model.graph.input.append(x) + model.graph.output.append(cast_output) + model.graph.node.append( + helper.make_node( + "Cast", + ["x"], + ["cast_output"], + name="user_cast", + to=TensorProto.INT32, + ) + ) + + result = convert_to_fp16(model, keep_io_types=True) + + checker.check_model(result) + assert any(node.name == "user_cast" for node in result.graph.node) def test_preserves_model_structure(self) -> None: """FP16 conversion preserves graph structure (node count diff ≤ 2).""" @@ -110,8 +5756,8 @@ def test_none_op_block_list_uses_ort_defaults(self) -> None: result = convert_to_fp16(model, op_block_list=None) assert result is not None - def test_skips_already_fp16_model(self) -> None: - """If all floating-point initializers are already FP16, conversion is skipped.""" + def test_preserves_already_fp16_model_without_casts(self) -> None: + """Already-FP16 graph I/O and initializers remain FP16 without extra Casts.""" # Build a model with FP16 initializers directly x = helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 4]) out = helper.make_tensor_value_info("out", TensorProto.FLOAT16, [1, 4]) @@ -121,15 +5767,21 @@ def test_skips_already_fp16_model(self) -> None: graph = helper.make_graph([add], "fp16_model", [x], [out], [weight]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) - original_nodes = len(model.graph.node) result = convert_to_fp16(model) - # Should return the same model unchanged (no Cast nodes inserted) - assert len(result.graph.node) == original_nodes - assert result is model + assert all( + value.type.tensor_type.elem_type == TensorProto.FLOAT16 + for value in (*result.graph.input, *result.graph.output) + ) + assert all( + initializer.data_type == TensorProto.FLOAT16 for initializer in result.graph.initializer + ) + assert all(node.op_type != "Cast" for node in result.graph.node) + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True) - def test_skips_fp16_model_with_int_initializers(self) -> None: - """FP16 model with non-float initializers (e.g. INT64 shapes) should still skip.""" + def test_preserves_fp16_model_with_int_initializers_without_casts(self) -> None: + """FP16 graph with INT64 shape initializers remains FP16 without extra Casts.""" x = helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 4]) out = helper.make_tensor_value_info("out", TensorProto.FLOAT16, [1, 4]) weight_data = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float16) @@ -140,8 +5792,16 @@ def test_skips_fp16_model_with_int_initializers(self) -> None: graph = helper.make_graph([add], "fp16_mixed", [x], [out], [weight, shape_tensor]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) - original_nodes = len(model.graph.node) result = convert_to_fp16(model) - assert len(result.graph.node) == original_nodes - assert result is model + assert all( + value.type.tensor_type.elem_type == TensorProto.FLOAT16 + for value in (*result.graph.input, *result.graph.output) + ) + assert any( + initializer.name == "shape" and initializer.data_type == TensorProto.INT64 + for initializer in result.graph.initializer + ) + assert all(node.op_type != "Cast" for node in result.graph.node) + checker.check_model(result) + shape_inference.infer_shapes(result, strict_mode=True)