Fix bug in change_dimension with custom classes (#72)

* fix bug in change_dimnension

* fix bypass

* inputs2

* add second input

* change

* split

* fix

* disable a test

* fix

* ruff

* use the latest

* fix

Author: xadupre

.github/workflows/ci.yml

@@ -94,10 +94,16 @@ jobs:
       - name: run tests bypassed
         run: PYTHONPATH=. python _unittests/ut_torch_models/test_tiny_llms_bypassed.py
 
+      - name: test image_classification
+        run: PYTHONPATH=. python _unittests/ut_tasks/test_tasks_image_classification.py
+
+      - name: test zero_shot_image_classification
+        run: PYTHONPATH=. python _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py
+
       - name: run tests
         run: |
           pip install pytest
-          PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py
+          PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests --ignore _unittests/ut_reference/test_backend_extended_reference_evaluator.py --ignore _unittests/ut_reference/test_backend_onnxruntime_evaluator.py --ignore _unittests/ut_torch_models/test_tiny_llms_bypassed.py --ignore _unittests/ut_tasks/test_tasks_zero_shot_image_classification.py --ignore _unittests/ut_tasks/test_tasks_image_classification.py
 
       - name: run backend tests python
         run: PYTHONPATH=. UNITTEST_GOING=1 pytest --durations=10 _unittests/ut_reference/test_backend_extended_reference_evaluator.py

CHANGELOGS.rst

@@ -4,6 +4,7 @@ Change Logs
 0.4.1
 +++++
 
+* :pr:`72`: fix change_dynamic_dimension for custom classes
 * :pr:`70`: support models options in command lines
 
 0.4.0

_doc/api/tasks/index.rst

@@ -9,8 +9,11 @@ All submodules contains the three following functions:
 * ``random_input_kwargs(config) -> kwargs, get_inputs``:
   produces values ``get_inputs`` can take to generate dummy inputs
   suitable for a model defined by its configuration
-* ``get_inputs(model, config, *args, **kwargs) -> dict(inputs=..., dynamic_shapes=...)``:
-  generates the dummy inputs and dynamic shapes for a specific model and configuration.
+* ``get_inputs(model, config, *args, add_second_input=False, **kwargs) -> dict(inputs=..., dynamic_shapes=...)``:
+  generates the dummy inputs and dynamic shapes for a specific model and configuration,
+  if ``add_second_input`` is True, the function should return a different set of inputs,
+  with different values for the dynamic dimension. This is usually better to
+  rely on the function as the dynamic dimensions may be correlated.
 
 For a specific task, you would write:
 

_unittests/ut_export/test_dynamic_shapes.py

@@ -1,5 +1,6 @@
 import unittest
 import torch
+import transformers
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_transformers
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
@@ -742,6 +743,18 @@ def test_couple_input_ds_change_dynamic_dimensions_fixed(self):
         self.assertEqual((1, 5, 8), new_input["A"].shape)
         self.assertEqual((1, 50), new_input["B"].shape)
 
+    def test_couple_input_ds_change_dynamic_dimensions_dynamic_cache(self):
+        inst = CoupleInputsDynamicShapes(
+            (),
+            {"A": make_dynamic_cache([(torch.ones((2, 2, 2, 2)), torch.ones((2, 2, 2, 2)))])},
+            {"A": [[{0: "batch", 2: "last"}], [{0: "batch", 2: "last"}]]},
+        )
+        with bypass_export_some_errors(patch_transformers=True):
+            new_inputs = inst.change_dynamic_dimensions()
+        self.assertIsInstance(new_inputs["A"], transformers.cache_utils.DynamicCache)
+        self.assertEqual((3, 2, 3, 2), new_inputs["A"].key_cache[0].shape)
+        self.assertEqual((3, 2, 3, 2), new_inputs["A"].value_cache[0].shape)
+
     @requires_transformers("4.51")
     def test_dynamic_cache_replace_by_string(self):
         n_layers = 2

_unittests/ut_helpers/test_helper.py

@@ -39,6 +39,8 @@
     dtype_to_tensor_dtype,
 )
 from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache, make_encoder_decoder_cache
