to_onnx return ONNXProgram

[GdoongMathew]

What does this PR do?

Fixes #20810

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--20811.org.readthedocs.build/en/20811/

[GdoongMathew]

Hi @Borda , a friendly ping for code review~

BTW, I'm wondering if there's any way to update the typing-extensions requirement during testing so that it could fit with onnxscript? Thanks.

[deependujha]

seems this to be the cause of issue:

pytorch-lightning/requirements/pytorch/base.txt

Line 10 in dd2912a

typing-extensions >=4.4.0, <4.14.0

I wonder what could be the reason behind this version pinning.

error:

The conflict is caused by:
    pytorch-lightning 2.5.1rc2 depends on typing-extensions<4.14.0 and ==4.4.0
    pytorch-lightning[extra,strategies,test] 2.5.1rc2 depends on typing-extensions<4.14.0 and ==4.4.0
    lightning-utilities 0.10.0 depends on typing-extensions
    torch 2.1.0 depends on typing-extensions
    onnxscript 0.2.5 depends on typing_extensions>=4.10
    onnxscript 0.2.4 depends on typing_extensions>=4.10
    onnxscript 0.2.3 depends on typing_extensions>=4.10
    onnxscript 0.2.2 depends on typing_extensions>=4.10

[deependujha]

Do you think it would be worth adding a check that verifies the output of the exported ONNX model matches the PyTorch model’s output (within tolerance)?

[GdoongMathew]

I think we could use the existed test test_if_inference_output_is_valid for this purpose? I added an parametrized dynamo to that unittest in the last commit as well.

Thank you ~.

[deependujha]

It's technically correct, but since we've only one type, what do you think about Optional["ONNXProgram"] for readability?

[GdoongMathew]

Yes. I've changed it in the last commit. Thanks.

[GdoongMathew]

seems this to be the cause of issue:

pytorch-lightning/requirements/pytorch/base.txt

Line 10 in dd2912a

typing-extensions >=4.4.0, <4.14.0

I wonder what could be the reason behind this version pinning.

error:

The conflict is caused by:
    pytorch-lightning 2.5.1rc2 depends on typing-extensions<4.14.0 and ==4.4.0
    pytorch-lightning[extra,strategies,test] 2.5.1rc2 depends on typing-extensions<4.14.0 and ==4.4.0
    lightning-utilities 0.10.0 depends on typing-extensions
    torch 2.1.0 depends on typing-extensions
    onnxscript 0.2.5 depends on typing_extensions>=4.10
    onnxscript 0.2.4 depends on typing_extensions>=4.10
    onnxscript 0.2.3 depends on typing_extensions>=4.10
    onnxscript 0.2.2 depends on typing_extensions>=4.10

Yea, I think that's where the problem is, and there're two solutions for this.

1. Update the requirement.
2. Figure out a way to only change the version to fit with onnxscript when testing.

Personally, I'd prefer the second one, but this probably would need some change to the ci config which I not familiar with.

[GdoongMathew]

Hi @deependujha all the changes are completed, just left with the requirement part.

And to answer your question, I think it's for testing the compatibility of all the requirements with the lowest supported version.