first step for gemma3

Author: xadupre

_unittests/ut_helpers/test_torch_helper.py

@@ -181,6 +181,39 @@ def forward(self, x, y):
             set(restored),
         )
 
+    @hide_stdout()
+    def test_steal_forward_dump_file_steal_append_drop(self):
+        class SubModel(torch.nn.Module):
+            def forward(self, x):
+                return x * x
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.s1 = SubModel()
+                self.s2 = SubModel()
+
+            def forward(self, x, y):
+                sx = self.s1(x)
+                steal_append("sx", sx)
+                return sx + self.s2(y)
+
+        inputs = dict(x=torch.rand(3, 4), y=torch.rand(3, 4))
+        model = Model()
+        dump_file = self.get_dump_file("test_steal_forward_dump_file_drop.onnx")
+        with steal_forward(model, dump_file=dump_file, dump_drop={"x"}):
+            model(**inputs)
+            model(**inputs)
+        self.assertExists(dump_file)
+        restored = create_input_tensors_from_onnx_model(dump_file)
+        self.assertEqual(
+            {("", 1, "I"), ("", 1, "O"), "sx", ("", 0, "O"), "sx_1", ("", 0, "I")},
+            set(restored),
+        )
+        first = restored[("", 0, "I")]
+        _a, kws = first
+        self.assertNotIn("x", kws)
+
     @hide_stdout()
     def test_steal_forward_submodules(self):
         class SubModel(torch.nn.Module):

_unittests/ut_tasks/test_data.py

@@ -0,0 +1,14 @@
+import unittest
+from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.tasks.data import get_data
+
+
+class TestTasks(ExtTestCase):
+    def test_get_data(self):
+        name = "dummies_imagetext2text_generation_gemma3.onnx"
+        data = get_data(name)
+        print(data)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_tasks/try_tasks.py

@@ -878,6 +878,7 @@ def test_imagetext2text_generation_gemma3_4b_it(self):
             model,
             dump_file=self.get_dump_file("test_imagetext2text_generation_gemma3_4b_it.onnx"),
             dump_drop={"attention_mask", "past_key_values", "pixel_values"},
+            save_as_external_data=False,
         ):
             generated_ids = model.generate(
                 **inputs, max_new_tokens=282, do_sample=False, cache_implementation="static"

onnx_diagnostic/helpers/torch_helper.py

@@ -287,6 +287,7 @@ def steal_forward(
     submodules: bool = False,
     verbose: int = 0,
     storage_limit: int = 2**27,
+    save_as_external_data: bool = True,
     **kwargs,
 ):
     """
@@ -305,6 +306,8 @@ def steal_forward(
         they can be restored with :func:`create_input_tensors_from_onnx_model
         <onnx_diagnostic.helpers.mini_onnx_builder.create_input_tensors_from_onnx_model>`
     :param dump_drop: to drop some inputs too big (only if dump_file is specified)
+    :param save_as_external_data: True by default, but maybe better to have everything
+        in a single file if possible
     :param submodules: if True and model is a module, the list extended with all the submodules
         the module contains
     :param verbose: verbosity
@@ -414,8 +417,14 @@ def forward(self, x, y):
                 size = torch_tensor_size(storage)
                 print(f"-- gather stored {len(storage)} objects, size={size // 2 ** 20} Mb")
             if dump_drop:
-                print(string_type(dump_drop))
-                stop
+                for k, v in storage.items():
+                    if k[-1] == "I":
+                        _args, kwargs = v
+                        ii = set(kwargs) & dump_drop
+                        if ii:
+                            for i in ii:
+                                print("---", i)
+                                del kwargs[i]
             proto = create_onnx_model_from_input_tensors(storage)
             if verbose:
                 print("-- dumps stored objects")
@@ -425,7 +434,7 @@ def forward(self, x, y):
             onnx.save(
                 proto,
                 dump_file,
-                save_as_external_data=True,
+                save_as_external_data=save_as_external_data,
                 all_tensors_to_one_file=True,
                 location=location,
             )

onnx_diagnostic/tasks/data/__init__.py

@@ -0,0 +1,13 @@
+import os
+
+
+def get_data(name: str):
+    """Returns data stored in this folder."""
+    filename = os.path.join(os.path.dirname(__file__), name)
+    assert os.path.exists(
+        filename
+    ), f"Unable to find a file with {name!r}, looked for {filename!r}"
+
+    from ...helpers.mini_onnx_builder import create_input_tensors_from_onnx_model
+
+    return create_input_tensors_from_onnx_model(filename)

pyproject.toml

@@ -1,3 +1,6 @@
+[tool.setuptools.package-data]
+onnx_diagnostic = ["tasks/data/*.onnx"]
+
 [tool.black]
 line-length = 95
 extend-exclude = '''.*clones.*'''

setup.py

@@ -10,7 +10,7 @@
 here = os.path.dirname(__file__)
 if here == "":
     here = "."
-package_data = {"onnx_diagnostic.validation": ["*.css", "*.js"]}
+package_data = {"onnx_diagnostic.tasks.data": ["*.onnx"]}
 
 try:
     with open(os.path.join(here, "requirements.txt"), "r") as f: