[AmitMY]

Stacked on #82.

Validate the labels and the latent mask independently

#82 derives the labels (next_word_labels) and the latent mask (latent_block_mask) in the model. These tests check both against oracles written from the words alone, in plain Python loops, sharing no code with what they test:

• oracle_labels: each word's label is the next word of its sequence, an empty word after the last one, and none for PAD words or for the words of a shift block before its ShiftIn.
• oracle_mask: each word sees itself and the earlier words of its sequence, plus every word of its shift block (ShiftOut to ShiftIn).

The texts are packed as in training: translation-style, plain, two blocks in one text, and SignWriting.

New tests:

• test_labels_match_an_oracle_of_the_words: packed at 16, 40 and 128 words.
• test_latent_mask_matches_an_oracle_of_the_words: the dense FlexAttention mask equals the oracle.
• test_masked_attention_matches_sdpa: now 200 words (over one 128-word FlexAttention block) against SDPA with the oracle's mask. Before, it built its expected mask with the same formula as the code.
• test_packing_is_invisible: each text's per-byte losses are the same alone as packed with the others and with PAD words.

Do the tests catch bugs?

Each of these deliberate bugs fails at least one new test (7 of 7):

Deliberate bug	Caught
Shift blocks attend causally only	✓
Sequences see each other	✓
Words outside blocks form one block (block_ids >= 0)	✓
A sequence's last word predicts the next sequence's first	✓
No empty word after a sequence's last word	✓
ShiftIn has no label	✓
ShiftIn outside its block	✓

Tests only: no code change.

🤖 Generated with Claude Code

https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f

---

Note

Low Risk
Test-only PR that adds regression coverage for label and mask behavior; no runtime or training code paths change.

Overview
Adds independent oracle helpers in conftest (oracle_labels, oracle_mask) plus shared ORACLE_TEXTS (translation shift blocks, plain text, multi-block, SignWriting) so tests can assert packing, labels, and attention without reusing the implementation under test.

New coverage: test_labels_match_an_oracle_of_the_words (packed lengths 16/40/128) checks label_mask and next_word_labels against the oracle; test_latent_mask_matches_an_oracle_of_the_words densifies FlexAttention’s latent_block_mask and compares to oracle_mask; test_packing_is_invisible checks per-byte losses match when texts are run alone vs one packed batch with PAD.

Strengthened: test_masked_attention_matches_sdpa now builds batches via pack_words/collate_fn at 200 words and uses the oracle dense mask for SDPA instead of mirroring the mask formula in the test.

Tests only—no production changes.

^{Reviewed by Cursor Bugbot for commit 8aa3598. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 8aa3598. Configure here.}

[cursor]

Packing test ignores RoPE positions

Medium Severity

test_packing_is_invisible expects per-word losses to match when texts run alone versus concatenated, but the latent transformer applies RoPE at absolute packed indices. Later texts therefore receive different rotary phases than they do alone, so the comparison can fail on a correct model or pass only by loose tolerance rather than proving sequence isolation.

 

^{Reviewed by Cursor Bugbot for commit 8aa3598. Configure here.}