[Feature request] Make OpSchema writable from Python

[justinchuby]

What is the problem that this feature solves?

Not being able to auto create OpSchema for functions in Python; not being able to do shape inference for custom ops.

Alternatives considered

Create a Python class that mirrors OpSchema. We can do that and expose a similar interface, but for onnx shape inference to work, we need to be able to create an OpSchema, register it and write shape inference functions in Python.

Describe the feature

Extend the OpSchema to be writable, and support registering this schema with ONNX.

• Allow new OpSchemas to be initialized in Python #5020
• Update documentation for OpSchema #5292
• repr methods for the classes Create repr functions for OpSchema #5117
• OpSchema registration Support register custom OpSchema by python #5906
• Shape inference binding: create inference code with Python

Will this influence the current api (Y/N)?

Y. It will be an extension of the current OpSchema api.

Feature Area

API

Are you willing to contribute it (Y/N)

Yes

Notes

assigned to @justinchuby

[jbachurski]

I think this is a really interesting proposal towards recognising custom operators. Spox currently uses ONNX OpSchemas to generate opset modules, and as such custom operator libraries could register more and we could use Spox utilities directly to generate opset modules.
Maybe we could take this a step further and consider allowing defining/overriding type and shape inference in a similar manner? Then we wouldn’t be locked to custom operators always failing shape inference internally.

[justinchuby]

I think that’s a great idea

[justinchuby]

We will create a single constructor and allow only shape inference overrides instead

[justinchuby]

https://pybind11-jagerman.readthedocs.io/en/stable/advanced.html#custom-constructors

[justinchuby]

Documentation: FormalParameter