onnx-diagnostic 0.8.7__py3-none-any.whl → 0.8.9__py3-none-any.whl

onnx_diagnostic/export/dynamic_shapes.py

@@ -329,7 +329,7 @@ class CoupleInputsDynamicShapes:
         if type(inputs) in (tuple, list, dict):
             # Type must be strict, some custom classes can inherit from those.
             assert type(inputs) is type(ds), (
-                f"Input type and dynamic shape type mush match but "
+                f"Input type and dynamic shapes type mush match but "
                 f"type(inputs)={type(inputs)}, type(ds)={type(ds)}, "
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
             )
@@ -352,6 +352,19 @@ class CoupleInputsDynamicShapes:
                     else None
                 )
             assert type(inputs) is dict, f"Unexpected type for inputs {type(inputs)}"
+            if set(inputs) != set(ds):
+                not_in_ds = {k for k in inputs if k not in ds}
+                not_in_inputs = {k for k in ds if k not in inputs}
+                assert not_in_inputs == {"kwargs"} and set(ds["kwargs"]) == not_in_ds, (
+                    f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
+                    f"inputs={string_type(inputs, with_shape=True)}, ds={ds}, "
+                    f"not_in_ds={not_in_ds}, not_in_inputs={not_in_inputs}"
+                )
+                # Tweak...
+                kws = ds["kwargs"]
+                del ds["kwargs"]
+                ds.update(kws)
+
             assert set(inputs) == set(ds), (
                 f"Keys mismatch between inputs {set(inputs)} and ds={set(ds)}, "
                 f"inputs={string_type(inputs, with_shape=True)}, ds={ds}"
@@ -366,13 +379,15 @@ class CoupleInputsDynamicShapes:
             return dvalue if dvalue else None
 
         # A custom class.
-        assert inputs.__class__ in torch.utils._pytree.SUPPORTED_NODES, (
+        assert inputs is None or inputs.__class__ in torch.utils._pytree.SUPPORTED_NODES, (
             f"Class {inputs.__class__.__name__!r} was not registered using "
             f"torch.utils._pytree.register_pytree_node, it is not possible to "
             f"map this class with the given dynamic shapes."
         )
         if flatten_unflatten:
             flatunflat = flatten_unflatten_for_dynamic_shapes(inputs)
+            if isinstance(flatunflat, (list, tuple, dict)) and len(flatunflat) == 0:
+                return flatunflat
             res = cls._generic_walker_step(
                 processor, flatunflat, ds, flatten_unflatten=flatten_unflatten
             )
@@ -667,6 +682,11 @@ class ModelInputs:
             if self.signature
             else None
         )
+        self.forward_parameters_kinds = (
+            {p.name: p.kind for p in self.signature.parameters.values()}
+            if self.signature
+            else None
+        )
         self.forward_ordered_parameter_names = (
             list(self.signature.parameters) if self.signature else None
         )
@@ -947,6 +967,8 @@ class ModelInputs:
         """
         Guesses the dynamic shapes for that module from two execution.
         If there is only one execution, then that would be static dimensions.
+        If the model signature is available, the kwargs are reordered following
+        the signature order, otherwise it follows the order given in the inputs.
 
         :param auto: if auto is True, use ``torch.export.Dim.AUTO`` for any
             dimension if the number of inputs is one,
@@ -973,7 +995,13 @@ class ModelInputs:
             len(s1) == 1
         ), f"Different numbers of positional arguments {s1} for {self.full_name}"
         s2 = set(tuple(sorted(set(i[1]))) for i in self.inputs)
-        assert len(s2) == 1, f"Different named arguments {s2} for {self.full_name}"
+        assert len(s2) > 0, f"empty {s2} for {self.full_name}"
+        if len(s2) > 1:
+            # We need to keep the largest set of inputs, the one including all the others.
+            sum_s2 = set()
+            for s in s2:
+                sum_s2 |= set(s)
+            s2 = {tuple(sum_s2)}
         args = []
         kwargs = {}
         for i in range(s1.pop()):
@@ -993,12 +1021,31 @@ class ModelInputs:
                 f"\ninputs[1]={string_type(self.inputs[1], with_shape=True)}"
             )
 
-            objs = [_[1][name] for _ in self.inputs]
+            objs = [_[1][name] for _ in self.inputs if name in _[1]]
             kwargs[name] = self.guess_dynamic_shape_object(
                 *objs,
                 auto=auto if isinstance(auto, bool) else f"{auto}_{i}I",
                 msg=lambda name=name: f" failing input {name!r}",
             )
+        # reordering
+        if kwargs:
+            if self.forward_ordered_parameter_names:
+                kwargs1 = {
+                    p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs
+                }
+                kwargs = {**kwargs1, **{k: v for k, v in kwargs.items() if k not in kwargs1}}
+            else:
+                # We reorder the same the way the input were given.
+                use = None
+                params = set(kwargs)
+                for _args, kws in self.inputs:
+                    if set(kws) == params:
+                        use = kws
+                        break
+                if use:
+                    ordered = list(use)
+                    kwargs = {k: kwargs[k] for k in ordered}
+
         return tuple(args), kwargs
 
     def move_to_kwargs(
@@ -1061,6 +1108,16 @@ class ModelInputs:
                 f"and kwargs={set(kwargs)}, "
                 f"forward_ordered_parameter_names={self.forward_ordered_parameter_names}"
             )
+        if kwargs is not None and self.forward_ordered_parameter_names:
+            kwargs1 = {
+                p: kwargs[p] for p in self.forward_ordered_parameter_names if p in kwargs
+            }
+            kwargs = {**kwargs1, **{k: v for k, v in kwargs.items() if k not in kwargs1}}
+        if kw_dyn is not None and self.forward_ordered_parameter_names:
+            kw_dyn1 = {
+                p: kw_dyn[p] for p in self.forward_ordered_parameter_names if p in kw_dyn
+            }
+            kw_dyn = {**kw_dyn1, **{k: v for k, v in kw_dyn.items() if k not in kw_dyn1}}
         return args, kwargs, (tuple(), kw_dyn)
 
     def validate_inputs_for_export(

onnx_diagnostic/export/shape_helper.py

@@ -47,7 +47,6 @@ def all_dynamic_shapes_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
             make_dynamic_cache,
             make_encoder_decoder_cache,
             make_mamba_cache,
-            make_sliding_window_cache,
             make_static_cache,
         )
         from onnx_diagnostic.export.shape_helper import all_dynamic_shapes_from_inputs
@@ -77,13 +76,6 @@ def all_dynamic_shapes_from_inputs(inputs: Any, dim_prefix: Any = "d") -> Any:
                     ]
                 ),
             ),
-            make_sliding_window_cache(
-                [
-                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
-                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
-                    (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
-                ]
-            ),
             make_static_cache(
                 [
                     (torch.rand((4, 5, 6, 7)), torch.rand((4, 5, 6, 7))),
@@ -210,6 +202,7 @@ def make_fake_with_dynamic_dimensions(
     This uses function :func:`onnx_diagnostic.helpers.fake_tensor_helper.make_fake`.
     Parameter ``existing`` is used to reused the same object when the dynamic
     dimension is given the same name as another one.
+    This function works with caches only if ``transformers>=4.57``.
 
     A simple tensor:
 

onnx_diagnostic/helpers/cache_helper.py

@@ -19,7 +19,7 @@ class CacheKeyValue:
         capi.value_cache
     """
 
-    def __init__(self, cache=None):
+    def __init__(self, cache=None, cls_layers=None):
         if hasattr(cache, "layers"):
             layers = [
                 layer
@@ -28,32 +28,52 @@ class CacheKeyValue:
             ]
             self.key_cache = [layer.keys for layer in layers]
             self.value_cache = [layer.values for layer in layers]
-            if None in self.key_cache or None in self.value_cache:
-                from .helper import string_type
-
-                raise AssertionError(
-                    f"issue with key_cache={string_type(self.key_cache)}, "
-                    f"or value_cache={string_type(self.value_cache)}, "
-                    f"cache.layers={string_type(cache.layers)}"
-                )
+            assert (
+                cls_layers is None
+            ), f"cache is {type(cache)}, cannot specify cls_layers={cls_layers}"
+            self.cls_layers = [type(lay) for lay in cache.layers]
         elif cache is not None and hasattr(cache, "key_cache"):
             self.key_cache = cache.key_cache
             self.value_cache = cache.value_cache
+            self.cls_layers = cls_layers
+        elif (
+            cache is not None
+            and isinstance(cache, list)
+            and all(isinstance(t, torch.Tensor) for t in cache)
+        ):
+            self.key_cache = cache[::2]
+            self.value_cache = cache[1::2]
+            self.cls_layers = cls_layers
         elif cache is None:
             self.key_cache = None
             self.value_cache = None
+            self.cls_layers = cls_layers
         else:
             raise NotImplementedError(f"type(cache)={type(cache)}")
 
     def make_dynamic_cache(self):
         """Does the reverse operation."""
-        return make_dynamic_cache(list(zip(self.key_cache, self.value_cache)))
+        return make_dynamic_cache(
+            list(zip(self.key_cache, self.value_cache)), cls_layers=self.cls_layers
+        )
 
     @property
     def n_layers(self) -> int:
         """Returns the number of layers."""
         return len(self.key_cache) if self.key_cache else 0
 
+    def __len__(self) -> int:
+        "Returns the number of tensors."
+        return len(self.key_cache) + len(self.value_cache)
+
+    def aslist(self) -> List[torch.Tensor]:
+        "Returns tensors in a list."
+        res = []
+        for i in range(self.n_layers):
+            res.append(self.key_cache[i])
+            res.append(self.value_cache[i])
+        return res
+
 
 def flatten_unflatten_for_dynamic_shapes(
     obj: Any,
@@ -164,12 +184,16 @@ if pv.Version(transformers.__version__) > pv.Version("4.49.99999"):
 
     def make_dynamic_cache(
         key_value_pairs: Union[List[torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]],
+        cls_layers: Optional[Union[str, List[type]]] = None,
     ) -> transformers.cache_utils.DynamicCache:
         """
         Creates an instance of :class:`transformers.cache_utils.DynamicCache`.
         This version is valid for ``transformers >= 4.50``.
 
         :param key_value_pairs: list of pairs of (key, values)
+        :param cls_layers: to select the appropriate class to use on each layer,
+            if specified, sliding_window is ignored, it can be a string
+            if all layers are expected to follow the same class
         :return: :class:`transformers.cache_utils.DynamicCache`
 
         Example:
@@ -200,15 +224,49 @@ if pv.Version(transformers.__version__) > pv.Version("4.49.99999"):
         are supported.
         """
         key_value_pairs = _preprocess_key_value_pairs(key_value_pairs)
+        cls_kwargs = {}
+        if isinstance(cls_layers, str):
+            assert hasattr(
+                transformers.cache_utils, cls_layers
+            ), f"Unable to find class {cls_layers!r} in transformers.cache_utils"
+            cls_layer = getattr(transformers.cache_utils, cls_layers)
+            if cls_layers == "DynamicSlidingWindowLayer":
+                cls_kwargs["sliding_window"] = key_value_pairs[0][0].shape[2]
+                assert isinstance(
+                    cls_kwargs["sliding_window"], int
+                ), f"sliding_window must be an integer but shape={key_value_pairs[0][0].shape}"
+        elif cls_layers is not None:
+            unique = set(cls_layers)
+            assert len(unique) == 1, f"Not implemented when cls_layers={cls_layers}"
+            cls_layer = unique.pop()
+            if (
+                hasattr(transformers.cache_utils, "DynamicSlidingWindowLayer")
+                and cls_layer == transformers.cache_utils.DynamicSlidingWindowLayer
+            ):
+                from .helper import string_type
+
+                assert key_value_pairs and key_value_pairs[0], (
+                    f"not implemented for key_value_pairs="
+                    f"{string_type(key_value_pairs, with_shape=True)}"
+                )
+                cls_kwargs["sliding_window"] = key_value_pairs[0][0].shape[2]
+                assert isinstance(
+                    cls_kwargs["sliding_window"], int
+                ), f"sliding_window must be an integer but shape={key_value_pairs[0][0].shape}"
+        else:
+            cls_layer = (
+                transformers.cache_utils.DynamicLayer
+                if hasattr(transformers.cache_utils, "DynamicLayer")
+                else None
+            )
+
         if (
             key_value_pairs
             and isinstance(key_value_pairs[0][0], torch._subclasses.fake_tensor.FakeTensor)
             and pv.Version(transformers.__version__) >= pv.Version("4.56")
         ):
             cache = transformers.cache_utils.DynamicCache()
-            cache.layers.extend(
-                [transformers.cache_utils.DynamicLayer() for _ in key_value_pairs]
-            )
+            cache.layers.extend([cls_layer(**cls_kwargs) for _ in key_value_pairs])
             for i, layer in enumerate(cache.layers):
                 k, v = key_value_pairs[i][0], key_value_pairs[i][1]
                 layer.dtype = k.dtype
@@ -222,14 +280,21 @@ if pv.Version(transformers.__version__) > pv.Version("4.49.99999"):
             )
             return finalize_cache(cache)
 
-        cache = transformers.cache_utils.DynamicCache(key_value_pairs)
-        if hasattr(cache, "layers") and len(key_value_pairs) < len(cache.layers):
-            # The cache constructor contains the two following lines
-            # (in cache_utils.py) which append empty layers when the cache is
-            # initialized. We need to remove them.
-            # self.num_hidden_layers = getattr(config, "num_hidden_layers", 1)
-            # self.append_new_layers(self.num_hidden_layers - 1)
-            cache.layers[:] = cache.layers[-len(key_value_pairs) :]
+        cache = transformers.cache_utils.DynamicCache()
+        if hasattr(cache, "layers") and cls_layer != transformers.cache_utils.DynamicLayer:
+            cache.layers.extend([cls_layer(**cls_kwargs) for _ in key_value_pairs])
+            for i, layer in enumerate(cache.layers):
+                layer.keys, layer.values = key_value_pairs[i][0], key_value_pairs[i][1]
+                layer.is_initialized = True
+        else:
+            cache = transformers.cache_utils.DynamicCache(key_value_pairs)
+            if hasattr(cache, "layers") and len(key_value_pairs) < len(cache.layers):
+                # The cache constructor contains the two following lines
+                # (in cache_utils.py) which append empty layers when the cache is
+                # initialized. We need to remove them.
+                # self.num_hidden_layers = getattr(config, "num_hidden_layers", 1)
+                # self.append_new_layers(self.num_hidden_layers - 1)
+                cache.layers[:] = cache.layers[-len(key_value_pairs) :]
         assert not hasattr(cache, "layers") or len(key_value_pairs) == len(cache.layers), (
             f"Unexpected number of layers in the cache ({len(cache.layers)}), "
             f"{len(key_value_pairs)} expected."
@@ -240,6 +305,7 @@ else:
 
     def make_dynamic_cache(
         key_value_pairs: Union[List[torch.Tensor], List[Tuple[torch.Tensor, torch.Tensor]]],
+        cls_layers: Optional[Union[str, List[type]]] = None,
     ) -> transformers.cache_utils.DynamicCache:
         """
         Creates an instance of :class:`transformers.cache_utils.DynamicCache`.
@@ -271,6 +337,7 @@ else:
             )
             print(string_type(past_key_values, with_shape=True))
         """
+        assert not cls_layers, "cls_layers cannot be used for transformers<5."
         key_value_pairs = _preprocess_key_value_pairs(key_value_pairs)
         cache = transformers.cache_utils.DynamicCache(len(key_value_pairs))  # type: ignore
         for i, (key, value) in enumerate(key_value_pairs):
@@ -516,9 +583,13 @@ if hasattr(transformers.cache_utils, "SlidingWindowCache"):
         )
         return finalize_cache(cache)
 
+    def get_make_hybrid_cache():
+        return make_sliding_window_cache
+
 else:
     make_sliding_window_cache = None  # type: ignore[assignment]
 
+
 if hasattr(transformers.cache_utils, "HybridCache"):
 
     def make_hybrid_cache(
@@ -680,9 +751,15 @@ if hasattr(transformers.cache_utils, "HybridCache"):
         )
         return finalize_cache(cache)
 
+    def get_make_hybrid_cache():
+        return make_hybrid_cache
+
 else:
     make_hybrid_cache = None  # type: ignore[assignment]
 
+    def get_make_hybrid_cache():
+        return None
+
 
 def finalize_cache(cache: transformers.cache_utils.Cache) -> transformers.cache_utils.Cache:
     """

onnx_diagnostic/helpers/fake_tensor_helper.py

@@ -105,6 +105,8 @@ class FakeTensorContext:
         reduced_tensor = self.from_tensor(true_tensor, static_shapes=True).sum(
             axis=tuple(sorted(sh)), keepdim=True
         )
+        if len(reduced_tensor.shape) == 0 == len(new_shape):
+            return reduced_tensor
         return reduced_tensor.expand(*new_shape)
 
     def make_fake(self, x: Any) -> Optional["FakeTensor"]:  # noqa: F821
@@ -144,19 +146,22 @@ class FakeTensorContext:
         """
         See
         :func:`onnx_diagnostic.export.shape_helper.make_fake_with_dynamic_dimensions`.
+        If caches are used, it requires ``transformers>=4.57``.
         """
         if x is None:
             return None, None
-        if isinstance(x, (list, tuple)):
+        if type(x) in (list, tuple):
             return x.__class__(
                 [
                     self.make_fake_with_dynamic_dimensions(i, dynamic_shapes=ds)
                     for i, ds in zip(x, dynamic_shapes)
                 ]
             )
-        if isinstance(x, dict):
+        if type(x) is dict:
             return {
-                k: self.make_fake_with_dynamic_dimensions(v, dynamic_shapes=dynamic_shapes[k])
+                k: self.make_fake_with_dynamic_dimensions(
+                    v, dynamic_shapes=dynamic_shapes[k] if dynamic_shapes else None
+                )
                 for k, v in x.items()
             }
         if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
@@ -187,6 +192,17 @@ class FakeTensorContext:
                 x.cross_attention_cache, dynamic_shapes=dynamic_shapes[1]
             )
             return x
+        if x.__class__.__name__ == "BaseModelOutput":
+            assert (
+                list(x.keys()) == ["last_hidden_state"] and x.last_hidden_state is not None
+            ), (
+                f"Field 'last_hidden_state' is empty for {type(x)} or other fields "
+                f"{list(x.keys())} are used."
+            )
+            x.last_hidden_state = self.make_fake_with_dynamic_dimensions(
+                x.last_hidden_state, dynamic_shapes=dynamic_shapes[0]
+            )
+            return x
         if hasattr(x, "shape"):
             assert dynamic_shapes is None or isinstance(dynamic_shapes, dict), (
                 f"dynamic_shapes must be a dictionary at this stage but "
@@ -197,9 +213,11 @@ class FakeTensorContext:
             for idim, dim in enumerate(x.shape):
                 if dynamic_shapes is not None and idim in dynamic_shapes:
                     s = dynamic_shapes[idim]
+                    if s.__class__.__name__ == "Dim":
+                        s = s.__name__
                     assert isinstance(s, str), (
                         f"Unexpected type {type(s)} in dynamic_shapes={dynamic_shapes} "
-                        f"at index {idim}"
+                        f"at index {idim}, self._mapping_str={self._mapping_str}"
                     )
                     if s in self._mapping_str:
                         dim = self._mapping_str[s]
@@ -217,10 +235,13 @@ class FakeTensorContext:
 
                 x = torch.empty(tuple(new_shape), dtype=x.dtype, device=x.device)
 
-            t = self.fake_reshape(x, dynamic_shapes)  # type: ignore[arg-type]
+            t = self.fake_reshape(x, dynamic_shapes) if dynamic_shapes else x  # type: ignore[arg-type]
             assert t.device == x.device, f"device mismatch {x.device} -> {t.device}"
             assert t.dtype == x.dtype, f"dtype mismatch {x.dtype} -> {t.dtype}"
             return t
+        if isinstance(x, (int, bool, float)):
+            # It is a constant, we don't change that.
+            return x
         from ..helpers import string_type
 
         raise TypeError(

onnx_diagnostic/helpers/helper.py

@@ -1,7 +1,6 @@
 import ast
 import enum
 import inspect
-import itertools
 import json
 from dataclasses import is_dataclass, fields
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@@ -704,9 +703,35 @@ def string_type(
     if obj.__class__.__name__ == "VirtualTensor":
         if verbose:
             print(f"[string_type] TT4:{type(obj)}")
+
+        def _torch_sym_int_to_str(value: "torch.SymInt") -> Union[int, str]:  #  noqa: F821
+            if isinstance(value, str):
+                return value
+            if hasattr(value, "node") and isinstance(value.node, str):
+                return f"{value.node}"
+
+            from torch.fx.experimental.sym_node import SymNode
+
+            if hasattr(value, "node") and isinstance(value.node, SymNode):
+                # '_expr' is safer than expr
+                return str(value.node._expr).replace(" ", "")
+
+            try:
+                val_int = int(value)
+                return val_int
+            except (
+                TypeError,
+                ValueError,
+                AttributeError,
+                torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode,
+            ):
+                pass
+
+            raise AssertionError(f"Unable to convert {value!r} into string")
+
         return (
             f"{obj.__class__.__name__}(name={obj.name!r}, "
-            f"dtype={obj.dtype}, shape={obj.shape})"
+            f"dtype={obj.dtype}, shape={tuple(_torch_sym_int_to_str(_) for _ in obj.shape)})"
         )
 
     if obj.__class__.__name__ == "KeyValuesWrapper":
@@ -775,6 +800,9 @@ def string_type(
             print(f"[string_type] TT8:{type(obj)}")
         return repr(obj).replace(" ", "").replace("\n", " ")
 
+    if isinstance(obj, torch.fx.proxy.Proxy):
+        return repr(obj)
+
     if ignore:
         if verbose:
             print(f"[string_type] CACHE4:{type(obj)}")
@@ -962,15 +990,17 @@ def flatten_object(x: Any, drop_keys: bool = False) -> Any:
     if x.__class__.__name__ in {"DynamicCache", "StaticCache", "HybridCache"}:
         from .cache_helper import CacheKeyValue
 
-        kc = CacheKeyValue(x)
-        return list(itertools.chain.from_iterable(zip(kc.key_cache, kc.value_cache)))
+        return CacheKeyValue(x).aslist()
 
     if x.__class__.__name__ == "EncoderDecoderCache":
-        res = flatten_object(x.self_attention_cache) + flatten_object(x.cross_attention_cache)
+        res = [
+            *flatten_object(x.self_attention_cache),
+            *flatten_object(x.cross_attention_cache),
+        ]
         return tuple(res)
     if x.__class__.__name__ == "MambaCache":
         if isinstance(x.conv_states, list):
-            res = flatten_object(x.conv_states) + flatten_object(x.ssm_states)
+            res = [*flatten_object(x.conv_states), *flatten_object(x.ssm_states)]
             return tuple(res)
         return (x.conv_states, x.ssm_states)
     if hasattr(x, "to_tuple"):

onnx_diagnostic/helpers/onnx_helper.py

@@ -28,6 +28,7 @@ from onnx import (
     NodeProto,
     OperatorSetIdProto,
     TensorProto,
+    TypeProto,
     ValueInfoProto,
     load as onnx_load,
 )
@@ -385,6 +386,12 @@ def pretty_onnx(
         shape_str = ",".join(map(str, shape))
         return f"{onnx_dtype_name(itype, exc=False)}[{shape_str}] {name}"
 
+    if isinstance(onx, TypeProto):
+        itype = onx.tensor_type.elem_type
+        shape = tuple((d.dim_param or d.dim_value) for d in onx.tensor_type.shape.dim)
+        shape_str = ",".join(map(str, shape))
+        return f"{onnx_dtype_name(itype, exc=False)}[{shape_str}]"
+
     if isinstance(onx, AttributeProto):
         att = onx
         if att.type == AttributeProto.INT:

onnx_diagnostic/helpers/ort_session.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import onnx
 import numpy as np
@@ -76,6 +77,10 @@ class _InferenceSession:
                     session_options.enable_profiling = enable_profiling
                 if optimized_model_filepath:
                     session_options.optimized_model_filepath = optimized_model_filepath
+                    session_options.add_session_config_entry(
+                        "session.optimized_model_external_initializers_file_name",
+                        f"{os.path.splitext(os.path.split(optimized_model_filepath)[-1])[0]}.data",
+                    )
                 if log_severity_level is not None:
                     session_options.log_severity_level = log_severity_level
                 if log_verbosity_level is not None:

onnx_diagnostic/helpers/rt_helper.py

@@ -41,7 +41,20 @@ def make_feeds(
     """
     # NOTE: position_ids is a special case because ModelBuilder does not usually use it,
     # because it's fued into rotary embedding in GQA.
-    if is_modelbuilder and isinstance(inputs, dict):
+    if is_modelbuilder and isinstance(inputs, dict) and "position_ids" in inputs:
+        position_ids = inputs["position_ids"]  # type: ignore[valid-type]
+        # We just check position_ids are contiguous.
+        assert isinstance(position_ids, torch.Tensor) and (
+            (
+                (position_ids - position_ids.min())
+                == torch.tensor(list(range(position_ids.shape[-1]))).unsqueeze(0)
+            )
+            .max()
+            .item()
+        ), (
+            f"ModelBuilder does not support position_ids={position_ids}, "
+            f"inputs={string_type(inputs, with_shape=True)}"
+        )
         inputs.pop("position_ids", None)  # Ensure 'position_ids' absent before removing.
 
     flat = flatten_object(inputs, drop_keys=True)

onnx_diagnostic/helpers/torch_helper.py

@@ -15,9 +15,6 @@ from .helper import string_type, size_type
 from .cache_helper import (
     make_dynamic_cache,
     make_encoder_decoder_cache,
-    make_hybrid_cache,
-    make_sliding_window_cache,
-    make_mamba_cache,
     make_static_cache,
     CacheKeyValue,
 )
@@ -769,10 +766,22 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
         return {to_any(t, to_value) for t in value}
     if type(value) is dict:
         return {k: to_any(t, to_value) for k, t in value.items()}
-    if value.__class__.__name__ in {"DynamicCache", "HybridCache"}:
-        make = dict(DynamicCache=make_dynamic_cache, HybridCache=make_hybrid_cache)
+    if value.__class__.__name__ == "DynamicCache":
         cc = CacheKeyValue(value)
-        return make[value.__class__.__name__](  # type: ignore[operator]
+        return make_dynamic_cache(
+            list(
+                zip(
+                    [t.to(to_value) if t is not None else t for t in cc.key_cache],
+                    [t.to(to_value) if t is not None else t for t in cc.value_cache],
+                )
+            ),
+            cls_layers=cc.cls_layers,
+        )
+    if value.__class__.__name__ == "HybridCache":
+        from .cache_helper import make_hybrid_cache
+
+        cc = CacheKeyValue(value)
+        return make_hybrid_cache(
             list(
                 zip(
                     [t.to(to_value) if t is not None else t for t in cc.key_cache],
@@ -843,7 +852,9 @@ def torch_deepcopy(value: Any) -> Any:
         from .cache_helper import CacheKeyValue
 
         ca = CacheKeyValue(value)
-        return make_dynamic_cache(torch_deepcopy(list(zip(ca.key_cache, ca.value_cache))))
+        return make_dynamic_cache(
+            torch_deepcopy(list(zip(ca.key_cache, ca.value_cache))), cls_layers=ca.cls_layers
+        )
     if value.__class__.__name__ == "StaticCache":
         from .cache_helper import CacheKeyValue
 
@@ -858,12 +869,12 @@ def torch_deepcopy(value: Any) -> Any:
             max_cache_len=max([value.max_cache_len, *[t.shape[2] for t in ca.key_cache]]),
         )
     if value.__class__.__name__ == "HybridCache":
-        from .cache_helper import CacheKeyValue
+        from .cache_helper import CacheKeyValue, make_hybrid_cache
 
         ca = CacheKeyValue(value)
         return make_hybrid_cache(torch_deepcopy(list(zip(ca.key_cache, ca.value_cache))))
     if value.__class__.__name__ == "SlidingWindowCache":
-        from .cache_helper import CacheKeyValue
+        from .cache_helper import CacheKeyValue, make_sliding_window_cache
 
         ca = CacheKeyValue(value)
         return make_sliding_window_cache(
@@ -875,6 +886,8 @@ def torch_deepcopy(value: Any) -> Any:
             torch_deepcopy(value.cross_attention_cache),
         )
     if value.__class__.__name__ == "MambaCache":
+        from .cache_helper import make_mamba_cache
+
         return make_mamba_cache(list(zip(value.conv_states, value.ssm_states)))
 
     if value.__class__ in torch.utils._pytree.SUPPORTED_NODES: