Write scores as a score directory

[luciaquirke]

Stacked on #415 (base feat/magic-per-token-per-query). Merge that first and this retargets to main cleanly.

Problem

MAGIC wrote a bare tensor to scores.pt, so every consumer had to recover from the shape what a score directory simply states. A 2-D tensor is [docs, seq_len] or [docs, queries] depending only on how the run was configured — which is why load_attribution_scores had to read config.yaml, fall back to the rank when there was none, and why scores_are_per_token needed the same treatment.

And the tensor was never self-contained anyway. It already leaned on config.yaml to disambiguate its own axes and on a doc_ids.pt sidecar to map shuffled chunks back to documents. It wasn't a simpler format than a score directory — it was an undocumented one.

Fix

Write the same score directory the scoring pipeline writes. info.json records attribute_tokens and num_scores, so both questions are answered by the store, and the reader collapses to: load it, negate if higher_is_better, report num_scores > 1.

Deleted: the .pt branch of load_attribution_scores, the config lookup and shape fallback behind it, scores_are_per_token's two non-directory branches, and read_first_step_config, whose only caller was that lookup.

Why the ragged store fits exactly

The one thing that made this non-obvious. MAGIC's weight grid is dense [rows, seq_len]; the token store is ragged, length - 1 values per row. They coincide, and not by luck:

• weighted_causal_lm_ce reads example_weight[:, :-1], so a row's columns past length - 1 never receive gradient — precisely the rows compute_num_token_grads allocates.
• to_grid restores width max(length - 1) + 1, which is the seq_len the grid started with, for both chunk_length > 0 and unpacked data.

So packing drops only structural zeros and the round-trip is exact in shape and value. test_save_magic_scores_round_trips_the_grid asserts that for single- and multi-query.

doc_ids

Moves into the directory as scores/doc_ids.npy, in the grid's shape, so a per-token store carries everything needed to aggregate per document rather than depending on a sibling file.

Output layout

run_path/scores/
  info.json      attribute_tokens, num_scores, dtype
  scores.bin     memmapped scores
  offsets.npy    per-token only
  doc_ids.npy    per-token only

Verified end to end on CPU for all three modes — per-token (5, 8), per-doc (5, 1), per-token per-query (5, 8, 2) — each loading back at the shape MAGIC used to write, with multi_query and scores_are_per_token reported from info.json.

Testing

110 passed, 38 skipped across every test file touching magic, validate or scores. Tests that covered the removed .pt path are gone; those that read scores.pt now read the directory, including test_distributed_magic.py. docs/magic.rst "Output files" rewritten, sphinx clean. pre-commit --all-files clean.

🤖 Generated with Claude Code