[6771663] Preserve ONNX API output types when wiring casts

[ajrasane]

What does this PR do?

Type of change: Bug fix

Preserves the public ONNX graph I/O types captured at the API boundary when PrecisionConverter wires output casts. Type inference can change the working graph's output declaration before conversion; consulting that mutated declaration caused the required cast back to the original public type to be discarded and metadata restoration to fail.

The converter now derives its I/O type map from the preserved boundary metadata and uses that map when deciding whether a cast should become a public graph output. A regression test covers an FP32 output whose working declaration is inferred as FP16, and the changelog records the corrected behavior.

Usage

# No API changes are required. Existing conversions now preserve the original
# public I/O declarations when keep_io_types=True.
converted = convert_to_f16(model, keep_io_types=True)

Testing

• Ran pytest tests/unit/onnx/autocast/test_precisionconverter.py (186 passed).
• Ran pytest tests/unit/onnx/autocast (249 passed).
• Ran Ruff check and format validation on the changed Python files.
• Verified the original minimal end-to-end reproduction with Python 3.12 and TensorRT 10.16.1.11; conversion now completes without the output-metadata restoration error.
• Verified the full CLI quantization path proceeds through the formerly failing one-Q/DQ scheme and successfully benchmarks the generated 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?: ✅
• Did you get Claude approval on this PR?: N/A

Additional Information

Tracking: [6771663]

Summary by CodeRabbit

• Bug Fixes

  • Fixed ONNX FP16 conversion to preserve public graph output types when type inference changes internal declarations.
  • Ensured output casts are inserted correctly when preserving input/output types is enabled, maintaining compatibility for downstream consumers.
  • Validated that converted models retain the expected output types and remain compliant with ONNX validation.

• Tests

  • Added regression coverage confirming preserved output types, correct cast insertion, and valid ONNX model generation.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

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: 0101cfcf-d5cf-415b-9729-13ebd139ea8b

📥 Commits

Reviewing files that changed from the base of the PR and between 5476e66 and ad24f7f.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/onnx/autocast/precisionconverter.py
• tests/unit/onnx/autocast/test_precisionconverter.py

🚧 Files skipped from review as they are similar to previous changes (1)

• CHANGELOG.rst

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

📝 Walkthrough

Walkthrough

FP16 conversion now uses preserved API-boundary I/O types when inserting output casts. A regression test verifies that a declared FLOAT output remains FLOAT and that the converted model passes ONNX validation.

Changes

ONNX output type preservation

Layer / File(s)	Summary
Preserve API-boundary types modelopt/onnx/autocast/precisionconverter.py	PrecisionConverter uses preserved original I/O metadata instead of modified working-model declarations when handling output casts.
Validate output cast behavior tests/unit/onnx/autocast/test_precisionconverter.py, CHANGELOG.rst	A regression test checks the declared FLOAT output, one output cast, and full ONNX validation. The changelog records the fix.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Suggested reviewers: gcunhase

Merge Risk: ⚪ Minimal · up to ad24f

The conversion change preserves declared ONNX output types and includes regression coverage for the reported inference scenario. No unresolved merge-blocking issue remains.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: preserving ONNX API output types when inserting casts.
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	PASS. The authoritative PR diff changes one modelopt Python file, one test file, and CHANGELOG.rst. The added implementation and test lines contain no torch.load(..., weights_only=False), numpy.load(.…

Full details: Docstring Coverage

Explanation

Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 5 functions across 2 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 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-6771663-triage

---

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

[cjluo-nv]

Bot review (claude-opus-5) — DM the bot to share feedback.

Approving: a small, correctly scoped fix that sources original_network_io from the boundary metadata instead of the post-inference working declaration, with a regression test and changelog entry.

No action needed.

• Checked the interaction with _add_cast's existing must_preserve_io_type guard and _sanity_check; both now read the same boundary map, and the fallback when original_network_io_metadata is None (i.e. keep_io_types=False) reproduces the previous behavior exactly.
• Cast-up/cast-down ordering for an output present in both lists still resolves to the boundary type, which is the intended outcome.

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2451/
Built to branch gh-pages at 2026-09-17 03:18 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.50%. Comparing base (b9cfdce) to head (ad24f7f).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2451      +/-   ##
==========================================
- Coverage   71.50%   71.50%   -0.01%     
==========================================
  Files         590      590              
  Lines       64749    64748       -1     
==========================================
- Hits        46297    46296       -1     
  Misses      18452    18452              

Flag	Coverage Δ
unit	57.85% <100.00%> (-0.01%)	⬇️

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.