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

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,24 +28,52 @@ class CacheKeyValue:
             ]
             self.key_cache = [layer.keys for layer in layers]
             self.value_cache = [layer.values for layer in 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,
@@ -156,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:
@@ -192,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
@@ -214,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."
@@ -232,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`.
@@ -263,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):
@@ -508,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(
@@ -672,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/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
@@ -991,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/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:

onnx_diagnostic/tasks/image_text_to_text.py

@@ -1,7 +1,7 @@
 import itertools
 from typing import Any, Callable, Dict, Optional, Tuple
 import torch
-from ..helpers.cache_helper import make_dynamic_cache, make_hybrid_cache
+from ..helpers.cache_helper import make_dynamic_cache, get_make_hybrid_cache
 from ..helpers.config_helper import (
     update_config,
     check_hasattr,
@@ -200,6 +200,9 @@ def _get_inputs_gemma3(
 
     _check_()
 
+    make_hybrid_cache = get_make_hybrid_cache()
+    assert make_hybrid_cache is not None, "not implemented when make_hybrid_cache is missing"
+
     inputs = dict(
         input_ids=dummies["input_ids"],
         token_type_ids=dummies["token_type_ids"],

onnx_diagnostic/tasks/text_generation.py

@@ -1,11 +1,6 @@
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 import torch
-from ..helpers.cache_helper import (
-    make_dynamic_cache,
-    make_mamba_cache,
-    make_sliding_window_cache,
-    make_static_cache,
-)
+from ..helpers.cache_helper import make_dynamic_cache, make_mamba_cache, make_static_cache
 from ..helpers.config_helper import (
     update_config,
     check_hasattr,
@@ -187,17 +182,22 @@ def get_inputs(
             if cls_cache is None or isinstance(cls_cache, str)
             else cls_cache.__name__
         )
-        make_caches = {
-            "DynamicCache": make_dynamic_cache,
-            "SlidingWindowCache": make_sliding_window_cache,
-            "StaticCache": make_static_cache,
-        }
-        assert cache_name is None or cache_name in make_caches, (
-            f"Unable to handle cls_cache={cache_name!r}, it should be in "
-            f"{sorted(make_caches)}"
-        )
-        make_cache = make_dynamic_cache if cache_name is None else make_caches[cache_name]
-        is_static = cache_name == "StaticCache"
+        if cache_name == "DynamicSlidingWindowCache":
+            from ..helpers.cache_helper import make_sliding_window_cache
+
+            make_cache = make_sliding_window_cache
+            is_static = False
+        else:
+            make_caches = {
+                "DynamicCache": make_dynamic_cache,
+                "StaticCache": make_static_cache,
+            }
+            assert cache_name is None or cache_name in make_caches, (
+                f"Unable to handle cls_cache={cache_name!r}, it should be in "
+                f"{sorted(make_caches)}"
+            )
+            make_cache = make_dynamic_cache if cache_name is None else make_caches[cache_name]  # type: ignore[assignment]
+            is_static = cache_name == "StaticCache"
 
         if is_static:
             # static

onnx_diagnostic/torch_export_patches/eval/__init__.py

@@ -521,7 +521,7 @@ def run_exporter(
     :param exporter: exporter
     :param cls_model: model class to create
     :param inputs: list of inputs to try
-    :param dynamic: use dynamic shape or not
+    :param dynamic: use dynamic shapes or not
     :param quiet: raise exception or not
     :param verbose: verbosity
     :return: results

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -7,15 +7,9 @@ import transformers
 from transformers.cache_utils import DynamicCache, StaticCache
 
 try:
-    from transformers.cache_utils import (
-        EncoderDecoderCache,
-        HybridCache,
-        SlidingWindowCache,
-    )
+    from transformers.cache_utils import EncoderDecoderCache
 except ImportError:
     EncoderDecoderCache = None
-    HybridCache = None
-    SlidingWindowCache = None
 from ..helpers import string_type
 from .serialization import _lower_name_with_
 
@@ -36,6 +30,24 @@ def get_mamba_cache_cls() -> type:
             return None
 
 
+def get_hybrid_cache_cls() -> type:
+    try:
+        from transformers.cache_utils import HybridCache
+
+        return HybridCache
+    except ImportError:
+        return None
+
+
+def get_sliding_window_cache_cls() -> type:
+    try:
+        from transformers.cache_utils import SlidingWindowCache
+
+        return SlidingWindowCache
+    except ImportError:
+        return None
+
+
 def register_class_serialization(
     cls,
     f_flatten: Callable,
@@ -179,18 +191,9 @@ def serialization_functions(
             flatten_dynamic_cache,
             unflatten_dynamic_cache,
             flatten_with_keys_dynamic_cache,
-            flatten_hybrid_cache,
-            unflatten_hybrid_cache,
-            flatten_with_keys_hybrid_cache,
-            flatten_mamba_cache,
-            unflatten_mamba_cache,
-            flatten_with_keys_mamba_cache,
             flatten_encoder_decoder_cache,
             unflatten_encoder_decoder_cache,
             flatten_with_keys_encoder_decoder_cache,
-            flatten_sliding_window_cache,
-            unflatten_sliding_window_cache,
-            flatten_with_keys_sliding_window_cache,
             flatten_static_cache,
             unflatten_static_cache,
             flatten_with_keys_static_cache,
@@ -208,14 +211,6 @@ def serialization_functions(
                 # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
                 verbose=verbose,
             ),
-            HybridCache: lambda verbose=verbose: register_class_serialization(
-                HybridCache,
-                flatten_hybrid_cache,
-                unflatten_hybrid_cache,
-                flatten_with_keys_hybrid_cache,
-                # f_check=make_dynamic_cache([(torch.rand((4, 4, 4)), torch.rand((4, 4, 4)))]),
-                verbose=verbose,
-            ),
             EncoderDecoderCache: lambda verbose=verbose: register_class_serialization(
                 EncoderDecoderCache,
                 flatten_encoder_decoder_cache,
@@ -223,13 +218,6 @@ def serialization_functions(
                 flatten_with_keys_encoder_decoder_cache,
                 verbose=verbose,
             ),
-            SlidingWindowCache: lambda verbose=verbose: register_class_serialization(
-                SlidingWindowCache,
-                flatten_sliding_window_cache,
-                unflatten_sliding_window_cache,
-                flatten_with_keys_sliding_window_cache,
-                verbose=verbose,
-            ),
             StaticCache: lambda verbose=verbose: register_class_serialization(
                 StaticCache,
                 flatten_static_cache,
@@ -240,6 +228,12 @@ def serialization_functions(
         }
         MambaCache = get_mamba_cache_cls()
         if MambaCache:
+            from .serialization.transformers_impl import (
+                flatten_mamba_cache,
+                unflatten_mamba_cache,
+                flatten_with_keys_mamba_cache,
+            )
+
             transformers_classes[MambaCache] = (
                 lambda verbose=verbose: register_class_serialization(
                     MambaCache,
@@ -249,6 +243,42 @@ def serialization_functions(
                     verbose=verbose,
                 )
             )
+        HybridCache = get_hybrid_cache_cls()
+        if HybridCache:
+            from .serialization.transformers_impl import (
+                flatten_hybrid_cache,
+                unflatten_hybrid_cache,
+                flatten_with_keys_hybrid_cache,
+            )
+
+            transformers_classes[HybridCache] = (
+                lambda verbose=verbose: register_class_serialization(
+                    HybridCache,
+                    flatten_hybrid_cache,
+                    unflatten_hybrid_cache,
+                    flatten_with_keys_hybrid_cache,
+                    verbose=verbose,
+                )
+            )
+
+        SlidingWindowCache = get_sliding_window_cache_cls()
+        if SlidingWindowCache:
+            from .serialization.transformers_impl import (
+                flatten_sliding_window_cache,
+                unflatten_sliding_window_cache,
+                flatten_with_keys_sliding_window_cache,
+            )
+
+            transformers_classes[SlidingWindowCache] = (
+                lambda verbose=verbose: register_class_serialization(
+                    SlidingWindowCache,
+                    flatten_sliding_window_cache,
+                    unflatten_sliding_window_cache,
+                    flatten_with_keys_sliding_window_cache,
+                    verbose=verbose,
+                )
+            )
+
         classes.update(transformers_classes)
 
     if patch_diffusers:
@@ -303,13 +333,7 @@ def unregister_class_serialization(cls: type, verbose: int = 0):
 
 
 def unregister_cache_serialization(undo: Dict[str, bool], verbose: int = 0):
-    """Undo all registrations."""
-    MambaCache = get_mamba_cache_cls()
-    cls_ensemble = (
-        {DynamicCache, EncoderDecoderCache}
-        | set(undo)
-        | ({MambaCache} if MambaCache else set())
-    )
+    cls_ensemble = {DynamicCache, EncoderDecoderCache} | set(undo)
     for cls in cls_ensemble:
         if undo.get(cls.__name__, False):
             unregister_class_serialization(cls, verbose)

onnx_diagnostic/torch_export_patches/patches/patch_torch.py

@@ -541,14 +541,17 @@ class patched_ShapeEnv:
                     # oblivious_var_to_val will be defined iff we have sizes
                     # with DimDynamic.OBLIVIOUS_SIZE type.
                     # See https://github.com/pytorch/pytorch/issues/137100#issuecomment-2495778113
+                    var_to_val = getattr(
+                        self,
+                        "unbacked_var_to_val",
+                        getattr(self, "oblivious_var_to_val", False),
+                    )
                     if (
-                        self.oblivious_var_to_val
-                        and not (
-                            correct_hint := orig_expr.xreplace(self.oblivious_var_to_val)
-                        ).free_symbols
+                        var_to_val
+                        and not (correct_hint := orig_expr.xreplace(var_to_val)).free_symbols
                         and not (
                             counterfactual_hint := orig_expr.xreplace(
-                                {k: max(2, v) for k, v in self.oblivious_var_to_val.items()}
+                                {k: max(2, v) for k, v in var_to_val.items()}
                             )
                         ).free_symbols
                         and correct_hint == counterfactual_hint
@@ -571,11 +574,11 @@ class patched_ShapeEnv:
                     # and if they pass we add a runtime assertions and continue.
                     if (
                         not ok
-                        and self.unbacked_var_to_val
+                        and var_to_val
                         and not (
-                            unsound_result := orig_expr.xreplace(
-                                self.unbacked_var_to_val
-                            ).xreplace(self.var_to_val)
+                            unsound_result := orig_expr.xreplace(var_to_val).xreplace(
+                                var_to_val
+                            )
                         ).free_symbols
                     ):
                         # pyrefly: ignore  # unbound-name