From ea8f814231bfff6acc373c97b0c5f56f3838d5c1 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Fri, 17 Jul 2026 01:06:59 +0000 Subject: [PATCH 1/4] Fix ONNX autocast metadata propagation Update ONNX autocast cleanup to keep folded Constant metadata in sync and infer GatherND output shapes in custom-op mode without running global ONNX shape inference. This preserves valid ONNX metadata after redundant Cast folding and avoids incorrect GatherND rank propagation when TensorRT plugin custom ops are present. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- modelopt/onnx/autocast/precisionconverter.py | 31 +++++++- modelopt/onnx/utils.py | 19 +++++ .../onnx/autocast/test_precisionconverter.py | 77 +++++++++++++++++++ 3 files changed, 126 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index 5474b57a062..25996daf49c 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -289,6 +289,33 @@ def _ensure_types_are_defined(self): def _propagate_types_shapes_custom_ops(self, model): """Propagate types and shapes after insertion of 'Cast' nodes or other graph modifications.""" logger.info("Propagating tensor shapes and types in model with custom ops.") + + def _get_shape(tensor): + if isinstance(tensor, gs.Constant): + return list(tensor.values.shape) + if not tensor.shape: + return None + return list(tensor.shape) + + def _infer_gathernd_op_shape(node): + if node.op != "GatherND" or len(node.inputs) < 2: + return None + + data_shape = _get_shape(node.inputs[0]) + indices_shape = _get_shape(node.inputs[1]) + if not data_shape or not indices_shape: + return None + + index_rank = indices_shape[-1] + batch_dims = node.attrs.get("batch_dims", 0) + if not isinstance(index_rank, int) or not isinstance(batch_dims, int): + return None + + suffix_start = batch_dims + index_rank + if suffix_start > len(data_shape): + return None + return indices_shape[:-1] + data_shape[suffix_start:] + graph = gs.import_onnx(model) traversed_tensors = [] @@ -398,7 +425,9 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1): # Set the output shape if not out.shape: - if isinstance(inp, gs.Constant): + if shape := _infer_gathernd_op_shape(node): + out.shape = shape + elif isinstance(inp, gs.Constant): out.shape = inp.values.shape elif inp.inputs and inp.inputs[0].op == "Constant": out.shape = inp.inputs[0].attrs["value"].values.shape diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 3b8edf76a84..c3efdcb63be 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1532,6 +1532,22 @@ def _convert_constant_values(constant_node: onnx.NodeProto, cast_node: onnx.Node break +def _sync_value_info_elem_type(graph: onnx.GraphProto, tensor_name: str, elem_type: int) -> None: + """Synchronize declarations for a tensor whose producer dtype changed.""" + for value_info in list(graph.value_info) + list(graph.input) + list(graph.output): + tensor_type = value_info.type.tensor_type + if value_info.name == tensor_name and tensor_type.elem_type: + tensor_type.elem_type = elem_type + + for node in graph.node: + for attr in node.attribute: + if attr.type == onnx.AttributeProto.GRAPH: + _sync_value_info_elem_type(attr.g, tensor_name, elem_type) + elif attr.type == onnx.AttributeProto.GRAPHS: + for subgraph in attr.graphs: + _sync_value_info_elem_type(subgraph, tensor_name, elem_type) + + def remove_redundant_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Removes both sequential casts and casts that don't change precision. @@ -1571,6 +1587,9 @@ def remove_redundant_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: assert len(cast_producers) == 1 and cast_producers[0].op_type == "Constant" constant_producer = cast_producers[0] _convert_constant_values(constant_producer, node) + _sync_value_info_elem_type( + onnx_model.graph, constant_producer.output[0], get_cast_to_type(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 4c0f19759e3..a7334743732 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -1854,3 +1854,80 @@ 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 + + +def test_folded_constant_cast_updates_value_info_type(): + const_tensor = numpy_helper.from_array( + np.array([1.0, 2.0], dtype=np.float32), name="const_value" + ) + const_node = helper.make_node( + "Constant", [], ["const_out"], name="const_node", value=const_tensor + ) + cast_node = helper.make_node( + "Cast", ["const_out"], ["cast_out"], name="cast_to_fp16", to=TensorProto.FLOAT16 + ) + identity_node = helper.make_node("Identity", ["cast_out"], ["Y"], name="identity") + + graph = helper.make_graph( + [const_node, cast_node, identity_node], + "constant_cast_value_info", + [], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2])], + [], + value_info=[helper.make_tensor_value_info("const_out", TensorProto.FLOAT, [2])], + ) + model = helper.make_model(graph, producer_name="constant_cast_value_info") + model.opset_import[0].version = 19 + model.ir_version = 10 + + folded = onnx_utils.remove_redundant_casts(model) + + assert [node.op_type for node in folded.graph.node] == ["Constant", "Identity"] + const_out = next(vi for vi in folded.graph.value_info if vi.name == "const_out") + assert const_out.type.tensor_type.elem_type == TensorProto.FLOAT16 + onnx.shape_inference.infer_shapes(folded, strict_mode=True, check_type=True) + + +def test_custom_op_mode_uses_schema_shape_for_standard_gathernd(): + data = helper.make_tensor_value_info("data", TensorProto.FLOAT, [1, 4, 2]) + plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [1, 4, 2]) + indices_init = numpy_helper.from_array( + np.array([[[0, 0], [3, 1]]], dtype=np.int64), name="indices" + ) + custom_node = helper.make_node( + "FakeTensorRTPlugin", ["plugin_in"], ["plugin_out"], name="fake_plugin" + ) + gather_node = helper.make_node( + "GatherND", + ["data", "indices"], + ["last_token_embed"], + name="shape_changing_gathernd", + batch_dims=1, + ) + graph = helper.make_graph( + [custom_node, gather_node], + "custom_op_gathernd_shape", + [data, plugin_in], + [ + helper.make_tensor_value_info("plugin_out", TensorProto.FLOAT, [1, 4, 2]), + helper.make_tensor_value_info("last_token_embed", TensorProto.FLOAT, None), + ], + [indices_init], + ) + model = helper.make_model(graph, producer_name="custom_op_gathernd_shape") + model.opset_import[0].version = 19 + model.ir_version = 10 + value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + custom_ops={"FakeTensorRTPlugin"}, + ) + propagated = converter._propagate_types_shapes_custom_ops(model) + + output = next(vi for vi in propagated.graph.output if vi.name == "last_token_embed") + assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == [1, 2] From ddd5b03f7089c181d5b59aa1bb1bcc70b19a52d3 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:45:59 +0000 Subject: [PATCH 2/4] Handle scalar GatherND shape propagation Treat an empty inferred GatherND shape as a valid scalar shape instead of falling through to the generic input-shape fallback in custom-op propagation. Add a regression test for the rank-0 GatherND output case. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- modelopt/onnx/autocast/precisionconverter.py | 3 +- .../onnx/autocast/test_precisionconverter.py | 42 +++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index 25996daf49c..faa85c252d5 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -425,7 +425,8 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1): # Set the output shape if not out.shape: - if shape := _infer_gathernd_op_shape(node): + shape = _infer_gathernd_op_shape(node) + if shape is not None: out.shape = shape elif isinstance(inp, gs.Constant): out.shape = inp.values.shape diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index a7334743732..9468a9e6d1c 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -1931,3 +1931,45 @@ def test_custom_op_mode_uses_schema_shape_for_standard_gathernd(): output = next(vi for vi in propagated.graph.output if vi.name == "last_token_embed") assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == [1, 2] + + +def test_custom_op_mode_preserves_scalar_gathernd_shape(): + data = helper.make_tensor_value_info("data", TensorProto.FLOAT, [4]) + plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [4]) + indices_init = numpy_helper.from_array(np.array([2], dtype=np.int64), name="indices") + custom_node = helper.make_node( + "FakeTensorRTPlugin", ["plugin_in"], ["plugin_out"], name="fake_plugin" + ) + gather_node = helper.make_node( + "GatherND", + ["data", "indices"], + ["selected_scalar"], + name="scalar_gathernd", + ) + graph = helper.make_graph( + [custom_node, gather_node], + "custom_op_scalar_gathernd_shape", + [data, plugin_in], + [ + helper.make_tensor_value_info("plugin_out", TensorProto.FLOAT, [4]), + helper.make_tensor_value_info("selected_scalar", TensorProto.FLOAT, None), + ], + [indices_init], + ) + model = helper.make_model(graph, producer_name="custom_op_scalar_gathernd_shape") + model.opset_import[0].version = 19 + model.ir_version = 10 + value_info_map, initializer_map, node_to_init_map = utils.setup_mappings(model) + + converter = PrecisionConverter( + model, + value_info_map, + initializer_map, + node_to_init_map, + keep_io_types=True, + custom_ops={"FakeTensorRTPlugin"}, + ) + propagated = converter._propagate_types_shapes_custom_ops(model) + + output = next(vi for vi in propagated.graph.output if vi.name == "selected_scalar") + assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == [] From c60259f83638ff381d99ba9418fb0c21fdce7ea0 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:51:39 +0000 Subject: [PATCH 3/4] Simplify GatherND shape check Use a walrus expression with an explicit None check so scalar GatherND shapes remain valid while keeping the custom-op propagation branch concise. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- modelopt/onnx/autocast/precisionconverter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py index faa85c252d5..3396045def4 100644 --- a/modelopt/onnx/autocast/precisionconverter.py +++ b/modelopt/onnx/autocast/precisionconverter.py @@ -425,8 +425,7 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1): # Set the output shape if not out.shape: - shape = _infer_gathernd_op_shape(node) - if shape is not None: + if (shape := _infer_gathernd_op_shape(node)) is not None: out.shape = shape elif isinstance(inp, gs.Constant): out.shape = inp.values.shape From 879ce4911c7459aaf40d3208e118142e08eaa091 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:56:09 +0000 Subject: [PATCH 4/4] Handle undefined Constant metadata declarations Synchronize matching ONNX tensor declarations even when the previous element type is TensorProto.UNDEFINED, and cover that Constant Cast folding case in the regression test. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> --- modelopt/onnx/utils.py | 5 ++--- tests/unit/onnx/autocast/test_precisionconverter.py | 5 +++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index c3efdcb63be..3f62376923c 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1535,9 +1535,8 @@ def _convert_constant_values(constant_node: onnx.NodeProto, cast_node: onnx.Node def _sync_value_info_elem_type(graph: onnx.GraphProto, tensor_name: str, elem_type: int) -> None: """Synchronize declarations for a tensor whose producer dtype changed.""" for value_info in list(graph.value_info) + list(graph.input) + list(graph.output): - tensor_type = value_info.type.tensor_type - if value_info.name == tensor_name and tensor_type.elem_type: - tensor_type.elem_type = elem_type + if value_info.name == tensor_name and value_info.type.HasField("tensor_type"): + value_info.type.tensor_type.elem_type = elem_type for node in graph.node: for attr in node.attribute: diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py index 9468a9e6d1c..856e3b42366 100644 --- a/tests/unit/onnx/autocast/test_precisionconverter.py +++ b/tests/unit/onnx/autocast/test_precisionconverter.py @@ -1856,7 +1856,8 @@ def test_if_subgraph_outer_scope_type_preservation( assert else_x_info[0].type.tensor_type.elem_type != onnx.TensorProto.UNDEFINED -def test_folded_constant_cast_updates_value_info_type(): +@pytest.mark.parametrize("value_info_elem_type", [TensorProto.FLOAT, TensorProto.UNDEFINED]) +def test_folded_constant_cast_updates_value_info_type(value_info_elem_type): const_tensor = numpy_helper.from_array( np.array([1.0, 2.0], dtype=np.float32), name="const_value" ) @@ -1874,7 +1875,7 @@ def test_folded_constant_cast_updates_value_info_type(): [], [helper.make_tensor_value_info("Y", TensorProto.FLOAT16, [2])], [], - value_info=[helper.make_tensor_value_info("const_out", TensorProto.FLOAT, [2])], + value_info=[helper.make_tensor_value_info("const_out", value_info_elem_type, [2])], ) model = helper.make_model(graph, producer_name="constant_cast_value_info") model.opset_import[0].version = 19