🐛[BUG]: AeroJEPA ignores unbatched TokenSet masks and decode_field_chunked accepts invalid arguments #1998

Version

2.3.0a0 (main at 3bbbe86)

On which installation method(s) does this occur?

Source

Describe the issue

Two groups of input-handling problems in the experimental AeroJEPA model (physicsnemo.experimental.models.aerojepa). None of them raise an error; they give wrong results or cryptic failures.

1. The mask of an unbatched TokenSet is ignored. QueryTokenDecoder.forward drops masked-out target tokens, but three other components treat every token of an unbatched TokenSet as valid:

• PrototypeTokenJEPAHead builds an all-True context mask for unbatched input, so masked-out context tokens still take part in cross-attention. The batched path uses context_tokens.mask correctly.
• pad_token_sets sets mask[i, :count] = True and computes the fallback global token with masked_mean(features, None). Tokens that the input set marks invalid become valid after packing, and they also enter the global token.
• flatten_valid_token_features returns rank-2 features unchanged even when a mask is passed.

Expected: masked-out tokens are excluded in every case, as they already are in the batched path.

2. AeroJEPA.decode_field_chunked does not check its arguments.

• An unknown precision (for example "bfloat16" or "amp") falls back to fp32 without any warning. The docstring says "Anything other than "fp32" enables torch.autocast", but the code only enables autocast for "fp16" and "bf16".
• chunk_size <= 0 fails inside the decoder with torch.cat(): expected a non-empty list of Tensors.
• query_sdf is required, even though decode_field and the decoder accept None when use_sdf=False. Omitting it raises TypeError: 'NoneType' object is not subscriptable.
• The same goes for QueryTokenDecoder(query_chunk_size=0): the constructor accepts it, and forward then fails with range() arg 3 must not be zero.

Expected: invalid arguments raise a clear ValueError, and query_sdf is optional when the decoder does not use it.

Minimum reproducible example

import torch
from physicsnemo.experimental.models.aerojepa import (
    PrototypeTokenJEPAHead, QueryTokenDecoder, TokenSet,
)
from physicsnemo.experimental.models.aerojepa.layers import (
    flatten_valid_token_features, pad_token_sets,
)

head = PrototypeTokenJEPAHead(
    token_dim=16, cond_dim=0, depth=1, num_heads=2, neighbor_k=4, query_pe_bands=2
).eval()
feats, coords, targets = torch.randn(10, 16), torch.randn(10, 3), torch.randn(5, 3)
mask = torch.arange(10) < 6

with torch.no_grad():
    trimmed = head(context_tokens=TokenSet(feats[:6], coords[:6]), target_positions=targets)
    masked = head(context_tokens=TokenSet(feats, coords, mask=mask), target_positions=targets)
print(torch.allclose(trimmed, masked, atol=1e-5))                        # False

packed = pad_token_sets([TokenSet(feats, coords, mask=mask)])
print(int(packed.mask.sum()))                                             # 10, expected 6
print(flatten_valid_token_features(feats, mask).shape[0])                 # 10, expected 6

QueryTokenDecoder(token_dim=16, use_sdf=False, query_chunk_size=0)        # accepted

For decode_field_chunked, run test_decode_field_chunked_* from the linked PR against main.

Relevant log output

precision='bfloat16' : accepted, autocast enabled=False
chunk_size=0         : ValueError: torch.cat(): expected a non-empty list of Tensors
query_sdf=None       : TypeError: 'NoneType' object is not subscriptable
query_chunk_size=0   : ValueError: range() arg 3 must not be zero

Environment details

Bare-metal, source install (uv, editable), Python 3.12.3, torch 2.14.0+cpu, Linux aarch64