[BangBang-afk]

Summary

Flags a known-invalid layout conversion in the atom cross-attention encoder (see #730) without changing behavior.

keys_ref_space_uid is converted from batch.ref_structure.ref_space_uid (token layout, shape (num_tokens, num_dense)) using batch.atom_cross_att.queries_to_keys, whose gather_idxs index the queries flat space (num_subsets, num_queries). The gathered values — and hence offsets_valid — are therefore not meaningful (in practice almost all False).

The two-step token -> queries -> keys conversion used for keys_ref_pos just above would be the correct pattern, but it is deliberately not applied here: the model was trained with this conversion, and changing it without retraining could deteriorate performance (per review discussion). The code now documents this with a NOTE(#730) comment at the conversion site so the next reader doesn't fix it blindly.

# NOTE(#730): batch.ref_structure.ref_space_uid is in token layout
# (num_tokens, num_dense) here, while queries_to_keys indexes the queries
# flat space (num_subsets, num_queries), so the gathered keys_ref_space_uid
# values -- and hence offsets_valid below -- are not meaningful (in practice
# almost all False). The two-step token -> queries -> keys conversion used
# for keys_ref_pos above would be the correct pattern, but this is left
# as-is intentionally: the model was trained with this conversion, and
# changing it without retraining could deteriorate performance.

Verification (no behavior change)

• Diff is comment-only: +9/-1, zero functional change, so inference outputs are bit-identical by construction.
• Evidence for the underlying claim, checked against source:
  • Batch.ref_structure.ref_space_uid shape (num_tokens, num_dense) (features.py:1685,1709).
  • AtomCrossAtt.queries_to_keys.input_shape is (num_subsets, num_queries); neighboring keys_mask/keys_single_cond/pair_cond_keys_input all gather from queries_* sources — this call is the sole inconsistency (confirmed by grep).
  • atom_layout.convert validates layout_shape[1:] == input_shape[1:] (atom_layout.py:955-961).
• No existing unit test covers this path (no atom_cross_attention test in repo); nothing to regress since behavior is unchanged.

Related to #730.

AI disclosure

Analysis assisted by AI and manually verified against atom_layout.convert and neighboring conversions.

[kilianmandon]

Given that the model was trained on the conversion as it was I think it might be problematic to just change the code. Without retraining, the change might deteriorate model performance. Maybe it would be a better idea to just flag the invalid conversion in the code?

[BangBang-afk]

Agreed -- good catch. The mask feeds straight into pair_act, so flipping it would shift activations away from what the weights saw in training. I reverted the behavior change; the PR now only adds a NOTE(#730) comment at the conversion site documenting the mismatch and why it is intentionally left as-is. Diff is comment-only (+9/-1), outputs bit-identical by construction.