Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

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

Open
  • Is something wrong?
#1,998 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
68/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start with PrototypeTokenJEPAHead, QueryTokenDecoder, pad_token_sets, and flatten_valid_token_features in physicsnemo.experimental.models.aerojepa and its layers module; trace how unbatched masks and chunked decoder arguments are handled. Run the linked PR's test_decode_field_chunked_* tests, then add or run coverage for masked TokenSets, invalid precision and chunk sizes, and query_sdf=None until invalid inputs fail clearly and valid masks are preserved.

Written by the indexing model from the issue text.

Description

Version

2.3.0a0 (main at 3bbbe86a)

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
Dominant language
Python
Stars
3.3k
Forks
787
Avg merge
3d 4h
Merged PRs (30d)
28

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/physicsnemo

All issues in NVIDIA/physicsnemo

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.