From 3ad0dbca07422f921e5cd732015942efbf334315 Mon Sep 17 00:00:00 2001
From: Lucia Quirke <luciarosequirke@gmail.com>
Date: Thu, 6 Aug 2026 18:54:28 +0900
Subject: [PATCH 1/5] feat(magic): support per-token per-query MAGIC
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

attribute_tokens=True with query_method="none" produced an unusable score
tensor. Per-token weights are [rows, seq_len] and the per-query stack used
dim=1, giving [rows, num_queries, seq_len] — query axis in the middle,
which nothing downstream reads. validate_scores takes shape[-1] as the
query count, so it compared seq_len against the query document count and
died naming the wrong dimension:

    ValueError: scores has 8 query columns but the query dataset has 2
    documents

on a run with 2 queries and seq_len 8.

Stack on dim=-1 instead. The query axis then comes last in both modes —
[rows, num_queries] per-doc, [rows, seq_len, num_queries] per-token —
matching the layout Scores.to_grid already produces for multi-query token
score directories, which load_attribution_scores already flags multi_query.
validate_scores needed no change: shape[-1] is the query count and
reshape(-1, num_queries) flattens the leading axes into leave-out units,
documents or token positions as appropriate. dim=-1 is identical to dim=1
for 1-D inputs, so per-doc per-query scores are unchanged.

Fix the padding trim in the same path, which applied weight_pad_count
regardless of rank while the main scoring path picks by rank. The two
differ once doc_ids are present (pad rows route to one synthetic doc id),
so a 5-doc dataset at batch_size 4 kept 7 of its 8 padded rows instead of
trimming to 5, leaving pad rows in the saved scores.

Teach both .pt classifiers about 3-D: scores_are_per_token so a reloaded
run sizes its weights per-token, and _pt_scores_are_per_query so it is
recognised as multi-query. 3-D needs no config lookup to disambiguate,
unlike 2-D.

The aggregation test is the numerical gate: per-token per-query scores
summed over each document's tokens reproduce the per-doc per-query run.
Both new end-to-end tests fail on the parent commit with shape (7, 2, 8)
against the expected (5, 8, 2) — wrong axis order and untrimmed padding
together.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 bergson/magic/cli.py          | 20 ++++-----
 bergson/validate.py           | 15 ++++---
 tests/test_per_query_magic.py | 82 +++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+), 17 deletions(-)

diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py
index 3bdcc130..9fef4aff 100644
--- a/bergson/magic/cli.py
+++ b/bergson/magic/cli.py
@@ -130,9 +130,8 @@ def compute_per_query_magic_scores(
     yields a single aggregate-query score. This scores each query separately:
     for each query document it takes that document's gradient at the final model
     as the backward cotangent and runs ``Trainer.backward`` over the saved
-    trajectory, producing a ``[num_train_docs, num_query_docs]`` matrix (the
-    layout ``validate_scores`` consumes: rows are leave-out docs, columns
-    queries).
+    trajectory, producing ``[num_train_docs, num_query_docs]`` — or
+    ``[num_train_docs, seq_len, num_query_docs]`` when attributing tokens.
 
     The backward is linear in the cotangent, so this is exact; the forward runs
     once and every query reuses its checkpoints (``cleanup=False``). Per-query
@@ -217,7 +216,7 @@ def compute_per_query_magic_scores(
 
         s = bwd_state.weight_grads.detach().cpu()
         if pad_count:
-            s = s[:-weight_pad_count]
+            s = s[:-weight_pad_count] if s.ndim == 1 else s[:-pad_count]
         if main:
             torch.save(s, qpath)
         per_query.append(s)
@@ -241,9 +240,7 @@ def compute_per_query_magic_scores(
     fwd_state.copy_(restored)
     del restored
 
-    # [num_train_docs, num_query_docs] — the layout validate_scores expects
-    # (rows are leave-out docs; columns are queries).
-    return torch.stack(per_query, dim=1)
+    return torch.stack(per_query, dim=-1)
 
 
 def scores_are_per_token(score_path: str) -> bool:
@@ -256,7 +253,9 @@ def scores_are_per_token(score_path: str) -> bool:
     if score_path.endswith(".npy"):
         return False
     scores = torch.load(score_path, map_location="cpu")
-    return isinstance(scores, torch.Tensor) and scores.ndim == 2 and scores.shape[1] > 1
+    return isinstance(scores, torch.Tensor) and (
+        scores.ndim == 3 or (scores.ndim == 2 and scores.shape[1] > 1)
+    )
 
 
 def attach_doc_ids_if_missing(dataset: Dataset) -> Dataset:
@@ -487,9 +486,8 @@ def worker(
 
     multi_query = False
     if not score_path and run_cfg.query_method == "none":
-        # Per-query MAGIC: one backward per query, sharing the forward. Yields a
-        # [num_query_docs, num_train_docs] score matrix (multi_query), the unit
-        # for a per-query LDS.
+        # Per-query MAGIC: one backward per query, sharing the forward. Yields
+        # the unit for a per-query LDS.
         if not isinstance(run_cfg, MagicConfig):
             raise RuntimeError("run_cfg must be a MagicConfig to compute scores")
         assert query_dataset is not None
diff --git a/bergson/validate.py b/bergson/validate.py
index 2d61a575..d8049983 100644
--- a/bergson/validate.py
+++ b/bergson/validate.py
@@ -46,7 +46,8 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
     Returns ``(scores, multi_query)``. Token score directories are per-token,
     with a query dimension when ``num_scores > 1`` (``[docs, seq_len,
     queries]``); plain score directories and ``.npy`` arrays are per-document,
-    with one column per query. A 2-D ``.pt`` tensor is ambiguous — per-token
+    with one column per query. A 3-D ``.pt`` tensor is per-token per-query
+    ``[docs, seq_len, queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token
     MAGIC scores are ``[docs, seq_len]`` and per-query MAGIC scores are
     ``[docs, queries]`` — so the run config next to it decides: it is
     per-query iff the run used ``query_method: none``.
@@ -84,12 +85,14 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
 
 
 def _pt_scores_are_per_query(score_path: str, scores: torch.Tensor) -> bool:
-    """Whether a 2-D ``.pt`` tensor is per-query ``[docs, queries]`` rather
-    than per-token ``[docs, seq_len]``. The shape alone cannot tell them
-    apart; the run config written next to ``scores.pt`` records which one
-    ``run_magic`` saved."""
-    if not (isinstance(scores, torch.Tensor) and scores.ndim == 2):
+    """Whether a ``.pt`` tensor carries a query axis. 3-D is unambiguously
+    ``[docs, seq_len, queries]``; a 2-D tensor is per-query ``[docs, queries]``
+    or per-token ``[docs, seq_len]`` and the shape alone cannot tell them
+    apart, so the run config next to ``scores.pt`` decides."""
+    if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3):
         return False
+    if scores.ndim == 3:
+        return True
     cfg_path = Path(score_path).parent / CONFIG_FILENAME
     if not cfg_path.is_file():
         return False
diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py
index 7f05e221..42f19010 100644
--- a/tests/test_per_query_magic.py
+++ b/tests/test_per_query_magic.py
@@ -197,3 +197,85 @@ def test_per_query_scores_only_real_queries_when_padded():
         import os
 
         assert not os.path.exists(f"{run_path}/per_query/q3.pt")
+
+
+def _per_query_run(tmp_path, attribute_tokens: bool, num_docs=5, seq_len=8, n_query=2):
+    """Run worker() in per-query mode; return (scores, doc_ids).
+
+    num_docs=5 at batch_size 4 pads by 3 rows, where weight_pad_count is 1 but
+    pad_count is 3 — the gap a rank-blind trim falls into.
+    """
+    from bergson.config.config import DataConfig
+    from bergson.magic.cli import worker
+
+    def ds(n):
+        toks = [[(d * seq_len + t) % 50 + 1 for t in range(seq_len)] for d in range(n)]
+        return Dataset.from_dict(
+            {
+                "input_ids": toks,
+                "labels": toks,
+                "doc_ids": [[d] * seq_len for d in range(n)],
+                "length": [seq_len] * n,
+            }
+        )
+
+    run_path = tmp_path / ("tok" if attribute_tokens else "doc")
+    run_cfg = MagicConfig(
+        run_path=str(run_path),
+        model="EleutherAI/pythia-14m",
+        data=DataConfig(dataset="unused", chunk_length=seq_len),
+        query=DataConfig(dataset="unused", chunk_length=seq_len),
+        batch_size=4,
+        attribute_tokens=attribute_tokens,
+        query_method="none",
+        skip_validation=True,
+    )
+    worker(0, 0, 1, ds(num_docs), ds(n_query), num_docs, n_query, run_cfg)
+
+    doc_ids = run_path / "doc_ids.pt"
+    return (
+        torch.load(run_path / "scores.pt"),
+        torch.load(doc_ids) if doc_ids.exists() else None,
+    )
+
+
+def test_per_query_per_token_layout(tmp_path):
+    """Per-token per-query scores are [docs, seq_len, queries], so doc_ids
+    indexes the leading two axes and validate_scores reads shape[-1]."""
+    scores, doc_ids = _per_query_run(tmp_path, attribute_tokens=True)
+
+    assert scores.shape == (5, 8, 2), f"got {tuple(scores.shape)}"
+    assert doc_ids is not None, "per-token run must write doc_ids.pt"
+    assert doc_ids.shape == scores.shape[:2]
+
+
+def test_per_query_per_token_aggregates_to_per_doc(tmp_path):
+    """Summing per-token per-query scores over each document's tokens
+    reproduces the per-doc per-query run."""
+    per_tok, doc_ids = _per_query_run(tmp_path, attribute_tokens=True)
+    per_doc, _ = _per_query_run(tmp_path, attribute_tokens=False)
+    num_docs, n_query = per_doc.shape
+    assert doc_ids is not None
+
+    agg = torch.zeros(num_docs, n_query, dtype=torch.float64)
+    agg.scatter_add_(
+        0,
+        doc_ids.reshape(-1, 1).expand(-1, n_query),
+        per_tok.reshape(-1, n_query).to(torch.float64),
+    )
+
+    torch.testing.assert_close(agg, per_doc.to(torch.float64), atol=1e-5, rtol=1e-4)
+
+
+def test_three_dim_scores_load_as_per_token_multi_query(tmp_path):
+    """A 3-D scores.pt is unambiguous: neither classifier needs the run config."""
+    from bergson.magic.cli import scores_are_per_token
+    from bergson.validate import load_attribution_scores
+
+    path = tmp_path / "scores.pt"
+    torch.save(torch.randn(4, 6, 3), path)
+
+    assert scores_are_per_token(str(path))
+    loaded, multi_query = load_attribution_scores(str(path))
+    assert multi_query
+    assert loaded.shape == (4, 6, 3)

From e01c6e5925cb3edd9d81f1c618886f67e6cc1fbb Mon Sep 17 00:00:00 2001
From: Lucia Quirke <luciarosequirke@gmail.com>
Date: Thu, 6 Aug 2026 19:04:00 +0900
Subject: [PATCH 2/5] refactor(magic): decide score format from the config, not
 the shape
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Score layout was inferred from tensor rank in several places, which is
guesswork: a 2-D .pt is [docs, seq_len] or [docs, queries] depending only
on how the run was configured. #407 established the fix for one of those
call sites — read the config.yaml that save_run_config writes next to
scores.pt — but scores_are_per_token was left sniffing shapes, and the
parsing lived inline rather than beside the other config readers.

Add read_first_step_config to config_io, next to read_config and
load_subconfig, and route both classifiers through it. The flags say
everything the rank could:

    query_method  attribute_tokens  layout
    none          yes               [docs, seq_len, queries]
    none          no                [docs, queries]
    mean/sum      yes               [docs, seq_len]
    mean/sum      no                [docs]

so a run is per-query iff query_method is none, whatever rank results, and
per-token iff attribute_tokens is set. Neither needs the shape. Score
directories keep reading info.json, and load_attribution_scores now takes
the query count from the store's num_scores rather than re-deriving it from
the grid it just built.

Shape survives in exactly one place: a .pt with no config beside it, where
nothing else is knowable and only rank 3 is unambiguous.

cfg_attributes_tokens reads attribute_tokens with the deprecated per_token
as an alias, in one place, so a run written with the current field name is
no longer missed.

test_load_attribution_scores_pt_per_query asserted that a 2-D tensor whose
config said query_method: none and per_token: true was single-query. No run
produces that pair — attributing tokens per query yields rank 3 — so the
case was describing an unreachable artifact and pinning the shape-derived
answer for it. Repointed at the 3-D tensor such a run does produce.

Drop six tests that asserted torch's own view()/reshape() indexing
semantics. They called no bergson code, so they could not fail unless
PyTorch itself changed, and the behaviour they stood in for is covered
end to end by the per-query aggregation test.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 bergson/config/config_io.py        |  25 +++++++
 bergson/magic/cli.py               |  19 ++++-
 bergson/validate.py                |  47 ++++--------
 tests/test_multi_query_validate.py |   9 ++-
 tests/test_per_query_magic.py      |  22 ++----
 tests/test_per_token_lds.py        | 114 ++++-------------------------
 6 files changed, 85 insertions(+), 151 deletions(-)

diff --git a/bergson/config/config_io.py b/bergson/config/config_io.py
index da50616d..105accc7 100644
--- a/bergson/config/config_io.py
+++ b/bergson/config/config_io.py
@@ -120,6 +120,31 @@ def read_config(path: str | Path) -> dict[str, Any]:
     return config
 
 
+def read_first_step_config(path: str | Path) -> dict[str, Any] | None:
+    """The first step's configuration from a run's ``config.yaml``.
+
+    ``path`` may be the run directory or an artifact inside it (``scores.pt``),
+    so callers holding an output path need not resolve it. Returns ``None``
+    when there is no readable config, which is the signal that a caller must
+    fall back to inspecting the artifact itself.
+    """
+    p = Path(path)
+    cfg_path = (p if p.is_dir() else p.parent) / CONFIG_FILENAME
+    if not cfg_path.is_file():
+        return None
+
+    try:
+        steps = read_config(cfg_path)["steps"]
+    except (ValueError, yaml.YAMLError):
+        return None
+
+    for step in steps:
+        for step_cfg in step.values():
+            if isinstance(step_cfg, dict):
+                return step_cfg
+    return None
+
+
 T = TypeVar("T", bound="FromDict")
 
 
diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py
index 9fef4aff..f5fe651c 100644
--- a/bergson/magic/cli.py
+++ b/bergson/magic/cli.py
@@ -28,7 +28,7 @@
 )
 
 from ..config.config import TrainingConfig, ValidationConfig
-from ..config.config_io import save_run_config
+from ..config.config_io import read_first_step_config, save_run_config
 from ..distributed import launch_distributed_run
 from ..utils.load_from_optimizer import (
     save_second_moments_as_optimizer_pt,
@@ -36,7 +36,11 @@
 from ..utils.logging import wandb_log_fn
 from ..utils.utils import get_device, get_device_index
 from ..utils.worker_utils import setup_data_pipeline
-from ..validate import load_attribution_scores, validate_scores
+from ..validate import (
+    cfg_attributes_tokens,
+    load_attribution_scores,
+    validate_scores,
+)
 from .config import MagicConfig
 from .data_stream import DataStream, pad_dataset_to_batch_size
 from .grad_accum import accumulate_grads
@@ -244,6 +248,12 @@ def compute_per_query_magic_scores(
 
 
 def scores_are_per_token(score_path: str) -> bool:
+    """Whether ``score_path`` holds per-token scores.
+
+    Score directories record it in ``info.json`` and ``.pt`` files in the run
+    config beside them. Without either -- an externally produced tensor --
+    fall back to the shape, where only 3-D is unambiguous.
+    """
     if os.path.isdir(score_path):
         info_path = os.path.join(score_path, "info.json")
         if not os.path.isfile(info_path):
@@ -252,6 +262,11 @@ def scores_are_per_token(score_path: str) -> bool:
             return bool(json.load(f).get("attribute_tokens", False))
     if score_path.endswith(".npy"):
         return False
+
+    step_cfg = read_first_step_config(score_path)
+    if step_cfg is not None:
+        return cfg_attributes_tokens(step_cfg)
+
     scores = torch.load(score_path, map_location="cpu")
     return isinstance(scores, torch.Tensor) and (
         scores.ndim == 3 or (scores.ndim == 2 and scores.shape[1] > 1)
diff --git a/bergson/validate.py b/bergson/validate.py
index d8049983..b681c12a 100644
--- a/bergson/validate.py
+++ b/bergson/validate.py
@@ -18,7 +18,6 @@
 import torch
 import torch.distributed as dist
 import torch.nn.functional as F
-import yaml
 from peft import PeftModel
 from scipy.stats import pearsonr, spearmanr
 from tqdm import tqdm
@@ -31,7 +30,7 @@
 )
 
 from .config.config import ScoreConfig, ValidationConfig
-from .config.config_io import CONFIG_FILENAME, load_subconfig, save_run_config
+from .config.config_io import load_subconfig, read_first_step_config, save_run_config
 from .data import Scores, load_scores, pad_and_tensor
 from .magic.data_stream import DataStream, pad_dataset_to_batch_size
 from .magic.trainer import TrainerState, prepare_trainer
@@ -50,7 +49,8 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
     ``[docs, seq_len, queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token
     MAGIC scores are ``[docs, seq_len]`` and per-query MAGIC scores are
     ``[docs, queries]`` — so the run config next to it decides: it is
-    per-query iff the run used ``query_method: none``.
+    per-query iff the run used ``query_method: none`` without attributing
+    tokens.
 
     Score directories are negated when their ``score_cfg.higher_is_better`` is
     set, aligning them with the loss-diff convention. ``.npy`` files carry no
@@ -70,7 +70,7 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
             scores = loaded.to_grid()
             if negate:
                 scores = -scores
-            return scores, scores.ndim == 3
+            return scores, loaded.num_scores > 1
 
         arr = np.asarray(loaded[:])
         # Copy: the slice is a read-only view onto the memmap.
@@ -78,37 +78,22 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
         scores = torch.from_numpy(arr.astype(out_dtype, copy=True))
         if negate:
             scores = -scores
-        return scores, scores.ndim == 2 and scores.shape[1] > 1
+        return scores, loaded.num_scores > 1
 
     scores = torch.load(score_path, map_location="cpu")
-    return scores, _pt_scores_are_per_query(score_path, scores)
+    if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3):
+        return scores, False
 
+    step_cfg = read_first_step_config(score_path)
+    if step_cfg is None:
+        return scores, scores.ndim == 3
+    return scores, step_cfg.get("query_method") == "none"
 
-def _pt_scores_are_per_query(score_path: str, scores: torch.Tensor) -> bool:
-    """Whether a ``.pt`` tensor carries a query axis. 3-D is unambiguously
-    ``[docs, seq_len, queries]``; a 2-D tensor is per-query ``[docs, queries]``
-    or per-token ``[docs, seq_len]`` and the shape alone cannot tell them
-    apart, so the run config next to ``scores.pt`` decides."""
-    if not isinstance(scores, torch.Tensor) or scores.ndim not in (2, 3):
-        return False
-    if scores.ndim == 3:
-        return True
-    cfg_path = Path(score_path).parent / CONFIG_FILENAME
-    if not cfg_path.is_file():
-        return False
-    with open(cfg_path) as f:
-        doc = yaml.safe_load(f)
-    if not isinstance(doc, dict):
-        return False
-    steps = doc.get("steps")
-    payload = (
-        next(iter(steps[0].values())) if isinstance(steps, list) and steps else doc
-    )
-    return (
-        isinstance(payload, dict)
-        and payload.get("query_method") == "none"
-        and not payload.get("per_token")
-    )
+
+def cfg_attributes_tokens(step_cfg: dict) -> bool:
+    """``attribute_tokens`` from a serialized run config, honouring the
+    deprecated ``per_token`` alias that ``MagicConfig`` still accepts."""
+    return bool(step_cfg.get("attribute_tokens") or step_cfg.get("per_token"))
 
 
 def bank_loss_cache_key(
diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py
index 7778bfbd..da3b70fc 100644
--- a/tests/test_multi_query_validate.py
+++ b/tests/test_multi_query_validate.py
@@ -174,11 +174,12 @@ def test_load_attribution_scores_pt_per_query(tmp_path):
     assert multi_query
     torch.testing.assert_close(scores, values)
 
-    # Per-token runs save [docs, seq_len]; never multi-query.
-    cfg = {"steps": [{"magic": {"query_method": "none", "per_token": True}}]}
+    # Attributing tokens as well makes it 3-D, and still per-query.
+    cfg = {"steps": [{"magic": {"query_method": "none", "attribute_tokens": True}}]}
     (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg))
-    _, multi_query = load_attribution_scores(str(tmp_path / "scores.pt"))
-    assert not multi_query
+    torch.save(torch.zeros(4, 3, 2), tmp_path / "tok.pt")
+    _, multi_query = load_attribution_scores(str(tmp_path / "tok.pt"))
+    assert multi_query
 
     cfg = {"steps": [{"magic": {"query_method": "mean"}}]}
     (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg))
diff --git a/tests/test_per_query_magic.py b/tests/test_per_query_magic.py
index 42f19010..a8910f8e 100644
--- a/tests/test_per_query_magic.py
+++ b/tests/test_per_query_magic.py
@@ -239,23 +239,16 @@ def ds(n):
     )
 
 
-def test_per_query_per_token_layout(tmp_path):
-    """Per-token per-query scores are [docs, seq_len, queries], so doc_ids
-    indexes the leading two axes and validate_scores reads shape[-1]."""
-    scores, doc_ids = _per_query_run(tmp_path, attribute_tokens=True)
-
-    assert scores.shape == (5, 8, 2), f"got {tuple(scores.shape)}"
-    assert doc_ids is not None, "per-token run must write doc_ids.pt"
-    assert doc_ids.shape == scores.shape[:2]
-
-
 def test_per_query_per_token_aggregates_to_per_doc(tmp_path):
-    """Summing per-token per-query scores over each document's tokens
-    reproduces the per-doc per-query run."""
+    """Per-token per-query scores are [docs, seq_len, queries], and summing
+    them over each document's tokens reproduces the per-doc per-query run."""
     per_tok, doc_ids = _per_query_run(tmp_path, attribute_tokens=True)
     per_doc, _ = _per_query_run(tmp_path, attribute_tokens=False)
     num_docs, n_query = per_doc.shape
-    assert doc_ids is not None
+
+    assert per_tok.shape == (num_docs, 8, n_query), f"got {tuple(per_tok.shape)}"
+    assert doc_ids is not None, "per-token run must write doc_ids.pt"
+    assert doc_ids.shape == per_tok.shape[:2]
 
     agg = torch.zeros(num_docs, n_query, dtype=torch.float64)
     agg.scatter_add_(
@@ -276,6 +269,5 @@ def test_three_dim_scores_load_as_per_token_multi_query(tmp_path):
     torch.save(torch.randn(4, 6, 3), path)
 
     assert scores_are_per_token(str(path))
-    loaded, multi_query = load_attribution_scores(str(path))
+    _, multi_query = load_attribution_scores(str(path))
     assert multi_query
-    assert loaded.shape == (4, 6, 3)
diff --git a/tests/test_per_token_lds.py b/tests/test_per_token_lds.py
index 188c5bf5..920bea5d 100644
--- a/tests/test_per_token_lds.py
+++ b/tests/test_per_token_lds.py
@@ -16,105 +16,6 @@
 from bergson.validate import load_attribution_scores
 
 
-def test_flat_reweight_per_doc_matches_direct_indexing():
-    """1-D (per-doc) weights: ``view(-1)[subset]`` equals direct indexing, so
-    the per-document leave-k-out path is unchanged by the flat generalization."""
-    weights = torch.ones(8)
-    subset = torch.tensor([1, 4, 6])
-
-    flat = weights.clone()
-    flat.view(-1)[subset] = 0.0
-    direct = torch.ones(8)
-    direct[subset] = 0.0
-
-    torch.testing.assert_close(flat, direct)
-
-
-def test_flat_reweight_per_token_hits_grid_positions():
-    """2-D (per-token) weights: a flat subset re-weights exactly the intended
-    ``doc * seq_len + token`` grid positions and nothing else."""
-    n_docs, seq_len = 4, 5
-    weights = torch.ones(n_docs, seq_len)
-    # doc 1 token 2, doc 3 token 0, doc 3 token 4
-    positions = [(1, 2), (3, 0), (3, 4)]
-    subset = torch.tensor([d * seq_len + t for d, t in positions])
-
-    weights.view(-1)[subset] = 2.0
-
-    expected = torch.ones(n_docs, seq_len)
-    for d, t in positions:
-        expected[d, t] = 2.0
-    torch.testing.assert_close(weights, expected)
-
-
-def test_whole_doc_token_subset_reweights_full_row():
-    """Selecting every token position of a document is equivalent to selecting
-    that document in the per-document formulation (its whole row is set)."""
-    n_docs, seq_len = 3, 6
-    weights = torch.ones(n_docs, seq_len)
-    doc = 1
-    subset = torch.arange(doc * seq_len, (doc + 1) * seq_len)
-
-    weights.view(-1)[subset] = 0.0
-
-    expected = torch.ones(n_docs, seq_len)
-    expected[doc] = 0.0
-    torch.testing.assert_close(weights, expected)
-
-
-def test_reweight_sequence_preserves_padding_rows():
-    """The retrain loop's exact sequence -- fill with 1.0, zero the padding
-    rows, then flat-reweight the subset -- leaves padding rows at zero, so
-    batch-size padding docs never re-enter training."""
-    n_docs, seq_len, pad_count = 4, 5, 1
-    weights = torch.ones(n_docs, seq_len)
-    weights.data[-pad_count:] = 0.0
-    subset = torch.tensor([0 * seq_len + 2, 1 * seq_len + 4])
-
-    weights.view(-1)[subset] = 0.0
-
-    assert weights[-pad_count:].eq(0).all()
-    expected_real = torch.ones(n_docs - pad_count, seq_len)
-    expected_real[0, 2] = 0.0
-    expected_real[1, 4] = 0.0
-    torch.testing.assert_close(weights[:-pad_count], expected_real)
-
-
-def test_score_sum_flat_per_token():
-    """``scores.reshape(-1)[subset].sum()`` sums the selected per-token scores
-    regardless of the score tensor's shape (1-D per doc or 2-D per token)."""
-    scores = torch.arange(12, dtype=torch.float32).reshape(3, 4)
-    subset = torch.tensor([0 * 4 + 1, 2 * 4 + 3])  # scores 1 and 11
-
-    got = scores.reshape(-1)[subset].sum()
-
-    torch.testing.assert_close(got, torch.tensor(12.0))
-    # Same expression on a flat (per-doc) score vector.
-    flat = torch.arange(6, dtype=torch.float32)
-    torch.testing.assert_close(
-        flat.reshape(-1)[torch.tensor([0, 5])].sum(), torch.tensor(5.0)
-    )
-
-
-def test_score_sum_flat_multi_query_per_token():
-    """``scores.reshape(-1, num_queries)[subset].sum(dim=0)`` sums each query's
-    scores over the selected token positions of a ``[docs, seq_len, queries]``
-    grid -- the same expression that serves per-doc multi-query scores."""
-    n_docs, seq_len, num_queries = 2, 3, 2
-    scores = torch.arange(12, dtype=torch.float32).reshape(n_docs, seq_len, num_queries)
-    # doc 0 token 1, doc 1 token 2
-    subset = torch.tensor([0 * seq_len + 1, 1 * seq_len + 2])
-
-    got = scores.reshape(-1, num_queries)[subset].sum(dim=0)
-
-    torch.testing.assert_close(got, scores[0, 1] + scores[1, 2])
-    # Per-doc multi-query [docs, queries] is served by the same expression.
-    per_doc = torch.arange(8, dtype=torch.float32).reshape(4, 2)
-    torch.testing.assert_close(
-        per_doc.reshape(-1, 2)[torch.tensor([1, 3])].sum(dim=0), per_doc[1] + per_doc[3]
-    )
-
-
 def test_load_attribution_scores_per_token_pt(tmp_path):
     """A 2-D ``.pt`` tensor is treated as per-token MAGIC scores, never
     multi-query -- so ``validate_scores`` keeps it 2-D for token re-weighting."""
@@ -225,3 +126,18 @@ def test_load_token_dir_keeps_query_dim_when_multiscore(tmp_path):
     torch.testing.assert_close(scores[0, 0], torch.tensor([1.0, 2.0]))
     torch.testing.assert_close(scores[1, 0], torch.tensor([5.0, 6.0]))
     torch.testing.assert_close(scores[1, 1], torch.tensor([0.0, 0.0]))
+
+
+def test_scores_are_per_token_from_run_config(tmp_path):
+    """With a run config beside it, the flag decides rather than the shape --
+    including ``attribute_tokens``, not just the deprecated ``per_token``."""
+    import yaml
+
+    from bergson.magic.cli import scores_are_per_token
+
+    # A shape the fallback above would read as per-doc.
+    torch.save(torch.zeros(4, 1), tmp_path / "scores.pt")
+    cfg = {"steps": [{"magic": {"query_method": "mean", "attribute_tokens": True}}]}
+    (tmp_path / "config.yaml").write_text(yaml.safe_dump(cfg))
+
+    assert scores_are_per_token(str(tmp_path / "scores.pt"))

From 2b51f7d15ec1a50e3604a0705d78b18608f9278d Mon Sep 17 00:00:00 2001
From: Lucia Quirke <luciarosequirke@gmail.com>
Date: Thu, 6 Aug 2026 19:28:13 +0900
Subject: [PATCH 3/5] refactor(score): stop duplicating the token score store
 format

save_sequence_scores delegates to MemmapSequenceScoreWriter, but its token
counterpart reimplemented MemmapTokenScoreWriter inline: the memmap
creation, offsets.npy, and an info.json payload identical field for field.
So the on-disk token score format was written from two places that had to
be kept in step by hand.

That matters more now that scores_are_per_token reads
info.json["attribute_tokens"] as authoritative: a drift between the two
writers stops being a cosmetic inconsistency and becomes a
misclassification.

Delegating needs the writer to accept what it actually uses. It only ever
took a Dataset to call compute_num_token_grads on it, while
save_token_scores already holds the offsets those counts came from, so
__init__ now takes num_token_grads and a from_dataset classmethod covers
the callers that hold a dataset.

Also gives the token writer the overwrite flag its sequence twin already
had, which delegation needs: save_token_scores wrote with mode="w+"
unconditionally, and without overwrite the writer would silently reuse a
stale scores.bin instead of replacing it.

Net 19 lines out of score_writer.py, and one place left that knows the
format.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 bergson/score/score.py         |  2 +-
 bergson/score/score_writer.py  | 67 ++++++++++++----------------------
 tests/test_attribute_tokens.py |  6 +--
 3 files changed, 28 insertions(+), 47 deletions(-)

diff --git a/bergson/score/score.py b/bergson/score/score.py
index 02ee6f8f..dac0c7c4 100644
--- a/bergson/score/score.py
+++ b/bergson/score/score.py
@@ -231,7 +231,7 @@ def create_scorer(
     else:
         num_scores = num_queries_total if num_queries_total is not None else num_queries
     if attribute_tokens:
-        writer = MemmapTokenScoreWriter(
+        writer = MemmapTokenScoreWriter.from_dataset(
             path,
             data,
             num_scores,
diff --git a/bergson/score/score_writer.py b/bergson/score/score_writer.py
index 3f7a06b6..d51dcc82 100644
--- a/bergson/score/score_writer.py
+++ b/bergson/score/score_writer.py
@@ -156,11 +156,12 @@ class MemmapTokenScoreWriter(ScoreWriter):
     def __init__(
         self,
         path: Path,
-        data: Dataset,
+        num_token_grads: np.ndarray,
         num_scores: int,
         *,
         dtype: torch.dtype = torch.float32,
         flush_interval: int = 64,
+        overwrite: bool = False,
     ):
         self.path = path
         self.num_scores = num_scores
@@ -168,10 +169,9 @@ def __init__(
         self.flush_interval = flush_interval
         self.num_batches_since_flush = 0
 
-        num_token_grads = compute_num_token_grads(data)
-        num_items = len(data)
+        num_items = len(num_token_grads)
         self.num_token_grads = num_token_grads
-        self.offsets = np.zeros(len(num_token_grads) + 1, dtype=np.int64)
+        self.offsets = np.zeros(num_items + 1, dtype=np.int64)
         np.cumsum(num_token_grads, out=self.offsets[1:])
         total_tokens = int(self.offsets[-1])
 
@@ -181,7 +181,7 @@ def __init__(
         struct_dtype, struct_dtype_json = _score_struct_dtype(num_scores, np_dtype)
 
         rank = dist.get_rank() if dist.is_initialized() else 0
-        if rank == 0 and not scores_file_path.exists():
+        if rank == 0 and (overwrite or not scores_file_path.exists()):
             print(f"Creating new token scores file: {scores_file_path}")
 
             self.scores = np.memmap(
@@ -216,6 +216,11 @@ def __init__(
             shape=(total_tokens,),
         )
 
+    @classmethod
+    def from_dataset(cls, path: Path, data: Dataset, num_scores: int, **kwargs):
+        """For callers holding the dataset the scores were computed over."""
+        return cls(path, compute_num_token_grads(data), num_scores, **kwargs)
+
     def __call__(self, indices: list[int], scores: torch.Tensor, query_offset: int = 0):
         # scores: [total_valid_in_batch, num_scores]
         scores = scores.to(dtype=self.dtype)
@@ -248,48 +253,24 @@ def save_token_scores(
     *,
     dtype: torch.dtype = torch.float32,
 ) -> None:
-    """One-shot equivalent of :class:`MemmapTokenScoreWriter`'s on-disk
-    layout, for callers that already hold the full flat ``(total_tokens,
-    num_scores)`` array in memory (e.g. summed across upstream per-token
-    score dirs sharing the same ``offsets``), rather than streaming batches
-    through ``__call__``.
+    """One-shot equivalent of :class:`MemmapTokenScoreWriter`, for callers that
+    already hold the full flat ``(total_tokens, num_scores)`` array in memory
+    (e.g. summed across upstream per-token score dirs sharing the same
+    ``offsets``), rather than streaming batches through ``__call__``.
     """
     if scores.ndim == 1:
         scores = scores[:, None]
-    total_tokens, num_scores = scores.shape
-    num_items = len(offsets) - 1
-
-    path.mkdir(parents=True, exist_ok=True)
-    np_dtype = convert_dtype_to_np(dtype)
-    struct_dtype, struct_dtype_json = _score_struct_dtype(num_scores, np_dtype)
-
-    mmap = np.memmap(
-        path / "scores.bin",
-        dtype=np.dtype(struct_dtype),  # type: ignore
-        mode="w+",
-        shape=(total_tokens,),
+    _, num_scores = scores.shape
+    num_token_grads = np.diff(offsets).astype(np.int64)
+
+    writer = MemmapTokenScoreWriter(
+        path, num_token_grads, num_scores, dtype=dtype, overwrite=True
     )
-    # numpy_to_tensor handles bf16, which torch.from_numpy cannot.
-    scores_np = tensor_to_numpy(numpy_to_tensor(np.ascontiguousarray(scores)).to(dtype))
-    for i in range(num_scores):
-        mmap[f"score_{i}"] = scores_np[:, i]
-        mmap[f"written_{i}"] = True
-    mmap.flush()
-
-    np.save(path / "offsets.npy", offsets)
-
-    with (path / "info.json").open("w") as f:
-        json.dump(
-            {
-                "attribute_tokens": True,
-                "num_items": num_items,
-                "num_rows": total_tokens,
-                "num_scores": num_scores,
-                "dtype": struct_dtype_json,
-            },
-            f,
-            indent=2,
-        )
+    writer(
+        list(range(len(num_token_grads))),
+        numpy_to_tensor(np.ascontiguousarray(scores)),
+    )
+    writer.flush()
 
 
 class MemmapSequenceScoreWriter(ScoreWriter):
diff --git a/tests/test_attribute_tokens.py b/tests/test_attribute_tokens.py
index 916c3954..38f4d2eb 100644
--- a/tests/test_attribute_tokens.py
+++ b/tests/test_attribute_tokens.py
@@ -236,7 +236,7 @@ def test_token_score_writer(tmp_path: Path):
     # lengths [4, 3] → num_token_grads [3, 2]
     ds = Dataset.from_dict({"input_ids": [[1, 2, 3, 4], [5, 6, 7]], "length": [4, 3]})
 
-    writer = MemmapTokenScoreWriter(
+    writer = MemmapTokenScoreWriter.from_dataset(
         tmp_path,
         data=ds,
         num_scores=2,
@@ -457,7 +457,7 @@ def test_token_score_e2e(tmp_path: Path, model, dataset):
     query_grads = {m: torch.randn(1, math.prod(shapes[m])) for m in modules}
 
     score_dtype = get_gradient_dtype(model)
-    writer = MemmapTokenScoreWriter(
+    writer = MemmapTokenScoreWriter.from_dataset(
         tmp_path / "scores",
         data=dataset,
         num_scores=1,
@@ -877,7 +877,7 @@ def test_trackstar_token_scores_sum_to_sequence_scores_on_disk(
 
     # --- Per-token scores via MemmapTokenScoreWriter ---
     tok_path = tmp_path / "tok_scores"
-    tok_writer = MemmapTokenScoreWriter(tok_path, dataset, 1, dtype=dtype)
+    tok_writer = MemmapTokenScoreWriter.from_dataset(tok_path, dataset, 1, dtype=dtype)
     tok_scorer = Scorer(
         query_grads=query_grads,
         modules=sorted_modules,

From 9048416da1b0287fd9f2ca81a270f98e55a30d4f Mon Sep 17 00:00:00 2001
From: Lucia Quirke <luciarosequirke@gmail.com>
Date: Thu, 6 Aug 2026 19:44:06 +0900
Subject: [PATCH 4/5] refactor(score): drop .npy attribution score support
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

bergson never writes a .npy score file — every writer emits a score
directory — so .npy was an ingest-only path for arrays produced outside
the scoring pipeline, and nothing in the repo feeds one: no config sets
scores: to a .npy, and the examples that save scores.npy read it straight
back with np.load rather than through load_attribution_scores.

Remove the branch from load_attribution_scores, scores_are_per_token and
worker's score-path dispatch, and with it ArrayScores, which existed only
to give a bare array the Scores interface.

It also carried its own rules. A .npy could not have a score_cfg, so it
alone skipped the higher_is_better negation and had to be supplied in the
loss-diff convention already; and it was the one input whose multi-query
flag came from a raw column count. Score directories record num_scores in
info.json, so the surviving formats all describe themselves.

The bank-loss-cache tests used .npy as a convenient way to hand a score
matrix to evaluate_retrained. They now write a score directory via
save_sequence_scores, which is what a caller would reach for, and the
multi_query parametrization still passes both ways.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 bergson/data.py                    | 35 ++++--------------------------
 bergson/magic/cli.py               |  5 +----
 bergson/validate.py                | 27 +++++++++--------------
 tests/test_bank_loss_cache.py      | 13 ++++++-----
 tests/test_multi_query_validate.py |  9 +-------
 tests/test_per_token_lds.py        | 10 ++-------
 6 files changed, 25 insertions(+), 74 deletions(-)

diff --git a/bergson/data.py b/bergson/data.py
index bf5f3418..caa21a58 100644
--- a/bergson/data.py
+++ b/bergson/data.py
@@ -640,40 +640,13 @@ def to_grid(self) -> torch.Tensor:
         return scores[..., 0] if self.num_scores == 1 else scores
 
 
-class ArrayScores:
-    """Dense ``[num_items, num_scores]`` score matrix with the same interface
-    as :class:`Scores`, for scores saved as a plain ``.npy`` (e.g. an ad hoc
-    array from outside the standard scoring pipeline)."""
-
-    def __init__(self, arr: np.ndarray):
-        if arr.ndim == 1:
-            arr = arr[:, None]
-        self.arr = arr
-        self.num_scores = arr.shape[1]
-
-    def __len__(self) -> int:
-        return self.arr.shape[0]
-
-    def __getitem__(self, key: Any) -> NDArray:
-        return self.arr[key]
-
-    def get(self, key: Any, score_idx: int = 0) -> NDArray:
-        return self.arr[key, score_idx]
-
-    def is_written(self) -> bool:
-        return True
-
-
-def load_scores(path: Path) -> Scores | ArrayScores:
+def load_scores(path: Path) -> Scores:
     """Load a score store written by the standard scoring pipeline.
 
-    A plain ``.npy`` array loads as :class:`ArrayScores`. Any other score
-    directory loads its ``scores.bin`` as :class:`Scores`, with ``offsets``
-    set when ``info["attribute_tokens"]`` marks it as a per-token store.
+    A score directory loads its ``scores.bin`` as :class:`Scores`, with
+    ``offsets`` set when ``info["attribute_tokens"]`` marks it as a per-token
+    store.
     """
-    if path.suffix == ".npy":
-        return ArrayScores(np.load(path))
-
     info_path = path / "info.json"
     with open(info_path, "r") as f:
         info = json.load(f)
diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py
index f5fe651c..78f33c44 100644
--- a/bergson/magic/cli.py
+++ b/bergson/magic/cli.py
@@ -260,9 +260,6 @@ def scores_are_per_token(score_path: str) -> bool:
             return False
         with open(info_path) as f:
             return bool(json.load(f).get("attribute_tokens", False))
-    if score_path.endswith(".npy"):
-        return False
-
     step_cfg = read_first_step_config(score_path)
     if step_cfg is not None:
         return cfg_attributes_tokens(step_cfg)
@@ -582,7 +579,7 @@ def worker(
             score_path = os.path.join(run_cfg.run_path, "scores.pt")
             torch.save(scores, score_path)
             print(f"Saved attribution scores to {score_path}")
-    elif os.path.isdir(score_path) or score_path.endswith(".npy"):
+    elif os.path.isdir(score_path):
         scores, multi_query = load_attribution_scores(score_path)
     else:
         scores = torch.load(score_path, map_location="cpu")
diff --git a/bergson/validate.py b/bergson/validate.py
index b681c12a..08cb3350 100644
--- a/bergson/validate.py
+++ b/bergson/validate.py
@@ -40,30 +40,23 @@
 
 
 def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
-    """Load attribution scores from a score directory, ``.npy``, or ``.pt`` file.
+    """Load attribution scores from a score directory or ``.pt`` file.
 
     Returns ``(scores, multi_query)``. Token score directories are per-token,
     with a query dimension when ``num_scores > 1`` (``[docs, seq_len,
-    queries]``); plain score directories and ``.npy`` arrays are per-document,
-    with one column per query. A 3-D ``.pt`` tensor is per-token per-query
-    ``[docs, seq_len, queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token
-    MAGIC scores are ``[docs, seq_len]`` and per-query MAGIC scores are
-    ``[docs, queries]`` — so the run config next to it decides: it is
-    per-query iff the run used ``query_method: none`` without attributing
-    tokens.
+    queries]``); plain score directories are per-document, with one column per
+    query. A 3-D ``.pt`` tensor is per-token per-query ``[docs, seq_len,
+    queries]``. A 2-D ``.pt`` tensor is ambiguous — per-token MAGIC scores are
+    ``[docs, seq_len]`` and per-query MAGIC scores are ``[docs, queries]`` —
+    so the run config next to it decides: it is per-query iff the run used
+    ``query_method: none`` without attributing tokens.
 
     Score directories are negated when their ``score_cfg.higher_is_better`` is
-    set, aligning them with the loss-diff convention. ``.npy`` files carry no
-    ``score_cfg`` and are loaded as-is: they must already be in the loss-diff
-    convention.
+    set, aligning them with the loss-diff convention.
     """
-    if os.path.isdir(score_path) or score_path.endswith(".npy"):
+    if os.path.isdir(score_path):
         loaded = load_scores(Path(score_path))
-        score_cfg = (
-            load_subconfig(score_path, "score_cfg", ScoreConfig)
-            if os.path.isdir(score_path)
-            else None
-        )
+        score_cfg = load_subconfig(score_path, "score_cfg", ScoreConfig)
         negate = score_cfg is not None and score_cfg.higher_is_better
 
         if isinstance(loaded, Scores) and loaded.offsets is not None:
diff --git a/tests/test_bank_loss_cache.py b/tests/test_bank_loss_cache.py
index f72985a4..a955154f 100644
--- a/tests/test_bank_loss_cache.py
+++ b/tests/test_bank_loss_cache.py
@@ -18,6 +18,7 @@
 import bergson.validate as validate
 from bergson.config.config import DataConfig
 from bergson.magic.config import MagicConfig
+from bergson.score.score_writer import save_sequence_scores
 from bergson.validate import bank_loss_cache_key, evaluate_retrained
 
 MODEL = "trl-internal-testing/tiny-Phi3ForCausalLM"
@@ -111,8 +112,8 @@ def test_evaluate_retrained_reuses_cached_bank_losses(
     cols = n_query if multi_query else 1
     rng = np.random.default_rng(0)
     scores = rng.standard_normal((NUM_DOCS, cols)).astype(np.float32)
-    score_path = tmp_path / "scores.npy"
-    np.save(score_path, scores)
+    score_path = tmp_path / "scores"
+    save_sequence_scores(score_path, scores)
 
     # First run: cold cache, evaluates the bank and writes the cache.
     cfg1 = _run_cfg(tmp_path, "run1", query_path)
@@ -150,8 +151,8 @@ def test_different_query_does_not_reuse_losses(tmp_path, model):
     root = _build_bank(tmp_path, model)
     query_path, n_query = _query_dataset(tmp_path)
     scores = np.random.default_rng(2).standard_normal((NUM_DOCS, n_query)).astype("f4")
-    score_path = tmp_path / "scores.npy"
-    np.save(score_path, scores)
+    score_path = tmp_path / "scores"
+    save_sequence_scores(score_path, scores)
 
     cfg_a = _run_cfg(tmp_path, "run_a", query_path)
     evaluate_retrained(cfg_a, str(root), score_path=str(score_path))
@@ -170,8 +171,8 @@ def test_evaluate_retrained_averages_over_dirs(tmp_path, model):
     query_path, _ = _query_dataset(tmp_path)
 
     scores = np.random.default_rng(0).standard_normal((NUM_DOCS, 1)).astype(np.float32)
-    score_path = tmp_path / "scores.npy"
-    np.save(score_path, scores)
+    score_path = tmp_path / "scores"
+    save_sequence_scores(score_path, scores)
 
     for name, dirs in [
         ("one", str(bank_a)),
diff --git a/tests/test_multi_query_validate.py b/tests/test_multi_query_validate.py
index da3b70fc..ad81b007 100644
--- a/tests/test_multi_query_validate.py
+++ b/tests/test_multi_query_validate.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 import torch.nn.functional as F
 from datasets import Dataset
@@ -72,13 +71,7 @@ def test_load_attribution_scores_single_column_dir(tmp_path):
     assert scores.shape == (4, 1)
 
 
-def test_load_attribution_scores_npy_and_pt(tmp_path):
-    npy_path = tmp_path / "scores.npy"
-    np.save(npy_path, np.zeros((5, 2), dtype=np.float32))
-    scores, multi_query = load_attribution_scores(str(npy_path))
-    assert multi_query
-    assert scores.shape == (5, 2)
-
+def test_load_attribution_scores_pt_without_config(tmp_path):
     # 2D tensors in .pt files are per-token MAGIC scores, never multi-query.
     pt_path = tmp_path / "scores.pt"
     torch.save(torch.zeros(5, 2), pt_path)
diff --git a/tests/test_per_token_lds.py b/tests/test_per_token_lds.py
index 920bea5d..f5fe33ec 100644
--- a/tests/test_per_token_lds.py
+++ b/tests/test_per_token_lds.py
@@ -82,8 +82,8 @@ def test_load_token_dir_scatters_packed_to_grid(tmp_path):
 
 def test_scores_are_per_token_inference(tmp_path):
     """``bergson validate`` infers per-token-ness from the scores themselves:
-    token dirs and 2-D ``.pt`` tensors are per-token; 1-D ``.pt``, ``.npy``,
-    and plain score dirs are per-document."""
+    token dirs and 2-D ``.pt`` tensors are per-token; 1-D ``.pt`` and plain
+    score dirs are per-document."""
     from bergson.magic.cli import scores_are_per_token
 
     token_dir = tmp_path / "token_dir"
@@ -103,12 +103,6 @@ def test_scores_are_per_token_inference(tmp_path):
     torch.save(torch.randn(6, 1), pt_col)
     assert not scores_are_per_token(str(pt_col))
 
-    import numpy as np
-
-    npy = tmp_path / "multi_query.npy"
-    np.save(npy, np.random.randn(6, 5))
-    assert not scores_are_per_token(str(npy))
-
     plain_dir = tmp_path / "plain_dir"
     plain_dir.mkdir()
     assert not scores_are_per_token(str(plain_dir))

From 3fdc70fb802f8e8c615924dea3f9b24290322191 Mon Sep 17 00:00:00 2001
From: Lucia Quirke <luciarosequirke@gmail.com>
Date: Thu, 6 Aug 2026 19:51:50 +0900
Subject: [PATCH 5/5] refactor(magic): inline the per-token config lookup

cfg_attributes_tokens had one caller left once _pt_scores_are_per_query
was inlined, and reaching across from magic.cli into validate for a two
key dict lookup bought nothing.

Drop the scores_are_per_token docstring with it: the branches say what
they read, and the function had none before this PR.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 bergson/magic/cli.py | 14 ++------------
 bergson/validate.py  |  6 ------
 2 files changed, 2 insertions(+), 18 deletions(-)

diff --git a/bergson/magic/cli.py b/bergson/magic/cli.py
index 78f33c44..47434c78 100644
--- a/bergson/magic/cli.py
+++ b/bergson/magic/cli.py
@@ -36,11 +36,7 @@
 from ..utils.logging import wandb_log_fn
 from ..utils.utils import get_device, get_device_index
 from ..utils.worker_utils import setup_data_pipeline
-from ..validate import (
-    cfg_attributes_tokens,
-    load_attribution_scores,
-    validate_scores,
-)
+from ..validate import load_attribution_scores, validate_scores
 from .config import MagicConfig
 from .data_stream import DataStream, pad_dataset_to_batch_size
 from .grad_accum import accumulate_grads
@@ -248,12 +244,6 @@ def compute_per_query_magic_scores(
 
 
 def scores_are_per_token(score_path: str) -> bool:
-    """Whether ``score_path`` holds per-token scores.
-
-    Score directories record it in ``info.json`` and ``.pt`` files in the run
-    config beside them. Without either -- an externally produced tensor --
-    fall back to the shape, where only 3-D is unambiguous.
-    """
     if os.path.isdir(score_path):
         info_path = os.path.join(score_path, "info.json")
         if not os.path.isfile(info_path):
@@ -262,7 +252,7 @@ def scores_are_per_token(score_path: str) -> bool:
             return bool(json.load(f).get("attribute_tokens", False))
     step_cfg = read_first_step_config(score_path)
     if step_cfg is not None:
-        return cfg_attributes_tokens(step_cfg)
+        return bool(step_cfg.get("attribute_tokens") or step_cfg.get("per_token"))
 
     scores = torch.load(score_path, map_location="cpu")
     return isinstance(scores, torch.Tensor) and (
diff --git a/bergson/validate.py b/bergson/validate.py
index 08cb3350..71c3805f 100644
--- a/bergson/validate.py
+++ b/bergson/validate.py
@@ -83,12 +83,6 @@ def load_attribution_scores(score_path: str) -> tuple[torch.Tensor, bool]:
     return scores, step_cfg.get("query_method") == "none"
 
 
-def cfg_attributes_tokens(step_cfg: dict) -> bool:
-    """``attribute_tokens`` from a serialized run config, honouring the
-    deprecated ``per_token`` alias that ``MagicConfig`` still accepts."""
-    return bool(step_cfg.get("attribute_tokens") or step_cfg.get("per_token"))
-
-
 def bank_loss_cache_key(
     run_cfg: ValidationConfig, multi_query: bool, num_subsets: int
 ) -> str: