[BUG] RFD3 unindexed motif loss uses an atom-level mask as a token-level mask

Describe the bug

In the RFD3 unindexed motif path, ground_truth["is_original_unindexed_token"] appears to be produced as an atom-level mask, but DiffusionLoss.forward consumes it as a token-level mask by indexing it with atom_to_token_map.

This can silently create an incorrect atom mask when the tensor length is neither the token count nor the current atom count.

To Reproduce

Observed in a debug/failing run with unindexed motif loss enabled:

loss_input["is_original_unindexed_token"]
# tensor([False, False, False, ..., False, False, False], device='cuda:0')
loss_input["is_original_unindexed_token"].shape
# torch.Size([1032])

tok_idx = network_input["f"]["atom_to_token_map"]
tok_idx.shape
# torch.Size([1799])
tok_idx
# tensor([0, 0, 0, ..., 145, 145, 145], device='cuda:0', dtype=torch.int32)

Here tok_idx is atom-level: length 1799 is the atom count, and tok_idx.max() + 1 == 146 is the token count. Therefore is_original_unindexed_token should be shape (146,) if it is token-level, or shape (1799,) if it is current atom-level. Its observed shape (1032,) is neither.

Relevant code on production at e412591:

models/rfd3/src/rfd3/transforms/conditioning_base.py appears to build the value as an atom-wise expanded mask:

is_unindexed_token = apply_and_spread_token_wise(
    atom_array,
    atom_array.is_motif_atom_unindexed.copy(),
    function=lambda x: np.any(x),
)

atom_array_expanded = self.expand_unindexed_motifs(...)
n_expanded_atoms = atom_array_expanded.array_length() - atom_array.array_length()
mask = np.concatenate([is_unindexed_token, np.zeros(n_expanded_atoms)])
data["ground_truth"]["is_original_unindexed_token"] = mask.astype(bool)

Then models/rfd3/src/rfd3/metrics/losses.py treats the tensor as token-level:

tok_idx = network_input["f"]["atom_to_token_map"]
is_original_unindexed_token = loss_input["is_original_unindexed_token"][tok_idx]

Because tok_idx contains token ids, this line expects loss_input["is_original_unindexed_token"] to be indexed by token id. If the provided mask is atom-level with length 1032, the loss silently interprets the first 146 atom positions as token flags.

Expected behavior

is_original_unindexed_token should have a single clear level:

• token-level shape (n_tokens,), then expanded in the loss with [tok_idx]; or
• atom-level shape (n_atoms,), then the loss should not index it with tok_idx.

Given the name and current loss code, token-level seems intended.

Screenshots

Not applicable.

Additional context

This affects unindexed motif loss weighting/masking, including w_L, the coordinate mask normalization, and the extra unindexed diffused LP norm term. The failure mode is especially easy to miss because the bad shape can still be large enough to avoid an IndexError.

A possible fix is to store a true token-level mask in the conditioning output, for example by sampling the spread atom mask at token starts:

is_unindexed_atom = apply_and_spread_token_wise(
    atom_array,
    atom_array.is_motif_atom_unindexed.copy(),
    function=lambda x: np.any(x),
)
is_original_unindexed_token = is_unindexed_atom[get_token_starts(atom_array)].astype(bool)
data["ground_truth"]["is_original_unindexed_token"] = is_original_unindexed_token

It would also help to add a guard in DiffusionLoss.forward:

n_tokens = network_input["f"]["is_polar"].shape[0]
assert loss_input["is_original_unindexed_token"].shape == (n_tokens,)

and a regression test that checks this mask length equals the token count after unindexed motif expansion and virtual atom handling.

Note: I initially filed this as #329 with the blank issue template by mistake; reposting here using the repository's bug report template.