API reference

The public surface of the engram package, generated from the source docstrings.

from engram import (
    EditorConfig, EngramEditor, edit_llm, get_engram, apply_engram,
    CovarianceCollector, Statistics,
    EngramResult, LayerScaleInfo,
    count_ratio, weight_norm, effective_rank, uniform, compose,
    LayerHandler, LinearHandler, Conv1DHandler,
)

EditorConfig dataclass

EditorConfig(
    storage_device: Optional[Union[str, device]] = None,
    absorb_bias: bool = True,
)

Configuration for engram extraction.

Attributes:

Name	Type	Description
storage_device	Optional[Union[str, device]]	Device for accumulating/holding covariance matrices. None (default) follows the model's device — fastest, with no per-batch GPU->CPU transfer. Set to "cpu" when the D x D covariances do not fit in VRAM (large/wide models): collection is slower (each batch's D x D is copied to CPU) but feasible.
absorb_bias	bool	When True (default, automatic), layers that have a bias are treated as the affine map y = Wx + b via homogeneous coordinates: the input is augmented with a constant 1 and the weight with the bias column, so the engram also corrects b. Bias-free layers (e.g. Llama/Mistral projections) are unaffected. Set False to edit W only (the original behavior).

EngramEditor

EngramEditor(
    model: Module,
    config: Optional[EditorConfig] = None,
    adapters: Optional[List[Any]] = None,
)

Covariance-based engram extractor for PyTorch / HuggingFace models.

The editor hooks every module whose type is in self.registry (by default nn.Linear and, when transformers is importable, HF Conv1D for the GPT-2 family). To restrict covariance to answer tokens, pass a mask_fn to collect_statistics (e.g. mask_fn=lambda b: b["labels"] != -100).

Example::

editor = EngramEditor(model, EditorConfig())
target = editor.collect_statistics(forget_loader, batch_fn=bf)   # Statistics
total = editor.collect_statistics(all_loader, batch_fn=bf)
edited = editor.edit(target, total, alpha=1.0)                   # paper edit (count_ratio)
# or, to reuse / re-scale the projection without recollecting:
#   engram = editor.compute_engram_weights(target, total)        # EngramResult
#   edited = editor.apply(engram, alpha=0.6, scale=weight_norm(1.0))

Source code in src/engram/editor.py

def __init__(
    self,
    model: nn.Module,
    config: Optional[EditorConfig] = None,
    adapters: Optional[List[Any]] = None,
) -> None:
    self.model = model
    self.config = config or EditorConfig()
    self.registry: Dict[type, Any] = {nn.Linear: LinearHandler()}
    conv1d = get_conv1d_class()
    if conv1d is not None:
        self.registry[conv1d] = Conv1DHandler()
    # opt-in extensions for layers the registry can't hook (e.g. fused MoE
    # experts via engram.moe.FusedExpertAdapter). Empty -> core is MoE-unaware.
    self.adapters = list(adapters or [])

collect_statistics

collect_statistics(
    dataloader: Iterable[Any],
    target_modules: Optional[Union[str, List[str]]] = None,
    batch_fn: Optional[Callable[[Any], Any]] = None,
    mask_fn: Optional[Callable[[Any], Tensor]] = None,
    layers_to_transform: Optional[
        Union[int, List[int]]
    ] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    target_layers: Optional[List[str]] = None,
) -> Statistics

Accumulate the mean input covariance mean(x^T x) + counts per supported layer.

Parameters:

