Support set schema inference function in python

.gitmodules

@@ -2,3 +2,6 @@
     path = third_party/pybind11
     url = https://github.com/pybind/pybind11.git
     branch = master
+[submodule "third_party/pybind11_protobuf"]
+	path = third_party/pybind11_protobuf
+	url = https://github.com/pybind/pybind11_protobuf.git

CMakeLists.txt

@@ -541,7 +541,25 @@ if(BUILD_ONNX_PYTHON)
     endif()
   endif()
 
-  add_library(onnx_cpp2py_export MODULE "${ONNX_ROOT}/onnx/cpp2py_export.cc")
+  set(SOURCE_FILE "${ONNX_ROOT}/onnx/cpp2py_export.cc")
+
+  # In order to keep the versions of the third-party libraries consistent,
+  # use the source files directly.
+  set(PYBIND11_PROTOBUF_DIR "${ONNX_ROOT}/third_party/pybind11_protobuf/")
+  if(EXISTS "${PYBIND11_PROTOBUF_DIR}")
+    set(SOURCE_FILE
+      "${SOURCE_FILE}"
+      ${PYBIND11_PROTOBUF_DIR}/pybind11_protobuf/native_proto_caster.h
+      ${PYBIND11_PROTOBUF_DIR}/pybind11_protobuf/enum_type_caster.h
+      ${PYBIND11_PROTOBUF_DIR}/pybind11_protobuf/proto_cast_util.cc
+      ${PYBIND11_PROTOBUF_DIR}/pybind11_protobuf/proto_cast_util.h
+      ${PYBIND11_PROTOBUF_DIR}/pybind11_protobuf/proto_caster_impl.h
+    )
+  else()
+    message(FATAL_ERROR "cannot find pybind11_protobuf")
+  endif()
+
+  add_library(onnx_cpp2py_export MODULE "${SOURCE_FILE}")
   set_target_properties(onnx_cpp2py_export PROPERTIES PREFIX "")
   set_target_properties(onnx_cpp2py_export
                         PROPERTIES COMPILE_FLAGS "-fvisibility=hidden")
@@ -565,7 +583,8 @@ if(BUILD_ONNX_PYTHON)
 
   target_include_directories(onnx_cpp2py_export PUBLIC
     "${pybind11_INCLUDE_DIRS}"
-    "${PYTHON_INCLUDE_DIRS}")
+    "${PYTHON_INCLUDE_DIRS}"
+    "${PYBIND11_PROTOBUF_DIR}")
 
   if(APPLE)
     set_target_properties(onnx_cpp2py_export

onnx/cpp2py_export.cc

@@ -2,8 +2,10 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
+#include <pybind11_protobuf/native_proto_caster.h>
 
 #include <climits>
 #include <limits>
@@ -15,6 +17,7 @@
 #include "onnx/defs/parser.h"
 #include "onnx/defs/printer.h"
 #include "onnx/defs/schema.h"
+#include "onnx/defs/shape_inference.h"
 #include "onnx/inliner/inliner.h"
 #include "onnx/py_utils.h"
 #include "onnx/shape_inference/implementation.h"
@@ -104,6 +107,8 @@ std::unordered_map<std::string, py::bytes> CallNodeInferenceFunction(
 }
 
 PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
+  pybind11_protobuf::ImportNativeProtoCasters();
+
   onnx_cpp2py_export.doc() = "Python interface to ONNX";
 
   onnx_cpp2py_export.attr("ONNX_ML") = py::bool_(
@@ -114,6 +119,9 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
 #endif // ONNX_ML
   );
 
+  // Avoid Segmentation fault if we not free the python function in Custom Schema
+  onnx_cpp2py_export.add_object("_cleanup", py::capsule([] { OpSchemaRegistry::OpSchemaDeregisterAll(); }));
+
   // Submodule `schema`
   auto defs = onnx_cpp2py_export.def_submodule("defs");
   defs.doc() = "Schema submodule";
@@ -394,7 +402,14 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
               func_proto.SerializeToString(&func_bytes);
             }
             return py::bytes(func_bytes);
