onnx-diagnostic 0.7.1__py3-none-any.whl → 0.7.2__py3-none-any.whl

onnx_diagnostic/__init__.py

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

onnx_diagnostic/_command_lines_parser.py

@@ -718,13 +718,13 @@ def get_parser_agg() -> ArgumentParser:
         "peak_gpu_torch,peak_gpu_nvidia,n_node_control_flow,"
         "n_node_constant,n_node_shape,n_node_expand,"
         "n_node_function,n_node_initializer,n_node_scatter,"
-        "time_export_unbiased",
+        "time_export_unbiased,onnx_n_nodes_no_cst,n_node_initializer_small",
         help="Columns to compute after the aggregation was done.",
     )
     parser.add_argument(
         "--views",
         default="agg-suite,agg-all,disc,speedup,time,time_export,err,cmd,"
-        "bucket-speedup,raw-short,counts,peak-gpu",
+        "bucket-speedup,raw-short,counts,peak-gpu,onnx",
         help="Views to add to the output files.",
     )
     parser.add_argument(
@@ -733,11 +733,28 @@ def get_parser_agg() -> ArgumentParser:
         help="Views to dump as csv files.",
     )
     parser.add_argument("-v", "--verbose", type=int, default=0, help="verbosity")
+    parser.add_argument(
+        "--filter-in",
+        default="",
+        help="adds a filter to filter in data, syntax is\n"
+        '``"<column1>:<value1>;<value2>/<column2>:<value3>"`` ...',
+    )
+    parser.add_argument(
+        "--filter-out",
+        default="",
+        help="adds a filter to filter out data, syntax is\n"
+        '``"<column1>:<value1>;<value2>/<column2>:<value3>"`` ...',
+    )
     return parser
 
 
 def _cmd_agg(argv: List[Any]):
-    from .helpers.log_helper import CubeLogsPerformance, open_dataframe, enumerate_csv_files
+    from .helpers.log_helper import (
+        CubeLogsPerformance,
+        open_dataframe,
+        enumerate_csv_files,
+        filter_data,
+    )
 
     parser = get_parser_agg()
     args = parser.parse_args(argv[1:])
@@ -748,7 +765,7 @@ def _cmd_agg(argv: List[Any]):
             args.inputs, verbose=args.verbose, filtering=lambda name: bool(reg.search(name))
         )
     )
-    assert csv, f"No csv files in {args.inputs}, csv={csv}"
+    assert csv, f"No csv files in {args.inputs}, args.filter={args.filter!r}, csv={csv}"
     if args.verbose:
         from tqdm import tqdm
 
@@ -761,7 +778,7 @@ def _cmd_agg(argv: List[Any]):
         assert (
             args.time in df.columns
         ), f"Missing time column {args.time!r} in {c!r}\n{df.head()}\n{sorted(df.columns)}"
-        dfs.append(df)
+        dfs.append(filter_data(df, filter_in=args.filter_in, filter_out=args.filter_out))
 
     drop_keys = set(args.drop_keys.split(","))
     cube = CubeLogsPerformance(

onnx_diagnostic/ext_test_case.py

@@ -756,6 +756,18 @@ class ExtTestCase(unittest.TestCase):
         "Adds a todo printed when all test are run."
         cls._todos.append((f, msg))
 
+    @classmethod
+    def ort(cls):
+        import onnxruntime
+
+        return onnxruntime
+
+    @classmethod
+    def to_onnx(self, *args, **kwargs):
+        from experimental_experiment.torch_interpreter import to_onnx
+
+        return to_onnx(*args, **kwargs)
+
     def print_model(self, model: "ModelProto"):  # noqa: F821
         "Prints a ModelProto"
         from onnx_diagnostic.helpers.onnx_helper import pretty_onnx
@@ -917,6 +929,15 @@ class ExtTestCase(unittest.TestCase):
             ]
             raise AssertionError("\n".join(rows))  # noqa: B904
 
+    def assertEqualDataFrame(self, d1, d2, **kwargs):
+        """
+        Checks that two dataframes are equal.
+        Calls :func:`pandas.testing.assert_frame_equal`.
+        """
+        from pandas.testing import assert_frame_equal
+
+        assert_frame_equal(d1, d2, **kwargs)
+
     def assertEqualTrue(self, value: Any, msg: str = ""):
         if value is True:
             return
@@ -967,6 +988,16 @@ class ExtTestCase(unittest.TestCase):
                 atol=atol,
                 rtol=rtol,
             )
