diff --git a/modelopt/onnx/autocast/convert.py b/modelopt/onnx/autocast/convert.py
index 0288c729dd4..d9bc8d68c99 100644
--- a/modelopt/onnx/autocast/convert.py
+++ b/modelopt/onnx/autocast/convert.py
@@ -23,6 +23,8 @@
 nodes.
 """
 
+from copy import deepcopy
+
 import numpy as np
 import onnx
 
@@ -44,6 +46,17 @@
 LATEST_IR_VERSION_SUPPORTED_BY_ORT = 10
 
 
+def _capture_network_io_metadata(
+    model: onnx.ModelProto, keep_io_types: bool
+) -> dict[str, list[onnx.ValueInfoProto]] | None:
+    if not keep_io_types:
+        return None
+    return {
+        "input": [deepcopy(io) for io in model.graph.input],
+        "output": [deepcopy(io) for io in model.graph.output],
+    }
+
+
 def convert_to_mixed_precision(
     onnx_path: str,
     low_precision_type: str = "fp16",
@@ -97,6 +110,7 @@ def convert_to_mixed_precision(
     # Load and process model
     model = onnx.load(onnx_path, load_external_data=True)
     assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16"
+    original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types)
 
     # Get original model's opset version
     original_opset = onnx_utils.get_opset_version(model)
@@ -172,6 +186,7 @@ def convert_to_mixed_precision(
         init_conversion_max_bytes=init_conversion_max_bytes,
         custom_ops=graph_sanitizer.custom_ops,
         use_standalone_type_inference=use_standalone_type_inference,
+        original_network_io_metadata=original_network_io_metadata,
     )
 
     # Obtain reference data
@@ -227,6 +242,7 @@ def convert_to_f16(
                INT4 requires 21, NVFP4 requires 23).
     """
     assert low_precision_type in ["fp16", "bf16"], "low_precision_type must be either fp16 or bf16"
+    original_network_io_metadata = _capture_network_io_metadata(model, keep_io_types)
 
     # Check Q/DQ precision types in the model and determine required opset
     qdq_precisions = get_qdq_precisions(model)
