Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

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

Aperta
  • Qualcosa non va?
#1,998 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

@harshaa765 ci sta già lavorando.

Dal 17/9/2026.

  • #1999 di @harshaa765 — aperta

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
68/100
Tipo di issue
Bug
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
machine-learning

Direzione di ricerca

Inizia con PrototypeTokenJEPAHead, QueryTokenDecoder, pad_token_sets e flatten_valid_token_features in physicsnemo.experimental.models.aerojepa e nel suo modulo layers; traccia come vengono gestite le maschere non batchate e gli argomenti del decoder chunked. Esegui i test test_decode_field_chunked_* del PR collegato, quindi aggiungi o esegui la coverage per TokenSets mascherati, valori di precision e dimensioni dei chunk non valide, e query_sdf=None, finché gli input non validi non falliscono chiaramente e le maschere valide vengono preservate.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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
Lingua principale
Python
Stelle
3.4k
Fork
824
Merge medio
3g 4h
PR unite (30g)
70

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/physicsnemo

Tutte le issue di NVIDIA/physicsnemo

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.