Skip to content

Support set schema inference function in python #5940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 8 commits into from
Apr 28, 2025
Merged

Support set schema inference function in python #5940

merged 8 commits into from
Apr 28, 2025
+430 −2

Conversation

OYCN
Copy link
Contributor

@OYCN OYCN commented Feb 16, 2024 •
edited
Loading

Description

Support set schema inference callback function for custion OpSchema before register in python.

example:

def func(ctx: onnx.shape_inference.InferenceContext):
    # get some info
    assert ctx.get_num_inputs() == 2
    value = ctx.get_input_type(0)
    ...
    # get or create output proto object
    output = ctx.get_output_type(0)
    # set type or shape
    ...
    # set the result proto
    ctx.set_output_type(0, output)

schema.set_type_and_shape_inference_function(func)

Note

Depends on #5906

Motivation and Context

Follow up of #5019

@OYCN
Copy link
Contributor Author

OYCN commented Feb 16, 2024

We're spending a considerable amount of code on passing proto objects between C++ and Python. However, the repository seems well-prepared for this feature.

@codecov Codecov
Copy link

codecov bot commented Feb 16, 2024 •
edited
Loading

Codecov Report

Attention: Patch coverage is 99.29577% with 1 line in your changes missing coverage. Please review.

Project coverage is 56.50%. Comparing base (c4a4d97) to head (66a5274).
Report is 6 commits behind head on main.

Files with missing lines Patch % Lines
onnx/test/shape_inference_test.py 99.28% 1 Missing ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##             main    #5940      +/-   ##
==========================================
+ Coverage   56.31%   56.50%   +0.18%     
==========================================
  Files         509      509              
  Lines       32607    32732     +125     
  Branches     3099     3099              
==========================================
+ Hits        18364    18495     +131     
+ Misses      13385    13379       -6     
  Partials      858      858

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

github-advanced-security[bot]
github-advanced-security bot found potential problems Feb 26, 2024
View reviewed changes
@OYCN
Copy link
Contributor Author

OYCN commented Mar 29, 2024

Hi @justinchuby ,

In this PR, we can implement shape inference on the Python side, similar to how it's done on the C++ side. If you have any suggestions for this implementation, I'm open to making adjustments accordingly.

@justinchuby justinchuby marked this pull request as ready for review March 29, 2024 18:38
@justinchuby justinchuby requested a review from a team as a code owner March 29, 2024 18:38
@justinchuby
Copy link
Member

Thank you! Is it ready to be reviewed?

@OYCN OYCN changed the title [WIP] Support set schema inference function in python Support set schema inference function in python Mar 31, 2024
@OYCN
Copy link
Contributor Author

OYCN commented Mar 31, 2024

Thank you! Is it ready to be reviewed?

Yes, I have removed the 'WIP' prefix from the title. Please feel free to leave any comments. :D

@justinchuby justinchuby self-assigned this Apr 1, 2024
xadupre
xadupre reviewed Apr 9, 2024
View reviewed changes
xadupre
xadupre reviewed Apr 9, 2024
View reviewed changes
@xadupre
Copy link
Contributor

xadupre commented Apr 9, 2024 •
edited by justinchuby
Loading

It would be a nice feature to have.

gramalingam
gramalingam reviewed Apr 17, 2024
View reviewed changes
gramalingam
gramalingam reviewed Apr 17, 2024
View reviewed changes
gramalingam
gramalingam reviewed Apr 17, 2024
View reviewed changes
gramalingam
gramalingam reviewed Apr 17, 2024
View reviewed changes
@gramalingam
Copy link
Contributor

Thanks for creating the PR! It would be great to add this functionality. My comments above.

OYCN
OYCN commented May 24, 2024
View reviewed changes
Copy link
Contributor Author

@OYCN OYCN left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Forgive my mistake, I didn't know that the response need to be commit manually. My reply is always in pending status.

@akmaru akmaru mentioned this pull request May 25, 2024
Open
@xadupre
Copy link
Contributor

xadupre commented May 27, 2024

I would do this in a different way. I think we should export in python all functions doing shape inference for every operator. We should have a python version of the algorithm doing shape inference. That way, adding new shape inference function would be easier and we could do local shape inference. This feature is needed when building a model. Shapes can be used to write efficient sequences of onnx operators.

@OYCN
Copy link
Contributor Author

OYCN commented May 30, 2024

Hi Xavier,

After thinking about it for a while, please correct me if my understanding is wrong. Implementing a Python version of shape inference may not eliminate the need to pass ONNX protos between C++ and Python. In the current architecture, if we want to obtain any proto in the shape inference function, we need to export InferenceContext from C++ to Python. Perhaps a more efficient implementation for interacting with information carried by proto is key to solving the problem. This could be done using native Python types (though it’s challenging to deliver graph protos), type pointers, or other methods.

Best Regards.

@gramalingam gramalingam added this to the 1.17 milestone Jul 3, 2024
@gramalingam
Copy link
Contributor

The code-freeze for 1.17 is coming up soon. We may need to move this to a later release.

@gramalingam gramalingam removed this from the 1.17 milestone Jul 25, 2024
@OYCN OYCN force-pushed the dev-py-schema-inference branch from 6cd35d0 to 2315678 Compare July 26, 2024 11:16
@OYCN
Copy link
Contributor Author

