[6425069][ONNX][Autocast] Fix autocast metadata propagation

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)) is not None:
+                        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

modelopt/onnx/utils.py

@@ -1532,6 +1532,21 @@ 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):
+        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:
+            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 +1586,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}")
 

tests/unit/onnx/autocast/test_precisionconverter.py

@@ -1854,3 +1854,123 @@ 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
+
+
+@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"
+    )
+    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", value_info_elem_type, [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]
+
+
+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] == []