onnx 1.18.0__cp310-cp310-win32.whl → 1.19.1rc1__cp310-cp310-win32.whl

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>
 
@@ -17,6 +18,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"
@@ -26,6 +28,61 @@
 #pragma error "Pybind11 must be >= 2.12 to be compatible with numpy 2.0."
 #endif
 
+#ifndef PYBIND11_DETAILED_ERROR_MESSAGES
+#define PYBIND11_DETAILED_ERROR_MESSAGES
+#endif
+
+template <typename T>
+struct PythonProtoTypeMap {};
+
+#define DEFINE_PROTO_TYPE_MAP(_ProtoType, PY_MODULE_NAME, PY_TYPE_NAME)                             \
+  template <>                                                                                       \
+  struct PythonProtoTypeMap<_ProtoType> {                                                           \
+    static constexpr auto FullName = pybind11::detail::const_name(PY_MODULE_NAME "." PY_TYPE_NAME); \
+    static constexpr auto TypeName = pybind11::detail::const_name(PY_TYPE_NAME);                    \
+    static constexpr auto ModuleName = pybind11::detail::const_name(PY_MODULE_NAME);                \
+  };
+
+#ifdef ONNX_USE_LITE_PROTO
+using BASE_PROTO_TYPE = ::google::protobuf::MessageLite;
+#else
+using BASE_PROTO_TYPE = ::google::protobuf::Message;
+#endif
+
+template <typename _ProtoType>
+class pybind11::detail::type_caster<_ProtoType, std::enable_if_t<std::is_base_of<BASE_PROTO_TYPE, _ProtoType>::value>> {
+ public:
+  PYBIND11_TYPE_CASTER(_ProtoType, PythonProtoTypeMap<_ProtoType>::FullName);
+  bool load(handle py_proto, bool) {
+    try {
+      if (!pybind11::hasattr(py_proto, "SerializeToString")) {
+        return false;
+      }
+      pybind11::bytes serialized = py_proto.attr("SerializeToString")();
+      std::string serialized_str = serialized;
+      if (!value.ParseFromString(serialized_str)) {
+        return false;
+      }
+      return true;
+    } catch (const pybind11::error_already_set&) {
+      return false;
+    }
+  }
+  static handle cast(const _ProtoType& cpp_proto, return_value_policy /* policy */, handle /* parent */) {
+    auto py_proto = pybind11::module::import(PythonProtoTypeMap<_ProtoType>::ModuleName.text)
+                        .attr(PythonProtoTypeMap<_ProtoType>::TypeName.text)();
+    std::string serialized = cpp_proto.SerializeAsString();
+    py_proto.attr("ParseFromString")(pybind11::bytes(serialized));
+    return py_proto.release();
+  }
+};
+
+DEFINE_PROTO_TYPE_MAP(ONNX_NAMESPACE::AttributeProto, "onnx", "AttributeProto");
+DEFINE_PROTO_TYPE_MAP(ONNX_NAMESPACE::TypeProto, "onnx", "TypeProto");
+DEFINE_PROTO_TYPE_MAP(ONNX_NAMESPACE::TensorProto, "onnx", "TensorProto");
+DEFINE_PROTO_TYPE_MAP(ONNX_NAMESPACE::SparseTensorProto, "onnx", "SparseTensorProto");
+DEFINE_PROTO_TYPE_MAP(ONNX_NAMESPACE::TensorShapeProto, "onnx", "TensorShapeProto");
+
 namespace ONNX_NAMESPACE {
 namespace py = pybind11;
 using namespace pybind11::literals;
@@ -110,6 +167,23 @@ static std::unordered_map<std::string, py::bytes> CallNodeInferenceFunction(
   return typeProtoBytes;
 }
 
+template <typename T>
+static std::tuple<std::vector<T>, std::vector<const T*>> ConvertPyObjToPtr(const std::vector<py::object>& pyObjs) {
+  std::vector<T> objs;
+  std::vector<const T*> ptrs;
+  objs.reserve(pyObjs.size());
+  ptrs.reserve(pyObjs.size());
+  for (const auto& obj : pyObjs) {
+    if (obj.is_none()) {
+      ptrs.push_back(nullptr);
+      continue;
+    }
+    objs.emplace_back(obj.cast<T>());
+    ptrs.push_back(&objs.back());
+  }
+  return std::make_tuple(std::move(objs), std::move(ptrs));
+}
+
 PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
   onnx_cpp2py_export.doc() = "Python interface to ONNX";
 
@@ -121,6 +195,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";
@@ -401,7 +478,15 @@ 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)
+      .def("get_type_and_shape_inference_function", &OpSchema::GetTypeAndShapeInferenceFunction);
 
   defs.def(
           "has_schema",
@@ -623,7 +708,19 @@ PYBIND11_MODULE(onnx_cpp2py_export, onnx_cpp2py_export) {
         ModelProto model{};
         ParseProtoFromPyBytes(&model, bytes);
         auto function_id_set = inliner::FunctionIdSet::Create(std::move(function_ids), exclude);
-        inliner::InlineSelectedFunctions(model, *function_id_set);
+        inliner::InlineSelectedLocalFunctions(model, *function_id_set);
+        std::string out;
+        model.SerializeToString(&out);
+        return py::bytes(out);
+      });
+
+  inliner.def(
+      "inline_selected_functions2",
+      [](const py::bytes& bytes, std::vector<std::pair<std::string, std::string>> function_ids, bool exclude) {
+        ModelProto model{};
+        ParseProtoFromPyBytes(&model, bytes);
+        auto function_id_set = inliner::FunctionIdSet::Create(std::move(function_ids), exclude);
+        inliner::InlineSelectedFunctions(model, *function_id_set, nullptr);
         std::string out;
         model.SerializeToString(&out);
         return py::bytes(out);
@@ -634,6 +731,48 @@ 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_context(shape_inference, "InferenceContext", "Inference context");
+
+  inference_context.def("get_attribute", &InferenceContext::getAttribute);
+  inference_context.def("get_num_inputs", &InferenceContext::getNumInputs);
+  inference_context.def("get_input_type", &InferenceContext::getInputType);
+  inference_context.def("has_input", &InferenceContext::hasInput);
+  inference_context.def("get_input_data", &InferenceContext::getInputData);
+  inference_context.def("get_num_outputs", &InferenceContext::getNumOutputs);
+  inference_context.def("get_output_type", &InferenceContext::getOutputType);
+  inference_context.def("set_output_type", [](InferenceContext& self, size_t idx, const TypeProto& src) {
+    auto* dst = self.getOutputType(idx);
+    if (dst == nullptr) {
+      return false;
+    }
+    dst->CopyFrom(src);
+    return true;
+  });
+  inference_context.def("has_output", &InferenceContext::hasOutput);
+  inference_context.def(
+      "get_graph_attribute_inferencer",
+      &InferenceContext::getGraphAttributeInferencer,
+      py::return_value_policy::reference_internal);
+  inference_context.def("get_input_sparse_data", &InferenceContext::getInputSparseData);
+  inference_context.def("get_symbolic_input", &InferenceContext::getSymbolicInput);
+  inference_context.def("get_display_name", &InferenceContext::getDisplayName);
+
+  py::class_<GraphInferencer> graph_inferencer(shape_inference, "GraphInferencer", "Graph Inferencer");
+  graph_inferencer.def(
+      "do_inferencing",
+      [](GraphInferencer& self,
+         const std::vector<py::object>& inputTypesObj,
+         const std::vector<py::object>& inputDataObj) {
+        auto inputTypesTuple = ConvertPyObjToPtr<ONNX_NAMESPACE::TypeProto>(inputTypesObj);
+        auto inputDataTuple = ConvertPyObjToPtr<ONNX_NAMESPACE::TensorProto>(inputDataObj);
+        auto ret = self.doInferencing(std::get<1>(inputTypesTuple), std::get<1>(inputDataTuple));
+        std::vector<py::object> ret_obj(ret.size());
+        for (size_t i = 0; i < ret.size(); ++i) {
+          ret_obj[i] = py::cast(ret[i]);
+        }
+        return ret_obj;
+      });
+
   shape_inference.def(
       "infer_shapes",
       [](const py::bytes& bytes, bool check_type, bool strict_mode, bool data_prop) {

onnx/defs/__init__.py

@@ -20,9 +20,8 @@ __all__ = [
     "SchemaError",
 ]
 
-
+import onnx
 import onnx.onnx_cpp2py_export.defs as C  # noqa: N812
-from onnx import AttributeProto, FunctionProto
 
 ONNX_DOMAIN = ""
 ONNX_ML_DOMAIN = "ai.onnx.ml"
@@ -48,7 +47,7 @@ def onnx_ml_opset_version() -> int:
 
 @property  # type: ignore
 def _function_proto(self):
-    func_proto = FunctionProto()
+    func_proto = onnx.FunctionProto()
     func_proto.ParseFromString(self._function_body)
     return func_proto
 
@@ -59,7 +58,7 @@ OpSchema.function_body = _function_proto  # type: ignore
 
 @property  # type: ignore
 def _attribute_default_value(self):
-    attr = AttributeProto()
+    attr = onnx.AttributeProto()
     attr.ParseFromString(self._default_value)
     return attr
 

onnx/defs/attr_proto_util.cc

@@ -7,46 +7,49 @@
 #include "attr_proto_util.h"
 
 #include <string>
+#include <utility>
 #include <vector>
 
 namespace ONNX_NAMESPACE {
 
-#define ADD_BASIC_ATTR_IMPL(type, enumType, field)                                \
-  AttributeProto MakeAttribute(const std::string& attr_name, const type& value) { \
-    AttributeProto a;                                                             \
-    a.set_name(attr_name);                                                        \
-    a.set_type(enumType);                                                         \
-    a.set_##field(value);                                                         \
-    return a;                                                                     \
+#define ADD_BASIC_ATTR_IMPL(type, enumType, field)                  \
+  AttributeProto MakeAttribute(std::string attr_name, type value) { \
+    AttributeProto a;                                               \
+    a.set_name(std::move(attr_name));                               \
+    a.set_type(enumType);                                           \
+    a.set_##field(value);                                           \
+    return a;                                                       \
   }
 
-#define ADD_ATTR_IMPL(type, enumType, field)                                      \
-  AttributeProto MakeAttribute(const std::string& attr_name, const type& value) { \
+#define ADD_ATTR_IMPL(type, enumType, field)                        \
+  AttributeProto MakeAttribute(std::string attr_name, type value) { \
+    AttributeProto a;                                               \
+    a.set_name(std::move(attr_name));                               \
+    a.set_type(enumType);                                           \
+    *(a.mutable_##field()) = std::move(value);                      \
+    return a;                                                       \
+  }
+
+#define ADD_LIST_ATTR_IMPL(type, enumType, field)                                 \
+  AttributeProto MakeAttribute(std::string attr_name, std::vector<type> values) { \
     AttributeProto a;                                                             \
-    a.set_name(attr_name);                                                        \
+    a.set_name(std::move(attr_name));                                             \
     a.set_type(enumType);                                                         \
-    *(a.mutable_##field()) = value;                                               \
+    for (auto&& val : std::move(values)) {                                        \
+      *(a.mutable_##field()->Add()) = std::move(val);                             \
+    }                                                                             \
     return a;                                                                     \
   }
 
-#define ADD_LIST_ATTR_IMPL(type, enumType, field)                                               \
-  AttributeProto MakeAttribute(const std::string& attr_name, const std::vector<type>& values) { \
-    AttributeProto a;                                                                           \
-    a.set_name(attr_name);                                                                      \
-    a.set_type(enumType);                                                                       \
-    for (const auto& val : values) {                                                            \
-      *(a.mutable_##field()->Add()) = val;                                                      \
-    }                                                                                           \
-    return a;                                                                                   \
-  }
-
 ADD_BASIC_ATTR_IMPL(float, AttributeProto_AttributeType_FLOAT, f)
 ADD_BASIC_ATTR_IMPL(int64_t, AttributeProto_AttributeType_INT, i)
 ADD_BASIC_ATTR_IMPL(std::string, AttributeProto_AttributeType_STRING, s)
 ADD_ATTR_IMPL(TensorProto, AttributeProto_AttributeType_TENSOR, t)
 ADD_ATTR_IMPL(GraphProto, AttributeProto_AttributeType_GRAPH, g)
 ADD_ATTR_IMPL(TypeProto, AttributeProto_AttributeType_TYPE_PROTO, tp)
+// NOLINTNEXTLINE(performance-unnecessary-value-param)
 ADD_LIST_ATTR_IMPL(float, AttributeProto_AttributeType_FLOATS, floats)
+// NOLINTNEXTLINE(performance-unnecessary-value-param)
 ADD_LIST_ATTR_IMPL(int64_t, AttributeProto_AttributeType_INTS, ints)
 ADD_LIST_ATTR_IMPL(std::string, AttributeProto_AttributeType_STRINGS, strings)
 ADD_LIST_ATTR_IMPL(TensorProto, AttributeProto_AttributeType_TENSORS, tensors)

onnx/defs/attr_proto_util.h

@@ -13,16 +13,18 @@
 
 namespace ONNX_NAMESPACE {
 
-AttributeProto MakeAttribute(const std::string& attr_name, const float& value);
-AttributeProto MakeAttribute(const std::string& attr_name, const int64_t& value);
-AttributeProto MakeAttribute(const std::string& attr_name, const std::string& value);
-AttributeProto MakeAttribute(const std::string& attr_name, const TensorProto& value);
-AttributeProto MakeAttribute(const std::string& attr_name, const GraphProto& value);
-AttributeProto MakeAttribute(const std::string& attr_name, const std::vector<float>& values);
-AttributeProto MakeAttribute(const std::string& attr_name, const std::vector<int64_t>& values);
-AttributeProto MakeAttribute(const std::string& attr_name, const std::vector<std::string>& values);
-AttributeProto MakeAttribute(const std::string& attr_name, const std::vector<TensorProto>& values);
-AttributeProto MakeAttribute(const std::string& attr_name, const std::vector<GraphProto>& values);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, float value);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, int64_t value);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::string value);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, TensorProto value);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, GraphProto value);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, TypeProto value);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::vector<float> values);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::vector<int64_t> values);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::vector<std::string> values);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::vector<TensorProto> values);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::vector<GraphProto> values);
+ONNX_API AttributeProto MakeAttribute(std::string attr_name, std::vector<TypeProto> values);
 
 // Make a "reference" attribute for a node in a function body.
 // <attr_name> specifies the attribute name of both the function node and its

onnx/defs/controlflow/defs.cc

@@ -12,10 +12,10 @@
 namespace ONNX_NAMESPACE {
 using SupportType = OpSchema::SupportType;
 
-static std::vector<std::string> control_flow_types_ir11() {
-  auto t = OpSchema::all_tensor_types_ir11();
-  auto s = OpSchema::all_tensor_sequence_types_ir11();
-  auto o = OpSchema::all_optional_types_ir11();
+static std::vector<std::string> control_flow_types_ir12() {
+  auto t = OpSchema::all_tensor_types_ir12();
+  auto s = OpSchema::all_tensor_sequence_types_ir12();
+  auto o = OpSchema::all_optional_types_ir12();
   t.insert(t.end(), s.begin(), s.end());
   t.insert(t.end(), o.begin(), o.end());
   return t;
@@ -23,7 +23,7 @@ static std::vector<std::string> control_flow_types_ir11() {
 
 ONNX_OPERATOR_SET_SCHEMA(
     If,
-    23,
+    24,
     OpSchema()
         .SetDoc("If conditional")
         .Input(0, "cond", "Condition for the if. The tensor must contain a single element.", "B")
@@ -63,12 +63,12 @@ ONNX_OPERATOR_SET_SCHEMA(
             AttributeProto::GRAPH)
         .TypeConstraint(
             "V",
-            control_flow_types_ir11(),
+            control_flow_types_ir12(),
             "All Tensor, Sequence(Tensor), Optional(Tensor), and Optional(Sequence(Tensor)) types up to IRv11.")
         .TypeConstraint("B", {"tensor(bool)"}, "Only bool")
         .TypeAndShapeInferenceFunction(IfInferenceFunction));
 
-static const char* Loop_ver16_doc = R"DOC(
+static const char* Loop_ver23_doc = R"DOC(
 Generic Looping construct. This loop has multiple termination conditions:
 
 1) Trip count. Iteration count specified at runtime. Set by
@@ -208,9 +208,9 @@ The input/output of subgraph (produced by loop node) matching is based on order
 
 ONNX_OPERATOR_SET_SCHEMA(
     Loop,
-    23,
+    24,
     OpSchema()
-        .SetDoc(Loop_ver16_doc)
+        .SetDoc(Loop_ver23_doc)
         .Input(
             0,
             "M",
@@ -253,13 +253,13 @@ ONNX_OPERATOR_SET_SCHEMA(
             AttributeProto::GRAPH)
         .TypeConstraint(
             "V",
-            control_flow_types_ir11(),
+            control_flow_types_ir12(),
             "All Tensor, Sequence(Tensor), Optional(Tensor), and Optional(Sequence(Tensor)) types up to IRv11.")
         .TypeConstraint("I", {"tensor(int64)"}, "tensor of int64, which should be a scalar.")
         .TypeConstraint("B", {"tensor(bool)"}, "tensor of bool, which should be a scalar.")
         .TypeAndShapeInferenceFunction(LoopInferenceFunction));
 
-static const char* scan_16_doc = R"DOC(
+static const char* scan_24_doc = R"DOC(
 Scan can be used to iterate over one or more scan_input tensors,
 constructing zero or more scan_output tensors. It combines ideas from general recurrences,
 functional programming constructs such as scan, fold, map, and zip, and is intended to enable
@@ -385,9 +385,9 @@ values are computed in the outer graph, they need to be passed in as extra state
 
 ONNX_OPERATOR_SET_SCHEMA(
     Scan,
-    23,
+    24,
     OpSchema()
-        .SetDoc(scan_16_doc)
+        .SetDoc(scan_24_doc)
         .Input(
             0,
             "initial_state_and_scan_inputs",
@@ -448,7 +448,7 @@ ONNX_OPERATOR_SET_SCHEMA(
             "range is [-r, r-1].",
             AttributeProto::INTS,
             false)
-        .TypeConstraint("V", OpSchema::all_tensor_types_ir11(), "All Tensor types up to IRv11.")
+        .TypeConstraint("V", OpSchema::all_tensor_types_ir12(), "All Tensor types up to IRv12.")
         .TypeAndShapeInferenceFunction(ScanInferenceFunction)); // Shares same shape inference as opset 11
 
 } // namespace ONNX_NAMESPACE