OYCN commented Jul 26, 2024

Hi @gramalingam

Sorry for the delay. Please let me briefly describe the changes I made since my last commit.

  • Because pybind_protobuf depends on absl::strings, I moved the FetchContent_* to the root scope in CMake and controlled them by the Build_* variable.
  • Added a check for the protobuf minor version, ensuring that it is at least x.22.y. This restriction comes from pybind_protobuf, so some CI will fatal. I think this is a bit radical, but I'm having some difficulty trying to compile protobuf as a fallback option.

It looks like there are plans to upgrade protobuf? (#6167)

@OYCN OYCN force-pushed the dev-py-schema-inference branch from 8070a53 to 2e7f99e Compare December 8, 2024 06:29
@andife
Copy link
Member

andife commented Mar 23, 2025

@OYCN #6725 Minimum protobuf version was upgraded to "protobuf version of 25.1" for upcoming 1.18 release

@justinchuby
Copy link
Member

Sorry for the delayed response. We are interested in getting this PR in if you are still interested

@OYCN
Copy link
Contributor Author

OYCN commented Apr 20, 2025 •
edited
Loading

Consider possible compatibility issues with zero copy (In fact, we have not enable this feature at present based on pybind11_protobuf) and the number of changes to the project, I removed pybind11_protobuf in the latest commit and add the object transport based on serialization and deserialization. Further optimization may be carried out in the future. If anyone has any suggestions, please feel free to comment.

For new types, it is expected to use these to add support. And it will be check at build-time (Only for the type which inherited from Message/Lite Type)

@OYCN OYCN marked this pull request as ready for review April 20, 2025 07:33
@OYCN
Copy link
Contributor Author

OYCN commented Apr 20, 2025

Maybe I should add more unit tests.

@andife andife added the run release CIs Use this label to trigger release tests in CI label Apr 20, 2025
@OYCN
Add more unittest
16ac26a 
Signed-off-by: oPluss <opluss@qq.com>
@OYCN OYCN force-pushed the dev-py-schema-inference branch from 20106ab to 16ac26a Compare April 20, 2025 12:58
@OYCN
Fix build error
32b0f0e 
Signed-off-by: oPluss <opluss@qq.com>
@OYCN
Copy link
Contributor Author

OYCN commented Apr 21, 2025

Hi @justinchuby , the PR has been ready for review. If you have any suggestion, please feel free to comment.

@justinchuby justinchuby added the module: shape inference Issues related to shape inference label Apr 21, 2025
justinchuby
justinchuby reviewed Apr 24, 2025
View reviewed changes
justinchuby
justinchuby reviewed Apr 24, 2025
View reviewed changes
@justinchuby
Copy link
Member

I think the main changes look good! Sorry for the delay. Just some minor nits

@github-project-automation github-project-automation bot moved this from In progress to Reviewer approved in PR Tracker Apr 24, 2025
justinchuby
justinchuby reviewed Apr 24, 2025
View reviewed changes
@justinchuby
Copy link
Member

cc @gramalingam @xadupre

@OYCN
rename
66a5274 
Signed-off-by: oPluss <opluss@qq.com>
@OYCN OYCN force-pushed the dev-py-schema-inference branch from ffd61fe to 66a5274 Compare April 26, 2025 03:39
@justinchuby
Copy link
Member

Need and additional review @xadupre @gramalingam

@xadupre
Copy link
Contributor

xadupre commented Apr 28, 2025

It would be great to add an example in the documentation. Let's merge this one but a standalone example would be great.

@xadupre xadupre added this pull request to the merge queue Apr 28, 2025
Merged via the queue into onnx:main with commit c0107db Apr 28, 2025
39 checks passed
@github-project-automation github-project-automation bot moved this from Reviewer approved to Done in PR Tracker Apr 28, 2025
gramalingam pushed a commit to gramalingam/onnx that referenced this pull request May 19, 2025
@OYCN @gramalingam
Support set schema inference function in python (onnx#5940)
a622946 
### Description

Support set schema inference callback function for custion OpSchema
before register in python.

example:

```
def func(ctx: onnx.shape_inference.InferenceContext):
    # get some info
    assert ctx.get_num_inputs() == 2
    value = ctx.get_input_type(0)
    ...
    # get or create output proto object
    output = ctx.get_output_type(0)
    # set type or shape
    ...
    # set the result proto
    ctx.set_output_type(0, output)

schema.set_type_and_shape_inference_function(func)
```

### Note

Depends on onnx#5906

### Motivation and Context

Follow up of onnx#5019

---------

Signed-off-by: oPluss <opluss@qq.com>
Signed-off-by: gramalingam <grama@microsoft.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@justinchuby justinchuby justinchuby approved these changes

@xadupre xadupre Awaiting requested review from xadupre

@gramalingam gramalingam Awaiting requested review from gramalingam

Assignees

@gramalingam gramalingam

@xadupre xadupre

Labels
module: shape inference Issues related to shape inference run release CIs Use this label to trigger release tests in CI
Projects
PR Tracker
Status: Done
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@OYCN @justinchuby @xadupre @gramalingam @andife