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"
        ):
            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)

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
) -> 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

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,
) -> 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.

    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())
    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).
        rtol = c_total.shape[-1] * torch.finfo(prec).eps
        pinv_total = torch.linalg.pinv(c_total, rtol=rtol)
        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=w.to(dev, dtype=prec) @ c_target @ pinv_total,
                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 = handler.weight_matrix(module, absorb_bias=absorbed).to(dev, dtype=prec) @ c_target @ pinv_total

        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
) -> 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,
) -> 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
    )
    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.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] = []

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._ref.clear()
    self._routed_mask = None

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).

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)."""
    return Statistics({k: v.to(device) for k, v in self.cov.items()}, 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.

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.
    """
    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)

save

save(path: Union[str, Path]) -> None

Save with torch.save (tagged format=2).

Source code in src/engram/stats.py

def save(self, path: Union[str, Path]) -> None:
    """Save with ``torch.save`` (tagged ``format=2``)."""
    torch.save({"format": _FORMAT, "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 not (isinstance(obj, dict) and obj.get("format") == _FORMAT and "cov" in obj):
        raise ValueError(
            f"{path!r} is not a Statistics file (format={_FORMAT}). It looks like a "
            "legacy raw-covariance dict; re-collect with EngramEditor.collect_statistics() "
            "— the mean+count format cannot be reconstructed from summed covariances."
        )
    return Statistics(obj["cov"], obj["count"])

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)