onnx 1.18.0__cp311-cp311-win32.whl → 1.19.0__cp311-cp311-win32.whl

onnx/serialization.py

@@ -57,7 +57,7 @@ class _Registry:
     def register(self, serializer: ProtoSerializer) -> None:
         self._serializers[serializer.supported_format] = serializer
         self._extension_to_format.update(
-            {ext: serializer.supported_format for ext in serializer.file_extensions}
+            dict.fromkeys(serializer.file_extensions, serializer.supported_format)
         )
 
     def get(self, fmt: str) -> ProtoSerializer:
@@ -118,7 +118,7 @@ class _ProtobufSerializer(ProtoSerializer):
             raise TypeError(
                 f"Parameter 'serialized' must be bytes, but got type: {type(serialized)}"
             )
-        decoded = typing.cast(Optional[int], proto.ParseFromString(serialized))
+        decoded = typing.cast("Optional[int]", proto.ParseFromString(serialized))
         if decoded is not None and decoded != len(serialized):
             raise google.protobuf.message.DecodeError(
                 f"Protobuf decoding consumed too few bytes: {decoded} out of {len(serialized)}"
@@ -130,7 +130,7 @@ class _TextProtoSerializer(ProtoSerializer):
     """Serialize and deserialize text proto."""
 
     supported_format = "textproto"
-    file_extensions = frozenset({".textproto", ".prototxt", ".pbtxt"})
+    file_extensions = frozenset({".txtpb", ".textproto", ".prototxt", ".pbtxt"})
 
     def serialize_proto(self, proto: _Proto) -> bytes:
         textproto = google.protobuf.text_format.MessageToString(proto)

onnx/shape_inference/implementation.cc

@@ -710,7 +710,7 @@ class ShapeInferenceImplBase {
   void FinalizeShapeInference() {
     auto& errors = getErrors();
     // Throw shape inference error if any. Error mode right now only supports 0 and 1.
-    // When set to 0, any node level shape inference errors are not thrown. This is to support backward compatiblity
+    // When set to 0, any node level shape inference errors are not thrown. This is to support backward compatibility
     // with 1.7 and earlier releases. When set to 1 it will throw all exceptions.
     // TODO: Add a more granular way for exception handling.
     if (!errors.empty() && options.error_mode > 0) {

onnx/shape_inference.py

@@ -14,11 +14,14 @@ from typing import TYPE_CHECKING
 
 import onnx
 import onnx.onnx_cpp2py_export.shape_inference as C  # noqa: N812
-from onnx import AttributeProto, FunctionProto, ModelProto, TypeProto
+from onnx.onnx_pb import AttributeProto, FunctionProto, ModelProto, TypeProto
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+GraphInferencer = C.GraphInferencer
+InferenceContext = C.InferenceContext
+
 
 def infer_shapes(
     model: ModelProto | bytes,

onnx/subbyte.py

@@ -22,13 +22,7 @@ UINT4_MAX = 15
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32_to_4bit_unpacked(*args, **kwargs):
-    return _float32_to_4bit_unpacked(*args, **kwargs)
-
-
-def _float32_to_4bit_unpacked(
-    x: np.ndarray | np.dtype | float, signed: bool
-) -> np.ndarray:
+def float32_to_4bit_unpacked(x: np.ndarray | float, signed: bool) -> np.ndarray:
     """Cast to 4bit via rounding and clipping (without packing).
 
     Args:
@@ -52,11 +46,7 @@ def _float32_to_4bit_unpacked(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32x2_to_4bitx2(*args, **kwargs):
-    return _float32x2_to_4bitx2(*args, **kwargs)
-
-
-def _float32x2_to_4bitx2(
+def float32x2_to_4bitx2(
     val_low: np.dtype, val_high: np.dtype, signed: bool
 ) -> np.ndarray:
     """Cast two elements to 4bit (via rounding and clipping) and pack
@@ -69,8 +59,8 @@ def _float32x2_to_4bitx2(
     Returns:
         An ndarray with a single int8/uint8 element, containing both int4 elements
     """
-    i8_high = _float32_to_4bit_unpacked(val_high, signed)
-    i8_low = _float32_to_4bit_unpacked(val_low, signed)
+    i8_high = float32_to_4bit_unpacked(val_high, signed)  # type: ignore[arg-type]
+    i8_low = float32_to_4bit_unpacked(val_low, signed)  # type: ignore[arg-type]
     return i8_high << 4 | i8_low & 0x0F
 
 
@@ -78,11 +68,7 @@ def _float32x2_to_4bitx2(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def unpack_4bitx2(*args, **kwargs):
-    return _unpack_4bitx2(*args, **kwargs)
-
-
-def _unpack_4bitx2(
+def unpack_4bitx2(
     x: npt.NDArray[np.uint8], dims: int | Sequence[int]
 ) -> npt.NDArray[np.uint8]:
     """Unpack an array of packed uint8 elements (4bitx2) into individual elements
@@ -113,11 +99,7 @@ def _unpack_4bitx2(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def unpack_single_4bitx2(*args, **kwargs):
-    return _unpack_single_4bitx2(*args, **kwargs)
-
-
-def _unpack_single_4bitx2(
+def unpack_single_4bitx2(
     x: np.ndarray | np.dtype | float, signed: bool
 ) -> tuple[np.ndarray, np.ndarray]:
     def unpack_signed(x):
@@ -144,11 +126,7 @@ def _unpack_single_4bitx2(
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32_to_float4e2m1_unpacked(*args, **kwargs):
-    return _float32_to_float4e2m1_unpacked(*args, **kwargs)
-
-
-def _float32_to_float4e2m1_unpacked(values: np.ndarray) -> np.ndarray:
+def float32_to_float4e2m1_unpacked(values: np.ndarray) -> np.ndarray:
     """Cast float32 to float4e2m1 (without packing).
 
     Args:
@@ -176,11 +154,7 @@ def _float32_to_float4e2m1_unpacked(values: np.ndarray) -> np.ndarray:
     "Deprecated since 1.18. Scheduled to remove in 1.20. Consider using libraries like ml_dtypes for dtype conversion",
     category=DeprecationWarning,
 )
-def float32x2_to_float4e2m1x2(*args, **kwargs):
-    return _float32x2_to_float4e2m1x2(*args, **kwargs)
-
-
-def _float32x2_to_float4e2m1x2(val_low: np.ndarray, val_high: np.ndarray) -> np.ndarray:
+def float32x2_to_float4e2m1x2(val_low: np.ndarray, val_high: np.ndarray) -> np.ndarray:
     """Cast two elements to float4e2m1 and pack to a single byte
     Args:
         val_low: element to be packed in the 4 LSB
@@ -189,6 +163,6 @@ def _float32x2_to_float4e2m1x2(val_low: np.ndarray, val_high: np.ndarray) -> np.
     Returns:
         An ndarray with uint8 elements, containing both float4e2m1 elements
     """
-    i8_high = _float32_to_float4e2m1_unpacked(val_high)
-    i8_low = _float32_to_float4e2m1_unpacked(val_low)
+    i8_high = float32_to_float4e2m1_unpacked(val_high)
+    i8_low = float32_to_float4e2m1_unpacked(val_low)
     return i8_high << 4 | i8_low & 0x0F

onnx/test/compose_test.py

@@ -393,7 +393,96 @@ class TestComposeFunctions(unittest.TestCase):
         m3 = compose.merge_models(m1, m2, io_map=io_map)
         checker.check_model(m3)
 
-    # FIXME: This function should be removed, as tests should not contain a copy of the tested logic.
+    def test_add_prefix_to_inputs_outputs(self) -> None:
+        """Tests prefixing inputs and outputs nodes."""
+        input_graph = """
+            <
+                ir_version: 6,
+                opset_import: [ "": 13]
+            >
+            agraph (float[N, 128] X, float[128, 10] W, float[10] B) => (float[N, 10] C)
+            {
+                T = MatMul(X, W)
+                S = Add(T, B)
+                C = Softmax(S)
+            }
+            """
+        input_model = _load_model(input_graph)
+        prefix = "pre_"
+        prefixed_model = compose.add_prefix(
+            model=input_model,
+            prefix=prefix,
+            rename_inputs=True,
+            rename_outputs=True,
+            rename_edges=True,
+        )
+
+        checker.check_graph(prefixed_model.graph)
+
+        for i in prefixed_model.graph.input:
+            self.assertTrue(i.name.startswith(prefix))
+        for o in prefixed_model.graph.output:
+            self.assertTrue(o.name.startswith(prefix))
+
+    def test_add_prefix_wo_inputs_outputs(self) -> None:
+        """Tests prefixing input and output nodes, when renaming of inputs/outputs is deactivated."""
+        input_graph = """
+            <
+                ir_version: 6,
+                opset_import: [ "": 13]
+            >
+            agraph (float[2,2] A) => (float[2,2] B)
+            {
+            B = Identity(A)
+            }
+            """
+        input_model = _load_model(input_graph)
+        prefix = "pre_"
+        prefixed_model = compose.add_prefix(
+            model=input_model,
+            prefix=prefix,
+            rename_inputs=False,
+            rename_outputs=False,
+            rename_edges=True,
+        )
+
+        checker.check_graph(prefixed_model.graph)
+        self.assertEqual(input_model.graph.input, prefixed_model.graph.input)
+        self.assertEqual(input_model.graph.output, prefixed_model.graph.output)
+
+    def test_add_prefix_with_loose_inputs_outputs(self) -> None:
+        """Tests prefixing of graphs with loose inputs."""
+        input_graph = """
+            <
+                ir_version: 7,
+                opset_import: [ "" : 13 ]
+            >
+            agraph (bool b, float[128] X, float[128] Y, bool unused_input) => (float[128] Z)
+            {
+            Z = If (b) <
+                then_branch = g1 () => (float[128] z_then) { z_then = Add(X, X) },
+                else_branch = g2 () => (float[128] z_else) { z_else = Sub(Y, Y) }
+                >
+            }
+            """
+        input_model = _load_model(input_graph)
+        prefix = "pre_"
+        prefixed_model = compose.add_prefix(
+            model=input_model,
+            prefix=prefix,
+            rename_inputs=True,
+            rename_outputs=True,
+            rename_edges=True,
+        )
+        checker.check_graph(prefixed_model.graph)
+
+        # Check if all inputs are prefixed even if input is not connected to node. e.g., 'unused_input'.
+        for i in prefixed_model.graph.input:
+            self.assertTrue(i.name.startswith(prefix))
+
+        for graph_output in prefixed_model.graph.output:
+            self.assertTrue(graph_output.name.startswith(prefix))
+
     def _test_add_prefix(
         self,
         rename_nodes: bool = False,
@@ -516,7 +605,9 @@ class TestComposeFunctions(unittest.TestCase):
 
     def test_add_prefix_edges(self) -> None:
         """Tests prefixing nodes edges. This will also rename inputs/outputs, since the names are shared"""
-        self._test_add_prefix(rename_edges=True)
+        self._test_add_prefix(
+            rename_edges=True, rename_inputs=True, rename_outputs=True
+        )
 
     def test_add_prefix_inputs(self) -> None:
         """Tests prefixing graph inputs only. Relevant node edges should be renamed as well"""

onnx/test/cpp/function_context_test.cc

@@ -71,9 +71,10 @@ BuildFloatFunctionBody(const FunctionBodyBuildContext& /*ctx*/, const OpSchema&
   // Create a scalar-tensor constant 2.0 of float type:
   auto two_as_tensor = ToTensor(2.0, TensorProto_DataType::TensorProto_DataType_FLOAT);
 
-  std::vector<FunctionBodyHelper::NodeDef> body{// nodes: {outputs, op, inputs, attributes}
-                                                {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
-                                                {{"Y"}, "Mul", {"X", "Two"}}};
+  std::vector<FunctionBodyHelper::NodeDef> body{
+      // nodes: {outputs, op, inputs, attributes}
+      {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
+      {{"Y"}, "Mul", {"X", "Two"}}};
 
   return BuildFunctionProto(functionProto, schema, body);
 }
@@ -92,7 +93,7 @@ static void RegisterCustomFuncFloatSchema() {
   (void)unused;
 }
 
-// Test for Context dependant function without type context
+// Test for Context dependent function without type context
 TEST(FunctionAPITest, ContextDependentFunctionTest) {
   RegisterCustomFuncFloatSchema();
 
@@ -130,9 +131,10 @@ BuildFunctionBody(const FunctionBodyBuildContext& ctx, const OpSchema& schema, F
   auto elem_type = (TensorProto_DataType)tp->tensor_type().elem_type();
   auto two_as_tensor = ToTensor(2.0, elem_type);
 
-  std::vector<FunctionBodyHelper::NodeDef> body{// nodes: {outputs, op, inputs, attributes}
-                                                {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
-                                                {{"Y"}, "Mul", {"X", "Two"}}};
+  std::vector<FunctionBodyHelper::NodeDef> body{
+      // nodes: {outputs, op, inputs, attributes}
+      {{"Two"}, "Constant", {}, {{"value", two_as_tensor}}},
+      {{"Y"}, "Mul", {"X", "Two"}}};
 
   return BuildFunctionProto(functionProto, schema, body);
 }

onnx/test/cpp/function_verify_test.cc

@@ -229,7 +229,7 @@ struct FunctionTypeChecker {
   // forTypeVar: This is used to iterate through all possible bindings of type-values
   // to all type-variables used in the op schema, and invoke the type-checker for
   // each possible instantiation.
-  void forTypeVar(int i) {
+  void forTypeVar(size_t i) {
     auto& typeConstraintVector = schema.typeConstraintParams();
     if (i < typeConstraintVector.size()) {
       std::string typeVar = typeConstraintVector[i].type_param_str;

onnx/test/cpp/inliner_test.cc

@@ -9,6 +9,7 @@
 #include "gtest/gtest.h"
 #include "onnx/checker.h"
 #include "onnx/common/constants.h"
+#include "onnx/defs/function.h"
 #include "onnx/defs/parser.h"
 #include "onnx/defs/printer.h"
 #include "onnx/defs/schema.h"
@@ -18,7 +19,11 @@
 namespace ONNX_NAMESPACE {
 namespace Test {
 
-static void InlineFunctions(ModelProto& model, const char* input, const inliner::FunctionIdSet* to_inline = nullptr) {
+static void InlineFunctions(
+    ModelProto& model,
+    const char* input,
+    const inliner::FunctionIdSet* to_inline = nullptr,
+    const ISchemaRegistry* schema_registry = nullptr) {
   OnnxParser parser(input);
   auto status = parser.Parse(model);
   EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
@@ -28,7 +33,9 @@ static void InlineFunctions(ModelProto& model, const char* input, const inliner:
   shape_inference::InferShapes(model);
 
   // std::cout << "Before inlining:\n" << ProtoToString(model) << "\n";
-  if (to_inline != nullptr)
+  if (schema_registry != nullptr)
+    inliner::InlineSelectedFunctions(model, *to_inline, schema_registry);
+  else if (to_inline != nullptr)
     inliner::InlineSelectedFunctions(model, *to_inline);
   else
     inliner::InlineLocalFunctions(model, true);
@@ -191,7 +198,7 @@ foo (x) => (y)
 
   ModelProto model;
   InlineFunctions(model, code);
-  // Check that valueinfo is propagated fron function to main graph.
+  // Check that valueinfo is propagated from function to main graph.
   auto& graph = model.graph();
   auto& temp_new_name = graph.node(0).output(0);
   auto& valueinfos = graph.value_info();
@@ -372,5 +379,127 @@ bar (x) => (y) {
   ASSERT_EQ(node2.attribute_size(), 0);
 }
 
+TEST(SchemaFunctionInliner, BasicTest) {
+  const char* code = R"ONNX(
+<ir_version: 8, opset_import: ["" : 18]>
+agraph (float[N, 128] X) => (float[N, 128] Y)
+{
+  Y = Softmax (X)
+}
+)ONNX";
+
+  ModelProto model;
+  inliner::FunctionIdVector to_inline = {{"", "Softmax"}};
+  auto to_inline_set = inliner::FunctionIdSet::Create(std::move(to_inline));
+  InlineFunctions(model, code, to_inline_set.get(), OpSchemaRegistry::Instance());
+  auto num_nodes = model.graph().node_size();
+  ASSERT_GT(num_nodes, 1);
+}
+
+static bool ContainsOp(const ModelProto& model, const char* op_type) {
+  for (const auto& node : model.graph().node()) {
+    if (node.op_type() == op_type) {
+      return true;
+    }
+  }
+  return false;
+}
+
+TEST(SchemaFunctionInliner, NestedTest) {
+  const char* code = R"ONNX(
+<ir_version: 8, opset_import: ["" : 18]>
+agraph (float[N, C] X, int32[N] expected) => (float Y)
+{
+  Y, log_prob = SoftmaxCrossEntropyLoss (X, expected)
+}
+)ONNX";
+
+  ModelProto model;
+  inliner::FunctionIdVector to_inline = {{"", "SoftmaxCrossEntropyLoss"}};
+  auto to_inline_set = inliner::FunctionIdSet::Create(std::move(to_inline));
+  InlineFunctions(model, code, to_inline_set.get(), OpSchemaRegistry::Instance());
+  auto num_nodes = model.graph().node_size();
+  ASSERT_GT(num_nodes, 1);
+  // Nested call to LogSoftmax should not be inlined.
+  ASSERT_TRUE(ContainsOp(model, "LogSoftmax"));
+
+  inliner::FunctionIdVector to_inline2 = {{"", "SoftmaxCrossEntropyLoss"}, {"", "LogSoftmax"}};
+  to_inline_set = inliner::FunctionIdSet::Create(std::move(to_inline2));
+  InlineFunctions(model, code, to_inline_set.get(), OpSchemaRegistry::Instance());
+  num_nodes = model.graph().node_size();
+  ASSERT_GT(num_nodes, 1);
+  // Nested call to LogSoftmax should be inlined.
+  ASSERT_FALSE(ContainsOp(model, "LogSoftmax"));
+}
+
+TEST(FunctionBuilder, AddInlinedCallBasic) {
+  // Test the AddInlinedCall functionality
+  GraphProto graph;
+
+  // Create a simple graph using parser for better readability
+  const char* graph_text = R"ONNX(
+test_graph (float x) => (float y)
+<float const_val = {2.0}>
+{
+    y = Add(x, const_val)
+}
+)ONNX";
+
+  auto status = OnnxParser::Parse(graph, graph_text);
+  EXPECT_TRUE(status.IsOK()) << status.ErrorMessage();
+
+  // Create a function and use AddInlinedCall
+  FunctionProto function;
+  FunctionBuilder builder(function);
+
+  builder.AddInlinedCall({"result"}, graph, {"input_x"}, "test");
+
+  // Verify the function has the expected structure
+  ASSERT_EQ(function.node_size(), 2); // One Constant node + one Add node
+
+  // Check the first node is a Constant
+  ASSERT_EQ(function.node(0).op_type(), "Constant");
+  ASSERT_EQ(function.node(0).output_size(), 1);
+  ASSERT_TRUE(function.node(0).output(0).find("test") != std::string::npos);
+
+  // Check the second node is an Add
+  ASSERT_EQ(function.node(1).op_type(), "Add");
+  ASSERT_EQ(function.node(1).input_size(), 2);
+  ASSERT_EQ(function.node(1).output_size(), 1);
+  ASSERT_EQ(function.node(1).input(0), "input_x"); // Should be renamed to actual input
+  ASSERT_EQ(function.node(1).output(0), "result"); // Should be renamed to actual output
+}
+
+TEST(Renamer, BasicFunctionality) {
+  // Test the Renamer class functionality
+  GraphProto graph;
+
+  // Add input to graph
+  auto* input = graph.add_input();
+  input->set_name("input");
+
+  // Create a Renamer instance
+  inliner::Renamer renamer("test", graph);
+
+  // Test binding names
+  renamer.BindName("formal_input", "actual_input");
+
+  // Test creating unique names and binding
+  std::string unique_name = renamer.BindToUniqueName("temp");
+  ASSERT_TRUE(unique_name.find("test") != std::string::npos);
+
+  // Test renaming a node
+  NodeProto node;
+  node.set_op_type("Add");
+  node.add_input("formal_input");
+  node.add_output("temp_output");
+
+  renamer.RenameNode(node);
+
+  // Verify renaming worked correctly
+  ASSERT_EQ(node.input(0), "actual_input"); // Should be bound to actual name
+  ASSERT_TRUE(node.output(0).find("test") != std::string::npos); // Should have prefix
+}
+
 } // namespace Test
 } // namespace ONNX_NAMESPACE

onnx/test/cpp/schema_registration_test.cc

@@ -127,7 +127,7 @@ TEST(SchemaRegistrationTest, RegisterSpecifiedOpsetSchemaVersion) {
 #endif
 }
 
-// Regsiter opset-11, then opset-14
+// Register opset-11, then opset-14
 // Expects Reg(11, 14) == Reg(11) U Reg(14)
 TEST(SchemaRegistrationTest, RegisterMultipleOpsetSchemaVersions_UpgradeVersion) {
 #ifdef __ONNX_DISABLE_STATIC_REGISTRATION
@@ -170,7 +170,7 @@ TEST(SchemaRegistrationTest, RegisterMultipleOpsetSchemaVersions_UpgradeVersion)
 #endif
 }
 
-// Regsiter opset-14, then opset-11
+// Register opset-14, then opset-11
 // Expects Reg(14, 11) == Reg(11) U Reg(14)
 TEST(SchemaRegistrationTest, RegisterMultipleOpsetSchemaVersions_DowngradeVersion) {
 #ifdef __ONNX_DISABLE_STATIC_REGISTRATION

onnx/test/cpp/test_main.cc

@@ -8,7 +8,7 @@
 
 #include "gtest/gtest.h"
 
-GTEST_API_ int main(int argc, char** argv) {
+int main(int argc, char** argv) {
   std::cout << "Running main() from test_main.cc" << '\n';
   ::testing::InitGoogleTest(&argc, argv);
   return RUN_ALL_TESTS();