fix MiniOnnxBuilder

Author: xadupre

_unittests/ut_helpers/test_mini_onnx_builder.py

@@ -1,5 +1,6 @@
 import unittest
 import numpy as np
+import ml_dtypes
 import onnx
 import torch
 from onnx_diagnostic.ext_test_case import ExtTestCase
@@ -59,7 +60,7 @@ def test_mini_onnx_builder_sequence_ort_randomize(self):
         self.assertEqual((2,), got[1].shape)
         self.assertEqual(np.float32, got[1].dtype)
 
-    def test_mini_onnx_builder(self):
+    def test_mini_onnx_builder1(self):
         data = [
             (
                 np.array([1, 2], dtype=np.int64),
@@ -153,6 +154,61 @@ def test_mini_onnx_builder(self):
                 restored = create_input_tensors_from_onnx_model(model)
                 self.assertEqualAny(inputs, restored)
 
+    def test_mini_onnx_builder2(self):
+        data = [
+            [],
+            {},
+            tuple(),
+            dict(one=np.array([1, 2], dtype=np.int64)),
+            [np.array([1, 2], dtype=np.int64)],
+            (np.array([1, 2], dtype=np.int64),),
+            [np.array([1, 2], dtype=np.int64), 1],
+            dict(one=np.array([1, 2], dtype=np.int64), two=1),
+            (np.array([1, 2], dtype=np.int64), 1),
+        ]
+
+        for inputs in data:
+            with self.subTest(types=string_type(inputs)):
+                model = create_onnx_model_from_input_tensors(inputs)
+                restored = create_input_tensors_from_onnx_model(model)
+                self.assertEqualAny(inputs, restored)
+
+    def test_mini_onnx_builder_type1(self):
+        data = [
+            np.array([1], dtype=dtype)
+            for dtype in [
+                np.float32,
+                np.double,
+                np.float16,
+                np.int32,
+                np.int64,
+                np.int16,
+                np.int8,
+                np.uint8,
+                np.uint16,
+                np.uint32,
+                np.uint64,
+                np.bool_,
+            ]
+        ]
+
+        for inputs in data:
+            with self.subTest(types=string_type(inputs)):
+                model = create_onnx_model_from_input_tensors(inputs)
+                restored = create_input_tensors_from_onnx_model(model)
+                self.assertEqualAny(inputs, restored)
+
+    def test_mini_onnx_builder_type2(self):
+        data = [
+            np.array([1], dtype=dtype) for dtype in [ml_dtypes.bfloat16, ml_dtypes.float8_e3m4]
+        ]
+
+        for inputs in data:
+            with self.subTest(types=string_type(inputs)):
+                model = create_onnx_model_from_input_tensors(inputs)
+                restored = create_input_tensors_from_onnx_model(model)
+                self.assertEqualAny(inputs, restored)
+
     def test_mini_onnx_builder_transformers(self):
         cache = make_dynamic_cache([(torch.ones((3, 3)), torch.ones((3, 3)) * 2)])
         dc = CacheKeyValue(cache)

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -563,6 +563,12 @@ def _make(ty: type, res: Any) -> Any:
             return d
         return ty(res)
 
+    if end and len(res) == 1:
+        if res[0] is None:
+            return next_pos, ty()
+        if isinstance(res[0], tuple) and len(res[0]) == 2 and res[0] == ("dict.", None):
+            return next_pos, ty()
+        return next_pos, ty(res)
     return next_pos, (
         ty() if len(res) == 1 and res[0] in (("dict.", None), None) else _make(ty, res)
     )
@@ -639,6 +645,8 @@ def create_input_tensors_from_onnx_model(
             return float(output[0])
         if name == "tensor":
             return torch.from_numpy(output).to(device)
-        raise AssertionError(f"Unexpected name {name!r} in {names}")
+        assert name.startswith(
+            ("list_", "list.", "dict.", "tuple_", "tuple.")
+        ), f"Unexpected name {name!r} in {names}"
 
     return _unflatten(sep, names, got, device=device)[1]