@@ -285,6 +301,7 @@ def convert_to_f16(
         custom_ops=sanitizer.custom_ops,
         tensor_block_dict=tensor_block_dict,
         use_standalone_type_inference=use_standalone_type_inference,
+        original_network_io_metadata=original_network_io_metadata,
     )
     high_precision_nodes = [node.name for node in model.graph.node if node.op_type in op_block_list]
     low_precision_nodes = [
diff --git a/modelopt/onnx/autocast/precisionconverter.py b/modelopt/onnx/autocast/precisionconverter.py
index da45cd3ecde..4cb96148692 100644
--- a/modelopt/onnx/autocast/precisionconverter.py
+++ b/modelopt/onnx/autocast/precisionconverter.py
@@ -98,6 +98,7 @@ def __init__(
         trt_plugins: list[str] | None = [],
         tensor_block_dict: dict[str, dict[str, list[int]]] = {},
         use_standalone_type_inference: bool = False,
+        original_network_io_metadata: dict[str, list[onnx.ValueInfoProto]] | None = None,
     ) -> None:
         """Initialize PrecisionConverter.
 
@@ -116,6 +117,7 @@ def __init__(
             trt_plugins: List of custom TensorRT plugin library paths in .so format (compiled shared library).
             tensor_block_dict: Dictionary of tensors (operation type and I/O indices) that should remain in FP32.
             use_standalone_type_inference: Use standalone type inference instead of ONNX's infer_shapes.
+            original_network_io_metadata: Original public input/output metadata captured at the API boundary.
         """
         self.model = deepcopy(model)
         self.value_info_map = value_info_map
@@ -139,6 +141,17 @@ def __init__(
         self.original_network_io.update(
             {io.name: io.type.tensor_type.elem_type for io in self.model.graph.output}
         )
+        self.original_network_io_metadata = (
+            {
+                "input": [deepcopy(io) for io in self.model.graph.input],
+                "output": [deepcopy(io) for io in self.model.graph.output],
+            }
+            if original_network_io_metadata is None
+            else {
+                field: [deepcopy(value) for value in values]
+                for field, values in original_network_io_metadata.items()
+            }
+        )
         self.min_opset = min_opset
         self.max_ir_version = max_ir_version
         self.trt_plugins = trt_plugins
@@ -276,6 +289,8 @@ def convert(
         # Remove redundant casts
         self._cleanup()
 
+        self._restore_original_io_metadata()
+
         self._sanity_check()
 
         return self.model
@@ -289,6 +304,120 @@ 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 tensor.shape is None:
+                return None
+            return list(tensor.shape)
+
+        def _get_const_values(tensor):
+            if isinstance(tensor, gs.Constant):
+                return tensor.values
+            if tensor.inputs and tensor.inputs[0].op == "Constant":
+                return tensor.inputs[0].attrs["value"].values
+            return None
+
+        def _get_int_attr(node, attr_name, default):
+            value = node.attrs.get(attr_name, default)
+            return value if isinstance(value, int) else None
+
+        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:]
+
+        def _infer_gather_op_shape(node):
+            if node.op != "Gather" or len(node.inputs) < 2:
+                return None
+
+            data_shape = _get_shape(node.inputs[0])
+            indices_shape = _get_shape(node.inputs[1])
+            if data_shape is None or indices_shape is None:
+                return None
+
+            axis = _get_int_attr(node, "axis", 0)
+            if axis is None:
+                return None
+            if axis < 0:
+                axis += len(data_shape)
+            if axis < 0 or axis >= len(data_shape):
+                return None
+
+            return data_shape[:axis] + indices_shape + data_shape[axis + 1 :]
+
+        def _infer_unsqueeze_op_shape(node):
+            if node.op != "Unsqueeze" or len(node.inputs) < 2:
+                return None
+
+            data_shape = _get_shape(node.inputs[0])
+            axes = _get_const_values(node.inputs[1])
+            if data_shape is None or axes is None:
+                return None
+
+            axes = [int(axis) for axis in np.asarray(axes).flatten()]
+            output_rank = len(data_shape) + len(axes)
+            normalized_axes = []
+            for axis in axes:
+                if axis < 0:
+                    axis += output_rank
+                if axis < 0 or axis >= output_rank:
+                    return None
+                normalized_axes.append(axis)
+
+            output_shape = list(data_shape)
+            for axis in sorted(normalized_axes):
+                output_shape.insert(axis, 1)
+            return output_shape
+
+        def _infer_shape_op_shape(node):
+            if node.op != "Shape" or not node.inputs:
+                return None
+
+            data_shape = _get_shape(node.inputs[0])
+            if data_shape is None:
+                return None
+
+            rank = len(data_shape)
+            start = _get_int_attr(node, "start", 0)
+            end = _get_int_attr(node, "end", rank)
+            if start is None or end is None:
+                return None
+            if start < 0:
+                start += rank
+            if end < 0:
+                end += rank
+            start = min(max(start, 0), rank)
+            end = min(max(end, 0), rank)
+            return [max(end - start, 0)]
+
+        def _infer_standard_op_shape(node):
+            for infer_shape in (
+                _infer_gathernd_op_shape,
+                _infer_gather_op_shape,
+                _infer_unsqueeze_op_shape,
+                _infer_shape_op_shape,
+            ):
+                shape = infer_shape(node)
+                if shape is not None:
+                    return shape
+            return None
+
         graph = gs.import_onnx(model)
         traversed_tensors = []
 
@@ -397,12 +526,14 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1):
                     out.dtype = np_type
 
                 # Set the output shape
-                if not out.shape:
-                    if isinstance(inp, gs.Constant):
+                if out.shape is None:
+                    if (shape := _infer_standard_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
-                    elif inp.shape:
+                    elif node.op in self.custom_ops and inp.shape:
                         out.shape = inp.shape
 
             # Propagate tensor types to the children nodes (until another Cast or Q node is met)
@@ -410,6 +541,32 @@ def _propagate_cast_type_through_nodes(node, np_type, iter=1):
 
         return gs.export_onnx(graph)
 
+    def _restore_original_io_metadata(self) -> None:
+        """Preserve complete public I/O metadata when keep_io_types=True."""
+        if not self.keep_io_types:
+            return
+
+        for io_field in ("input", "output"):
+            current_values = list(getattr(self.model.graph, io_field))
+            original_values = self.original_network_io_metadata[io_field]
+            current_names = [value.name for value in current_values]
+            original_names = [value.name for value in original_values]
+            if current_names != original_names:
+                raise RuntimeError(
+                    f"Cannot restore public graph {io_field} metadata because names changed: "
+                    f"{original_names} -> {current_names}"
+                )
+            for current_value, original_value in zip(current_values, original_values, strict=True):
+                if (
+                    current_value.type.tensor_type.elem_type
+                    != original_value.type.tensor_type.elem_type
+                ):
+                    raise RuntimeError(
+                        f"Cannot restore public graph {io_field} metadata for {current_value.name}: "
+                        "element type changed"
+                    )
+                current_value.CopyFrom(original_value)
+
     def _is_bf16(self, type: PrecisionTypes = None) -> bool:
         if type is None:
             type = self.low_precision_type
@@ -1043,6 +1200,7 @@ def _cleanup(self):
         # Remove redundant casts
         self._remove_redundant_casts()
         self._deduplicate_network_output_producers()
+        self._refresh_gathernd_pre_cast_declarations()
 
     def _cleanup_no_consumer_nodes(self):
         network_outputs = {o.name for o in self.model.graph.output}
@@ -1184,6 +1342,113 @@ def _fix_network_output_names(self):
                 self.model
             )
 
+    def _get_tensor_shape(self, tensor_name: str) -> list[int | str | None] | None:
+        initializer = next(
+            (value for value in self.model.graph.initializer if value.name == tensor_name), None
+        )
+        if initializer is not None:
+            return list(initializer.dims)
+
+        for value_info in (
+            *self.model.graph.input,
+            *self.model.graph.output,
+            *self.model.graph.value_info,
+        ):
+            if value_info.name != tensor_name:
+                continue
+            tensor_type = value_info.type.tensor_type
+            if not tensor_type.HasField("shape"):
+                continue
+            shape = []
+            for dim in tensor_type.shape.dim:
+                if dim.HasField("dim_value"):
+                    shape.append(dim.dim_value)
+                elif dim.HasField("dim_param"):
+                    shape.append(dim.dim_param)
+                else:
+                    shape.append(None)
+            return shape
+
+        producer_nodes = onnx_utils.get_producer_nodes(self.model, tensor_name)
+        if len(producer_nodes) == 1 and producer_nodes[0].op_type == "Cast":
+            return self._get_tensor_shape(producer_nodes[0].input[0])
+        return None
+
+    def _get_tensor_elem_type(self, tensor_name: str) -> int | None:
+        for value_info in (
+            *self.model.graph.input,
+            *self.model.graph.output,
+            *self.model.graph.value_info,
+        ):
+            if value_info.name == tensor_name and value_info.type.HasField("tensor_type"):
+                elem_type = value_info.type.tensor_type.elem_type
+                if elem_type != onnx.TensorProto.UNDEFINED:
+                    return elem_type
+        initializer = next(
+            (value for value in self.model.graph.initializer if value.name == tensor_name), None
+        )
+        return initializer.data_type if initializer is not None else None
+
+    def _refresh_gathernd_output_declaration(self, node_name: str, tensor_name: str) -> None:
+        node = next(
+            (
+                node
+                for node in self.model.graph.node
+                if node.name == node_name and tensor_name in node.output
+            ),
+            None,
+        )
+        if node is None or node.op_type != "GatherND" or len(node.input) < 2:
+            return
+
+        data_shape = self._get_tensor_shape(node.input[0])
+        indices_shape = self._get_tensor_shape(node.input[1])
+        if data_shape is None or indices_shape is None or not indices_shape:
+            return
+        index_rank = indices_shape[-1]
+        batch_dims = next(
+            (attr.i for attr in node.attribute if attr.name == "batch_dims"),
+            0,
+        )
+        if not isinstance(index_rank, int) or not isinstance(batch_dims, int):
+            return
+        suffix_start = batch_dims + index_rank
+        if suffix_start > len(data_shape):
+            return
+
+        shape = indices_shape[:-1] + data_shape[suffix_start:]
+        elem_type = (
+            self._get_tensor_elem_type(tensor_name)
+            or self._get_tensor_elem_type(node.input[0])
+            or self.low_precision_type.onnx_type
+        )
+        value_info = self.value_info_map.get(tensor_name)
+        updated_value_info = helper.make_tensor_value_info(tensor_name, elem_type, shape)
+        if value_info is None:
+            self.model.graph.value_info.append(updated_value_info)
+        else:
+            value_info.CopyFrom(updated_value_info)
+
+    def _refresh_gathernd_pre_cast_declarations(self) -> None:
+        gathernd_pre_cast_outputs = [
+            (node.name, output)
+            for node in self.model.graph.node
+            if node.op_type == "GatherND"
+            for output in node.output
+            if output.endswith("_pre_cast")
+        ]
+        if not gathernd_pre_cast_outputs:
+            return
+
+        self.value_info_map, self.initializer_map, self.node_to_init_map = utils.setup_mappings(
+            self.model
+        )
+        for node_name, tensor_name in gathernd_pre_cast_outputs:
+            self._refresh_gathernd_output_declaration(node_name, tensor_name)
+        self.value_info_map, self.initializer_map, self.node_to_init_map = utils.setup_mappings(
+            self.model
+        )
+
     def _deduplicate_network_output_producers(self):
         for output in self.model.graph.output:
             producer_nodes = onnx_utils.get_producer_nodes(self.model, output.name)
diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py
index bc37c4a9333..04be94a743c 100644
--- a/modelopt/onnx/utils.py
+++ b/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}")
 