+from onnx_diagnostic.torch_models.hghub.hub_api import get_pretrained_config
+
 
 TFLOAT = onnx.TensorProto.FLOAT
 
@@ -484,6 +486,11 @@ def test_flatten_encoder_decoder_cache(self):
         s = string_type(inputs)
         self.assertIn("EncoderDecoderCache", s)
 
+    def test_string_typeçconfig(self):
+        conf = get_pretrained_config("microsoft/phi-2")
+        s = string_type(conf)
+        self.assertStartsWith("PhiConfig(**{", s)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks.py

@@ -1,6 +1,6 @@
 import unittest
 import torch
-from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers, has_torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers
 from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
 from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
 from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
@@ -10,11 +10,27 @@ class TestTasks(ExtTestCase):
     @hide_stdout()
     def test_text2text_generation(self):
         mid = "sshleifer/tiny-marian-en-de"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "text2text-generation")
         self.assertIn((data["size"], data["n_weights"]), [(473928, 118482)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         raise unittest.SkipTest(f"not working for {mid!r}")
         model(**inputs)
+        model(**data["inputs2"])
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+    @hide_stdout()
+    def test_text_generation(self):
+        mid = "arnir0/Tiny-LLM"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "text-generation")
+        self.assertIn((data["size"], data["n_weights"]), [(51955968, 12988992)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
         with bypass_export_some_errors(patch_transformers=True, verbose=10):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
@@ -23,9 +39,11 @@ def test_text2text_generation(self):
     @hide_stdout()
     def test_automatic_speech_recognition(self):
         mid = "openai/whisper-tiny"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "automatic-speech-recognition")
         self.assertIn((data["size"], data["n_weights"]), [(132115968, 33028992)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**data["inputs2"])
         Dim = torch.export.Dim
         self.maxDiff = None
         self.assertIn("{0:Dim(batch),1:DYN(seq_length)}", self.string_type(ds))
@@ -90,27 +108,15 @@ def test_automatic_speech_recognition(self):
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
             )
 
-    @hide_stdout()
-    def test_imagetext2text_generation(self):
-        mid = "HuggingFaceM4/tiny-random-idefics"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
-        self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)])
-        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
-        model(**inputs)
-        if not has_torch("2.10"):
-            raise unittest.SkipTest("sym_max does not work with dynamic dimension")
-        with bypass_export_some_errors(patch_transformers=True, verbose=10):
-            torch.export.export(
-                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
-            )
-
     @hide_stdout()
     def test_fill_mask(self):
         mid = "google-bert/bert-base-multilingual-cased"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "fill-mask")
         self.assertIn((data["size"], data["n_weights"]), [(428383212, 107095803)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
+        model(**data["inputs2"])
         with bypass_export_some_errors(patch_transformers=True, verbose=10):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
@@ -119,10 +125,12 @@ def test_fill_mask(self):
     @hide_stdout()
     def test_feature_extraction(self):
         mid = "facebook/bart-base"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "feature-extraction")
         self.assertIn((data["size"], data["n_weights"]), [(557681664, 139420416)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
+        model(**data["inputs2"])
         with bypass_export_some_errors(patch_transformers=True, verbose=10):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
@@ -131,10 +139,12 @@ def test_feature_extraction(self):
     @hide_stdout()
     def test_text_classification(self):
         mid = "Intel/bert-base-uncased-mrpc"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "text-classification")
         self.assertIn((data["size"], data["n_weights"]), [(154420232, 38605058)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
+        model(**data["inputs2"])
         with bypass_export_some_errors(patch_transformers=True, verbose=10):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
@@ -143,10 +153,12 @@ def test_text_classification(self):
     @hide_stdout()
     def test_sentence_similary(self):
         mid = "sentence-transformers/all-MiniLM-L6-v1"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "sentence-similarity")
         self.assertIn((data["size"], data["n_weights"]), [(62461440, 15615360)])
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
+        model(**data["inputs2"])
         with bypass_export_some_errors(patch_transformers=True, verbose=10):
             torch.export.export(
                 model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
@@ -155,9 +167,11 @@ def test_sentence_similary(self):
     @hide_stdout()
     def test_falcon_mamba_dev(self):
         mid = "tiiuae/falcon-mamba-tiny-dev"
-        data = get_untrained_model_with_inputs(mid, verbose=1)
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "text-generation")
         model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
         model(**inputs)
+        model(**data["inputs2"])
         self.assertIn((data["size"], data["n_weights"]), [(138640384, 34660096)])
         if not has_transformers("4.55"):
             raise unittest.SkipTest("The model has control flow.")

_unittests/ut_tasks/test_tasks_image_classification.py

@@ -0,0 +1,28 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasks(ExtTestCase):
+    @hide_stdout()
+    def test_image_classification(self):
+        mid = "hf-internal-testing/tiny-random-BeitForImageClassification"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "image-classification")
+        self.assertIn((data["size"], data["n_weights"]), [(56880, 14220)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        if not has_transformers("4.51.999"):
+            raise unittest.SkipTest("Requires transformers>=4.52")
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks_image_text_to_text.py

@@ -0,0 +1,30 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, has_transformers, has_torch
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasks(ExtTestCase):
+    @hide_stdout()
+    def test_image_text_to_text(self):
+        mid = "HuggingFaceM4/tiny-random-idefics"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "image-text-to-text")
+        self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        if not has_transformers("4.55"):
+            raise unittest.SkipTest("The model has control flow.")
+        if not has_torch("2.7.99"):
+            raise unittest.SkipTest("sym_max does not work with dynamic dimension")
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_tasks/test_tasks_zero_shot_image_classification.py

@@ -0,0 +1,27 @@
+import unittest
+import torch
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout, requires_torch
+from onnx_diagnostic.torch_models.hghub.model_inputs import get_untrained_model_with_inputs
+from onnx_diagnostic.torch_export_patches import bypass_export_some_errors
+from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+class TestTasks(ExtTestCase):
+    @requires_torch("2.7.99")
+    @hide_stdout()
+    def test_zero_shot_image_classification(self):
+        mid = "openai/clip-vit-base-patch16"
+        data = get_untrained_model_with_inputs(mid, verbose=1, add_second_input=True)
+        self.assertEqual(data["task"], "zero-shot-image-classification")
+        self.assertIn((data["size"], data["n_weights"]), [(188872708, 47218177)])
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        model(**inputs)
+        model(**data["inputs2"])
+        with bypass_export_some_errors(patch_transformers=True, verbose=10):
+            torch.export.export(
+                model, (), kwargs=inputs, dynamic_shapes=use_dyn_not_str(ds), strict=False
+            )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

onnx_diagnostic/export/dynamic_shapes.py

@@ -363,16 +363,20 @@ def _generic_walker_step(
         )
         if flatten_unflatten:
             flatunflat = flatten_unflatten_for_dynamic_shapes(inputs)
-            return cls._generic_walker_step(
+            res = cls._generic_walker_step(
                 processor, flatunflat, ds, flatten_unflatten=flatten_unflatten
             )
-        flat, _spec = torch.utils._pytree.tree_flatten(inputs)
+            # Should we restore the original class?
+            return res
+        flat, spec = torch.utils._pytree.tree_flatten(inputs)
         if all(isinstance(t, torch.Tensor) for t in flat):
             # We need to flatten dynamic shapes as well
             ds = flatten_dynamic_shapes(ds)
-        return cls._generic_walker_step(
+        res = cls._generic_walker_step(
             processor, flat, ds, flatten_unflatten=flatten_unflatten
         )
+        # Then we restore the original class.
+        return torch.utils._pytree.tree_unflatten(res, spec)
 
     class ChangeDimensionProcessor:
         def __init__(self, desired_values):