src/more_nnsight/_saved_activation_values.py

from __future__ import annotations from typing import Any class SavedActivationValues: """Provides dynamic attribute and index access over a sparse activation tree.""" def __init__(self, values: dict[Any, Any], path: str = "saved_activation") -> None: """Initializes one view into the saved activation hierarchy.""" self._values = values self._path = path def __getattr__(self, name: str) -> Any: """Exposes saved attribute-style paths while rejecting unsaved ones.""" if name.startswith("_"): raise AttributeError(name) try: value = self._values[name] except KeyError as exc: raise AttributeError(f"Unknown saved activation path: {self._path}.{name}") from exc return self._wrap(value, f"{self._path}.{name}") def __getitem__(self, index: int) -> Any: """Exposes saved index-style paths while rejecting unsaved ones.""" try: value = self._values[index] except KeyError as exc: raise KeyError(f"Unknown saved activation path: {self._path}[{index}]") from exc return self._wrap(value, f"{self._path}[{index}]") @staticmethod def _wrap(value: Any, path: str) -> Any: """Returns child nodes for nested paths and tensors for saved leaves.""" if isinstance(value, dict): return SavedActivationValues(value, path) return value