Support set schema inference function in python

onnx/cpp2py_export.cc

@@ -2,6 +2,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#include <pybind11/functional.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -394,6 +395,12 @@
               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) {
+            auto wrapper = [=](InferenceContext& ctx) { func(&ctx); };
+            return op->TypeAndShapeInferenceFunction(wrapper);
           });
 
   defs.def(
@@ -625,6 +632,85 @@
   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* ctx, std::string name) {
+    auto attr = ctx->getAttribute(name);
+    std::string data;
+    attr->SerializeToString(&data);
+    return py::bytes(data);
+  });
+  inference_ctx.def("get_num_inputs", &InferenceContext::getNumInputs);
+  inference_ctx.def("has_input", &InferenceContext::hasInput);
+  inference_ctx.def("__get_input_type", [](InferenceContext* ctx, size_t index) {
+    auto type = ctx->getInputType(index);
+    std::string data;
+    type->SerializeToString(&data);
+    return py::bytes(data);
+  });
+  inference_ctx.def("__get_input_data", [](InferenceContext* ctx, size_t index) {
+    auto tensor = ctx->getInputData(index);
+    std::string data;
+    tensor->SerializeToString(&data);
+    return py::bytes(data);
+  });
+  inference_ctx.def("__get_input_sparse_data", [](InferenceContext* ctx, size_t index) {
+    auto stensor = ctx->getInputSparseData(index);
+    std::string data;
+    stensor->SerializeToString(&data);
+    return py::bytes(data);
+  });
+  inference_ctx.def("__get_symbolic_input", [](InferenceContext* ctx, size_t index) {
+    auto shape = ctx->getSymbolicInput(index);
+    std::string data;
+    shape->SerializeToString(&data);
+    return py::bytes(data);
+  });
+  inference_ctx.def("__get_graph_attribute_inferencer", &InferenceContext::getGraphAttributeInferencer);
+  inference_ctx.def("get_num_outputs", &InferenceContext::getNumOutputs);
+  inference_ctx.def("__get_output_type", [](InferenceContext* ctx, size_t index) {
+    auto type = ctx->getOutputType(index);
+    std::string data;
+    type->SerializeToString(&data);
+    return py::bytes(data);
+  });
+  inference_ctx.def("__set_output_type", [](InferenceContext* ctx, size_t index, py::bytes bytes) {
+    auto type = ctx->getOutputType(index);
+    ParseProtoFromPyBytes(type, bytes);
+  });
+
+  py::class_<GraphInferencer> graph_inferencer(shape_inference, "GraphInferencer", "Graph Inferencer");
+  graph_inferencer.def(
+      "__do_inferencing",
+      [](GraphInferencer* inferencer,
+         const std::vector<py::bytes>& input_types,
+         const std::vector<py::bytes>& input_data) {
+        std::vector<TypeProto> type_proto;
+        std::vector<TensorProto> tensor_proto;
+        std::vector<const TypeProto*> type_inputs;
+        std::vector<const TensorProto*> tensor_inputs;
+        for (const auto& bytes : input_types) {
+          TypeProto proto{};
+          ParseProtoFromPyBytes(&proto, bytes);
+          type_proto.emplace_back(proto);
+          type_inputs.emplace_back(&type_proto.back());
+        }
+        for (const auto& bytes : input_data) {
+          TensorProto proto{};
+          ParseProtoFromPyBytes(&proto, bytes);
+          tensor_proto.emplace_back(proto);
+          tensor_inputs.emplace_back(&tensor_proto.back());
+        }
+        auto ret = inferencer->doInferencing(type_inputs, tensor_inputs);
+        std::vector<py::bytes> out;
+        for (const auto& type : ret) {
+          std::string data;
+          type->SerializeToString(&data);
+          out.emplace_back(py::bytes(data));
+        }
+        return out;
+      });
+
   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,38 @@
+from types import List
+
+from onnx import AttributeProto, TypeProto, TensorProto, SparseTensorProto, TensorShapeProto
 
 class InferenceError(Exception): ...
 
+class GraphInferencer:
+    # Impl in cpp (onnx/cpp2py_export.cc)
+    def __do_inferencing(self, input_types: List[bytes], input_data: List[bytes]) -> List[bytes]: ...
+    # Impl in py (onnx/shape_inference.py)
+    def do_inferencing(self, input_types: List[TypeProto], input_data: List[TensorProto]) -> List[TypeProto]: ...
+
+class InferenceContext:
+    # Impl in cpp  (onnx/cpp2py_export.cc)
+    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) -> bytes: ...
+    def __get_input_type(self, idx: int) -> bytes: ...
+    def __get_input_data(self, idx: int) -> bytes: ...
+    def __get_input_sparse_data(self, idx: int) -> bytes: ...
+    def __get_symbolic_input(self, idx: int) -> bytes: ...
+    def __get_graph_attribute_inferencer(self) -> GraphInferencer: ...
+    def __get_output_type(self, idx: int) -> bytes: ...
+    def __set_output_type(self, idx: int, output: bytes) -> None: ...
+    # Impl in py (onnx/shape_inference.py)
+    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, output: TypeProto) -> None: ...
+
 def infer_shapes(
     b: bytes, check_type: bool, strict_mode: bool, data_prop: bool
 ) -> bytes: ...

onnx/shape_inference.py

@@ -14,7 +14,58 @@
 
 import onnx
 import onnx.onnx_cpp2py_export.shape_inference as C  # noqa: N812
-from onnx import AttributeProto, FunctionProto, ModelProto, TypeProto
+from onnx import (
+    AttributeProto,
+    FunctionProto,
+    ModelProto,
+    SparseTensorProto,
+    TensorProto,
+    TensorShapeProto,
+    TypeProto,
+)
+
+GraphInferencer = C.GraphInferencer
+
+
+def _do_inferencing(
+    self, input_types: list[TypeProto], input_data: list[TensorProto]
+) -> list[TypeProto]:
+    input_types_bytes = [proto.SerializeToString() for proto in input_types]
+    input_data_bytes = [proto.SerializeToString() for proto in input_data]
+    ret = self.__impl.__do_inferencing(input_types_bytes, input_data_bytes)
+    return [TypeProto.FromString(data) for data in ret]
+
+
+GraphInferencer.do_inferencing = _do_inferencing  # type: ignore
+
+
+def _parse_to_proto(attr, proto_type):
+    def impl(self, *args, **kwargs):
+        data = getattr(self, attr)(*args, **kwargs)
+        return proto_type.FromString(data)
+
+    return impl
+
+
+InferenceContext = C.InferenceContext
+InferenceContext.get_attribute = _parse_to_proto("__get_attribute", AttributeProto)  # type: ignore
+InferenceContext.get_input_type = _parse_to_proto("__get_input_type", TypeProto)  # type: ignore
+InferenceContext.get_input_data = _parse_to_proto("__get_input_data", TensorProto)  # type: ignore
+InferenceContext.get_input_sparse_data = _parse_to_proto(  # type: ignore
+    "__get_input_sparse_data", SparseTensorProto
+)
+InferenceContext.get_symbolic_input = _parse_to_proto(  # type: ignore
+    "__get_symbolic_input", TensorShapeProto
+)
+InferenceContext.get_output_type = _parse_to_proto("__get_output_type", TypeProto)  # type: ignore
+
+
+def _op_set_output_type(self, idx: int, output: TypeProto):
+    data = output.SerializeToString()
+    self.__set_output_type(idx, data)
+
+
+InferenceContext.set_output_type = _op_set_output_type  # type: ignore
 
 
 def infer_shapes(