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

[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.

[Borda]

not sure if the failing tests are related to this change or they are a bit flaky... 🤔

[GdoongMathew]

Hi @Borda
I think that's was coming from the updated version of pydantic. Normally during testing with the oldest pytorch, pydantic 1.x was installed, but in this PR, because of the onnxscript dependency, it needs typing-extensions 1.10 or later, which forces the ci to install pydantic 2.0, which causes the problem during importing deepspeed.
That's only from looking at the log. If there's anything I could change in the ci setting, I'd love to test it out~

[GdoongMathew]

Or I'm wondering if it's possible to install onnxscript only when torch version 2.4.0 or above is installed during testing 🤔

[Borda]

I'm wondering if it's possible to install onnxscript only when torch version 2.4.0 or above is installed during testing

Why this limitation? You can make some tests skip if PT is below a certain version...

[Borda]

So why do we need this bump?

[GdoongMathew]

So this is governed by onnxscript here, and I didn't think of a better solution to pass the tests with the oldest pytorch without changing the current ci.

Also mentioned in #20811 (comment)

[GdoongMathew]

I'm wondering if it's possible to install onnxscript only when torch version 2.4.0 or above is installed during testing

Why this limitation? You can make some tests skip if PT is below a certain version...

Yes, currently only torch 2.7.0 or higher supports running the newly added tests.
Also, onnxscript needs to be installed additionally as its not included in torch's standard requirements - similar to onnx. Therefore, I suggest installing onnxscript only when the appropriate Torch version is installed.