[6058870] Fix ONNX AutoCast keep_io_types for empty I/O tensors

[ajrasane]

What does this PR do?

Type of change: Bug fix

ONNX AutoCast (used by modelopt.onnx.quantization with --high_precision_dtype fp16/bf16) failed its internal sanity check with Unexpected type in I/O tensor <name>, keep_io_types=True, original type: 1, converted type: 10 when a network input or output is an empty tensor (a tensor with a dimension of size 0).

Root cause: PrecisionConverter._add_cast treats empty tensors specially — instead of inserting a Cast node, it "fake-casts" them by retyping their value_info in place (there is no data to convert). However, setup_mappings stores the same graph.input / graph.output ValueInfoProto objects in value_info_map, so retyping an empty tensor that is a network input/output silently changed the model's declared I/O type to the low-precision type — violating the keep_io_types=True contract and tripping the sanity check.

Fix (modelopt/onnx/autocast/precisionconverter.py): skip the metadata-only fake-cast for protected network I/O tensors (when keep_io_types=True and the target type differs from the original I/O type) and fall through to insert a real Cast node. This preserves the declared I/O type while still bridging precision into the low-precision consumer/producer. Intermediate empty tensors are unchanged and still fake-cast.

Usage

# Previously failed with "Sanity check failed: Unexpected type in I/O tensor ...";
# now completes and produces a valid mixed-precision model.
python -m modelopt.onnx.quantization \
    --quantize_mode int8 \
    --high_precision_dtype fp16 \
    --onnx_path model_with_empty_io_tensor.onnx \
    --output_path model_int8_fp16.onnx

Testing

• Added test_empty_tensor_network_input_keep_io_types (fp16 + bf16, both type-inference paths) in tests/unit/onnx/autocast/test_precisionconverter.py, asserting network I/O types stay FP32 and a real Cast bridges the empty input. Verified it fails before the fix and passes after.
• Full autocast unit suite passes: CUDA_VISIBLE_DEVICES="" pytest tests/unit/onnx/autocast/ → 202 passed.
• Reproduced the original failure on the reported model and confirmed the command now succeeds end-to-end; the resulting ONNX model passes onnx.checker, runs in ONNX Runtime with correct FP32 I/O, and builds a strongly-typed TensorRT engine.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • Fixed an ONNX AutoCast issue where empty tensors with keep_io_types=True could silently change model I/O types; now a proper Cast is inserted for protected inputs/outputs to preserve declared I/O types.

• Tests

  • Added a regression test to ensure empty network-input tensors are handled correctly in mixed-precision ONNX conversion when keep_io_types=True.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 1fd0509b-eee7-4188-8a87-6c18b8e74a4d

📥 Commits

Reviewing files that changed from the base of the PR and between 02f3513 and f224e41.

📒 Files selected for processing (1)

• tests/unit/onnx/autocast/test_precisionconverter.py

---

📝 Walkthrough

Walkthrough

This PR fixes an ONNX AutoCast bug where keep_io_types=True silently retyped empty network inputs/outputs via fake-casting. The fix inserts real Cast nodes instead, adds a regression test, and updates the changelog.

Changes

Empty tensor I/O type preservation fix

Layer / File(s)	Summary
Empty tensor type preservation logic modelopt/onnx/autocast/precisionconverter.py	PrecisionConverter._add_cast now detects empty tensors that are original network I/O; when keep_io_types=True and the target cast type differs from the original I/O element type, it skips fake-casting and inserts a real Cast node instead.
Regression test and changelog tests/unit/onnx/autocast/test_precisionconverter.py, CHANGELOG.rst	New test test_empty_tensor_network_input_keep_io_types validates that empty network inputs preserve FP32 type and receive real Cast nodes under mixed-precision conversion; changelog entry documents the fix for release 0.45.

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main fix: preventing fake-casting of empty I/O tensors in ONNX AutoCast's keep_io_types feature.
Docstring Coverage	✅ Passed	Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	PR contains no security anti-patterns: no unsafe torch.load/numpy.load, eval/exec, trust_remote_code, # nosec comments, or non-permissive dependencies. All changes use existing dependencies.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch ajrasane/nvbug_6058870

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-06-04 21:20 UTC

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.37%. Comparing base (0081473) to head (f224e41).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1626      +/-   ##
==========================================
- Coverage   77.38%   77.37%   -0.01%     
==========================================
  Files         482      482              
  Lines       52960    52961       +1     
==========================================
- Hits        40985    40981       -4     
- Misses      11975    11980       +5     

Flag	Coverage Δ
examples	43.12% <100.00%> (+1.13%)	⬆️
gpu	59.52% <100.00%> (-0.48%)	⬇️
unit	53.92% <100.00%> (-0.02%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cjluo-nv]

Bot review — DM the bot to share feedback.

Targeted bug fix for ONNX AutoCast keep_io_types=True with empty I/O tensors. The root-cause analysis (value_info_map aliasing graph.input/graph.output ValueInfoProto, so in-place "fake-cast" mutates declared I/O type) is correct, and the guard cleanly falls through to the real-Cast path only for protected I/O tensors whose target type differs from the original. original_network_io is captured in __init__ from the deepcopied model, so the comparison remains stable across _sanitize_model. Test covers both fp16/bf16 and both type-inference paths and asserts both the I/O-type invariant and the presence of a real Cast node. CHANGELOG updated. Minor gap: no test for the symmetric empty-network-output case, but the fix is symmetric and not worth blocking.

[galagam]

I'm amazed at the amount of possible corner cases... This time a corner case inside a a WAR for another corner case... Thanks for fixing.

[gcunhase]

Added a comment, LGTM after that's resolved.