+        elif expected.__class__.__name__ == "StaticCache":
+            self.assertEqual(type(expected), type(value), msg=msg)
+            self.assertEqual(expected.max_cache_len, value.max_cache_len)
+            atts = ["key_cache", "value_cache"]
+            self.assertEqualAny(
+                {k: expected.__dict__.get(k, None) for k in atts},
+                {k: value.__dict__.get(k, None) for k in atts},
+                atol=atol,
+                rtol=rtol,
+            )
         elif expected.__class__.__name__ == "EncoderDecoderCache":
             self.assertEqual(type(expected), type(value), msg=msg)
             atts = ["self_attention_cache", "cross_attention_cache"]

onnx_diagnostic/helpers/cache_helper.py

@@ -154,10 +154,12 @@ else:
 
 def make_static_cache(
     key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
+    max_cache_len: Optional[int] = None,
 ) -> transformers.cache_utils.DynamicCache:
     """
     Creates an instance of :class:`transformers.cache_utils.StaticCache`.
     :param key_value_pairs: list of pairs of (key, values)
+    :param max_cache_len: max_cache_length or something inferred from the vector
     :return: :class:`transformers.cache_utils.StaticCache`
 
     Example:
@@ -179,7 +181,8 @@ def make_static_cache(
                     torch.randn(bsize, nheads, slen, dim),
                 )
                 for i in range(n_layers)
-            ]
+            ],
+            max_cache_len=10,
         )
         print(string_type(past_key_values, with_shape=True))
     """
@@ -190,24 +193,32 @@ def make_static_cache(
             self.num_attention_heads = key_value_pairs[0][0].shape[1]
             self.num_hidden_layers = len(key_value_pairs)
 
+    assert max_cache_len is not None, (
+        f"max_cache_len={max_cache_len} cannot be setup "
+        f"automatically yet from shape {key_value_pairs[0][0].shape}"
+    )
+    torch._check(
+        max_cache_len >= key_value_pairs[0][0].shape[2],
+        (
+            f"max_cache_len={max_cache_len} cannot be smaller "
+            f"shape[2]={key_value_pairs[0][0].shape[2]} in shape "
+            f"{key_value_pairs[0][0].shape}"
+        ),
+    )
     cache = transformers.cache_utils.StaticCache(
         _config(),
         max_batch_size=key_value_pairs[0][0].shape[0],
         device=key_value_pairs[0][0].device,
         dtype=key_value_pairs[0][0].dtype,
-        max_cache_len=key_value_pairs[0][0].shape[2],
+        max_cache_len=max_cache_len,
     )
     for i in range(len(key_value_pairs)):
-        assert cache.key_cache[i].shape == key_value_pairs[i][0].shape, (
-            f"Shape mismatch, expected {cache.key_cache[i].shape}, "
-            f"got {key_value_pairs[i][0].shape}"
-        )
-        cache.key_cache[i][:, :, :, :] = key_value_pairs[i][0]
-        assert cache.value_cache[i].shape == key_value_pairs[i][1].shape, (
-            f"Shape mismatch, expected {cache.value_cache[i].shape}, "
-            f"got {key_value_pairs[i][1].shape}"
-        )
-        cache.value_cache[i][:, :, :, :] = key_value_pairs[i][1]
+        assert (
+            key_value_pairs[i][0].shape == key_value_pairs[i][1].shape
+        ), f"Shape mismatch {key_value_pairs[i][0].shape} != {key_value_pairs[i][1].shape}"
+        d = key_value_pairs[i][1].shape[2]
+        cache.key_cache[i][:, :, :d, :] = key_value_pairs[i][0]
+        cache.value_cache[i][:, :, :d, :] = key_value_pairs[i][1]
     return cache
 
 

onnx_diagnostic/helpers/config_helper.py

@@ -43,7 +43,10 @@ def update_config(config: Any, mkwargs: Dict[str, Any]):
             else:
                 update_config(getattr(config, k), v)
             continue
-        setattr(config, k, v)
+        if type(config) is dict:
+            config[k] = v
+        else:
+            setattr(config, k, v)
 
 
 def _pick(config, *atts, exceptions: Optional[Dict[str, Callable]] = None):
@@ -66,6 +69,18 @@ def _pick(config, *atts, exceptions: Optional[Dict[str, Callable]] = None):
     raise AssertionError(f"Unable to find any of these {atts!r} in {config}")
 
 
+def pick(config, name: str, default_value: Any) -> Any:
+    """
+    Returns the value of a attribute if config has it
+    otherwise the default value.
+    """
+    if not config:
+        return default_value
+    if type(config) is dict:
+        return config.get(name, default_value)
+    return getattr(config, name, default_value)
+
+
 @functools.cache
 def config_class_from_architecture(arch: str, exc: bool = False) -> Optional[type]:
     """