Replaces long initiliazer by rando values (#98)

xadupre · web-flow · commit a8b45f9a0fd7 · 2025-03-10T00:33:09.000+01:00
* Replaces long initiliazer by rando values

* fix display

* fix issues
diff --git a/_doc/api/translate_api.rst b/_doc/api/translate_api.rst
@@ -39,6 +39,12 @@ InnerEmitter
 .. autoclass:: onnx_array_api.translate_api.inner_emitter.InnerEmitter
     :members:
 
+InnerEmitterShortInitializer
+++++++++++++++++++++++++++++
+
+.. autoclass:: onnx_array_api.translate_api.inner_emitter.InnerEmitterShortInitializer
+    :members:
+
 LightEmitter
 ++++++++++++
 
diff --git a/_unittests/ut_ort/test_ort_profile.py b/_unittests/ut_ort/test_ort_profile.py
@@ -57,8 +57,6 @@ def myloss(x, y):
         prof = ort_profile(optimized, feeds)
         events = {
             "kernel_time",
-            "fence_before",
-            "fence_after",
             "SequentialExecutor::Execute",
             "model_run",
             "model_loading_array",
diff --git a/_unittests/ut_translate_api/test_translate_classic.py b/_unittests/ut_translate_api/test_translate_classic.py
@@ -178,6 +178,75 @@ def test_transpose(self):
         self.maxDiff = None
         self.assertEqual(expected, code)
 
+    def test_transpose_short(self):
+        onx = (
+            start(opset=19)
+            .vin("X")
+            .reshape((-1, 1))
+            .Transpose(perm=[1, 0])
+            .rename("Y")
+            .vout()
+            .to_onnx()
+        )
+        self.assertIsInstance(onx, ModelProto)
+        self.assertIn("Transpose", str(onx))
+        ref = ReferenceEvaluator(onx)
+        a = np.arange(10).astype(np.float32)
+        got = ref.run(None, {"X": a})[0]
+        self.assertEqualArray(a.reshape((-1, 1)).T, got)
+
+        code = translate(onx, api="onnx-short")
+        expected = dedent(
+            """
+            opset_imports = [
+                make_opsetid('', 19),
+            ]
+            inputs = []
+            outputs = []
+            nodes = []
+            initializers = []
+            sparse_initializers = []
+            functions = []
+            initializers.append(
+                from_array(
+                    np.array([-1, 1], dtype=np.int64),
+                    name='r'
+                )
+            )
+            inputs.append(make_tensor_value_info('X', TensorProto.FLOAT, shape=[]))
+            nodes.append(
+                make_node_extended(
+                    'Reshape',
+                    ['X', 'r'],
+                    ['r0_0']
+                )
+            )
+            nodes.append(
+                make_node_extended(
+                    'Transpose',
+                    ['r0_0'],
+                    ['Y'],
+                    perm=[1, 0]
+                )
+            )
+            outputs.append(make_tensor_value_info('Y', TensorProto.FLOAT, shape=[]))
+            graph = make_graph(
+                nodes,
+                'light_api',
+                inputs,
+                outputs,
+                initializers,
+                sparse_initializer=sparse_initializers,
+            )
+            model = make_model(
+                graph,
+                functions=functions,
+                opset_imports=opset_imports
+            )"""
+        ).strip("\n")
+        self.maxDiff = None
+        self.assertEqual(expected, code)
+
     def test_topk_reverse(self):
         onx = (
             start(opset=19)
diff --git a/onnx_array_api/translate_api/__init__.py b/onnx_array_api/translate_api/__init__.py
@@ -1,6 +1,6 @@
 from onnx import ModelProto
 from .translate import Translater
-from .inner_emitter import InnerEmitter
+from .inner_emitter import InnerEmitter, InnerEmitterShortInitializer
 from .builder_emitter import BuilderEmitter
 
 
@@ -16,7 +16,8 @@ def translate(proto: ModelProto, single_line: bool = False, api: str = "light")
         :class:`onnx_array_api.translate_api.light_emitter.LightEmitter`,
         another value is `"onnx"` which is the inner API implemented
         in onnx package, `"builder"` follows the syntax for the
-        class :class:`onnx_array_api.graph_api.GraphBuilder`
+        class :class:`onnx_array_api.graph_api.GraphBuilder`,
+        `"onnx-short"` replaces long initializer with random values
     :return: code
 
     .. runpython::
@@ -84,6 +85,9 @@ class :class:`onnx_array_api.graph_api.GraphBuilder`
     if api == "onnx":
         tr = Translater(proto, emitter=InnerEmitter())
         return tr.export(as_str=True)
+    if api == "onnx-short":
+        tr = Translater(proto, emitter=InnerEmitterShortInitializer())
+        return tr.export(as_str=True)
     if api == "builder":
         tr = Translater(proto, emitter=BuilderEmitter())
         return tr.export(as_str=True)
diff --git a/onnx_array_api/translate_api/inner_emitter.py b/onnx_array_api/translate_api/inner_emitter.py
@@ -106,6 +106,7 @@ def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]:
                 raise NotImplementedError(f"Unexpected dtype={sdtype}.")
         else:
             sdtype = f"np.{sdtype}"
+
         return [
             "initializers.append(",
             f"    {fra}(",
@@ -209,3 +210,57 @@ def _emit_end_function(self, **kwargs: Dict[str, Any]) -> List[str]:
             ")",
         ]
         return lines
+
+
+class InnerEmitterShortInitializer(InnerEmitter):
+    """
+    Converts event into proper code.
+    Initializer are replaced by random values if too big.
+    """
+
+    def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]:
+        name = kwargs["name"]
+        value = kwargs["value"]
+        repl = {"bool": "bool_", "object": "object_", "str": "str_"}
+        fra = "from_array"
+        sdtype = repl.get(str(value.dtype), str(value.dtype))
+        if sdtype.startswith("("):
+            from onnx.reference.custom_element_types import float8e4m3fn
+
+            if sdtype == str(float8e4m3fn):
+                sdtype = "float8e4m3fn"
+                fra = "from_array_extended"
+            else:
+                raise NotImplementedError(f"Unexpected dtype={sdtype}.")
+        else:
+            sdtype = f"np.{sdtype}"
+        if value.size <= 16:
+            return [
+                "initializers.append(",
+                f"    {fra}(",
+                f"        np.array({value.tolist()}, dtype={sdtype}),",
+                f"        name={name!r}",
+                "    )",
+                ")",
+            ]
+        if "int" in sdtype:
+            return [
+                f"value = np.random.randint(0, 10, size={value.shape})"
+                f".astype({sdtype})",
+                "initializers.append(",
+                f"    {fra}(",
+                f"        np.array(value, dtype={sdtype}),",
+                f"        name={name!r}",
+                "    )",
+                ")",
+            ]
+        return [
+            f"value = np.random.randn({', '.join(map(str,value.shape))})"
+            f".astype({sdtype})",
+            "initializers.append(",
+            f"    {fra}(",
+            f"        np.array(value, dtype={sdtype}),",
+            f"        name={name!r}",
+            "    )",
+            ")",
+        ]
