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

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();

onnx/test/helper_test.py

@@ -23,7 +23,6 @@ from onnx import (
     SequenceProto,
     TensorProto,
     TypeProto,
-    _custom_element_types,
     checker,
     defs,
     helper,
@@ -423,6 +422,7 @@ class TestHelperNodeFunctions(unittest.TestCase):
         test([("", 21)], 10)
         test([("", 22)], 10)
         test([("", 23)], 11)
+        test([("", 24)], 12)
         # standard opset can be referred to using empty-string or "ai.onnx"
         test([("ai.onnx", 9)], 4)
         test([("ai.onnx.ml", 2)], 6)
@@ -534,11 +534,8 @@ class TestHelperTensorFunctions(unittest.TestCase):
         np.testing.assert_allclose(numpy_helper.to_array(tensor).view(np.uint16), array)
 
     def test_make_float8e4m3fn_tensor_raw(self) -> None:
-        expected = np.array([0, 0.5, 1, 448, 10], dtype=np.float32)
-        f8 = np.array(
-            [helper.float32_to_float8e4m3(x) for x in expected], dtype=np.uint8
-        )
-        packed_values = f8.tobytes()
+        expected = np.array([0, 0.5, 1, 448, 10], dtype=ml_dtypes.float8_e4m3fn)
+        packed_values = expected.tobytes()
         y = helper.make_tensor(
             name="test",
             data_type=TensorProto.FLOAT8E4M3FN,
@@ -547,15 +544,11 @@ class TestHelperTensorFunctions(unittest.TestCase):
             raw=True,
         )
         ynp = numpy_helper.to_array(y)
-        np.testing.assert_equal(
-            ynp.view(np.uint8), expected.astype(ml_dtypes.float8_e4m3fn).view(np.uint8)
-        )
+        np.testing.assert_equal(ynp.view(np.uint8), expected.view(np.uint8))
 
     def test_make_float8e4m3fnuz_tensor_raw(self) -> None:
         expected = np.array([0, 0.5, 1, 240, 10], dtype=np.float32)
-        f8 = np.array(
-            [helper.float32_to_float8e4m3(x, uz=True) for x in expected], dtype=np.uint8
-        )
+        f8 = expected.astype(ml_dtypes.float8_e4m3fnuz)
         packed_values = f8.tobytes()
         y = helper.make_tensor(
             name="test",
@@ -571,11 +564,8 @@ class TestHelperTensorFunctions(unittest.TestCase):
         )
 
     def test_make_float8e5m2_tensor_raw(self) -> None:
-        expected = np.array([0, 0.5, 1, 49152, 10], dtype=np.float32)
-        f8 = np.array(
-            [helper.float32_to_float8e5m2(x) for x in expected], dtype=np.uint8
-        )
-        packed_values = f8.tobytes()
+        expected = np.array([0, 0.5, 1, 49152, 10], dtype=ml_dtypes.float8_e5m2)
+        packed_values = expected.tobytes()
         y = helper.make_tensor(
             name="test",
             data_type=TensorProto.FLOAT8E5M2,
@@ -584,17 +574,11 @@ class TestHelperTensorFunctions(unittest.TestCase):
             raw=True,
         )
         ynp = numpy_helper.to_array(y)
-        np.testing.assert_equal(
-            ynp.view(np.uint8), expected.astype(ml_dtypes.float8_e5m2).view(np.uint8)
-        )
+        np.testing.assert_equal(ynp.view(np.uint8), expected.view(np.uint8))
 
     def test_make_float8e5m2fnuz_tensor_raw(self) -> None:
-        expected = np.array([0, 0.5, 1, 49152, 10], dtype=np.float32)
-        f8 = np.array(
-            [helper.float32_to_float8e5m2(x, fn=True, uz=True) for x in expected],
-            dtype=np.uint8,
-        )
-        packed_values = f8.tobytes()
+        expected = np.array([0, 0.5, 1, 49152, 10], dtype=ml_dtypes.float8_e5m2fnuz)
+        packed_values = expected.tobytes()
         y = helper.make_tensor(
             name="test",
             data_type=TensorProto.FLOAT8E5M2FNUZ,
@@ -603,9 +587,36 @@ class TestHelperTensorFunctions(unittest.TestCase):
             raw=True,
         )
         ynp = numpy_helper.to_array(y)
+        np.testing.assert_equal(ynp.view(np.uint8), expected.view(np.uint8))
+
+    def test_make_float8e8m0_tensor(self) -> None:
+        y = helper.make_tensor(
+            "scale",
+            TensorProto.FLOAT8E8M0,
+            [7],
+            [0, 0.124, 1.0, 1.5, 2.0, 2.1, np.finfo(np.float32).max],
+        )
+        ynp = numpy_helper.to_array(y)
+        expected = np.array([0, 124, 127, 128, 128, 129, 254], dtype=np.uint8)
+        np.testing.assert_equal(ynp.view(np.uint8), expected)
+
+    def test_make_float8e8m0_tensor_raw(self) -> None:
+        e8m0_raw = np.array([120, 124, 127, 128, 129, 140], dtype=np.uint8)
+        packed_values = e8m0_raw.tobytes()
+        y = helper.make_tensor(
+            name="test",
+            data_type=TensorProto.FLOAT8E8M0,
+            dims=list(e8m0_raw.shape),
+            vals=packed_values,
+            raw=True,
+        )
+        ynp = numpy_helper.to_array(y)
+        expected = np.array(
+            [0.0078125, 0.125, 1.0, 2.0, 4.0, 8192], dtype=ml_dtypes.float8_e8m0fnu
+        )
         np.testing.assert_equal(
             ynp.view(np.uint8),
-            expected.astype(ml_dtypes.float8_e5m2fnuz).view(np.uint8),
+            expected.view(np.uint8),
         )
 
     @parameterized.parameterized.expand(
@@ -656,13 +667,7 @@ class TestHelperTensorFunctions(unittest.TestCase):
         )
     )
     def test_make_4bit_raw_tensor(self, dtype, dims) -> None:
-        type_range = {
-            TensorProto.UINT4: (0, 15),
-            TensorProto.INT4: (-8, 7),
-        }
-        data = np.random.randint(
-            type_range[dtype][0], high=type_range[dtype][1] + 1, size=dims
-        ).astype(np.uint8)
+        data = np.random.randint(0, high=16, size=dims, dtype=np.uint8)
         packed_data = _pack_4bit(data)
 
         y = helper.make_tensor(
@@ -696,7 +701,7 @@ class TestHelperTensorFunctions(unittest.TestCase):
         expected = np.array(
             [0, 0.5, 1, 6, -0.5, -6, -4], dtype=ml_dtypes.float4_e2m1fn
         ).view(np.uint8)
-        np.testing.assert_equal(ynp, expected)
+        np.testing.assert_equal(ynp.view(np.uint8), expected)
 
     def test_make_sparse_tensor(self) -> None:
         values = [1.1, 2.2, 3.3, 4.4, 5.5]
@@ -882,24 +887,7 @@ class TestPrintableGraph(unittest.TestCase):
 
 @pytest.mark.parametrize(
     "tensor_dtype",
-    [
-        t
-        for t in helper.get_all_tensor_dtypes()
-        if t
-        not in {
-            TensorProto.BFLOAT16,
-            TensorProto.FLOAT8E4M3FN,
-            TensorProto.FLOAT8E4M3FNUZ,
-            TensorProto.FLOAT8E5M2,
-            TensorProto.FLOAT8E5M2FNUZ,
-            TensorProto.UINT4,
-            TensorProto.INT4,
-            TensorProto.STRING,
-            TensorProto.COMPLEX64,
-            TensorProto.COMPLEX128,
-            TensorProto.FLOAT4E2M1,
-        }
-    ],
+    [t for t in helper.get_all_tensor_dtypes() if t != TensorProto.STRING],
     ids=lambda tensor_dtype: helper.tensor_dtype_to_string(tensor_dtype),
 )
 def test_make_tensor_vals(tensor_dtype: int) -> None:
@@ -909,41 +897,64 @@ def test_make_tensor_vals(tensor_dtype: int) -> None:
     tensor = helper.make_tensor(
         name="test", data_type=tensor_dtype, dims=np_array.shape, vals=np_array
     )
-    np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
+    roundtrip_array = numpy_helper.to_array(tensor)
+    if tensor_dtype in {
+        TensorProto.FLOAT8E5M2FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.BFLOAT16,
+        TensorProto.FLOAT8E8M0,
+    }:
+        # There is a bug in ml_dtypes that causes equality checks to fail for these dtypes
+        # See https://github.com/jax-ml/ml_dtypes/issues/301
+        assert roundtrip_array.shape == np_array.shape
+        assert roundtrip_array.dtype == np_array.dtype
+        assert roundtrip_array.tobytes() == np_array.tobytes()
+    else:
+        np.testing.assert_equal(np_array, roundtrip_array)
 
 
 @pytest.mark.parametrize(
     "tensor_dtype",
-    [
-        t
-        for t in helper.get_all_tensor_dtypes()
-        if t
-        not in {
-            TensorProto.BFLOAT16,
-            TensorProto.STRING,
-            TensorProto.FLOAT8E4M3FN,
-            TensorProto.FLOAT8E4M3FNUZ,
-            TensorProto.FLOAT8E5M2,
-            TensorProto.FLOAT8E5M2FNUZ,
-            TensorProto.UINT4,
-            TensorProto.INT4,
-            TensorProto.FLOAT4E2M1,
-        }
-    ],
+    [t for t in helper.get_all_tensor_dtypes() if t != TensorProto.STRING],
     ids=lambda tensor_dtype: helper.tensor_dtype_to_string(tensor_dtype),
 )
 def test_make_tensor_raw(tensor_dtype: int) -> None:
     np_array = np.random.randn(2, 3).astype(
         helper.tensor_dtype_to_np_dtype(tensor_dtype)
     )
+    if tensor_dtype in {
+        TensorProto.FLOAT4E2M1,
+        TensorProto.INT4,
+        TensorProto.UINT4,
+    }:
+        vals = _pack_4bit(np_array).tobytes()
+    else:
+        vals = np_array.tobytes()
     tensor = helper.make_tensor(
         name="test",
         data_type=tensor_dtype,
         dims=np_array.shape,
-        vals=np_array.tobytes(),
+        vals=vals,
         raw=True,
     )
-    np.testing.assert_equal(np_array, numpy_helper.to_array(tensor))
+    roundtrip_array = numpy_helper.to_array(tensor)
+    if tensor_dtype in {
+        TensorProto.FLOAT8E5M2FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.BFLOAT16,
+        TensorProto.FLOAT8E8M0,
+    }:
+        # There is a bug in ml_dtypes that causes equality checks to fail for these dtypes
+        # See https://github.com/jax-ml/ml_dtypes/issues/301
+        assert roundtrip_array.shape == np_array.shape
+        assert roundtrip_array.dtype == np_array.dtype
+        assert roundtrip_array.tobytes() == np_array.tobytes()
+    else:
+        np.testing.assert_equal(np_array, roundtrip_array)
 
 
 class TestHelperMappingFunctions(unittest.TestCase):
@@ -967,17 +978,15 @@ class TestHelperMappingFunctions(unittest.TestCase):
 
     def test_tensor_dtype_to_np_dtype_bfloat16(self) -> None:
         self.assertEqual(
-            helper.tensor_dtype_to_np_dtype(TensorProto.BFLOAT16), np.dtype("float32")
+            helper.tensor_dtype_to_np_dtype(TensorProto.BFLOAT16), ml_dtypes.bfloat16
         )
 
     def test_tensor_dtype_to_storage_tensor_dtype_bfloat16(self) -> None:
         self.assertEqual(
             helper.tensor_dtype_to_storage_tensor_dtype(TensorProto.BFLOAT16),
-            TensorProto.UINT16,
+            TensorProto.INT32,
         )
 
-    # BFloat16 tensor uses TensorProto.UINT16 as storage type;
-    # And the field name for TensorProto.UINT16 is int32_data
     def test_tensor_dtype_to_field_bfloat16(self) -> None:
         self.assertEqual(
             helper.tensor_dtype_to_field(TensorProto.BFLOAT16), "int32_data"
@@ -1011,17 +1020,6 @@ class TestAttrTypeToStr(unittest.TestCase):
         result = helper._attr_type_to_str(9999)
         self.assertEqual(result, "UNDEFINED")
 
-    def test_custom_types(self):
-        def _get(name):
-            if hasattr(_custom_element_types, name):
-                return getattr(_custom_element_types, name)
-            name = f"float8{name}"
-            return getattr(_custom_element_types, name)
-
-        for k, v in _custom_element_types.mapping_name_to_data_type.items():
-            self.assertEqual(helper.np_dtype_to_tensor_dtype(_get(k)), v)
-
 
 if __name__ == "__main__":
     unittest.main()
-    pytest.main([__file__])

onnx/test/inference_function_test.py

@@ -103,8 +103,8 @@ class TestInferenceFunctionCall(unittest.TestCase):
                 ADD_SCHEMA,
                 ["A", "B"],
                 ["C"],
-                _to_tensor_types(ins),  # type: ignore[arg-type]
-            ) == _to_tensor_types(outs)  # type: ignore[arg-type]
+                _to_tensor_types(ins),
+            ) == _to_tensor_types(outs)
 
     def test_clip_inference_with_optional_input(self) -> None:
         # Test case where the second input is optional

onnx/test/inliner_test.py

@@ -111,6 +111,22 @@ class InlinerTest(unittest.TestCase):
         self.assertEqual(function_nodes[0].op_type, "Add")
         self.assertEqual(function_nodes[1].op_type, "Mul")
 
+    def test_schema_function_inlining(self):
+        model = parser.parse_model(
+            """
+            <ir_version: 8, opset_import: [ "" : 20]>
+            agraph (float[N] X) => (float[N] Y)
+            {
+                Y = Softsign (X)
+            }
+        """
+        )
+        inlined = inliner.inline_selected_functions(
+            model, [], exclude=True, inline_schema_functions=True
+        )
+        inlined_nodes = inlined.graph.node
+        self.assertIn("Abs", [n.op_type for n in inlined_nodes])
+
 
 if __name__ == "__main__":
     unittest.main()

onnx/test/model_container_refeval_test.py

@@ -97,7 +97,7 @@ class TestLargeOnnxReferenceEvaluator(unittest.TestCase):
             ],
             dtype=np.float32,
         )
-        npt.assert_allclose(expected, got[0])  # type: ignore[index]
+        npt.assert_allclose(got[0], expected)
 
     def test_large_onnx_no_large_initializer(self):
         model_proto = _linear_regression()