extend coverage

Author: xadupre

_unittests/ut_torch_export_patches/test_onnx_export_errors.py

@@ -5,6 +5,7 @@
     requires_transformers,
     skipif_ci_windows,
     ignore_warnings,
+    hide_stdout,
 )
 from onnx_diagnostic.helpers import string_type
 from onnx_diagnostic.torch_export_patches.onnx_export_errors import (
@@ -16,6 +17,7 @@ class TestOnnxExportErrors(ExtTestCase):
     @requires_transformers("4.49.999")
     @skipif_ci_windows("not working on Windows")
     @ignore_warnings(UserWarning)
+    @hide_stdout()
     def test_pytree_flatten_mamba_cache(self):
         import torch
         import torch.utils._pytree as py_pytree
@@ -31,7 +33,7 @@ def __init__(self):
 
         cache = MambaCache(_config(), max_batch_size=1, device="cpu")
 
-        with bypass_export_some_errors():
+        with bypass_export_some_errors(verbose=1):
             values, spec = py_pytree.tree_flatten(cache)
             cache2 = py_pytree.tree_unflatten(values, spec)
             self.assertEqual(cache.dtype, cache2.dtype)
@@ -46,6 +48,7 @@ def __init__(self):
     @requires_torch("2.7")
     @skipif_ci_windows("not working on Windows")
     @ignore_warnings(UserWarning)
+    @hide_stdout()
     def test_exportable_mamba_cache(self):
         import torch
         from transformers.models.mamba.modeling_mamba import MambaCache
@@ -73,7 +76,7 @@ def forward(self, x: torch.Tensor, cache: MambaCache):
         model = Model()
         model(x, cache)
 
-        with bypass_export_some_errors():
+        with bypass_export_some_errors(replace_dynamic_cache=True, verbose=1):
             cache = MambaCache(_config(), max_batch_size=1, device="cpu")
             torch.export.export(Model(), (x, cache))
 

_unittests/ut_xrun_doc/test_args.py

@@ -25,6 +25,30 @@ def test_args(self):
         self.assertEqual(args.with_mask, 0)
         self.assertEqual(args.optim, "")
 
+    def test_args_expose(self):
+        try:
+            args = get_parsed_args(
+                "plot_custom_backend_llama",
+                config=(
+                    "medium",
+                    "large or medium depending, large means closer to the real model",
+                ),
+                num_hidden_layers=(1, "number of hidden layers"),
+                with_mask=(0, "tries with a mask as a secondary input"),
+                optim=("", "Optimization to apply, empty string for all"),
+                description="doc",
+                new_args=["--config", "m"],
+                expose="repeat,warmup",
+            )
+        except SystemExit as e:
+            raise AssertionError(f"SystemExist caught: {e}")
+        self.assertEqual(args.config, "m")
+        self.assertEqual(args.num_hidden_layers, 1)
+        self.assertEqual(args.with_mask, 0)
+        self.assertEqual(args.optim, "")
+        self.assertEqual(args.repeat, 10)
+        self.assertEqual(args.warmup, 5)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_xrun_doc/test_onnx_tools.py

@@ -4,7 +4,7 @@
 import onnx.numpy_helper as onh
 from onnx import TensorProto
 from onnx.checker import check_model
-from onnx_diagnostic.ext_test_case import ExtTestCase
+from onnx_diagnostic.ext_test_case import ExtTestCase, hide_stdout
 from onnx_diagnostic.onnx_tools import (
     onnx_lighten,
     onnx_unlighten,
@@ -72,6 +72,7 @@ def test_onnx_find(self):
         self.assertIn("xm2", res[0].output)
         self.assertIn("xm2", res[1].input)
 
+    @hide_stdout()
     def test__validate_function(self):
         new_domain = "custom"
 
@@ -87,7 +88,7 @@ def test__validate_function(self):
             [oh.make_opsetid("", 14)],
             [],
         )
-        _validate_function(linear_regression)
+        _validate_function(linear_regression, verbose=1)
 
 
 if __name__ == "__main__":

_unittests/ut_xrun_doc/test_unit_test.py

@@ -84,6 +84,24 @@ def test_measure_time(self):
             },
         )
 
+    def test_measure_time_max(self):
+        res = measure_time(lambda: math.cos(0.5), max_time=0.1)
+        self.assertIsInstance(res, dict)
+        self.assertEqual(
+            set(res),
+            {
+                "min_exec",
+                "max_exec",
+                "average",
+                "warmup_time",
+                "context_size",
+                "deviation",
+                "repeat",
+                "ttime",
+                "number",
+            },
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers.py

@@ -66,12 +66,8 @@ def size_type(dtype: Any) -> int:
         return 4
     if dtype == np.float16 or dtype == np.int16:
         return 2
-    if dtype == np.int16:
-        return 2
     if dtype == np.int32:
         return 4
-    if dtype == np.int64:
-        return 8
     if dtype == np.int8:
         return 1
     if hasattr(np, "uint64"):