fix test

Author: xadupre

.github/workflows/ci.yml

@@ -62,12 +62,12 @@ jobs:
       - name: tiny-llm torch.export.export
         run: |
           export PYTHONPATH=.
-          python _unittests/ut_torch_models/test_llms.py
+          python _unittests/ut_torch_models/test_tiny_llms.py
 
       - name: tiny-llm onnx
         run: |
           export PYTHONPATH=.
-          python _unittests/ut_torch_models/test_llms_onnx.py
+          python _unittests/ut_torch_models/test_tiny_llms_onnx.py
 
       - name: run tests
         run: |

.github/workflows/documentation.yml

@@ -56,12 +56,12 @@ jobs:
       - name: tiny-llm torch.export.export
         run: |
           export PYTHONPATH=.
-          python _unittests/ut_torch_models/test_llms.py
+          python _unittests/ut_torch_models/test_tiny_llms.py
 
       - name: tiny-llm onnx
         run: |
           export PYTHONPATH=.
-          python _unittests/ut_torch_models/test_llms_onnx.py
+          python _unittests/ut_torch_models/test_tiny_llms_onnx.py
 
       - name: Generate coverage report
         run: |

_unittests/ut_torch_export_patches/test_onnx_export_errors.py

@@ -6,6 +6,7 @@
     skipif_ci_windows,
     ignore_warnings,
     hide_stdout,
+    has_transformers,
 )
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import (
@@ -68,10 +69,12 @@ def forward(self, x: torch.Tensor, cache: MambaCache):
                 return x2
 
         cache = MambaCache(_config(), max_batch_size=1, device="cpu")
-        self.assertEqual(
-            string_type(cache),
-            "MambaCache(conv_states=#64[T10r3,...], ssm_states=#64[T10r3,...])",
-        )
+        if has_transformers("4.50"):
+            # MambaCache was updated in 4.50
+            self.assertEqual(
+                "MambaCache(conv_states=#64[T10r3,...], ssm_states=#64[T10r3,...])",
+                string_type(cache),
+            )
         x = torch.ones(2, 8, 16).to(torch.float16)
         model = Model()
         model(x, cache)

_unittests/ut_torch_models/test_llms.py renamed to _unittests/ut_torch_models/test_tiny_llms.py

@@ -189,6 +189,7 @@ def string_type(
     """
     if obj is None:
         return "None"
+    # tuple
     if isinstance(obj, tuple):
         if len(obj) == 1:
             s = string_type(
@@ -225,6 +226,7 @@ def string_type(
             mini, maxi, avg = min(obj), max(obj), sum(float(_) for _ in obj) / len(obj)
             return f"#{len(obj)}({tt},...)[{mini},{maxi}:A[{avg}]]"
         return f"#{len(obj)}({tt},...)"
+    # list
     if isinstance(obj, list):
         if len(obj) < limit:
             js = ",".join(
@@ -251,6 +253,7 @@ def string_type(
             mini, maxi, avg = min(obj), max(obj), sum(float(_) for _ in obj) / len(obj)
             return f"#{len(obj)}[{tt},...][{mini},{maxi}:{avg}]"
         return f"#{len(obj)}[{tt},...]"
+    # set
     if isinstance(obj, set):
         if len(obj) < 10:
             js = ",".join(
@@ -269,6 +272,7 @@ def string_type(
             mini, maxi, avg = min(obj), max(obj), sum(float(_) for _ in obj) / len(obj)
             return f"{{...}}#{len(obj)}[{mini},{maxi}:A{avg}]"
         return f"{{...}}#{len(obj)}" if with_shape else "{...}"
+    # dict
     if isinstance(obj, dict):
         if len(obj) == 0:
             return "{}"
@@ -281,6 +285,7 @@ def string_type(
         )
         s = ",".join(f"{kv[0]}:{string_type(kv[1],**kws)}" for kv in obj.items())
         return f"dict({s})"
+    # arrat
     if isinstance(obj, np.ndarray):
         if with_min_max:
             s = string_type(obj, with_shape=with_shape)
@@ -303,6 +308,7 @@ def string_type(
 
     import torch
 
+    # Dim, SymInt
     if isinstance(obj, torch.export.dynamic_shapes._DerivedDim):
         return "DerivedDim"
     if isinstance(obj, torch.export.dynamic_shapes._Dim):
@@ -311,13 +317,13 @@ def string_type(
         return "SymInt"
     if isinstance(obj, torch.SymFloat):
         return "SymFloat"
+    # Tensors
     if isinstance(obj, torch._subclasses.fake_tensor.FakeTensor):
         i = torch_dtype_to_onnx_dtype(obj.dtype)
         prefix = ("G" if obj.get_device() >= 0 else "C") if with_device else ""
         if not with_shape:
             return f"{prefix}F{i}r{len(obj.shape)}"
         return f"{prefix}F{i}s{'x'.join(map(str, obj.shape))}"
-
     if isinstance(obj, torch.Tensor):
         if with_min_max:
             s = string_type(obj, with_shape=with_shape, with_device=with_device)