diff --git a/tests/unit/onnx/autocast/test_precisionconverter.py b/tests/unit/onnx/autocast/test_precisionconverter.py
index 9872e0c4f2b..64fc423d737 100644
--- a/tests/unit/onnx/autocast/test_precisionconverter.py
+++ b/tests/unit/onnx/autocast/test_precisionconverter.py
@@ -20,7 +20,8 @@
 
 import modelopt.onnx.autocast.utils as utils
 import modelopt.onnx.utils as onnx_utils
-from modelopt.onnx.autocast.convert import convert_to_mixed_precision
+from modelopt.onnx.autocast.convert import convert_to_f16, convert_to_mixed_precision
+from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer
 from modelopt.onnx.autocast.logging_config import configure_logging
 from modelopt.onnx.autocast.precisionconverter import PrecisionConverter
 
@@ -1955,3 +1956,326 @@ 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] == []
+
+
+def test_custom_op_mode_uses_schema_shapes_for_standard_rank_changes():
+    gather_data = helper.make_tensor_value_info("gather_data", TensorProto.FLOAT, [4])
+    unsqueeze_data = helper.make_tensor_value_info("unsqueeze_data", TensorProto.FLOAT, [4])
+    plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [1])
+    gather_index = numpy_helper.from_array(np.array(0, dtype=np.int64), "gather_index")
+    unsqueeze_axes = numpy_helper.from_array(np.array([0], dtype=np.int64), "unsqueeze_axes")
+    gather_node = helper.make_node(
+        "Gather", ["gather_data", "gather_index"], ["gather_y_pre_cast"], name="gather"
+    )
+    gather_cast = helper.make_node(
+        "Cast", ["gather_y_pre_cast"], ["gather_y"], name="gather_cast", to=TensorProto.FLOAT
+    )
+    unsqueeze_node = helper.make_node(
+        "Unsqueeze",
+        ["unsqueeze_data", "unsqueeze_axes"],
+        ["unsqueeze_y_pre_cast"],
+        name="unsqueeze",
+    )
+    unsqueeze_cast = helper.make_node(
+        "Cast",
+        ["unsqueeze_y_pre_cast"],
+        ["unsqueeze_y"],
+        name="unsqueeze_cast",
+        to=TensorProto.FLOAT,
+    )
+    custom_node = helper.make_node(
+        "FakePlugin", ["plugin_in"], ["plugin_y"], name="plugin", domain="test.plugins"
+    )
+    graph = helper.make_graph(
+        [gather_node, gather_cast, unsqueeze_node, unsqueeze_cast, custom_node],
+        "custom_op_standard_rank_changes",
+        [gather_data, unsqueeze_data, plugin_in],
+        [
+            helper.make_tensor_value_info("gather_y", TensorProto.FLOAT, []),
+            helper.make_tensor_value_info("unsqueeze_y", TensorProto.FLOAT, [1, 4]),
+            helper.make_tensor_value_info("plugin_y", TensorProto.FLOAT, [1]),
+        ],
+        [gather_index, unsqueeze_axes],
+    )
+    model = helper.make_model(
+        graph,
+        producer_name="custom_op_standard_rank_changes",
+        opset_imports=[helper.make_opsetid("", 19), helper.make_opsetid("test.plugins", 1)],
+        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={"FakePlugin"},
+    )
+    propagated = converter._propagate_types_shapes_custom_ops(model)
+
+    value_infos = {vi.name: vi for vi in [*propagated.graph.value_info, *propagated.graph.output]}
+    assert [
+        dim.dim_value for dim in value_infos["gather_y_pre_cast"].type.tensor_type.shape.dim
+    ] == []
+    assert [
+        dim.dim_value for dim in value_infos["unsqueeze_y_pre_cast"].type.tensor_type.shape.dim
+    ] == [1, 4]
+    assert [dim.dim_value for dim in value_infos["gather_y"].type.tensor_type.shape.dim] == []
+    assert [dim.dim_value for dim in value_infos["unsqueeze_y"].type.tensor_type.shape.dim] == [
+        1,
+        4,
+    ]
+    onnx.checker.check_model(propagated, full_check=True)
+
+
+def test_custom_op_mode_preserves_known_scalar_custom_op_shape():
+    plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [4])
+    custom_node = helper.make_node(
+        "FakePlugin", ["plugin_in"], ["plugin_scalar"], name="plugin", domain="test.plugins"
+    )
+    graph = helper.make_graph(
+        [custom_node],
+        "custom_op_known_scalar_shape",
+        [plugin_in],
+        [helper.make_tensor_value_info("plugin_scalar", TensorProto.FLOAT, [])],
+        [],
+    )
+    model = helper.make_model(
+        graph,
+        producer_name="custom_op_known_scalar_shape",
+        opset_imports=[helper.make_opsetid("", 19), helper.make_opsetid("test.plugins", 1)],
+        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={"FakePlugin"},
+    )
+    propagated = converter._propagate_types_shapes_custom_ops(model)
+
+    output = next(vi for vi in propagated.graph.output if vi.name == "plugin_scalar")
+    assert [dim.dim_value for dim in output.type.tensor_type.shape.dim] == []
+
+
+def _shape_of(value):
+    shape = []
+    for dim in value.type.tensor_type.shape.dim:
+        if dim.HasField("dim_value"):
+            shape.append(dim.dim_value)
+        elif dim.HasField("dim_param"):
+            shape.append(dim.dim_param)
+        else:
+            shape.append(None)
+    return shape
+
+
+def test_convert_to_f16_restores_public_io_metadata_from_entry_boundary():
+    graph_input = helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 3])
+    graph_output = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [None, None])
+    graph_output.doc_string = "Public output dimensions intentionally unspecified"
+    node = helper.make_node("Identity", ["X"], ["Y"], name="Identity_0")
+    graph = helper.make_graph([node], "public_io_boundary", [graph_input], [graph_output])
+    model = helper.make_model(
+        graph,
+        producer_name="public_io_boundary",
+        opset_imports=[helper.make_opsetid("", 19)],
+        ir_version=10,
+    )
+
+    converted = convert_to_f16(
+        model, keep_io_types=True, op_block_list=[], trt_plugins=[], opset=19
+    )
+
+    output = next(vi for vi in converted.graph.output if vi.name == "Y")
+    assert output.type.tensor_type.elem_type == TensorProto.FLOAT
+    assert _shape_of(output) == [None, None]
+    assert output.doc_string == graph_output.doc_string
+    onnx.checker.check_model(converted, full_check=True)
+
+
+def test_convert_to_f16_refreshes_gathernd_pre_cast_declaration(monkeypatch):
+    def discover_test_plugins_without_trt(self):
+        self.custom_ops = {
+            node.op_type for node in self.model.graph.node if node.domain == "test.plugins"
+        }
+        self.custom_ops_low_precision_nodes = []
+
+    monkeypatch.setattr(GraphSanitizer, "find_custom_nodes", discover_test_plugins_without_trt)
+
+    data = helper.make_tensor_value_info("data", TensorProto.FLOAT, [1, 4, 2])
+    plugin_in = helper.make_tensor_value_info("plugin_in", TensorProto.FLOAT, [1])
+    last_token_embed = helper.make_tensor_value_info("last_token_embed", TensorProto.FLOAT, [1, 2])
+    plugin_y = helper.make_tensor_value_info("plugin_y", TensorProto.FLOAT, [1])
+    indices = numpy_helper.from_array(np.array([[3]], dtype=np.int64), name="indices")
+    gathernd = helper.make_node(
+        "GatherND",
+        ["data", "indices"],
+        ["last_token_embed"],
+        name="/lm_head/GatherND",
+        batch_dims=1,
+    )
+    plugin = helper.make_node(
+        "FakePlugin",
+        ["plugin_in"],
+        ["plugin_y"],
+        name="synthetic_plugin",
+        domain="test.plugins",
+    )
+    graph = helper.make_graph(
+        [gathernd, plugin],
+        "public_gathernd_pre_cast_declaration",
+        [data, plugin_in],
+        [last_token_embed, plugin_y],
+        initializer=[indices],
+    )
+    model = helper.make_model(
+        graph,
+        producer_name="public_gathernd_pre_cast_declaration",
+        opset_imports=[helper.make_opsetid("", 19), helper.make_opsetid("test.plugins", 1)],
+        ir_version=10,
+    )
+
+    converted = convert_to_f16(
+        model,
+        keep_io_types=True,
+        op_block_list=["FakePlugin"],
+        low_precision_type="fp16",
+    )
+
+    declarations = {
+        value.name: value
+        for value in (*converted.graph.input, *converted.graph.output, *converted.graph.value_info)
+    }
+    pre_cast = declarations["last_token_embed_pre_cast"]
+    assert pre_cast.type.tensor_type.elem_type == TensorProto.FLOAT16
+    assert _shape_of(pre_cast) == [1, 2]
+    assert declarations["last_token_embed"].type.tensor_type.elem_type == TensorProto.FLOAT
+    assert _shape_of(declarations["last_token_embed"]) == [1, 2]
+    onnx.checker.check_model(converted, full_check=True)