onnx-diagnostic 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.4.1"
+__version__ = "0.4.2"
 __author__ = "Xavier Dupré"

onnx_diagnostic/helpers/cache_helper.py

@@ -155,6 +155,7 @@ def make_mamba_cache(
     key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
 ) -> transformers.cache_utils.MambaCache:
     "Creates a :class:`transformers.cache_utils.MambaCache`."
+    dtype = key_value_pairs[0][0].dtype
 
     class _config:
         def __init__(self):
@@ -162,14 +163,23 @@ def make_mamba_cache(
             self.conv_kernel = key_value_pairs[0][0].shape[-1]
             self.state_size = key_value_pairs[0][1].shape[-1]
             self.num_hidden_layers = len(key_value_pairs)
-            self.dtype = key_value_pairs[0][0].dtype
+            self.dtype = dtype
 
     cache = transformers.cache_utils.MambaCache(
         _config(),
         max_batch_size=key_value_pairs[0][0].shape[0],
         device=key_value_pairs[0][0].device,
+        dtype=dtype,
     )
     for i in range(len(key_value_pairs)):
+        assert cache.conv_states[i].dtype == dtype, (
+            f"Type mismatch for cache.conv_states[{i}].dtype="
+            f"{cache.conv_states[i].dtype} != {dtype}"
+        )
+        assert cache.ssm_states[i].dtype == dtype, (
+            f"Type mismatch for cache.ssm_states[{i}].dtype="
+            f"{cache.ssm_states[i].dtype} != {dtype}"
+        )
         assert cache.conv_states[i].shape == key_value_pairs[i][0].shape, (
             f"Shape mismatch, expected {cache.conv_states[i].shape}, "
             f"got {key_value_pairs[i][0].shape}"

onnx_diagnostic/helpers/helper.py

@@ -1404,6 +1404,28 @@ def max_diff(
             f"level={level}"
         )
 
+    if expected.__class__.__name__ == "SlidingWindowCache":
+        if got.__class__.__name__ == "SlidingWindowCache":
+            if verbose >= 6:
+                print(f"[max_diff] DynamicCache: {string_type(expected)} ? {string_type(got)}")
+            return max_diff(
+                [expected.key_cache, expected.value_cache],
+                [got.key_cache, got.value_cache],
+                verbose=verbose,
+            )
+        if isinstance(got, tuple) and len(got) == 2:
+            return max_diff(
+                [expected.key_cache, expected.value_cache],
+                [got[0], got[1]],
+                verbose=verbose,
+            )
+        raise AssertionError(
+            f"SlidingWindowCache not fully implemented with classes "
+            f"{expected.__class__.__name__!r} and {got.__class__.__name__!r}, "
+            f"and expected={string_type(expected)}, got={string_type(got)},\n"
+            f"level={level}"
+        )
+
     if expected.__class__.__name__ == "EncoderDecoderCache":
         if got.__class__.__name__ == "EncoderDecoderCache":
             if verbose >= 6:

onnx_diagnostic/helpers/torch_test_helper.py

@@ -8,6 +8,7 @@ from .cache_helper import (
     make_dynamic_cache,
     make_encoder_decoder_cache,
     make_sliding_window_cache,
+    make_mamba_cache,
 )
 
 
@@ -346,6 +347,8 @@ def torch_deepcopy(value: Any) -> Any:
     """
     Makes a deepcopy.
     """
+    if value is None:
+        return None
     if isinstance(value, (int, float, str)):
         return value
     if isinstance(value, tuple):
@@ -376,6 +379,9 @@ def torch_deepcopy(value: Any) -> Any:
             torch_deepcopy(value.self_attention_cache),
             torch_deepcopy(value.cross_attention_cache),
         )
+    if value.__class__.__name__ == "MambaCache":
+        return make_mamba_cache(list(zip(value.conv_states, value.ssm_states)))
+
     if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
         args, spec = torch.utils._pytree.tree_flatten(value)
         new_args = torch_deepcopy(args)

{onnx_diagnostic-0.4.1.dist-info → onnx_diagnostic-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnx-diagnostic
-Version: 0.4.1
+Version: 0.4.2
 Summary: Investigate ONNX models
 Home-page: https://github.com/sdpython/onnx-diagnostic
 Author: Xavier Dupré

{onnx_diagnostic-0.4.1.dist-info → onnx_diagnostic-0.4.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-onnx_diagnostic/__init__.py,sha256=bh5VOJtQu31ncsPxImpF46_gmHuLSnb8C4SsLWLh1_k,164
+onnx_diagnostic/__init__.py,sha256=wVSctxhjG5jNBmX9oZ_oUVWt2QU4P1s8bsgeKXDj0YI,164
 onnx_diagnostic/__main__.py,sha256=YmyV_Aq_ianDlHyKLHMa6h8YK3ZmFPpLVHLKjM91aCk,79
 onnx_diagnostic/_command_lines_parser.py,sha256=kOECT1BccZc38vmVc3jF3xvXGDpcocvLuUGoPkzte08,14753
 onnx_diagnostic/doc.py,sha256=MTuT7Kxyvn7KEy84liQeFeqhugJrUQhjjpx21F72Uxw,926
@@ -9,14 +9,14 @@ onnx_diagnostic/export/validate.py,sha256=FBI7Sercyio_Ozw542z5Jh7u6mzKFyDxSRL2hS
 onnx_diagnostic/helpers/__init__.py,sha256=21OeajRAtfTaQ-IRa7UqQemJq5lh_7iYMTuZUcNcEqU,67
 onnx_diagnostic/helpers/args_helper.py,sha256=7pTrw1A1wuNvLdXJdpda5spPI140FylwSmxxZTGu_4E,4389
 onnx_diagnostic/helpers/bench_run.py,sha256=CGA6VMJZMH2gDhVueT9ypNm4PMcjGrrGFYp08nhWj9k,16539
-onnx_diagnostic/helpers/cache_helper.py,sha256=5AaZA3F0Tta99QGr8x0mpFJnqGMZauknjtlb1voetEo,7949
+onnx_diagnostic/helpers/cache_helper.py,sha256=soKjyIXa7EQgALd9PAUGIKYzXlJGoLevYiQDsxoqkQ4,8349
 onnx_diagnostic/helpers/config_helper.py,sha256=zmxKA54xYTNOQyf2MP0FnHsplDmYYtBRLIJlLBrdjUU,3039
-onnx_diagnostic/helpers/helper.py,sha256=8VdJqJfMXU2hB0IKEkhWtERJTOBur8GTxFMpKv1OxPM,53397
+onnx_diagnostic/helpers/helper.py,sha256=HpSws6esUq86jDQKj7jUM4g-Xdd3FqsvcXBX6QDHDlg,54349
 onnx_diagnostic/helpers/memory_peak.py,sha256=lgQm5DvjxfSw9nEPBYyZKMlaGVe-dcP5jqo7SxBOIS0,6380
 onnx_diagnostic/helpers/onnx_helper.py,sha256=4jX_BxbJ29jW-LaCZib9lJXNci4u8iBM7qitx4KrDQU,29336
 onnx_diagnostic/helpers/ort_session.py,sha256=hlM6h0Bn0b5adNVK-QBwyt0SPFzo2Z7AZ5c0Zmdy928,26774
 onnx_diagnostic/helpers/rt_helper.py,sha256=zFIA3HTvogok6EUjWfTAXUd-BfP_Sh4pwEmLcdqPbl0,1774
-onnx_diagnostic/helpers/torch_test_helper.py,sha256=EvhliihBUIff2CEGIXB2X15IfOn5svBg4H1HgUidGB4,13376
+onnx_diagnostic/helpers/torch_test_helper.py,sha256=qaXk1sqzx2kBvgAHYJ5Xzn3pdF1mn-N41tJBq59OejY,13570
 onnx_diagnostic/reference/__init__.py,sha256=0Al5kins8LlBICAsszEZ59thMwmaARBO6fMwtYpKOOQ,98
 onnx_diagnostic/reference/evaluator.py,sha256=gzfcgzc2oC99ynFJ4FF_JPlm-52_OKtpSrLBq7S-QR0,8804
 onnx_diagnostic/reference/ort_evaluator.py,sha256=aYMydWPh9oEAHf1exCu8XFfdpfXvDRY_3d4sKJe5ruo,16369
@@ -81,8 +81,8 @@ onnx_diagnostic/torch_models/untrained/llm_phi2.py,sha256=ynBTDHJHCk44NjLT_t6OiF
 onnx_diagnostic/torch_models/untrained/llm_tiny_llm.py,sha256=7N3fGvT_4Mn4NbIo0Qk57c6DMc3OXGWyvj_P41rjwSY,3513
 onnx_diagnostic/torch_onnx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnx_diagnostic/torch_onnx/sbs.py,sha256=HEGDHhV9pfXxpBQrpOWPNWGMsNfOebWewyAazi9poV8,16872
-onnx_diagnostic-0.4.1.dist-info/licenses/LICENSE.txt,sha256=Vv6TXglX6Rc0d-f8aREhayhT-6PMQXEyOmI2NKlUCMc,1045
-onnx_diagnostic-0.4.1.dist-info/METADATA,sha256=XSoBJwePsa6TtlJj_ok66Xf6RIqByU5e1oihbqMQt_s,5511
-onnx_diagnostic-0.4.1.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
-onnx_diagnostic-0.4.1.dist-info/top_level.txt,sha256=KwNkXewmcobM3ZT1DJLVWH6ebJzA5qKg7cWqKfpGNT4,16
-onnx_diagnostic-0.4.1.dist-info/RECORD,,
+onnx_diagnostic-0.4.2.dist-info/licenses/LICENSE.txt,sha256=Vv6TXglX6Rc0d-f8aREhayhT-6PMQXEyOmI2NKlUCMc,1045
+onnx_diagnostic-0.4.2.dist-info/METADATA,sha256=xkEoLhmlKpx91wddmUKXOpmDwdsEuW2y_c7bbk7cAVw,5511
+onnx_diagnostic-0.4.2.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+onnx_diagnostic-0.4.2.dist-info/top_level.txt,sha256=KwNkXewmcobM3ZT1DJLVWH6ebJzA5qKg7cWqKfpGNT4,16
+onnx_diagnostic-0.4.2.dist-info/RECORD,,