Name	Type	Description	Default
dataloader	Iterable[Any]	iterable of batches.	required
target_modules	Optional[Union[str, List[str]]]	which modules to collect, LoRA/PEFT convention. A list matches by name suffix (["down_proj", "q_proj"] hits every layer's down_proj/q_proj); a string is a regex matched against the full module path (r".*layers\.5\..*down_proj"). None (default) collects every supported layer ("all-linear").	None
layers_to_transform	Optional[Union[int, List[int]]]	restrict to these decoder-layer indices (an int or list of ints), like PEFT. Combined with target_modules as an AND filter; None (default) applies no index restriction.	None
layers_pattern	Optional[Union[str, List[str]]]	container name(s) holding the layer index ("layers", "h", …). None scans common patterns.	None
target_layers	Optional[List[str]]	deprecated alias of target_modules (exact module names still match, as they did before).	None
batch_fn	Optional[Callable[[Any], Any]]	maps a batch to model inputs (a tensor, a tuple of positional args, or a dict of keyword args). Defaults to batch[0] (vision-style loaders). For HF models pass e.g. lambda b: {"input_ids": b["input_ids"], "attention_mask": b["attention_mask"]}.	None
mask_fn	Optional[Callable[[Any], Tensor]]	optional batch -> bool tensor selecting which tokens enter the covariance (one entry per flattened token, e.g. lambda b: b["labels"] != -100 for answer-token-only LLM editing). Applied to every layer type (nn.Linear, Conv1D, …).	None

Returns:

Name	Type	Description
A	Statistics	class:engram.stats.Statistics: per-layer mean covariance [D, D]
	Statistics	on config.storage_device plus the sample counts. D is the layer's
	Statistics	input dim, or input dim + 1 when bias absorption applies
	Statistics	(config.absorb_bias and the layer has a bias).

Source code in src/engram/editor.py

def collect_statistics(
    self,
    dataloader: Iterable[Any],
    target_modules: Optional[Union[str, List[str]]] = None,
    batch_fn: Optional[Callable[[Any], Any]] = None,
    mask_fn: Optional[Callable[[Any], torch.Tensor]] = None,
    layers_to_transform: Optional[Union[int, List[int]]] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    target_layers: Optional[List[str]] = None,
) -> Statistics:
    """Accumulate the mean input covariance ``mean(x^T x)`` + counts per supported layer.

    Args:
        dataloader: iterable of batches.
        target_modules: which modules to collect, **LoRA/PEFT convention**. A
            list matches by name suffix (``["down_proj", "q_proj"]`` hits every
            layer's ``down_proj``/``q_proj``); a string is a regex matched
            against the full module path (``r".*layers\\.5\\..*down_proj"``).
            ``None`` (default) collects every supported layer ("all-linear").
        layers_to_transform: restrict to these decoder-layer indices (an int or
            list of ints), like PEFT. Combined with ``target_modules`` as an AND
            filter; ``None`` (default) applies no index restriction.
        layers_pattern: container name(s) holding the layer index
            (``"layers"``, ``"h"``, …). ``None`` scans common patterns.
        target_layers: **deprecated** alias of ``target_modules`` (exact module
            names still match, as they did before).
        batch_fn: maps a batch to model inputs (a tensor, a tuple of
            positional args, or a dict of keyword args). Defaults to
            ``batch[0]`` (vision-style loaders). For HF models pass e.g.
            ``lambda b: {"input_ids": b["input_ids"], "attention_mask": b["attention_mask"]}``.
        mask_fn: optional ``batch -> bool tensor`` selecting which tokens enter
            the covariance (one entry per flattened token, e.g.
            ``lambda b: b["labels"] != -100`` for answer-token-only LLM editing).
            Applied to every layer type (``nn.Linear``, ``Conv1D``, …).

    Returns:
        A :class:`engram.stats.Statistics`: per-layer **mean** covariance ``[D, D]``
        on ``config.storage_device`` plus the sample counts. ``D`` is the layer's
        input dim, or ``input dim + 1`` when bias absorption applies
        (``config.absorb_bias`` and the layer has a bias).
    """
    if target_layers is not None:
        if target_modules is None:
            target_modules = target_layers
        warnings.warn(
            "target_layers is deprecated; use target_modules "
            "(LoRA-style: list = name suffix, str = regex).",
            DeprecationWarning,
            stacklevel=2,
        )
    collector = CovarianceCollector(
        self.model,
        self.config,
        self.registry,
        target_modules=target_modules,
        layers_to_transform=layers_to_transform,
        layers_pattern=layers_pattern,
        adapters=self.adapters,
    )
    self.model.eval()

    with collector, torch.inference_mode():
        for batch in tqdm(
            dataloader, disable=None, desc="Collecting covariance"
        ):
            # Unconditionally, not only when a mask_fn exists: the input-sharing window is
            # only ever valid inside one forward pass. A loader that refills a preallocated
            # tensor would otherwise be served the previous batch's product.
            collector.begin_batch()
            if mask_fn is not None:
                collector.set_mask(mask_fn(batch))
            raw_inputs = batch_fn(batch) if batch_fn else batch[0]
            self._forward(self._move_to_device(raw_inputs))

    return Statistics(collector.covariance_matrices, collector.sample_counts).dedupe()

merge_statistics staticmethod

merge_statistics(*stats: Statistics) -> Statistics

Count-weighted merge of several :class:Statistics (used to build totals).

Source code in src/engram/editor.py

@staticmethod
def merge_statistics(*stats: Statistics) -> Statistics:
    """Count-weighted merge of several :class:`Statistics` (used to build totals)."""
    return Statistics.merge(*stats)

compute_engram_weights

compute_engram_weights(
    target_covariances: Union[Statistics, List[Statistics]],
    total_covariance: Statistics,
    *,
    compute_erank: bool = False,
    rank_fraction: Optional[float] = None,
    inverse_method: str = "eigh",
    inverse_solver: str = "exact",
    rank_floor: str = "rtol",
    inverse_precision: Optional[dtype] = float64,
    condition_cap: Optional[float] = None,
    cut: str = "rtol",
    energy_fraction: float = 0.99,
    ridge_delta: float = 1e-06
) -> EngramResult

Compute the per-layer engram projection P = W . C_target . pinv(C_total).

C_* are the mean covariances from :meth:collect_statistics. P is the pure projection — the paper's n / N sample-count factor is not folded in here; it is applied at edit time by the scaling function (default :func:engram.scaling.count_ratio), so the scaling can be swapped without recomputing. Whether a layer was bias-absorbed is inferred from the covariance size (D == in + 1), kept consistent with collection without re-passing a flag.

Parameters:

Name	Type	Description	Default
target_covariances	Union[Statistics, List[Statistics]]	statistics of the data to isolate (the "forget"/target set). A single :class:Statistics, or a list merged first (count-weighted) via :meth:merge_statistics.	required
total_covariance	Statistics	statistics of the total/reference set.	required
compute_erank	bool	also compute each layer's target/total effective rank and store them in the result (needed only by :func:engram.scaling.effective_rank). Default False — skips the extra eigendecompositions.	False
rank_fraction	Optional[float]	None (default) keeps the existing singular-value cut (rtol = D * eps). A float f in (0, 1] switches to a rank-based cut: invert only the top ceil(f * D) eigen-directions of C_total — a conditioning-independent criterion that needs no ridge/damping (see :mod:engram.inverse).	None
inverse_solver	str	"exact" (default) runs a full eigh; "randomized" solves only the top ceil(rank_fraction * D) directions (O(D^2 k)) — worth it for aggressive truncation (f <= 0.1).	'exact'
cut	str	which directions to invert — "rtol" (default, the historical relative singular-value cut), "mp"/"mp_n" (random-matrix noise-bulk edge, see :mod:engram.rmt), "energy" (keep energy_fraction of the trace), or "ridge" (no truncation; damp by ridge_delta * lambda_max).	'rtol'
condition_cap	Optional[float]	impose one condition-number cap on every layer (rtol = 1 / condition_cap) instead of the width-dependent default.	None
inverse_precision	Optional[dtype]	dtype the eigendecomposition is solved in — torch.float64 by default. Upcasting a float32 covariance is lossless, and it is what makes the truncation deterministic: at float32 the eigenvalues nearest the cut carry ~1e-7 error, so a direction whose 1/lambda is ~1/rtol drifts in and out of the kept set and moves the projection by ~2%. At float64 the keep-set is identical across solvers (measured: 0 mismatches, 0.000000 relative difference) for ~14% more time. Pass None to solve in the covariance's own dtype.	float64
rank_floor	str	"rtol" (default) also drops kept directions below rtol * lam_max; "none" makes rank_fraction the only cut.	'rtol'
inverse_method	str	"eigh" (default) computes the pseudo-inverse from the symmetric eigendecomposition — same operator as the SVD-based torch.linalg.pinv for these symmetric PSD matrices, severalfold faster. "svd" keeps the previous torch.linalg.pinv code path (only valid with rank_fraction=None).	'eigh'

Returns:

Name	Type	Description
An	EngramResult	class:engram.scaling.EngramResult: layers[name] holds the projection
	EngramResult	plus the per-layer scale inputs (counts n/N, weight, total covariance);
	EngramResult	bias[name] holds the bias projection for bias-absorbed layers.

Source code in src/engram/editor.py

@torch.no_grad()
def compute_engram_weights(
    self,
    target_covariances: Union[Statistics, List[Statistics]],
    total_covariance: Statistics,
    *,
    compute_erank: bool = False,
    rank_fraction: Optional[float] = None,
    inverse_method: str = "eigh",
    inverse_solver: str = "exact",
    rank_floor: str = "rtol",
    inverse_precision: Optional[torch.dtype] = torch.float64,
    condition_cap: Optional[float] = None,
    cut: str = "rtol",
    energy_fraction: float = 0.99,
    ridge_delta: float = 1e-6,
) -> EngramResult:
    """Compute the per-layer engram projection ``P = W . C_target . pinv(C_total)``.

    ``C_*`` are the **mean** covariances from :meth:`collect_statistics`. ``P`` is the
    *pure* projection — the paper's ``n / N`` sample-count factor is **not** folded in
    here; it is applied at edit time by the scaling function (default
    :func:`engram.scaling.count_ratio`), so the scaling can be swapped without
    recomputing. Whether a layer was bias-absorbed is inferred from the covariance
    size (``D == in + 1``), kept consistent with collection without re-passing a flag.

    Args:
        target_covariances: statistics of the data to isolate (the "forget"/target
            set). A single :class:`Statistics`, or a list merged first
            (count-weighted) via :meth:`merge_statistics`.
        total_covariance: statistics of the total/reference set.
        compute_erank: also compute each layer's target/total effective rank and store
            them in the result (needed only by :func:`engram.scaling.effective_rank`).
            Default ``False`` — skips the extra eigendecompositions.
        rank_fraction: ``None`` (default) keeps the existing singular-value cut
            (``rtol = D * eps``). A float ``f`` in ``(0, 1]`` switches to a
            **rank-based** cut: invert only the top ``ceil(f * D)`` eigen-directions
            of ``C_total`` — a conditioning-independent criterion that needs no
            ridge/damping (see :mod:`engram.inverse`).
        inverse_solver: ``"exact"`` (default) runs a full ``eigh``; ``"randomized"``
            solves only the top ``ceil(rank_fraction * D)`` directions
            (``O(D^2 k)``) — worth it for aggressive truncation (``f <= 0.1``).
        cut: which directions to invert — ``"rtol"`` (default, the historical relative
            singular-value cut), ``"mp"``/``"mp_n"`` (random-matrix noise-bulk edge, see
            :mod:`engram.rmt`), ``"energy"`` (keep ``energy_fraction`` of the trace), or
            ``"ridge"`` (no truncation; damp by ``ridge_delta * lambda_max``).
        condition_cap: impose one condition-number cap on every layer
            (``rtol = 1 / condition_cap``) instead of the width-dependent default.
        inverse_precision: dtype the eigendecomposition is solved in — ``torch.float64``
            by default. Upcasting a float32 covariance is lossless, and it is what makes
            the truncation deterministic: at float32 the eigenvalues nearest the cut carry
            ~1e-7 error, so a direction whose ``1/lambda`` is ~1/rtol drifts in and out of
            the kept set and moves the projection by ~2%. At float64 the keep-set is
            identical across solvers (measured: 0 mismatches, 0.000000 relative difference)
            for ~14% more time. Pass ``None`` to solve in the covariance's own dtype.
        rank_floor: ``"rtol"`` (default) also drops kept directions below
            ``rtol * lam_max``; ``"none"`` makes ``rank_fraction`` the only cut.
        inverse_method: ``"eigh"`` (default) computes the pseudo-inverse from the
            symmetric eigendecomposition — same operator as the SVD-based
            ``torch.linalg.pinv`` for these symmetric PSD matrices, severalfold
            faster. ``"svd"`` keeps the previous ``torch.linalg.pinv`` code path
            (only valid with ``rank_fraction=None``).

    Returns:
        An :class:`engram.scaling.EngramResult`: ``layers[name]`` holds the projection
        plus the per-layer scale inputs (counts ``n``/``N``, weight, total covariance);
        ``bias[name]`` holds the bias projection for bias-absorbed layers.
    """
    if isinstance(target_covariances, list):
        target_covariances = self.merge_statistics(*target_covariances)

    missing = [k for k in target_covariances if k not in total_covariance]
    if missing:
        warnings.warn(
            f"compute_engram_weights: {len(missing)} target layer(s) are absent from "
            f"total_covariance and were skipped (e.g. {missing[:3]}) — is total a superset "
            f"of target?",
            stacklevel=2,
        )

    result = EngramResult()
    modules = dict(self.model.named_modules())
    # Layers sharing an input share their covariance object, so the eigendecomposition is
    # computed once per distinct matrix rather than once per layer.
    factor_cache: Dict[Any, Any] = {}
    dev = self._model_device  # float32 throughout — float64's pinv is catastrophic on
    prec = torch.float32      # ill-conditioned C_total (see CovarianceCollector).

    for layer_name, c_target in tqdm(
        target_covariances.items(), disable=None, desc="Computing engram"
    ):
        if layer_name not in total_covariance:
            continue
        c_target = c_target.to(dev, dtype=prec)
        c_total = total_covariance[layer_name].to(dev, dtype=prec)
        n = int(target_covariances.count.get(layer_name, 0))
        N = int(total_covariance.count.get(layer_name, 0))
        # pinv's rcond is load-bearing: float32's coarse singular-value cut discards the
        # near-null directions of the ill-conditioned C_total (float64 keeps and
        # 1/sigma-amplifies them -> catastrophic edit). Pin rtol to torch's own default
        # formula (max(M,N) * eps) so this regularization is explicit and independent of
        # any future change to the library's default tolerance. C_total is square (D x D).
        # Dtype-independent cut (engram.inverse.default_rtol): identical to the historical
        # D * eps_float32 value, but pinned as a constant so precision changes do not change
        # the operator being solved.
        rtol = None if condition_cap is not None else default_rtol(c_total.shape[-1])
        u_k = inv_lam = None
        if inverse_method not in ("eigh", "svd"):
            raise ValueError(f"inverse_method must be 'eigh' or 'svd', got {inverse_method!r}")
        if condition_cap is not None and condition_cap <= 1.0:
            raise ValueError(f"condition_cap must be > 1, got {condition_cap}")
        if inverse_method == "svd":
            if rank_fraction is not None:
                raise ValueError("rank_fraction requires inverse_method='eigh'")
            pinv_total = torch.linalg.pinv(
                c_total, rtol=(1.0 / condition_cap if condition_cap else rtol)
            )
        else:
            # factored form: pinv = U_k diag(inv_lam) U_k^T, applied right-to-left so the
            # D x D inverse is never materialized (identical result, D^2 less memory).
            # Siblings sharing a covariance are consecutive in iteration order, so a
            # single-entry cache captures the reuse without retaining one U_k per distinct
            # covariance — which would pin tens of GB on a large model.
            ckey = (id(total_covariance[layer_name]), N)   # N only matters for cut="mp_n"
            cached = factor_cache.get(ckey)
            if cached is None:
                factor_cache.clear()
                cached = spectral_factors(
                    c_total, rank_fraction=rank_fraction, rtol=rtol,
                    method=inverse_solver, floor=rank_floor,
                    compute_dtype=inverse_precision, condition_cap=condition_cap,
                    cut=cut, energy_fraction=energy_fraction, ridge_delta=ridge_delta,
                    n_samples=N or None,
                )
                factor_cache[ckey] = cached
            u_k, inv_lam = cached
            pinv_total = None

        def _project(w_mat: torch.Tensor) -> torch.Tensor:
            wc = w_mat @ c_target
            if pinv_total is not None:
                return wc @ pinv_total
            return ((wc @ u_k) * inv_lam) @ u_k.mT
        er_t = _erank(c_target) if compute_erank else None  # only effective_rank needs these
        er_a = _erank(c_total) if compute_erank else None

        module = modules.get(layer_name)
        if module is None:
            # not a hooked module — an opt-in adapter (e.g. fused MoE experts) may own it
            adapter = next((a for a in self.adapters if a.owns(layer_name, self.model)), None)
            if adapter is None:
                continue
            w = adapter.weight_for(layer_name, self.model)  # [out, in] (a Parameter slice)
            result.layers[layer_name] = LayerScaleInfo(
                name=layer_name, weight_fro=float(w.float().norm()),  # scalar; full weight not retained
                projection=_project(w.to(dev, dtype=prec)),
                n=n, N=N, target_erank=er_t, total_erank=er_a,
            )
            continue

        handler = handler_for(self.registry, module)
        if handler is None:
            continue

        in_dim = handler.get_input_dim(module, absorb_bias=False)
        absorbed = c_total.shape[0] == in_dim + 1 and getattr(module, "bias", None) is not None
        full = _project(handler.weight_matrix(module, absorb_bias=absorbed).to(dev, dtype=prec))

        if absorbed:
            proj = handler.to_weight_shape(full[:, :in_dim], module)
            result.bias[layer_name] = full[:, in_dim:].reshape(-1)
        else:
            proj = handler.to_weight_shape(full, module)
        result.layers[layer_name] = LayerScaleInfo(
            name=layer_name, weight_fro=float(module.weight.float().norm()),  # scalar; immune to later in-place edits
            projection=proj, n=n, N=N, target_erank=er_t, total_erank=er_a,
        )

    return result

apply

apply(
    engram: EngramResult,
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    inplace: bool = False
) -> Module

Subtract the engram from the model: W <- W - alpha * f_l * P_l.

Parameters:

Name	Type	Description	Default
engram	EngramResult	result of :meth:compute_engram_weights (projections + scale inputs).	required
alpha	float	global edit strength (the paper's strength; 1.0 = full removal at the default scaling).	1.0
scale	Optional[ScaleFn]	a scaling function {name: LayerScaleInfo} -> {name: float} giving the per-layer factor f_l. Defaults to :func:engram.scaling.count_ratio (f_l = n / N), which reproduces the paper. See :mod:engram.scaling for weight_norm / effective_rank / uniform / compose.	None
inplace	bool	edit self.model in place; otherwise edit and return a deep copy.	False

Returns:

Type	Description
Module	the edited model. Fused-expert keys are written to their 3D-Parameter slices
Module	via the registered adapter.

Source code in src/engram/editor.py

@torch.no_grad()
def apply(
    self,
    engram: EngramResult,
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    inplace: bool = False,
) -> nn.Module:
    """Subtract the engram from the model: ``W <- W - alpha * f_l * P_l``.

    Args:
        engram: result of :meth:`compute_engram_weights` (projections + scale inputs).
        alpha: global edit strength (the paper's strength; ``1.0`` = full removal at
            the default scaling).
        scale: a scaling function ``{name: LayerScaleInfo} -> {name: float}`` giving the
            per-layer factor ``f_l``. Defaults to :func:`engram.scaling.count_ratio`
            (``f_l = n / N``), which reproduces the paper. See :mod:`engram.scaling`
            for ``weight_norm`` / ``effective_rank`` / ``uniform`` / ``compose``.
        inplace: edit ``self.model`` in place; otherwise edit and return a deep copy.

    Returns:
        the edited model. Fused-expert keys are written to their 3D-Parameter slices
        via the registered adapter.
    """
    if scale is None:
        scale = count_ratio(1.0)
    model = self.model if inplace else copy.deepcopy(self.model)
    modules = dict(model.named_modules())
    factors = scale(engram.layers)

    for name, info in engram.layers.items():
        s = alpha * float(factors.get(name, 0.0))
        if s == 0.0:
            continue
        module = modules.get(name)
        if module is not None:
            module.weight.data -= (s * info.projection).to(module.weight.device, module.weight.dtype)
            b = engram.bias.get(name)
            if b is not None and getattr(module, "bias", None) is not None:
                module.bias.data -= (s * b).to(module.bias.device, module.bias.dtype)
            continue
        adapter = next((a for a in self.adapters if a.owns(name, model)), None)
        if adapter is not None:
            adapter.apply_delta(model, name, s * info.projection)
    return model

edit

edit(
    target_covariances: Union[Statistics, List[Statistics]],
    total_covariance: Statistics,
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    inplace: bool = False,
    compute_erank: bool = False,
    rank_fraction: Optional[float] = None,
    inverse_method: str = "eigh"
) -> Module

One call: :meth:compute_engram_weights then :meth:apply; returns the edited model.

Source code in src/engram/editor.py

def edit(
    self,
    target_covariances: Union[Statistics, List[Statistics]],
    total_covariance: Statistics,
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    inplace: bool = False,
    compute_erank: bool = False,
    rank_fraction: Optional[float] = None,
    inverse_method: str = "eigh",
) -> nn.Module:
    """One call: :meth:`compute_engram_weights` then :meth:`apply`; returns the edited model."""
    engram = self.compute_engram_weights(
        target_covariances, total_covariance, compute_erank=compute_erank,
        rank_fraction=rank_fraction, inverse_method=inverse_method,
    )
    return self.apply(engram, alpha=alpha, scale=scale, inplace=inplace)

save_statistics

save_statistics(
    stats: Statistics, path: Union[str, Path]
) -> None

Save collected :class:Statistics (mean covariances + counts) with torch.save.

Source code in src/engram/editor.py

def save_statistics(self, stats: Statistics, path: Union[str, Path]) -> None:
    """Save collected :class:`Statistics` (mean covariances + counts) with ``torch.save``."""
    stats.save(path)
    logger.info("Saved statistics to %s", path)

load_statistics

load_statistics(path: Union[str, Path]) -> Statistics

Load :class:Statistics onto the storage device (the model's device by default).

Source code in src/engram/editor.py

def load_statistics(self, path: Union[str, Path]) -> Statistics:
    """Load :class:`Statistics` onto the storage device (the model's device by default)."""
    return Statistics.load(path, map_location=self._storage_device)

edit_llm

edit_llm(
    model: Module,
    tokenizer: Any,
    forget: Iterable[Item],
    total: Iterable[Item],
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    target_modules: Optional[Union[str, List[str]]] = None,
    layers_to_transform: Optional[
        Union[int, List[int]]
    ] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    max_length: int = 512,
    batch_size: int = 8,
    inplace: bool = False,
    config: Optional[EditorConfig] = None,
    adapters: Optional[List[Any]] = None
) -> Module

Collect forget/total covariances over text and apply the engram, in one call.

Exactly :func:get_engram followed by :func:apply_engram. To sweep alpha without recollecting, call those two directly: engram = get_engram(...) once, then apply_engram(model, engram, alpha=...) per setting.

Parameters:

Name	Type	Description	Default
model	Module	a HuggingFace causal LM (or any module whose forward takes input_ids / attention_mask).	required
tokenizer	Any	a HF tokenizer (used as tokenizer(text)["input_ids"]; no chat template is applied).	required
forget	Iterable[Item]	the set to unlearn — an iterable of str or (prompt, answer).	required
total	Iterable[Item]	the reference set (typically forget + retain or a broad corpus), same item types as forget.	required
alpha	float	global edit strength.	1.0
scale	Optional[ScaleFn]	per-layer scaling function (default :func:engram.scaling.count_ratio, the paper's n/N); see :mod:engram.scaling.	None
target_modules	Optional[Union[str, List[str]]]	restrict which layers are edited (LoRA/PEFT convention; list = name suffix, str = regex), forwarded to :meth:EngramEditor.collect_statistics.	None
layers_to_transform	Optional[Union[int, List[int]]]	restrict to these decoder-layer indices (forwarded).	None
layers_pattern	Optional[Union[str, List[str]]]	the layer-index container name, e.g. "layers" (forwarded).	None
max_length	int	tokenizer truncation length.	512
batch_size	int	covariance-collection batch size.	8
inplace	bool	edit model in place; otherwise edit and return a deep copy.	False
config	Optional[EditorConfig]	an :class:EditorConfig (covariance storage device, bias absorption).	None
adapters	Optional[List[Any]]	optional fused-MoE adapters (e.g. [engram.moe.FusedExpertAdapter()]).	None

Returns:

Type	Description
Module	the edited model.

Source code in src/engram/llm.py

def edit_llm(
    model: nn.Module,
    tokenizer: Any,
    forget: Iterable[Item],
    total: Iterable[Item],
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    target_modules: Optional[Union[str, List[str]]] = None,
    layers_to_transform: Optional[Union[int, List[int]]] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    max_length: int = 512,
    batch_size: int = 8,
    inplace: bool = False,
    config: Optional[EditorConfig] = None,
    adapters: Optional[List[Any]] = None,
) -> nn.Module:
    """Collect forget/total covariances over text and apply the engram, in one call.

    Exactly :func:`get_engram` followed by :func:`apply_engram`. To sweep ``alpha`` without
    recollecting, call those two directly: ``engram = get_engram(...)`` once, then
    ``apply_engram(model, engram, alpha=...)`` per setting.

    Args:
        model: a HuggingFace causal LM (or any module whose ``forward`` takes
            ``input_ids`` / ``attention_mask``).
        tokenizer: a HF tokenizer (used as ``tokenizer(text)["input_ids"]``; no chat
            template is applied).
        forget: the set to unlearn — an iterable of ``str`` or ``(prompt, answer)``.
        total: the reference set (typically ``forget + retain`` or a broad corpus), same
            item types as ``forget``.
        alpha: global edit strength.
        scale: per-layer scaling function (default :func:`engram.scaling.count_ratio`,
            the paper's ``n/N``); see :mod:`engram.scaling`.
        target_modules: restrict which layers are edited (LoRA/PEFT convention; list =
            name suffix, str = regex), forwarded to :meth:`EngramEditor.collect_statistics`.
        layers_to_transform: restrict to these decoder-layer indices (forwarded).
        layers_pattern: the layer-index container name, e.g. ``"layers"`` (forwarded).
        max_length: tokenizer truncation length.
        batch_size: covariance-collection batch size.
        inplace: edit ``model`` in place; otherwise edit and return a deep copy.
        config: an :class:`EditorConfig` (covariance storage device, bias absorption).
        adapters: optional fused-MoE adapters (e.g. ``[engram.moe.FusedExpertAdapter()]``).

    Returns:
        the edited model.
    """
    engram = get_engram(
        model, tokenizer, forget, total,
        target_modules=target_modules, layers_to_transform=layers_to_transform,
        layers_pattern=layers_pattern, max_length=max_length, batch_size=batch_size,
        config=config, adapters=adapters,
    )
    return apply_engram(model, engram, alpha=alpha, scale=scale, inplace=inplace, adapters=adapters)

get_engram

get_engram(
    model: Module,
    tokenizer: Any,
    forget: Iterable[Item],
    total: Iterable[Item],
    *,
    target_modules: Optional[Union[str, List[str]]] = None,
    layers_to_transform: Optional[
        Union[int, List[int]]
    ] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    max_length: int = 512,
    batch_size: int = 8,
    compute_erank: bool = False,
    config: Optional[EditorConfig] = None,
    adapters: Optional[List[Any]] = None
) -> EngramResult

Tokenize forget/total, collect covariances, and compute the engram (alpha-free).

The expensive half of :func:edit_llm — two covariance passes plus one pseudo-inverse per layer — separated out so you can run it once and then call :func:apply_engram cheaply for as many alpha / scale settings as you like (the projection P = W . C_target . pinv(C_total) carries no alpha). Item types and selection knobs are the same as :func:edit_llm.

Parameters:

Name	Type	Description	Default
model	Module	a HuggingFace causal LM (forward takes input_ids / attention_mask).	required
tokenizer	Any	a HF tokenizer (no chat template is applied).	required
forget	Iterable[Item]	the set to unlearn — an iterable of str or (prompt, answer).	required
total	Iterable[Item]	the reference set (typically forget + retain), same item types.	required
target_modules	Optional[Union[str, List[str]]]	restrict which layers are collected (LoRA/PEFT convention).	None
layers_to_transform	Optional[Union[int, List[int]]]	restrict to these decoder-layer indices.	None
layers_pattern	Optional[Union[str, List[str]]]	the layer-index container name, e.g. "layers".	None
max_length	int	tokenizer truncation length.	512
batch_size	int	covariance-collection batch size.	8
compute_erank	bool	also compute per-layer effective ranks (needed only by :func:engram.scaling.effective_rank).	False
config	Optional[EditorConfig]	an :class:EditorConfig (covariance storage device, bias absorption).	None
adapters	Optional[List[Any]]	optional fused-MoE adapters (e.g. [engram.moe.FusedExpertAdapter()]).	None

Returns:

Name	Type	Description
an	EngramResult	class:engram.scaling.EngramResult — feed it to :func:apply_engram.

Source code in src/engram/llm.py

def get_engram(
    model: nn.Module,
    tokenizer: Any,
    forget: Iterable[Item],
    total: Iterable[Item],
    *,
    target_modules: Optional[Union[str, List[str]]] = None,
    layers_to_transform: Optional[Union[int, List[int]]] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    max_length: int = 512,
    batch_size: int = 8,
    compute_erank: bool = False,
    config: Optional[EditorConfig] = None,
    adapters: Optional[List[Any]] = None,
) -> EngramResult:
    """Tokenize forget/total, collect covariances, and compute the engram (``alpha``-free).

    The **expensive** half of :func:`edit_llm` — two covariance passes plus one
    pseudo-inverse per layer — separated out so you can run it **once** and then call
    :func:`apply_engram` cheaply for as many ``alpha`` / ``scale`` settings as you like
    (the projection ``P = W . C_target . pinv(C_total)`` carries no ``alpha``). Item types
    and selection knobs are the same as :func:`edit_llm`.

    Args:
        model: a HuggingFace causal LM (forward takes ``input_ids`` / ``attention_mask``).
        tokenizer: a HF tokenizer (no chat template is applied).
        forget: the set to unlearn — an iterable of ``str`` or ``(prompt, answer)``.
        total: the reference set (typically ``forget + retain``), same item types.
        target_modules: restrict which layers are collected (LoRA/PEFT convention).
        layers_to_transform: restrict to these decoder-layer indices.
        layers_pattern: the layer-index container name, e.g. ``"layers"``.
        max_length: tokenizer truncation length.
        batch_size: covariance-collection batch size.
        compute_erank: also compute per-layer effective ranks (needed only by
            :func:`engram.scaling.effective_rank`).
        config: an :class:`EditorConfig` (covariance storage device, bias absorption).
        adapters: optional fused-MoE adapters (e.g. ``[engram.moe.FusedExpertAdapter()]``).

    Returns:
        an :class:`engram.scaling.EngramResult` — feed it to :func:`apply_engram`.
    """
    pad_id = tokenizer.pad_token_id
    if pad_id is None:
        pad_id = tokenizer.eos_token_id
    if pad_id is None:
        raise ValueError("tokenizer has neither pad_token_id nor eos_token_id; set tokenizer.pad_token.")

    editor = EngramEditor(model, config or EditorConfig(), adapters=adapters)
    dev = editor._model_device
    sel = dict(target_modules=target_modules, layers_to_transform=layers_to_transform, layers_pattern=layers_pattern)
    feats = lambda b: {"input_ids": b["input_ids"].to(dev), "attention_mask": b["attention_mask"].to(dev)}
    mask_fn = lambda b: b["labels"] != _IGNORE  # answer tokens (tuples) or all real tokens (str)

    g_forget = editor.collect_statistics(
        _loader(tokenizer, forget, max_length, batch_size, pad_id), batch_fn=feats, mask_fn=mask_fn, **sel
    )
    g_total = editor.collect_statistics(
        _loader(tokenizer, total, max_length, batch_size, pad_id), batch_fn=feats, mask_fn=mask_fn, **sel
    )
    return editor.compute_engram_weights(g_forget, g_total, compute_erank=compute_erank)

apply_engram

apply_engram(
    model: Module,
    engram: EngramResult,
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    inplace: bool = False,
    adapters: Optional[List[Any]] = None
) -> Module

Apply a precomputed engram (from :func:get_engram): W <- W - alpha * f_l * P_l.

Cheap — a model copy plus one subtraction per layer — so call it repeatedly with different alpha / scale to tune the forget/retain trade-off without re-collecting covariances. alpha=0 is a no-op (returns the unedited model).

Parameters:

Name	Type	Description	Default
model	Module	the same model :func:get_engram saw (the projection was computed from its weights).	required
engram	EngramResult	the :class:engram.scaling.EngramResult from :func:get_engram.	required
alpha	float	global edit strength.	1.0
scale	Optional[ScaleFn]	per-layer scaling function (default :func:engram.scaling.count_ratio, the paper's n/N); see :mod:engram.scaling.	None
inplace	bool	edit model in place; otherwise edit and return a deep copy (default).	False
adapters	Optional[List[Any]]	fused-MoE adapters, if any were passed to :func:get_engram.	None

Returns:

Type	Description
Module	the edited model.

Source code in src/engram/llm.py

def apply_engram(
    model: nn.Module,
    engram: EngramResult,
    *,
    alpha: float = 1.0,
    scale: Optional[ScaleFn] = None,
    inplace: bool = False,
    adapters: Optional[List[Any]] = None,
) -> nn.Module:
    """Apply a precomputed engram (from :func:`get_engram`): ``W <- W - alpha * f_l * P_l``.

    Cheap — a model copy plus one subtraction per layer — so call it repeatedly with
    different ``alpha`` / ``scale`` to tune the forget/retain trade-off **without
    re-collecting** covariances. ``alpha=0`` is a no-op (returns the unedited model).

    Args:
        model: the **same** model :func:`get_engram` saw (the projection was computed
            from its weights).
        engram: the :class:`engram.scaling.EngramResult` from :func:`get_engram`.
        alpha: global edit strength.
        scale: per-layer scaling function (default :func:`engram.scaling.count_ratio`,
            the paper's ``n/N``); see :mod:`engram.scaling`.
        inplace: edit ``model`` in place; otherwise edit and return a deep copy (default).
        adapters: fused-MoE adapters, if any were passed to :func:`get_engram`.

    Returns:
        the edited model.
    """
    return EngramEditor(model, adapters=adapters).apply(engram, alpha=alpha, scale=scale, inplace=inplace)

LinearHandler

Bases: LayerHandler

Handler for nn.Linear (weight already stored as [out, in]).

Conv1DHandler

Bases: LayerHandler

Handler for HuggingFace Conv1D (GPT-2). Weight is stored [in, out].

The effective linear operator is weight.T (shape [out, in]), so the canonical matrix is weight.t() and the engram weight part is transposed back to [in, out] to match module.weight.

LayerHandler

Bases: ABC

Abstract per-layer-type handler.

absorb_bias is threaded through so a single handler covers both the plain (y = Wx) and bias-absorbed (y = [W|b][x;1]) cases.

get_input_dim abstractmethod

get_input_dim(
    module: Module, absorb_bias: bool = False
) -> int

Covariance dimension D (in or in+1 when absorbing).

Source code in src/engram/handlers.py

@abstractmethod
def get_input_dim(self, module: nn.Module, absorb_bias: bool = False) -> int:
    """Covariance dimension ``D`` (``in`` or ``in+1`` when absorbing)."""

reshape_input abstractmethod

reshape_input(
    module: Module, inputs: Any, absorb_bias: bool = False
) -> Tensor

Flatten the forward-pre-hook input to [N, D] (augmented if absorbing).

Source code in src/engram/handlers.py

@abstractmethod
def reshape_input(
    self, module: nn.Module, inputs: Any, absorb_bias: bool = False
) -> torch.Tensor:
    """Flatten the forward-pre-hook input to ``[N, D]`` (augmented if absorbing)."""

weight_matrix abstractmethod

weight_matrix(
    module: Module, absorb_bias: bool = False
) -> Tensor

Return the weight as canonical [out, D] ([W | b] if absorbing).

Source code in src/engram/handlers.py

@abstractmethod
def weight_matrix(self, module: nn.Module, absorb_bias: bool = False) -> torch.Tensor:
    """Return the weight as canonical ``[out, D]`` (``[W | b]`` if absorbing)."""

to_weight_shape abstractmethod

to_weight_shape(w: Tensor, module: Module) -> Tensor

Map a canonical [out, in] weight part back to module.weight's shape.

Source code in src/engram/handlers.py

@abstractmethod
def to_weight_shape(self, w: torch.Tensor, module: nn.Module) -> torch.Tensor:
    """Map a canonical ``[out, in]`` weight part back to ``module.weight``'s shape."""

CovarianceCollector

CovarianceCollector(
    model: Module,
    config: EditorConfig,
    registry: Dict[Type[Module], LayerHandler],
    target_modules: Optional[Union[str, List[str]]] = None,
    layers_to_transform: Optional[
        Union[int, List[int]]
    ] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    adapters: Optional[List[Any]] = None,
)

Context manager accumulating per-layer mean input covariance + sample count.

Source code in src/engram/collectors.py

def __init__(
    self,
    model: nn.Module,
    config: EditorConfig,
    registry: Dict[Type[nn.Module], LayerHandler],
    target_modules: Optional[Union[str, List[str]]] = None,
    layers_to_transform: Optional[Union[int, List[int]]] = None,
    layers_pattern: Optional[Union[str, List[str]]] = None,
    adapters: Optional[List[Any]] = None,
) -> None:
    self.model = model
    self.config = config
    self.registry = registry
    self.target_modules = target_modules
    self.layers_to_transform = layers_to_transform
    self.layers_pattern = layers_pattern
    self.adapters = list(adapters or [])  # opt-in extensions (e.g. fused MoE experts)
    self.covariance_matrices: Dict[str, torch.Tensor] = {}  # per-layer MEAN of x^T x
    self.sample_counts: Dict[str, int] = {}                 # per-layer rows that entered the mean
    self._recent: "OrderedDict[Any, Any]" = OrderedDict()   # MRU of shared inputs (see _remember)
    self._shared_hits = 0                                   # x^T x calls skipped by sharing
    self.current_mask: Optional[torch.Tensor] = None
    # per-batch routing-alignment state (used only when a layer sees a subset)
    self._proj: Dict[int, torch.Tensor] = {}                       # H -> random projection [H, _FP_DIM]
    self._ref: Dict[int, Tuple[torch.Tensor, torch.Tensor]] = {}   # H -> (fingerprint[N,_FP_DIM], mask[N])
    self._routed_mask: Optional[torch.Tensor] = None
    self._storage_device: Optional[torch.device] = None  # resolved in __enter__
    self._hook_handles: List[Any] = []

shared_hits property

shared_hits: int

How many x^T x products were skipped because a sibling had already computed one.

set_mask

set_mask(mask: Optional[Tensor]) -> None

Set the per-batch token mask and reset routing-alignment state.

Source code in src/engram/collectors.py

def set_mask(self, mask: Optional[torch.Tensor]) -> None:
    """Set the per-batch token mask and reset routing-alignment state."""
    self.current_mask = mask
    self.begin_batch()
    self._ref.clear()
    self._routed_mask = None

begin_batch

begin_batch() -> None

Drop the input-sharing window. Sharing is only ever valid inside one forward pass: a caller that reuses the same batch tensor object across iterations (or re-masks it) must not be served a stale product.

Source code in src/engram/collectors.py

def begin_batch(self) -> None:
    """Drop the input-sharing window. Sharing is only ever valid inside one forward pass: a
    caller that reuses the same batch tensor object across iterations (or re-masks it) must
    not be served a stale product."""
    self._recent.clear()

Statistics & scaling

collect_statistics returns a Statistics (mean covariances + sample counts); compute_engram_weights returns an EngramResult of per-layer projections. The per-layer edit weighting is a pluggable scaling function — see Guide → Scaling.

Statistics dataclass

Statistics(
    cov: Dict[str, Tensor] = dict(),
    count: Dict[str, int] = dict(),
)

Per-layer mean input covariance (cov) and sample count (count).

cov[name] is the mean of x^T x over the count[name] rows that entered layer name during collection. Keys are module names, or "<experts>.gate_up_proj.<e>" / "....down_proj.<e>" for fused-MoE experts. Behaves like a read-only mapping over cov (stats[name], name in stats, iteration) with the parallel count dict alongside.

to

to(device: Union[str, device]) -> 'Statistics'

Move every covariance to device (counts are plain ints, copied as-is).

Layers that share one covariance (q/k/v, gate/up) keep sharing it after the move — moving each member separately would silently triple the memory the collector saved.

Source code in src/engram/stats.py

def to(self, device: Union[str, torch.device]) -> "Statistics":
    """Move every covariance to ``device`` (counts are plain ints, copied as-is).

    Layers that share one covariance (q/k/v, gate/up) keep sharing it after the move —
    moving each member separately would silently triple the memory the collector saved.
    """
    moved: Dict[int, torch.Tensor] = {}
    cov: Dict[str, torch.Tensor] = {}
    for k, v in self.cov.items():
        oid = id(v)
        if oid not in moved:
            moved[oid] = v.to(device)
        cov[k] = moved[oid]
    return Statistics(cov, dict(self.count))

merge staticmethod

merge(*stats: 'Statistics') -> 'Statistics'

Count-weighted merge: C = sum(n_i C_i) / sum(n_i), N = sum(n_i).

Equivalent to having collected over the concatenated token streams. Keys are unioned; a key present in only some inputs contributes only its own (count, mean). Combined incrementally (C += n_i/(N+n_i) (C_i - C)) so no large n_i * C_i intermediate is formed.

Note that a merge materializes one tensor per key: layers that shared a covariance during collection no longer do afterwards. That costs memory on a merged result but keeps the arithmetic obviously correct; to() and save() do preserve sharing.

Source code in src/engram/stats.py

@staticmethod
def merge(*stats: "Statistics") -> "Statistics":
    """Count-weighted merge: ``C = sum(n_i C_i) / sum(n_i)``, ``N = sum(n_i)``.

    Equivalent to having collected over the concatenated token streams. Keys are
    unioned; a key present in only some inputs contributes only its own
    ``(count, mean)``. Combined incrementally (``C += n_i/(N+n_i) (C_i - C)``) so
    no large ``n_i * C_i`` intermediate is formed.

    Note that a merge materializes one tensor per key: layers that shared a covariance
    during collection no longer do afterwards. That costs memory on a merged result but
    keeps the arithmetic obviously correct; ``to()`` and ``save()`` do preserve sharing.
    """
    cov: Dict[str, torch.Tensor] = {}
    count: Dict[str, int] = {}
    for s in stats:
        for k, c in s.cov.items():
            n = int(s.count.get(k, 0))
            if k not in cov:
                cov[k] = c.to(torch.float32).clone()
                count[k] = n
                continue
            prev = count[k]
            total = prev + n
            if total > 0:
                c = c.to(cov[k].device, torch.float32)
                cov[k] = cov[k] + (n / total) * (c - cov[k])
            count[k] = total
    return Statistics(cov, count)

dedupe

dedupe() -> 'Statistics'

Collapse bit-identical covariances onto one tensor each.

Layers fed by the same input — q/k/v off one LayerNorm, gate/up off the other — accumulate covariances that are equal to the last bit. Sharing one tensor between them costs nothing (they are read-only from here on) and saves that fraction of memory, file size and eigendecompositions.

This runs after collection rather than during it, so the accumulation itself stays exactly what it was: every layer folds every batch it saw, with its own count. Only tensors that are already identical, and whose counts agree, are merged.

The merged covariances are the same object, so writing into one in place writes into all of them. Treat them as read-only — or call :meth:merge on the result, which materializes one tensor per key.

Source code in src/engram/stats.py

def dedupe(self) -> "Statistics":
    """Collapse bit-identical covariances onto one tensor each.

    Layers fed by the same input — ``q``/``k``/``v`` off one LayerNorm, ``gate``/``up`` off
    the other — accumulate covariances that are equal to the last bit. Sharing one tensor
    between them costs nothing (they are read-only from here on) and saves that fraction of
    memory, file size and eigendecompositions.

    This runs *after* collection rather than during it, so the accumulation itself stays
    exactly what it was: every layer folds every batch it saw, with its own count. Only
    tensors that are already identical, and whose counts agree, are merged.

    The merged covariances are the *same object*, so writing into one in place writes into
    all of them. Treat them as read-only — or call :meth:`merge` on the result, which
    materializes one tensor per key.
    """
    by_shape: Dict[Any, List[str]] = {}
    for name, c in self.cov.items():
        by_shape.setdefault((tuple(c.shape), c.dtype, c.device, self.count.get(name)), []).append(name)
    cov = dict(self.cov)
    for names in by_shape.values():
        if len(names) < 2:
            continue
        canon: List[str] = []
        for name in names:
            for c0 in canon:
                if cov[name] is cov[c0] or torch.equal(cov[name], cov[c0]):
                    cov[name] = cov[c0]
                    break
            else:
                canon.append(name)
    return Statistics(cov, dict(self.count))

save

save(
    path: Union[str, Path], *, packed: bool = True
) -> None

Save with torch.save.

packed=True (default, tag format=3) stores each covariance as its upper triangle only — the matrices are symmetric, so this halves the file with the upper triangle preserved bit-exactly. packed=False writes the dense format=2 layout for compatibility with older readers.

Source code in src/engram/stats.py

def save(self, path: Union[str, Path], *, packed: bool = True) -> None:
    """Save with ``torch.save``.

    ``packed=True`` (default, tag ``format=3``) stores each covariance as its
    upper triangle only — the matrices are symmetric, so this halves the file
    with the upper triangle preserved bit-exactly. ``packed=False`` writes the
    dense ``format=2`` layout for compatibility with older readers.
    """
    if packed:
        # Layers fed by the same tensor (q/k/v, gate/up) hold the SAME accumulator object,
        # so the file stores it once and records who shares it. One packed copy is live at
        # a time.
        packed_cov, alias, owner_of = {}, {}, {}
        for k, v in self.cov.items():
            oid = id(v)
            if oid in owner_of:
                alias[k] = owner_of[oid]
                continue
            owner_of[oid] = k
            packed_cov[k] = _pack_symmetric(v)
        # A reader that does not know about aliases would silently lose the aliased layers,
        # so a file that has them gets its own tag and fails loudly on 0.9.x instead.
        obj = {"format": _FORMAT_ALIASED if alias else _FORMAT,
               "cov_packed": packed_cov, "count": self.count}
        if alias:
            obj["alias"] = alias
        torch.save(obj, path)
    else:
        torch.save({"format": _FORMAT_DENSE, "cov": self.cov, "count": self.count}, path)

load staticmethod

load(
    path: Union[str, Path],
    map_location: Union[str, device, None] = None,
) -> "Statistics"

Load a :class:Statistics. Rejects the legacy raw-covariance dict format.

Source code in src/engram/stats.py

@staticmethod
def load(
    path: Union[str, Path], map_location: Union[str, torch.device, None] = None
) -> "Statistics":
    """Load a :class:`Statistics`. Rejects the legacy raw-covariance dict format."""
    obj = torch.load(path, map_location=map_location, weights_only=True)
    if isinstance(obj, dict) and obj.get("format") in (_FORMAT, _FORMAT_ALIASED) and "cov_packed" in obj:
        # unpack on the host, then honour map_location
        cov = {k: _unpack_symmetric(v) for k, v in obj["cov_packed"].items()}
        if map_location is not None:
            cov = {k: v.to(map_location) for k, v in cov.items()}
        for member, owner in obj.get("alias", {}).items():
            cov[member] = cov[owner]      # restore the sharing, not a copy
        return Statistics(cov, obj["count"])
    if isinstance(obj, dict) and obj.get("format") == _FORMAT_DENSE and "cov" in obj:
        return Statistics(obj["cov"], obj["count"])
    raise ValueError(
        f"{path!r} is not a Statistics file (format={_FORMAT_ALIASED}/{_FORMAT} packed or "
        f"{_FORMAT_DENSE} dense). It looks like a legacy raw-covariance dict; "
        "re-collect with EngramEditor.collect_statistics() — the mean+count "
        "format cannot be reconstructed from summed covariances."
    )

EngramResult dataclass

EngramResult(
    layers: Dict[str, LayerScaleInfo] = dict(),
    bias: Dict[str, Tensor] = dict(),
)

Output of :meth:EngramEditor.compute_engram_weights: projections + scale inputs.

layers[name].projection is the pure engram projection (no n / N); bias[name] is the matching bias projection for bias-absorbed layers. Pass the whole result to :meth:EngramEditor.apply.

LayerScaleInfo dataclass

LayerScaleInfo(
    name: str,
    weight_fro: float,
    projection: Tensor,
    n: int,
    N: int,
    target_erank: Optional[float] = None,
    total_erank: Optional[float] = None,
)

Per-layer inputs handed to a scaling function.

count_ratio

count_ratio(power: float = 1.0) -> ScaleFn

f_l = (n_l / N_l) ** power. Default (power=1) reproduces the paper exactly.

n_l / N_l is the target/total sample-count ratio implicit in the paper's summed covariances. n == 0 or N == 0 -> 0 (a layer no target token reached is not edited).

Source code in src/engram/scaling.py

def count_ratio(power: float = 1.0) -> ScaleFn:
    """``f_l = (n_l / N_l) ** power``. Default (``power=1``) reproduces the paper exactly.

    ``n_l / N_l`` is the target/total sample-count ratio implicit in the paper's summed
    covariances. ``n == 0`` or ``N == 0`` -> 0 (a layer no target token reached is not edited).
    """

    def fn(infos: Dict[str, LayerScaleInfo]) -> Dict[str, float]:
        return {
            k: (i.n / i.N) ** power if (i.n > 0 and i.N > 0) else 0.0
            for k, i in infos.items()
        }

    return fn

weight_norm

weight_norm(power: float = 1.0) -> ScaleFn

f_l = (rel_l / max rel) ** power with rel_l = ||P_l|| / ||W_l||.

Edits each layer in proportion to how strongly the engram occupies it (the relative Frobenius norm). This is the per-layer "adaptive" weighting; compose with :func:count_ratio to also keep the paper's n / N factor.

Source code in src/engram/scaling.py

def weight_norm(power: float = 1.0) -> ScaleFn:
    """``f_l = (rel_l / max rel) ** power`` with ``rel_l = ||P_l|| / ||W_l||``.

    Edits each layer in proportion to how strongly the engram occupies it (the relative
    Frobenius norm). This is the per-layer "adaptive" weighting; compose with
    :func:`count_ratio` to also keep the paper's ``n / N`` factor.
    """

    def fn(infos: Dict[str, LayerScaleInfo]) -> Dict[str, float]:
        rel = {
            k: i.projection.norm().item() / (i.weight_fro + _EPS)
            for k, i in infos.items()
        }
        mx = max(rel.values(), default=1.0) or 1.0
        return {k: (r / mx) ** power for k, r in rel.items()}

    return fn

effective_rank

effective_rank(power: float = 1.0) -> ScaleFn

f_l = (er(C_target_l) / er(C_total_l)) ** power per layer.

The ratio of the effective rank (see :func:_erank) of the target covariance to that of the total — how dimensionally rich the forget inputs are relative to the whole at each layer. Needs the per-layer effective ranks, i.e. compute_engram_weights(..., compute_erank=True).

Source code in src/engram/scaling.py

def effective_rank(power: float = 1.0) -> ScaleFn:
    """``f_l = (er(C_target_l) / er(C_total_l)) ** power`` per layer.

    The ratio of the **effective rank** (see :func:`_erank`) of the target covariance to
    that of the total — how dimensionally rich the forget inputs are relative to the whole
    at each layer. Needs the per-layer effective ranks, i.e.
    ``compute_engram_weights(..., compute_erank=True)``.
    """

    def fn(infos: Dict[str, LayerScaleInfo]) -> Dict[str, float]:
        out: Dict[str, float] = {}
        for k, i in infos.items():
            if i.target_erank is None or i.total_erank is None:
                raise ValueError(
                    "effective_rank needs per-layer effective ranks; recompute with "
                    "compute_engram_weights(..., compute_erank=True)."
                )
            ratio = (i.target_erank / i.total_erank) if i.total_erank > 0 else 0.0
            out[k] = ratio ** power
        return out

    return fn

uniform

uniform() -> ScaleFn

f_l = 1 for every layer (subtract the bare projection, ignoring counts/norms).

Source code in src/engram/scaling.py

def uniform() -> ScaleFn:
    """``f_l = 1`` for every layer (subtract the bare projection, ignoring counts/norms)."""

    def fn(infos: Dict[str, LayerScaleInfo]) -> Dict[str, float]:
        return {k: 1.0 for k in infos}

    return fn

compose

compose(*fns: ScaleFn) -> ScaleFn

Multiply scaling functions per layer: f_l = prod_i fns[i](infos)[l].

e.g. compose(count_ratio(1.0), weight_norm(1.0)) is the paper's n / N weighting further modulated by relative weight-norm.

Source code in src/engram/scaling.py

def compose(*fns: ScaleFn) -> ScaleFn:
    """Multiply scaling functions per layer: ``f_l = prod_i fns[i](infos)[l]``.

    e.g. ``compose(count_ratio(1.0), weight_norm(1.0))`` is the paper's ``n / N``
    weighting further modulated by relative weight-norm.
    """

    def fn(infos: Dict[str, LayerScaleInfo]) -> Dict[str, float]:
        out: Dict[str, float] = {k: 1.0 for k in infos}
        for f in fns:
            d = f(infos)
            for k in out:
                out[k] *= d.get(k, 0.0)
        return out

    return fn

MoE (optional)

Support for transformers ≥5 fused-expert MoE layers lives in a separate, detachable module — import it explicitly; the core never depends on it:

from engram import EngramEditor
from engram.moe import FusedExpertAdapter, apply_engram_weights

editor = EngramEditor(model, adapters=[FusedExpertAdapter()])

See Guide → Mixture-of-experts.

FusedExpertAdapter

FusedExpertAdapter()

Collects per-expert input covariance for standard fused MoE experts (opt-in).

Plugs into the core via three calls: attach/detach (the collector adds its hooks during a covariance pass) and owns/weight_for (the editor asks it to resolve covariance keys it doesn't recognize as modules).

Source code in src/engram/moe.py

def __init__(self) -> None:
    self._handles: List[Any] = []
    self._modules: Dict[str, nn.Module] = {}     # experts-module name -> module
    self._warned: set = set()

owns

owns(key: str, model: Module) -> bool

True if key names a fused-expert slice present on model.

Stateless — resolved from model, so it works whether or not a covariance pass was run this session (e.g. applying engrams loaded from disk).

Source code in src/engram/moe.py

def owns(self, key: str, model: nn.Module) -> bool:
    """True if ``key`` names a fused-expert slice present on ``model``.

    Stateless — resolved from ``model``, so it works whether or not a covariance
    pass was run this session (e.g. applying engrams loaded from disk).
    """
    split = _split_key(key)
    if split is None:
        return False
    name, proj, e = split
    mod = dict(model.named_modules()).get(name)
    return mod is not None and _is_standard_fused(mod) and 0 <= e < getattr(mod, proj).shape[0]

apply_delta

apply_delta(model: Module, key: str, delta: Tensor) -> None

Subtract delta from the expert's 3D-Parameter slice, in place.

Source code in src/engram/moe.py

def apply_delta(self, model: nn.Module, key: str, delta: torch.Tensor) -> None:
    """Subtract ``delta`` from the expert's 3D-Parameter slice, in place."""
    name, proj, e = _split_key(key)  # type: ignore[misc]
    param = getattr(dict(model.named_modules())[name], proj)
    param.data[e] -= delta.to(param.dtype)

apply_engram_weights

apply_engram_weights(
    model: Module,
    weight_engrams: Dict[str, Tensor],
    alpha: float = 1.0,
) -> None

Low-level: subtract alpha * weight_engrams[key] from each weight, handling fused keys.

weight_engrams is a plain {key: delta} of final per-key tensors (already scaled by any per-layer factor). Real-module keys edit module.weight; fused-expert keys ("<experts>.gate_up_proj.<e>") edit param.data[e] of the 3D Parameter. The primary path is :meth:EngramEditor.apply, which dispatches fused keys here via the adapter and applies the pluggable scaling for you; use this only when you hold raw deltas.

Source code in src/engram/moe.py

def apply_engram_weights(model: nn.Module, weight_engrams: Dict[str, torch.Tensor], alpha: float = 1.0) -> None:
    """Low-level: subtract ``alpha * weight_engrams[key]`` from each weight, handling fused keys.

    ``weight_engrams`` is a plain ``{key: delta}`` of **final** per-key tensors (already
    scaled by any per-layer factor). Real-module keys edit ``module.weight``; fused-expert
    keys (``"<experts>.gate_up_proj.<e>"``) edit ``param.data[e]`` of the 3D Parameter. The
    primary path is :meth:`EngramEditor.apply`, which dispatches fused keys here via the
    adapter and applies the pluggable scaling for you; use this only when you hold raw deltas.
    """
    modules = dict(model.named_modules())
    with torch.no_grad():
        for key, w in weight_engrams.items():
            if key in modules:
                modules[key].weight.data -= (alpha * w).to(modules[key].weight.dtype)
                continue
            split = _split_key(key)
            if split is None:
                continue
            name, proj, e = split
            param = getattr(modules[name], proj)
            param.data[e] -= (alpha * w).to(param.dtype)

Benchmarks

run

run(
    model,
    tok,
    *,
    split: str = "forget10",
    alpha: Optional[float] = None,
    scale: str = "adaptive",
    reference: Union[str, Statistics, None] = "tofu",
    level: str = "quick",
    device: str = "cuda",
    splits: Optional[Dict[str, Any]] = None,
    cache: Optional[Dict[str, str]] = None,
    n_total: int = 4000
) -> Dict[str, Any]

Collect, edit and score in one call. Returns {"before", "after", "alpha", "scale"}.

reference: "tofu" uses the benchmark's own 4000-sample reference set; a :class:~engram.Statistics uses that instead (e.g. a self-generated one from :func:engram.generate_corpus).

Source code in src/engram/benchmarks/tofu.py

def run(
    model,
    tok,
    *,
    split: str = "forget10",
    alpha: Optional[float] = None,
    scale: str = "adaptive",
    reference: Union[str, Statistics, None] = "tofu",
    level: str = "quick",
    device: str = "cuda",
    splits: Optional[Dict[str, Any]] = None,
    cache: Optional[Dict[str, str]] = None,
    n_total: int = 4000,
) -> Dict[str, Any]:
    """Collect, edit and score in one call. Returns ``{"before", "after", "alpha", "scale"}``.

    ``reference``: ``"tofu"`` uses the benchmark's own 4000-sample reference set; a
    :class:`~engram.Statistics` uses that instead (e.g. a self-generated one from
    :func:`engram.generate_corpus`).
    """
    if scale not in SCALES:
        raise ValueError(f"scale must be one of {sorted(SCALES)}, got {scale!r}")
    alpha = DEFAULT_ALPHA[scale] if alpha is None else alpha
    splits = splits if splits is not None else load_splits(split, n_total=n_total)
    target, ref, editor = collect(model, tok, splits, device=device, cache=cache)
    if isinstance(reference, Statistics):
        ref = reference
    elif reference not in ("tofu", None):
        raise ValueError(f"reference must be 'tofu' or a Statistics, got {reference!r}")
    before = evaluate(model, tok, splits, level=level, device=device)
    engram = editor.compute_engram_weights(target, ref)
    edited = editor.apply(engram, alpha=alpha, scale=SCALES[scale]()).eval()
    after = evaluate(edited, tok, splits, level=level, device=device, baseline=before)
    return {"before": before, "after": after, "alpha": alpha, "scale": scale,
            "model": edited, "engram": engram}

search

search(
    model,
    tok,
    *,
    split: str = "forget10",
    alphas: Sequence[float] = (
        0.3,
        0.6,
        0.9,
        1.2,
        1.5,
        2.0,
    ),
    scales: Sequence[str] = ("plain", "adaptive"),
    objective: Union[
        str, Callable[[Report], float]
    ] = "overall",
    utility_floor: float = 0.9,
    coarse_level: str = "quick",
    final_level: str = "full",
    top_k: int = 3,
    reference: Union[str, Statistics, None] = "tofu",
    device: str = "cuda",
    splits: Optional[Dict[str, Any]] = None,
    cache: Optional[Dict[str, str]] = None,
    verbose: bool = True
) -> Dict[str, Any]

Find the best (alpha, scale), cheaply.

The engram is computed once: alpha only scales the subtraction, so a sweep costs evaluations, not extractions. That is why this is coarse-to-fine — every candidate is scored at coarse_level (seconds), and only the top_k are re-scored at final_level.

The default objective is the paper's Overall. It contains Utility, so maximizing it cannot be won by editing harder — unlike a raw forget-minus-retain score, which rises right up to the point where the model collapses. That raw score is available as objective="proxy" for quick-level-only sweeps; treat it as a ranking, not a verdict. Objective/level mismatches are rejected before any model is scored.

Source code in src/engram/benchmarks/tofu.py

def search(
    model,
    tok,
    *,
    split: str = "forget10",
    alphas: Sequence[float] = (0.3, 0.6, 0.9, 1.2, 1.5, 2.0),
    scales: Sequence[str] = ("plain", "adaptive"),
    objective: Union[str, Callable[[Report], float]] = "overall",
    utility_floor: float = 0.9,
    coarse_level: str = "quick",
    final_level: str = "full",
    top_k: int = 3,
    reference: Union[str, Statistics, None] = "tofu",
    device: str = "cuda",
    splits: Optional[Dict[str, Any]] = None,
    cache: Optional[Dict[str, str]] = None,
    verbose: bool = True,
) -> Dict[str, Any]:
    """Find the best ``(alpha, scale)``, cheaply.

    The engram is computed **once**: ``alpha`` only scales the subtraction, so a sweep costs
    evaluations, not extractions. That is why this is coarse-to-fine — every candidate is scored
    at ``coarse_level`` (seconds), and only the ``top_k`` are re-scored at ``final_level``.

    The default objective is the paper's ``Overall``. It contains Utility, so maximizing it
    cannot be won by editing harder — unlike a raw forget-minus-retain score, which rises right
    up to the point where the model collapses. That raw score is available as
    ``objective="proxy"`` for quick-level-only sweeps; treat it as a ranking, not a verdict.
    Objective/level mismatches are rejected before any model is scored.
    """
    _check_objective(objective, final_level)          # fail here, not after an hour of scoring
    for lvl in (coarse_level, final_level):
        if lvl not in ("quick", "utility", "full"):
            raise ValueError(f"level must be 'quick', 'utility' or 'full', got {lvl!r}")
    for sc in scales:
        if sc not in SCALES:
            raise ValueError(f"scale must be one of {sorted(SCALES)}, got {sc!r}")

    splits = splits if splits is not None else load_splits(split)
    target, ref, editor = collect(model, tok, splits, device=device, cache=cache)
    if isinstance(reference, Statistics):
        ref = reference
    base_coarse = evaluate(model, tok, splits, level=coarse_level, device=device)
    engram = editor.compute_engram_weights(target, ref)

    coarse: List[Dict[str, Any]] = []
    for sc in scales:
        for a in alphas:
            edited = editor.apply(engram, alpha=a, scale=SCALES[sc]()).eval()
            rep = evaluate(edited, tok, splits, level=coarse_level, device=device,
                           baseline=base_coarse)
            coarse.append({"alpha": a, "scale": sc, "report": rep})
            if verbose:
                print(f"[tofu.search] alpha={a:<5} scale={sc:<9} {rep}", flush=True)
            del edited
            torch.cuda.empty_cache()

    # rank the coarse pass by selective forgetting, then confirm the survivors properly
    coarse.sort(key=lambda r: (r["report"].forget_delta or 0) - (r["report"].retain_delta or 0),
                reverse=True)
    finalists = coarse[:max(1, top_k)]
    if final_level == coarse_level:
        scored = finalists
        base_final = base_coarse
    else:
        base_final = evaluate(model, tok, splits, level=final_level, device=device)
        scored = []
        for cand in finalists:
            edited = editor.apply(engram, alpha=cand["alpha"],
                                  scale=SCALES[cand["scale"]]()).eval()
            rep = evaluate(edited, tok, splits, level=final_level, device=device,
                           baseline=base_final)
            scored.append({"alpha": cand["alpha"], "scale": cand["scale"], "report": rep})
            if verbose:
                print(f"[tofu.search] final alpha={cand['alpha']} scale={cand['scale']} {rep}",
                      flush=True)
            del edited
            torch.cuda.empty_cache()

    best = max(scored, key=lambda r: _objective(r["report"], objective, utility_floor, base_final))
    return {"best": best, "finalists": scored, "coarse": coarse,
            "baseline": base_final, "engram": engram, "editor": editor}

evaluate

evaluate(
    model,
    tok,
    splits: Optional[Dict[str, Any]] = None,
    *,
    level: str = "quick",
    device: str = "cuda",
    baseline: Optional["Report"] = None,
    n_retain: int = 200,
    n_utility: int = 100,
    bs: int = 16
) -> Report

Score one model. level is "quick" | "utility" | "full" (see the module docstring).

Pass baseline (the same call on the unedited model) to get deltas and, at "full", the rescaled composite.

Source code in src/engram/benchmarks/tofu.py

def evaluate(
    model,
    tok,
    splits: Optional[Dict[str, Any]] = None,
    *,
    level: str = "quick",
    device: str = "cuda",
    baseline: Optional["Report"] = None,
    n_retain: int = 200,
    n_utility: int = 100,
    bs: int = 16,
) -> Report:
    """Score one model. ``level`` is ``"quick"`` | ``"utility"`` | ``"full"`` (see the module docstring).

    Pass ``baseline`` (the same call on the unedited model) to get deltas and, at ``"full"``,
    the rescaled composite.
    """
    if level not in ("quick", "utility", "full"):
        raise ValueError(f"level must be 'quick', 'utility' or 'full', got {level!r}")
    splits = splits if splits is not None else load_splits()
    retain_eval = splits["retain"].select(range(min(n_retain, len(splits["retain"]))))
    rep = Report(level=level,
                 forget_nll=answer_nll(model, splits["forget"], tok, device, bs),
                 retain_nll=answer_nll(model, retain_eval, tok, device, bs))
    if baseline is not None:
        rep.forget_delta = rep.forget_nll - baseline.forget_nll
        rep.retain_delta = rep.retain_nll - baseline.retain_nll
    if level in ("utility", "full"):
        u = _utility(model, tok, splits, n=n_utility, bs=bs)
        rep.utility, rep.utility_retain = u["utility"], u["utility_retain"]
        rep.extra.update(u)
    if level == "full":
        from ._tofu_official import compute_full_tofu

        raw = compute_full_tofu(model, tok, _official_splits(splits), bs=bs)
        rep.extra["raw"] = raw
        if baseline is not None and "raw" in baseline.extra:
            rep.extra["scores"] = _composite(raw, baseline.extra["raw"])
            rep.overall = rep.extra["scores"]["Overall"]
    return rep

collect

collect(
    model,
    tok,
    splits: Dict[str, Any],
    *,
    device: str = "cuda",
    bs: int = 8,
    cache: Optional[Dict[str, str]] = None
) -> Tuple[Statistics, Statistics, EngramEditor]

Answer-token-masked covariances for the forget set (target) and the full set (reference).

cache={"target": path, "reference": path} reuses saved statistics — the covariance does not depend on any edit hyper-parameter, so it is collected once and swept over.

Source code in src/engram/benchmarks/tofu.py

def collect(
    model,
    tok,
    splits: Dict[str, Any],
    *,
    device: str = "cuda",
    bs: int = 8,
    cache: Optional[Dict[str, str]] = None,
) -> Tuple[Statistics, Statistics, EngramEditor]:
    """Answer-token-masked covariances for the forget set (target) and the full set (reference).

    ``cache={"target": path, "reference": path}`` reuses saved statistics — the covariance does
    not depend on any edit hyper-parameter, so it is collected once and swept over.
    """
    editor = EngramEditor(model, EditorConfig(storage_device=torch.device(device)))
    feats = lambda b: {"input_ids": b["input_ids"].to(device),
                       "attention_mask": b["attention_mask"].to(device)}
    mask = lambda b: b["labels"] != IGNORE

    def one(key: str, rows) -> Statistics:
        path = (cache or {}).get(key)
        if path:
            import os
            if os.path.exists(path):
                return Statistics.load(path, map_location=device)
        dl = DataLoader(_QA(rows, tok), batch_size=bs, collate_fn=_collate(tok.pad_token_id))
        st = editor.collect_statistics(dl, batch_fn=feats, mask_fn=mask)
        if path:
            st.save(path)
        return st

    return one("target", splits["forget"]), one("reference", splits["total"]), editor

Report dataclass

Report(
    level: str,
    forget_nll: float,
    retain_nll: float,
    forget_delta: Optional[float] = None,
    retain_delta: Optional[float] = None,
    utility: Optional[float] = None,
    utility_retain: Optional[float] = None,
    overall: Optional[float] = None,
    extra: Dict[str, Any] = dict(),
)

What one edited model scored. extra carries the level's own sub-metrics.

selectivity property

selectivity: Optional[float]

How much of the damage landed on the forget set rather than the retain set.