From de3a34c519f16918f5bf84c4c94bdaf18b7240e5 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:13:35 +0000 Subject: [PATCH] fix(onnx): consistent types on If/Loop/Scan subgraphs during FP16/BF16 conversion ONNX `--high_precision_dtype fp16` (and the autocast PrecisionConverter) used to blindly convert every control-flow subgraph initializer to the parent node's precision. This produced inconsistent types and broke shape inference / TensorRT strongly-typed parsing on models with control-flow subgraphs: - A `Gemm` inside an `If` branch reading an outer-scope activation got fp16 weights but an fp32 activation -> "B has inconsistent type". - `Resize` `scales` (which must stay fp32) was converted to fp16 -> "ParseData type mismatch". The converter now keeps a subgraph node in low precision only when all of its float inputs are subgraph initializers eligible for low precision; any node with a float activation/outer-scope input, or an input that must remain high precision per the ONNX spec, stays in high precision so each node's inputs share one precision. Float outer-scope captures and low->high precision boundaries inside a subgraph are reconciled with `Cast` nodes, then the subgraph is topologically re-sorted so each cast lands after its producer and before its consumers (no assumption that the incoming subgraph is sorted). Captured-tensor `value_info` is synced to the capturing tensor's real precision, control-flow node outputs are treated as high precision (their subgraph bodies are kept high precision), and `Constant` folding refreshes (creating one when missing) the constant's `value_info` so a same-type-constrained consumer (e.g. `Greater`) is not left with a stale, conflicting type. Known limitation: a low-precision `Loop`/`Scan` whose body carries a float loop-carried or scan-input state var is not yet reconciled (the body's formal inputs are treated as high precision while the main graph casts the parent's inputs to low precision); such bodies are tracked separately. `If` branches and high-precision `Loop`/`Scan` bodies are handled and covered by regression tests. Validated on both reported models: both convert successfully, pass strict `onnx.shape_inference` type checking, load in ONNX Runtime, and match the FP32 reference outputs within FP16 tolerance. Fixes bug 6058841 Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/onnx/autocast/precisionconverter.py | 281 ++++++++++++++-- modelopt/onnx/utils.py | 19 +- .../onnx/autocast/test_precisionconverter.py | 303 ++++++++++++++++++ 4 files changed, 576 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aadd0f0563f..151ea49584c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -64,6 +64,7 @@ Experimental - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3. - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export. +- Fix ONNX FP16/BF16 conversion (``--high_precision_dtype fp16``) producing inconsistent tensor types on models with control-flow subgraphs (e.g. a ``Gemm`` inside an ``If`` branch reading an outer-scope activation alongside converted weights, or ``Resize`` ``scales`` that must stay FP32). Subgraph nodes now only run in low precision when all their float inputs are subgraph initializers; outer-scope captures and low-to-high-precision boundaries inside subgraphs are reconciled with ``Cast`` nodes, and ``Constant`` folding refreshes the constant's ``value_info`` so strongly-typed parsers (TensorRT) no longer reject the model. This is a behavioral change: previously a low-precision control-flow parent converted *every* float subgraph initializer, so a weight inside a branch that also reads an outer-scope FP32 activation could become FP16; such weights now stay FP32 so each node's inputs keep a single precision. 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index da45cd3ecde..2515597ea2b 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -693,11 +693,31 @@ def _convert_initializers( def _convert_initializers_recursive( self, low_precision_nodes: list[str], high_precision_nodes: list[str] ) -> None: - """Convert initializers in main graph and all subgraphs to appropriate precision. - - For the main graph, uses sophisticated consumer tracking to determine precision. - For subgraphs, inherits precision from the parent control flow node and converts - all initializers to that precision (no runtime casts). + """Convert initializers in the main graph and reconcile precision inside all subgraphs. + + For the main graph, uses consumer tracking to determine each initializer's precision + (see :meth:`_convert_initializers`). + + Control-flow subgraphs (e.g. If/Loop/Scan bodies) do not get the activation-bracketing casts + that the main graph uses, so a subgraph node is only converted to low precision when *all* of + its float inputs are subgraph initializers that may be in low precision (i.e. the node consumes + no float activation/outer-scope tensor and none of its inputs must stay high precision per the + ONNX spec, such as ``Resize`` ``scales``). Such a node's initializers are converted to low + precision; every other subgraph node and its initializers are kept in high precision so that + each node's inputs share a single precision. Float tensors captured from the enclosing scope, + and the outputs of any low-precision subgraph node feeding a high-precision one, are reconciled + with a ``Cast`` inserted inside the subgraph. + + This is a behavioral change: previously a low-precision control-flow parent converted *every* + float subgraph initializer, so a weight inside a branch that also reads an outer-scope float + activation could end up low precision while the activation stayed high precision (the + inconsistent-type bug this gating fixes). + + Known limitation: a low-precision ``Loop``/``Scan`` whose body carries a float loop-carried + or scan-input state var is not yet reconciled here (its formal inputs are treated as high + precision while the main graph casts the parent's inputs to low precision). Such bodies are + tracked separately; this method currently targets ``If`` branches and high-precision + ``Loop``/``Scan`` bodies. Args: low_precision_nodes: List of node names in main graph that are low precision. @@ -706,42 +726,251 @@ def _convert_initializers_recursive( # Convert main graph initializers with full consumer tracking self._convert_initializers(low_precision_nodes, high_precision_nodes) - # Convert subgraph initializers - walk all subgraphs and convert based on parent node precision + # Precompute, for each main-graph activation, the (raw) precision it has after main-graph + # conversion: a subgraph that captures it sees this raw precision, because the main-graph + # cast-up only rewires main-graph consumers (not subgraph captures). low_precision_nodes_set = set(low_precision_nodes) + main_producer_precision: dict[str, int] = {} + for node in self.model.graph.node: + # Control-flow nodes (If/Loop/Scan) execute a subgraph that is kept in high precision, so + # their outputs are high precision regardless of the node's low/high classification. + is_control_flow = any( + attr.type in (onnx.AttributeProto.GRAPH, onnx.AttributeProto.GRAPHS) + for attr in node.attribute + ) + producer_type = ( + self.low_precision_type.onnx_type + if node.name in low_precision_nodes_set and not is_control_flow + else self.high_precision_type.onnx_type + ) + for output_name in node.output: + main_producer_precision[output_name] = producer_type + + def _capture_type(name: str) -> int | None: + """Return the float precision (ONNX type) of an outer-scope tensor seen in a subgraph. + + Float-ness is read from the (pre-conversion) type info, since it is invariant under + precision conversion; the precision is then taken from the producing main-graph node. + Network inputs use their declared type in ``value_info_map`` (already updated in place + when ``keep_io_types`` is False). Returns None for non-float or unknown tensors. + """ + if name in self.initializer_map: + base_type = self.initializer_map[name].data_type + elif name in self.value_info_map: + base_type = self.value_info_map[name].type.tensor_type.elem_type + else: + return None + if base_type not in ONNX_TYPES: + return None + return main_producer_precision.get(name, base_type) def _convert_subgraph_callback( graph: onnx.GraphProto, parent: onnx.NodeProto, is_subgraph: bool ) -> None: if not is_subgraph or parent is None: return + parent_is_low_precision = parent.name in low_precision_nodes_set + self._convert_subgraph_precision(graph, parent_is_low_precision, _capture_type) + + utils.walk_subgraphs_recursive(self.model.graph, _convert_subgraph_callback) + + def _precision_type_from_onnx_type(self, onnx_type: int) -> PrecisionTypes | None: + """Return the converter precision metadata for a supported ONNX float type.""" + return next((p for p in PRECISION_MAP.values() if p.onnx_type == onnx_type), None) + + def _convert_subgraph_precision( + self, subgraph: onnx.GraphProto, parent_is_low_precision: bool, capture_type_fn + ) -> None: + """Convert a single control-flow subgraph to a consistent precision. - # Inherit precision from parent control flow node - target_type = ( - self.low_precision_type - if parent.name in low_precision_nodes_set - else self.high_precision_type + See :meth:`_convert_initializers_recursive` for the conversion policy. + + Args: + subgraph: The subgraph (e.g. an If branch or a Loop/Scan body) to convert. + parent_is_low_precision: Whether the parent control-flow node is low precision. + capture_type_fn: Maps an outer-scope tensor name to its float ONNX type (or None). + """ + target = self.low_precision_type + high = self.high_precision_type + + local_inits = {init.name: init for init in subgraph.initializer} + local_produced = {out for node in subgraph.node for out in node.output} + formal_inputs = {inp.name for inp in subgraph.input} + # Original (pre-conversion) element types of subgraph-local tensors. Whether a tensor is + # float is invariant under fp32<->fp16 conversion, so this is a reliable "is this float?" + # source that prevents casting non-float tensors (e.g. int axes/indices/shapes). + local_elem_type = {vi.name: vi.type.tensor_type.elem_type for vi in subgraph.value_info} + local_elem_type.update({vi.name: vi.type.tensor_type.elem_type for vi in subgraph.output}) + + # Map of tensor name -> list of (node, input index) consumers within this subgraph. + consumers: dict[str, list[InputIndexTracker]] = defaultdict(list) + for node in subgraph.node: + for idx, input_name in enumerate(node.input): + if input_name: + consumers[input_name].append(InputIndexTracker(node=node, node_index=idx)) + + def _is_low_precision_eligible_init(node: onnx.NodeProto, input_name: str) -> bool: + init = local_inits.get(input_name) + return ( + init is not None + and init.data_type in ONNX_TYPES + and not self._should_skip_low_precision_input_conversion(node, input_name) ) - # Convert all float initializers to target precision - for init in graph.initializer: - if init.data_type not in ONNX_TYPES or init.data_type == target_type.onnx_type: - continue + def _known_elem_type(input_name: str) -> int | None: + """Best-effort (pre-conversion) element type of a tensor visible here, or None.""" + if input_name in local_inits: + return local_inits[input_name].data_type + if input_name in local_elem_type: + return local_elem_type[input_name] + if input_name in self.value_info_map: + return self.value_info_map[input_name].type.tensor_type.elem_type + if input_name in self.initializer_map: + return self.initializer_map[input_name].data_type + return None - from_type = self._precision_type_from_onnx_type(init.data_type) - if from_type is None: - logger.debug( - f"Skipping subgraph initializer {init.name} with unsupported type {init.data_type}" - ) + # 1. Classify each subgraph node. A node is converted to low precision only if it has at least + # one low-precision-eligible float initializer input and every other input is a tensor we + # know is non-float (e.g. int axes/indices). Any float activation/outer-scope input (or an + # input of unknown type) keeps the node in high precision, since no bracketing casts are + # inserted around subgraph activations. + node_is_low: dict[str, bool] = {} + for node in subgraph.node: + low = parent_is_low_precision and ( + node.op_type not in self.op_types_not_supported_in_low_precision + ) + has_low_precision_init = False + if low: + for input_name in node.input: + if not input_name: + continue + if _is_low_precision_eligible_init(node, input_name): + has_low_precision_init = True + continue + elem_type = _known_elem_type(input_name) + if elem_type is None or elem_type in ONNX_TYPES: + # Float or unknown non-initializer input: keep the node in high precision. + low = False + break + node_is_low[node.name] = low and has_low_precision_init + + # 2. Convert the initializers consumed by low-precision nodes to low precision. An initializer + # shared by low- and high-precision nodes is duplicated so each consumer keeps one precision. + for init in list(subgraph.initializer): + if init.data_type not in ONNX_TYPES: + continue + from_type = self._precision_type_from_onnx_type(init.data_type) + assert from_type is not None + low_consumers: list[InputIndexTracker] = [] + high_consumers: list[InputIndexTracker] = [] + for c in consumers.get(init.name, []): + if node_is_low.get(c.node.name) and _is_low_precision_eligible_init( + c.node, init.name + ): + low_consumers.append(c) + else: + high_consumers.append(c) + + if not low_consumers: + # Keep in high precision (covers Resize-scales-style inputs and unused initializers). + if init.data_type != high.onnx_type: + init.CopyFrom(self._convert_initializer_data(init, from_type, high)) + elif not high_consumers: + # Convert the single-precision initializer in place. + if init.data_type != target.onnx_type: + init.CopyFrom(self._convert_initializer_data(init, from_type, target)) + else: + # Shared: keep the original high precision and add a low-precision duplicate. + low_init = self._convert_initializer_data(init, from_type, target) + low_init.name = f"{init.name}_{target.str_short}" + subgraph.initializer.extend([low_init]) + for consumer in low_consumers: + consumer.node.input[consumer.node_index] = low_init.name + if init.data_type != high.onnx_type: + init.CopyFrom(self._convert_initializer_data(init, from_type, high)) + + # 3. Reconcile float tensors whose precision does not match the consuming node: outer-scope + # captures and the outputs of low-precision nodes feeding high-precision ones. Casts are + # collected first and inserted afterwards to avoid mutating the node list while iterating. + local_produced_low = { + out for node in subgraph.node if node_is_low.get(node.name) for out in node.output + } + + def _current_float_type(input_name: str) -> int | None: + """Float precision (ONNX type) of a subgraph input tensor, or None if it is non-float. + + Non-float tensors (int indices/axes/shapes, bool conditions) must never be cast. + """ + if input_name in local_produced: + elem_type = local_elem_type.get(input_name) + if elem_type is not None and elem_type not in ONNX_TYPES: + return None # known non-float subgraph activation + if input_name in local_produced_low: + return target.onnx_type # output of a converted low-precision node + return high.onnx_type if elem_type in ONNX_TYPES else None + if input_name in formal_inputs: + elem_type = local_elem_type.get(input_name) + return high.onnx_type if elem_type in ONNX_TYPES else None + return capture_type_fn(input_name) # outer-scope capture (None if non-float) + + # (tensor name, target onnx type) -> cast output name + casts_to_insert: dict[tuple[str, int], str] = {} + rewrites: list[tuple[InputIndexTracker, str]] = [] + # Float outer-scope captures and the (current) precision they have in the enclosing scope. + captured_types: dict[str, int] = {} + for node in subgraph.node: + node_low = node_is_low.get(node.name, False) + for idx, input_name in enumerate(node.input): + if not input_name or input_name in local_inits: + continue + current = _current_float_type(input_name) + if current is None: + continue # non-float tensor: never cast + if input_name not in local_produced and input_name not in formal_inputs: + captured_types[input_name] = current + desired = ( + target.onnx_type + if node_low + and not self._should_skip_low_precision_input_conversion(node, input_name) + else high.onnx_type + ) + if current == desired: continue - new_init = self._convert_initializer_data(init, from_type, target_type) - init.CopyFrom(new_init) + key = (input_name, desired) + if key not in casts_to_insert: + desired_type = self._precision_type_from_onnx_type(desired) + assert desired_type is not None + short = desired_type.str_short + casts_to_insert[key] = f"{input_name}_subgraph_cast_to_{short}" + rewrites.append( + (InputIndexTracker(node=node, node_index=idx), casts_to_insert[key]) + ) - utils.walk_subgraphs_recursive(self.model.graph, _convert_subgraph_callback) + # Sync any preserved outer-scope value_info inside the subgraph with the capture's current + # main-graph precision. Otherwise a stale type (e.g. fp32 for a tensor the main graph now + # produces in fp16) makes strongly-typed parsers (and ORT) reject the If subgraph. + for vi in subgraph.value_info: + if vi.name in captured_types and vi.type.HasField("tensor_type"): + vi.type.tensor_type.elem_type = captured_types[vi.name] - def _precision_type_from_onnx_type(self, onnx_type: int) -> PrecisionTypes | None: - """Return the converter precision metadata for a supported ONNX float type.""" - return next((p for p in PRECISION_MAP.values() if p.onnx_type == onnx_type), None) + if not casts_to_insert: + return + + for tracker, cast_output in rewrites: + tracker.node.input[tracker.node_index] = cast_output + + # Insert the Cast nodes and topologically re-sort the subgraph so each cast lands after its + # producer and before its consumers. The sort makes no assumption about the incoming node + # order, so a hand-built (unsorted) subgraph is handled correctly too. + cast_nodes = [ + helper.make_node( + "Cast", inputs=[input_name], outputs=[cast_output], to=desired, name=cast_output + ) + for (input_name, desired), cast_output in casts_to_insert.items() + ] + subgraph.node.extend(cast_nodes) + onnx_utils.topologically_sort_graph_nodes(subgraph) def _convert_initializer_data( self, diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index bc37c4a9333..bfa0bdc25da 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1502,7 +1502,9 @@ def _is_foldable_constant_cast_pattern(model: onnx.ModelProto, node: onnx.NodePr return False -def _convert_constant_values(constant_node: onnx.NodeProto, cast_node: onnx.NodeProto) -> None: +def _convert_constant_values( + onnx_model: onnx.ModelProto, constant_node: onnx.NodeProto, cast_node: onnx.NodeProto +) -> None: """Convert the Constant node's values to the Cast node's target type.""" cast_to_type = get_cast_to_type(cast_node) for attr in constant_node.attribute: @@ -1529,6 +1531,19 @@ def _convert_constant_values(constant_node: onnx.NodeProto, cast_node: onnx.Node new_tensor = onnx.numpy_helper.from_array(new_array, attr.t.name) attr.t.CopyFrom(new_tensor) + + const_output = constant_node.output[0] + const_vi = next( + (vi for vi in onnx_model.graph.value_info if vi.name == const_output), None + ) + if const_vi is not None: + const_vi.type.tensor_type.elem_type = cast_to_type + else: + onnx_model.graph.value_info.append( + onnx.helper.make_tensor_value_info( + const_output, cast_to_type, list(new_tensor.dims) + ) + ) break @@ -1570,7 +1585,7 @@ def remove_redundant_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: cast_producers = get_producer_nodes(onnx_model, node.input[0]) assert len(cast_producers) == 1 and cast_producers[0].op_type == "Constant" constant_producer = cast_producers[0] - _convert_constant_values(constant_producer, node) + _convert_constant_values(onnx_model, constant_producer, node) _bypass_cast_node(onnx_model, node) logger.debug(f"Found foldable Constant->Cast pattern, removing {node.name}") diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index 9872e0c4f2b..4f2871bc5ed 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -1955,3 +1955,306 @@ def test_if_subgraph_outer_scope_type_preservation( assert len(else_x_info) > 0, "X value_info should be preserved in else branch" assert then_x_info[0].type.tensor_type.elem_type != onnx.TensorProto.UNDEFINED assert else_x_info[0].type.tensor_type.elem_type != onnx.TensorProto.UNDEFINED + + +#################################################################################################### +# Regression tests for bug 6058841: inconsistent tensor types on control-flow If nodes during +# ONNX FP16/BF16 conversion. +# +# Converting a model with control-flow subgraphs to FP16 used to blindly convert every subgraph +# initializer to the parent node's precision, which broke models where a subgraph node also consumes +# a float activation/outer-scope tensor (e.g. a Gemm reading a network input) or whose inputs must +# stay in high precision per the ONNX spec (e.g. Resize 'scales'). +#################################################################################################### +@pytest.fixture +def model_if_subgraph_gemm_outer_input(): + """If branches with a Gemm consuming an outer-scope input plus subgraph weight initializers.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3]) + + def _branch(name): + w = numpy_helper.from_array(np.random.randn(4, 3).astype(np.float32), name=f"w_{name}") + b = numpy_helper.from_array(np.random.randn(3).astype(np.float32), name=f"b_{name}") + out = helper.make_tensor_value_info(f"{name}_out", TensorProto.FLOAT, [1, 3]) + gemm = helper.make_node( + "Gemm", ["X", f"w_{name}", f"b_{name}"], [f"{name}_out"], name=f"{name}_gemm" + ) + return helper.make_graph([gemm], f"{name}_branch", [], [out], [w, b]) + + if_node = helper.make_node( + "If", + ["condition"], + ["Y"], + name="if_node", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + main_graph = helper.make_graph([if_node], "model_if_gemm", [x, condition], [y]) + model = helper.make_model(main_graph, producer_name="model_if_gemm") + model.opset_import[0].version = 20 + model.ir_version = 10 + onnx.checker.check_model(model) + return setup_mappings(model) + + +@pytest.mark.parametrize("keep_io_types", [True, False]) +@pytest.mark.parametrize("low_precision_type", ["fp16", "bf16"]) +@pytest.mark.parametrize("use_standalone_type_inference", [True, False]) +def test_if_subgraph_gemm_with_outer_scope_input( + model_if_subgraph_gemm_outer_input, + keep_io_types, + low_precision_type, + use_standalone_type_inference, +): + """A Gemm inside an If branch consuming an outer-scope input must not end up with fp16 weights + feeding alongside an fp32 activation (regression test for bug 6058841).""" + model, value_info_map, initializer_map, node_to_init_map = model_if_subgraph_gemm_outer_input + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=keep_io_types, + low_precision_type=low_precision_type, + use_standalone_type_inference=use_standalone_type_inference, + ) + converted_model = converter.convert(high_precision_nodes=[], low_precision_nodes=["if_node"]) + onnx.checker.check_model(converted_model) + # Strict type checking must pass; this is what failed before the fix. + onnx.shape_inference.infer_shapes(converted_model, strict_mode=True, check_type=True) + + +@pytest.fixture +def model_if_subgraph_resize(): + """If branches containing a Resize whose 'roi'/'scales' inputs must remain in high precision.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3, 8, 8]) + condition = helper.make_tensor_value_info("condition", TensorProto.BOOL, []) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3, 16, 16]) + + def _branch(name): + roi = numpy_helper.from_array(np.array([], dtype=np.float32), name=f"roi_{name}") + scales = numpy_helper.from_array( + np.array([1.0, 1.0, 2.0, 2.0], dtype=np.float32), name=f"scales_{name}" + ) + out = helper.make_tensor_value_info(f"{name}_out", TensorProto.FLOAT, [1, 3, 16, 16]) + resize = helper.make_node( + "Resize", + ["X", f"roi_{name}", f"scales_{name}"], + [f"{name}_out"], + name=f"{name}_resize", + mode="nearest", + ) + return helper.make_graph([resize], f"{name}_branch", [], [out], [roi, scales]) + + if_node = helper.make_node( + "If", + ["condition"], + ["Y"], + name="if_node", + then_branch=_branch("then"), + else_branch=_branch("else"), + ) + main_graph = helper.make_graph([if_node], "model_if_resize", [x, condition], [y]) + model = helper.make_model(main_graph, producer_name="model_if_resize") + model.opset_import[0].version = 20 + model.ir_version = 10 + onnx.checker.check_model(model) + return setup_mappings(model) + + +@pytest.mark.parametrize("low_precision_type", ["fp16", "bf16"]) +@pytest.mark.parametrize("use_standalone_type_inference", [True, False]) +def test_if_subgraph_resize_scales_stay_high_precision( + model_if_subgraph_resize, low_precision_type, use_standalone_type_inference +): + """Resize 'scales' inside an If branch must remain FP32 (regression test for bug 6058841).""" + model, value_info_map, initializer_map, node_to_init_map = model_if_subgraph_resize + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + low_precision_type=low_precision_type, + use_standalone_type_inference=use_standalone_type_inference, + ) + converted_model = converter.convert(high_precision_nodes=[], low_precision_nodes=["if_node"]) + onnx.checker.check_model(converted_model) + onnx.shape_inference.infer_shapes(converted_model, strict_mode=True, check_type=True) + + # The 'scales' (and 'roi') initializers must stay FP32 in both branches. + if_node = next(n for n in converted_model.graph.node if n.op_type == "If") + for attr in if_node.attribute: + if attr.type == onnx.AttributeProto.GRAPH: + scales = [init for init in attr.g.initializer if init.name.startswith("scales_")] + assert scales, "scales initializer should be present in the branch" + for init in scales: + assert init.data_type == TensorProto.FLOAT, ( + f"Resize scales must remain FP32, but '{init.name}' is {init.data_type}" + ) + + +@pytest.fixture +def model_chained_if_capture(): + """Two chained If nodes; the second's subgraph captures the first If node's output.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + cond1 = helper.make_tensor_value_info("cond1", TensorProto.BOOL, []) + cond2 = helper.make_tensor_value_info("cond2", TensorProto.BOOL, []) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 3]) + + def _gemm_branch(name, data, k, n): + w = numpy_helper.from_array(np.random.randn(k, n).astype(np.float32), name=f"w_{name}") + out = helper.make_tensor_value_info(f"{name}_out", TensorProto.FLOAT, [1, n]) + gemm = helper.make_node("Gemm", [data, f"w_{name}"], [f"{name}_out"], name=f"{name}_gemm") + return helper.make_graph([gemm], f"{name}_branch", [], [out], [w]) + + if1 = helper.make_node( + "If", + ["cond1"], + ["mid"], + name="if1", + then_branch=_gemm_branch("then1", "X", 4, 3), + else_branch=_gemm_branch("else1", "X", 4, 3), + ) + if2 = helper.make_node( + "If", + ["cond2"], + ["Y"], + name="if2", + then_branch=_gemm_branch("then2", "mid", 3, 3), + else_branch=_gemm_branch("else2", "mid", 3, 3), + ) + main_graph = helper.make_graph([if1, if2], "chained_if", [x, cond1, cond2], [y]) + model = helper.make_model(main_graph, producer_name="chained_if") + model.opset_import[0].version = 20 + model.ir_version = 10 + onnx.checker.check_model(model) + return setup_mappings(model) + + +@pytest.mark.parametrize("low_precision_type", ["fp16", "bf16"]) +def test_chained_if_subgraph_capture(model_chained_if_capture, low_precision_type): + """An If subgraph capturing another control-flow node's output must reconcile its precision + (regression test for bug 6058841; an If subgraph capturing another If node's output).""" + model, value_info_map, initializer_map, node_to_init_map = model_chained_if_capture + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + low_precision_type=low_precision_type, + ) + converted_model = converter.convert(high_precision_nodes=[], low_precision_nodes=["if1", "if2"]) + onnx.checker.check_model(converted_model) + onnx.shape_inference.infer_shapes(converted_model, strict_mode=True, check_type=True) + + +@pytest.mark.parametrize("low_precision_type", ["fp16", "bf16"]) +def test_constant_cast_fold_refreshes_value_info(low_precision_type): + """Folding a Constant->Cast must refresh the constant's value_info, otherwise a same-type + constrained consumer (e.g. Greater) sees a stale, conflicting type and strict type inference + fails (regression test for bug 6058841; Constant feeding a same-type-constrained Greater).""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [4]) + y = helper.make_tensor_value_info("Y", TensorProto.BOOL, [4]) + w = numpy_helper.from_array(np.ones([4], dtype=np.float32), name="w") + nodes = [ + helper.make_node("Mul", ["X", "w"], ["m0"], name="mul0"), + helper.make_node( + "Constant", + [], + ["c0"], + name="const0", + value=numpy_helper.from_array(np.array([0.5, 0.5, 0.5, 0.5], dtype=np.float32), "cv"), + ), + helper.make_node("Greater", ["m0", "c0"], ["Y"], name="greater0"), + ] + graph = helper.make_graph(nodes, "const_greater", [x], [y], [w]) + model = helper.make_model(graph, producer_name="const_greater") + model.opset_import[0].version = 20 + model.ir_version = 10 + onnx.checker.check_model(model) + model, value_info_map, initializer_map, node_to_init_map = setup_mappings(model) + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + low_precision_type=low_precision_type, + ) + converted_model = converter.convert( + high_precision_nodes=[], low_precision_nodes=["mul0", "greater0"] + ) + onnx.checker.check_model(converted_model) + onnx.shape_inference.infer_shapes(converted_model, strict_mode=True, check_type=True) + + +@pytest.fixture +def model_loop_subgraph_capture(): + """A Loop whose body captures a low-precision outer-scope activation. + + A main-graph ``Mul`` runs in low precision and produces ``pre``; the high-precision Loop body + reads ``pre`` (an outer-scope capture) alongside its float loop-carried state var, so the body + must reconcile the captured tensor's precision with a ``Cast``. + """ + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 3]) + trip_count = helper.make_tensor_value_info("M", TensorProto.INT64, []) + cond = helper.make_tensor_value_info("cond", TensorProto.BOOL, []) + acc_init = helper.make_tensor_value_info("acc_init", TensorProto.FLOAT, [1, 3]) + acc_final = helper.make_tensor_value_info("acc_final", TensorProto.FLOAT, [1, 3]) + + iter_num = helper.make_tensor_value_info("iter_num", TensorProto.INT64, []) + cond_in = helper.make_tensor_value_info("cond_in", TensorProto.BOOL, []) + acc_in = helper.make_tensor_value_info("acc_in", TensorProto.FLOAT, [1, 3]) + cond_out = helper.make_tensor_value_info("cond_out", TensorProto.BOOL, []) + acc_out = helper.make_tensor_value_info("acc_out", TensorProto.FLOAT, [1, 3]) + body = helper.make_graph( + [ + helper.make_node("Add", ["acc_in", "pre"], ["acc_out"], name="body_acc"), + helper.make_node("Identity", ["cond_in"], ["cond_out"], name="body_cond"), + ], + "loop_body", + [iter_num, cond_in, acc_in], + [cond_out, acc_out], + ) + + scale = numpy_helper.from_array(np.ones((1, 3), dtype=np.float32), name="scale") + pre = helper.make_node("Mul", ["X", "scale"], ["pre"], name="pre_mul") + loop = helper.make_node( + "Loop", ["M", "cond", "acc_init"], ["acc_final"], name="loop_node", body=body + ) + main_graph = helper.make_graph( + [pre, loop], "model_loop", [x, trip_count, cond, acc_init], [acc_final], [scale] + ) + model = helper.make_model(main_graph, producer_name="model_loop") + model.opset_import[0].version = 20 + model.ir_version = 10 + onnx.checker.check_model(model) + return setup_mappings(model) + + +@pytest.mark.parametrize("keep_io_types", [True, False]) +@pytest.mark.parametrize("low_precision_type", ["fp16", "bf16"]) +def test_loop_subgraph_high_precision_capture( + model_loop_subgraph_capture, keep_io_types, low_precision_type +): + """A high-precision Loop body capturing a low-precision outer-scope activation must reconcile it + with a ``Cast`` so the subgraph stays a single precision (regression test for bug 6058841; + control-flow subgraph capture). Low-precision Loop/Scan bodies with float loop-carried inputs are + tracked separately and not exercised here.""" + model, value_info_map, initializer_map, node_to_init_map = model_loop_subgraph_capture + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=keep_io_types, + low_precision_type=low_precision_type, + ) + converted_model = converter.convert( + high_precision_nodes=["loop_node"], low_precision_nodes=["pre_mul"] + ) + onnx.checker.check_model(converted_model) + onnx.shape_inference.infer_shapes(converted_model, strict_mode=True, check_type=True)