fix bug

Author: xadupre

_unittests/ut_helpers/test_mini_onnx_builder.py

@@ -151,6 +151,48 @@ def test_mini_onnx_builder_transformers_sep(self):
                 restored = create_input_tensors_from_onnx_model(model, sep="#")
                 self.assertEqualAny(inputs, restored)
 
+    def test_specific_data(self):
+        data = {
+            ("amain", 0, "I"): (
+                (
+                    torch.rand((2, 16, 3, 448, 448), dtype=torch.float16),
+                    torch.rand((2, 16, 32, 32), dtype=torch.float16),
+                    torch.rand((2, 2)).to(torch.int64),
+                ),
+                {},
+            ),
+        }
+        model = create_onnx_model_from_input_tensors(data)
+        shapes = [
+            tuple(d.dim_value for d in i.type.tensor_type.shape.dim)
+            for i in model.graph.output
+        ]
+        self.assertEqual(shapes, [(2, 16, 3, 448, 448), (2, 16, 32, 32), (2, 2), (0,)])
+        names = [i.name for i in model.graph.output]
+        self.assertEqual(
+            [
+                "dict._((amain,0,I))___tuple_0___tuple_0___tensor",
+                "dict._((amain,0,I))___tuple_0___tuple_1___tensor",
+                "dict._((amain,0,I))___tuple_0___tuple_2.___tensor",
+                "dict._((amain,0,I))___tuple_1.___dict.___empty",
+            ],
+            names,
+        )
+        shapes = [tuple(i.dims) for i in model.graph.initializer]
+        self.assertEqual(shapes, [(2, 16, 3, 448, 448), (2, 16, 32, 32), (2, 2), (0,)])
+        names = [i.name for i in model.graph.initializer]
+        self.assertEqual(
+            [
+                "t_dict._((amain,0,I))___tuple_0___tuple_0___tensor",
+                "t_dict._((amain,0,I))___tuple_0___tuple_1___tensor",
+                "t_dict._((amain,0,I))___tuple_0___tuple_2.___tensor",
+                "t_dict._((amain,0,I))___tuple_1.___dict.___empty",
+            ],
+            names,
+        )
+        restored = create_input_tensors_from_onnx_model(model)
+        self.assertEqualAny(data, restored)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -139,6 +139,9 @@ def append_output_initializer(
                 return
 
         init_name = f"t_{name}"
+        assert (
+            init_name not in self.initializers_dict
+        ), f"name={init_name!r} already in {sorted(self.initializers_dict)}"
         self.initializers_dict[init_name] = tensor
         shape = tuple(map(int, tensor.shape))
         self.outputs.append(
@@ -324,21 +327,21 @@ def _flatten_iterator(obj: Any, sep: str) -> Iterator:
                 for i, o in enumerate(obj):
                     if i == len(obj) - 1:
                         for p, oo in _flatten_iterator(o, sep):
-                            yield f"tuple.{sep}{p}", oo
+                            yield f"tuple_{i}.{sep}{p}", oo
                     else:
                         for p, oo in _flatten_iterator(o, sep):
-                            yield f"tuple{sep}{p}", oo
+                            yield f"tuple_{i}{sep}{p}", oo
         elif isinstance(obj, list):
             if not obj:
                 yield f"list.{sep}empty", None
             else:
                 for i, o in enumerate(obj):
                     if i == len(obj) - 1:
                         for p, oo in _flatten_iterator(o, sep):
-                            yield f"list.{sep}{p}", oo
+                            yield f"list_{i}.{sep}{p}", oo
                     else:
                         for p, oo in _flatten_iterator(o, sep):
-                            yield f"list{sep}{p}", oo
+                            yield f"list_{i}{sep}{p}", oo
         elif isinstance(obj, dict):
             if not obj:
                 yield f"dict.{sep}empty", None