-          });
+          })
+      .def(
+          "set_type_and_shape_inference_function",
+          [](OpSchema& op, const std::function<void(InferenceContext*)>& func) -> OpSchema& {
+            auto wrapper = [=](InferenceContext& ctx) { func(&ctx); };
+            return op.TypeAndShapeInferenceFunction(wrapper);
+          },
+          py::return_value_policy::reference_internal);
 
   defs.def(
           "has_schema",
@@ -625,6 +640,29 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
   shape_inference.doc() = "Shape Inference submodule";
   py::register_exception<InferenceError>(shape_inference, "InferenceError");
 
+  py::class_<InferenceContext> inference_ctx(shape_inference, "InferenceContext", "Inference context");
+
+  inference_ctx.def("get_attribute", &InferenceContext::getAttribute);
+  inference_ctx.def("get_num_inputs", &InferenceContext::getNumInputs);
+  inference_ctx.def("has_input", &InferenceContext::hasInput);
+  inference_ctx.def("get_input_type", &InferenceContext::getInputType);
+  inference_ctx.def("get_input_data", &InferenceContext::getInputData);
+  inference_ctx.def("get_input_sparse_data", &InferenceContext::getInputSparseData);
+  inference_ctx.def("get_symbolic_input", &InferenceContext::getSymbolicInput);
+  inference_ctx.def("get_graph_attribute_inferencer", &InferenceContext::getGraphAttributeInferencer);
+  inference_ctx.def("get_num_outputs", &InferenceContext::getNumOutputs);
+  inference_ctx.def("get_output_type", &InferenceContext::getOutputType, py::return_value_policy::reference);
+  inference_ctx.def("set_output_type", [](InferenceContext& self, size_t idx, const TypeProto& src) {
+    auto* dst = self.getOutputType(idx);
+    if (dst == nullptr || dst == &src) {
+      return;
+    }
+    dst->CopyFrom(src);
+  });
+
+  py::class_<GraphInferencer> graph_inferencer(shape_inference, "GraphInferencer", "Graph Inferencer");
+  graph_inferencer.def("do_inferencing", &GraphInferencer::doInferencing);
+
   shape_inference.def(
       "infer_shapes",
       [](const py::bytes& bytes, bool check_type, bool strict_mode, bool data_prop) {

onnx/onnx_cpp2py_export/defs.pyi

@@ -68,6 +68,7 @@ class OpSchema:
     ) -> dict[str, bytes]: ...
     @property
     def function_body(self) -> FunctionProto: ...
+    def set_type_and_shape_inference_function(self, func) -> None: ...
 
     class TypeConstraintParam:
         def __init__(

onnx/onnx_cpp2py_export/shape_inference.pyi

@@ -1,6 +1,25 @@
+from types import List
+
+from onnx import AttributeProto, TypeProto, TensorProto, SparseTensorProto, TensorShapeProto
 
 class InferenceError(Exception): ...
 
+class GraphInferencer:
+    def do_inferencing(self, input_types: List[TypeProto], input_data: List[TensorProto]) -> List[TypeProto]: ...
+
+class InferenceContext:
+    def get_num_inputs(self) -> int: ...
+    def has_input(self, idx: int) -> bool: ...
+    def get_num_outputs(self) -> int: ...
+    def get_attribute(self, name: str) -> AttributeProto: ...
+    def get_input_type(self, idx: int) -> TypeProto: ...
+    def get_input_data(self, idx: int) -> TensorProto: ...
+    def get_input_sparse_data(self, idx: int) -> SparseTensorProto: ...
+    def get_symbolic_input(self, idx: int) -> TensorShapeProto: ...
+    def get_graph_attribute_inferencer(self) -> GraphInferencer: ...
+    def get_output_type(self, idx: int) -> TypeProto: ...
+    def set_output_type(self, idx: int, type: TypeProto) -> None: ...
+
 def infer_shapes(
     b: bytes, check_type: bool, strict_mode: bool, data_prop: bool
 ) -> bytes: ...

onnx/shape_inference.py

@@ -16,6 +16,9 @@
 import onnx.onnx_cpp2py_export.shape_inference as C  # noqa: N812
 from onnx import AttributeProto, FunctionProto, ModelProto, TypeProto
 
+GraphInferencer = C.GraphInferencer
+InferenceContext = C.InferenceContext
+
 
 def infer_shapes(
     model: ModelProto | bytes,

onnx/test/shape_inference_test.py

@@ -10172,6 +10172,103 @@ def test_check_type_when_schema_has_empty_io(self):
             op_schema.name, op_schema.since_version, op_schema.domain
         )
 
+    def test_custom_schema_shape_inference(self) -> None:
+        # CustomOp schema:
+        #   attrs:
+        #       out_len: [L0, L1, ...]
+        #   inputs:
+        #       a[N, La]
+        #       b[N, Lb]
+        #   outputs:
+        #       out0[N, La * Lb, L0]
+        #       out1[N, La * Lb, L1]
+        #       ...
+        N = 3
+        La = 32
+        Lb = 64
+        out_len = [1, 2]
+        outs = [f"out{i}" for i in range(len(out_len))]
+        graph = self._make_graph(
+            [
+                ("a", TensorProto.FLOAT, (N, La)),
+                ("b", TensorProto.FLOAT, (N, Lb)),
+            ],
+            [make_node("CustomOp", ["a", "b"], outs, out_len=out_len)],
+            [],
+        )
+        with self.assertRaises(onnx.checker.ValidationError):
+            self._assert_inferred(
+                graph,
+                [
+                    make_tensor_value_info(
+                        f"out{i}", TensorProto.FLOAT, (N, La * Lb, Li)
+                    )
+                    for i, Li in enumerate(out_len)
+                ],
+            )
+
+        schema = OpSchema(
+            "CustomOp",
+            "",
+            1,
+            inputs=[
+                defs.OpSchema.FormalParameter("a", "float"),
+                defs.OpSchema.FormalParameter("b", "float"),
+            ],
+            outputs=[
+                defs.OpSchema.FormalParameter(
+                    "out", "float", param_option=OpSchema.FormalParameterOption.Variadic
+                ),
+            ],
+            attributes=[
+                defs.OpSchema.Attribute("out_len", defs.OpSchema.AttrType.INTS)
+            ],
+        )
+
+        def func(ctx: onnx.shape_inference.InferenceContext):
+            def parse_tensor_input(t: TypeProto):
+                assert isinstance(t, TypeProto)
+                return (
+                    t.tensor_type.elem_type,
+                    [
+                        d.dim_value if d.HasField("dim_value") else None
+                        for d in t.tensor_type.shape.dim
+                    ],
+                )
+
+            assert ctx.get_num_inputs() == 2
+            in0 = ctx.get_input_type(0)
+            in1 = ctx.get_input_type(1)
+            in0_type, in0_shape = parse_tensor_input(in0)
+            in1_type, in1_shape = parse_tensor_input(in1)
+            assert in0_type == in1_type == TensorProto.FLOAT
+            assert len(in0_shape) == len(in1_shape) == 2
+            assert in0_shape[0] == in1_shape[0]
+            N, La = in0_shape
+            _, Lb = in1_shape
+            attr = ctx.get_attribute("out_len")
+            out_len = attr.ints
+            assert len(out_len) == ctx.get_num_outputs()
+            for i in range(ctx.get_num_outputs()):
+                out = ctx.get_output_type(i)
+                out.tensor_type.elem_type = in0_type
+                out.tensor_type.shape.dim.add().dim_value = N
+                out.tensor_type.shape.dim.add().dim_value = La * Lb
+                out.tensor_type.shape.dim.add().dim_value = out_len[i]
+                ctx.set_output_type(i, out)
+
+        schema.set_type_and_shape_inference_function(func)
+        onnx.defs.register_schema(schema)
+
+        self._assert_inferred(
+            graph,
+            [
+                make_tensor_value_info(f"out{i}", TensorProto.FLOAT, (N, La * Lb, Li))
+                for i, Li in enumerate(out_len)
+            ],
+        )
+        onnx.defs.deregister_schema(schema.name, schema.since_version, schema.domain)
+
 
 if __name__ == "__main__":
     unittest.main()