A 4D (b, t, x, c) analog of physicsnemo.nn.DiTBlock for field sequences: keeps
the DiT template (pluggable spatial-attention backend via get_attention,
adaLN-Zero conditioning, drop-path, gated MLP) and adds two optional gated
sub-layers:
- temporal (video) attention across time, with a pluggable time<->space reshard
  for context parallelism (manual all-to-all or ShardTensor.redistribute);
- observation cross-attention (pixel latents attend to packed per-pixel obs
  tokens via the vendored PixelCrossAttention).

New modules: temporal_attention.py (TemporalAttention + RoPE + causal/window
mask), sharding.py (shard_x/shard_t all-to-all + ShardTensor variants),
obs_packing.py (AttentionPacking/PixelGroupMap, copied as-is). Tests cover the
CPU plain + temporal paths and the CUDA full spatial+obs+temporal path
(forward/backward + grad flow).

Commit message authored by AI

Replace the _pixel_attn_kernels() helper + _k alias with an inline
`from . import _pixel_attn_kernels as kernels` inside the kernel-launch
wrappers, matching PNM's lazy optional-backend pattern (cf.
mesh/visualization/draw_mesh.py importing _matplotlib_impl/_pyvista_impl).

Commit message authored by AI

Bundle the observation tokens and ragged packing metadata into one
ObsCrossAttention object (tokens + cu_seqlens_k + max_seqlen_k + group_map)
so VideoDiTBlock's obs sub-layer takes a single argument instead of separate
obs_tokens + packing args. Drop the redundant data-pipeline AttentionPacking
struct (its counts/npix/hpx_level/pixel_order/is_packed fields are unused by
the model). Add jaxtyping Float/Int shape hints across the block, temporal
attention, sharding, and obs-packing modules.

Commit message authored by AI

Use the released physicsnemo.models.dit shape names: the token feature axis is
'hidden_size' (not 'channels'/'dim') and conditioning is 'condition_embed_dim',
matching DiTBlock's annotations and our own hidden_size constructor arg.

Commit message authored by AI

Add the ragged pixel cross-attention test suite (16 tests) validating the
Triton kernel forward/backward against a readable PyTorch GQA reference
(_ragged_gqa_reference), plus packed-grid, small-pixel grouping (grouped ==
ungrouped bit-for-bit), nn.Module wiring, empty-tokens DDP-safety, and config
validation. Port build_pixel_group_map (the CSR small-pixel grouping helper,
pure function of cu_seqlens_k) into obs_packing.py to support the grouping test.

Commit message authored by AI

A HEALPix field-sequence diffusion transformer over (B, C, T, npix): reuses the
existing HEALPixPatchTokenizer/HEALPixPatchDetokenizer (which already fold the
time axis and add the calendar embedding) and EDM conditioning, reshapes the
flat token sequence to (B, T, X, hidden) for the VideoDiTBlock stack (spatial +
optional factorized temporal + optional observation cross-attention), then back
for the detokenizer. Observations enter as a prebuilt ObsCrossAttention bundle.

Tests: CPU dense+temporal and CUDA dense+temporal+obs forward/backward.

Commit message authored by AI

- VideoDiT no longer hardcodes HEALPix: it takes a pluggable tokenizer /
  detokenizer (the grid lives only in tokenization; the backbone is grid-
  agnostic) plus time_length, threading tokenizer_kwargs and reshaping the flat
  time-major token sequence to (B, T, X, hidden) for the blocks.
- VideoDiTBlock's spatial+MLP path now mirrors physicsnemo.nn.DiTBlock: renamed
  emb_channels -> condition_embed_dim and added the DiTBlock dropout args
  (attn_drop_rate, proj_drop_rate, mlp_drop_rate, final_mlp_dropout), wired into
  get_attention and the MLP. Its args are now a superset of DiTBlock's.

Commit message authored by AI

…pose VideoDiT

- VideoDiTBlock subclasses physicsnemo.nn.DiTBlock, reusing its spatial
  attention, gated MLP, pre-norms, 6-chunk adaLN-Zero and drop-path; adds
  optional gated temporal + observation cross-attention sub-layers.
- is_causal and obs/temporal config moved to __init__; obs args behind
  obs_kwargs, temporal behind temporal_kwargs (explicit Dicts, no **kwargs).
- VideoDiT inherits physicsnemo.Module and composes the conditioning embedder +
  pluggable tokenizer/detokenizer + blocks (no production DiT change).
- sharding: @torch._dynamo.disable on the ShardTensor reshard.

Commit message authored by AI

…-attention

Generalize the experimental video DiT into a grid-agnostic DiT with a time axis:
- add a shared ndim-agnostic AdaLayerNormZero with a zero_init toggle (using
  get_layer_norm) and SiLU inside the module;
- compose VideoDiTBlock from shared building blocks instead of subclassing
  DiTBlock, moving modulation into per-sub-layer AdaLayerNormZero;
- replace the obs-specific cross-attention with a generic pluggable
  CrossAttentionModuleBase slot + opaque context; PixelCrossAttention is the
  reference impl and now owns the fold/ragged-unpack;
- make VideoDiT a kwarg-superset of DiT (drop_path_rates, conditioning_embedder
  choice, attn_kwargs, block_kwargs, dit_initialization).

Commit message authored by AI

…tion

…oss-attn, adaLN naming

- VideoDiT/VideoDiTBlock: cross_attention is a per-block factory (no deepcopy);
  forward context typed Optional[Any]; conditioning resolution inlined.
- Time axis first-class: HEALPix tokenizer gains separate_time_axis; detokenizer
  infers flat-vs-time-first from input rank (backwards-compatible with the v1
  flat + time_length path).
- adaLN-Zero attrs renamed role-first: attn_norm / temporal_attn_norm /
  cross_attn_norm / mlp_norm.

Commit message authored by AI

… Triton kernel

Port the v2 FiLM-conditioned observation tokenizer as an initial drop-in:
- obs_film_tokenizer.py: ObsTokenizerFiLM module + pure-PyTorch reference,
  custom-op wrappers, and the fused_film_tokenizer_triton entry point. Forward
  dispatches to the Triton kernel on CUDA when triton is available, else the
  reference path.
- _film_kernels.py: private fused FiLM forward/backward Triton kernels, guarded
  by OptionalImport("triton") (no bare import) and imported lazily.
- test_obs_film_tokenizer.py: CPU reference smoke test plus CUDA Triton-vs-
  reference parity tests.

Settings are kept as-is for now (TODO(polish) markers on unused ones).

Commit message authored by AI

DiT.__init__ indexed input_size[1] unconditionally, which IndexErrors for
non-2D tokenizers (e.g. HEALPix, where HealDA passes a 1-tuple input_size).
The latent grid is only consumed by the NATTEN backends, so guard it.

Commit message authored by AI

Match the public module stem (obs_film_tokenizer.py), mirroring
pixel_cross_attention.py / _pixel_attn_kernels.py.

Commit message authored by AI

…kenizer

obs_film_tokenizer.py -> obs_tokenizer.py, _obs_film_kernels.py ->
_obs_tokenizer_kernels.py (class stays ObsTokenizerFiLM; "FiLM" distinguishes
the impl from the existing ObsTokenizer and leaves room to evolve).

Commit message authored by AI

Compose the upstreamed VideoDiT backbone with the FiLM ObsTokenizerFiLM and
per-block PixelCrossAttention into the production v2 video+obs data-assimilation
architecture (hidden 1536, 16 heads, 32 layers, time_length 8, linear causal
temporal attention, drop-path 0.0 for the first 4 blocks then 0.1), so the
existing healda checkpoint can be loaded. HealDAv2 hosts the FiLM obs tokenizer
and assembles the per-pixel ObsCrossAttention context the backbone consumes.

Commit message authored by AI

…tidy wording

- VideoDiT.set_context_parallel(mode, target) fans the temporal time<->space
  reshard config out to every block (the per-block setter had no caller, so CP
  was never actually enabled through the model).
- VideoDiT.forward asserts the tokenizer emits 4D (B,T,X,hidden).
- Drop the coined "field sequence(s)" phrasing from docstrings/comments.

Commit message authored by AI

…ntract

Make the base a general cross-attention sub-layer rather than one "injected
into a video DiT block": forward takes (*batch, hidden_size) latents and an
opaque context: Any. Drops the field-sequence wording and trims the docstring.

Commit message authored by AI

…ocks=2)

Replace the separate attn_norm + mlp_norm AdaLayerNormZero(n_blocks=1) pair with
a single norm1 = AdaLayerNormZero(n_blocks=2) that emits the spatial-attention
modulation plus the raw MLP shift/scale/gate, and a parameter-free LayerNorm
MLP pre-norm modulated by those. This matches the DiT/diffusers layout and the
production checkpoint's single norm1.linear. temporal_attn_norm/cross_attn_norm
stay one-block adaLNs; initialize_weights now zeroes norm1.

Commit message authored by AI

Rename the packed-observation container to ObsContext and make tokens optional
(unset until the observation tokenizer fills it). PixelCrossAttention.forward
now consumes an ObsContext, raising if tokens is unset; a None group_map keeps
using the ungrouped ragged path (the model layer does not build it). Updates
all references and tests.

Commit message authored by AI

…orward

ObsContext now also carries the raw per-observation arrays (values,
float_metadata, obs_type, channel, platform) alongside the ragged packing, so
HealDAv2.forward takes one obs: ObsContext instead of loose per-obs and packing
args. The forward runs the tokenizer, fills tokens via dataclasses.replace, and
passes the context through; it no longer builds the pixel group map (a None
group_map uses the ungrouped ragged path). build_pixel_group_map stays for
callers that precompute it.

Commit message authored by AI

…mporal RoPE

Replace the hand-rolled RotaryPositionEmbedding with PhysicsNeMo's
nn.module.rope.RotaryPositionEmbedding1D (math-verified equivalent: same
interleaved-pair rotation and theta^(-2k/d) schedule). The temporal q/k are
transposed so the time axis is the -2 dim the module rotates, then restored.
Its cos/sin are non-persistent buffers recomputed at init, so the old
persistent rope.freqs_cos/sin keys no longer exist.

Commit message authored by AI

…rim docstring

Remove the unused nchannel, nplatform, and use_global_channel_platform_ids
constructor params (the embedding tables are always sized GLOBAL_MAX_* and id
spaces are the caller's responsibility). Trim the editorial rationale and
conv/sat first-linear essays in the class docstring to terse facts.

Commit message authored by AI

…meter-free

HealDAv2 passes attn_kwargs={"qk_norm_type":"RMSNorm","qk_norm_affine":False}.
PhysicsNeMo's TimmSelfAttention declares those exact kwargs (translating them to
timm's qk_norm + RmsNorm norm_layer), so affine-free RMSNorm actually engages and
the names are not silently swallowed. Add a regression test that the spatial q/k
norm modules exist (not Identity) and have no learnable parameters.

Commit message authored by AI

Lightweight audit cleanups: jaxtyping annotations on the two HEALPix hpx
tokenizer forwards, a descriptive CalendarEmbedding ValueError message, an
accurate uneven-shards comment in sharding.py, a Literal["apex","torch"] hint on
AdaLayerNormZero.layernorm_backend, and removal of the dead VideoDiTBlock
self.hidden_size attribute. The hpx __init__ **kwargs safety nets are kept.

Commit message authored by AI

…n; tidy HealDAv2 wiring

Drop the per-block hot-path validation from PixelCrossAttention.forward (tokens-set
and cu_seqlens-vs-pixel-count) and instead validate the packing's structural shape
once in ObsContext.__post_init__; tokens-set is guaranteed by HealDAv2.forward.
Build the HEALPix (de)tokenizers as locals instead of inline in the VideoDiT call,
and note the drop-path zero-first stability rationale in the docstring.

Commit message authored by AI

…packing->obs_context

Extract the kernel-companion packing primitives into a single
pixel_attention_utils.py: sort_and_pack (inlined Triton counting-sort kernel,
guarded by triton availability, argsort fallback), counts_to_cu_seqlens, and
build_pixel_group_map. These operate on plain index/count tensors, so the data
pipeline builds the packing the model only consumes. obs_packing.py is now purely
the ObsContext + PixelGroupMap contract, renamed to obs_context.py. Also build
HealDAv2's cross-attention via functools.partial and keep the cross-attention
base docstring generic. Adds CPU parity tests for the packing utils.

Commit message authored by AI

…pers

Split the adaLN-Zero op into a projection module and two apply helpers so
every sub-layer composes the same way and the pieces can be reused across
DiT/video blocks (3D and 4D states):

- adaln.py: AdaLNModulation (c -> 3*n_blocks shift/scale/gate chunks) plus
  standalone modulate() and gated_residual(); the affine-free pre-norm now
  lives at the call site instead of inside the module.
- video_dit_block.py: rewire to the uniform norm -> modulate -> gated_residual
  pattern (no block-0 special case); projections named norm1_modulation /
  temporal_attn_modulation / cross_attn_modulation, each attention sub-layer
  owning its own parameter-free pre-norm (attn_norm / mlp_norm /
  temporal_attn_norm / cross_attn_norm).
- video_dit.py: update the adaLN docstring reference.
- test: track the norm1_modulation rename.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…te _kernels modules

Consolidate each op's Triton device code and its host-side glue (launch
dispatch, autograd Function, custom-op registration) into the op's private
_*_kernels module, leaving the model files as just the nn.Module + public API:

- _pixel_attn_kernels.py / _obs_tokenizer_kernels.py: now own the full Triton
  op stack instead of only the @triton.jit kernels.
- pixel_cross_attention.py: slimmed to the PixelCrossAttention module;
  pixel_attention_reference -> _pixel_attention_reference (now private).
- obs_tokenizer.py: slimmed to the ObsTokenizerFiLM module.
- __init__.py: add package docstring + __all__; export PixelCrossAttention and
  ObsTokenizerFiLM alongside the existing public classes.
- pixel_attention_utils.py: docstring trim.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The packing helpers (sort_and_pack, counts_to_cu_seqlens, build_pixel_group_map,
counting_sort_and_pack + its counting-sort Triton kernel) build the ragged
layout that pixel_attention consumes and are used only by pixel cross-attention,
so co-locate them in pixel_cross_attention.py and drop the separate module.
Repoint the two tests and update the obs_context doc reference.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…2-arch

…ntext fields

- Consolidate Triton backends (pixel-attention GQA + counting sort, FiLM
  tokenizer, autotune cache) into a new kernels/ subpackage. Move
  pixel_attention() dispatch and the counting-sort kernel out of
  pixel_cross_attention.py into kernels/pixel_attention.py; move the
  sort_and_pack/counts_to_cu_seqlens/build_pixel_group_map packing utilities
  into obs_context.py.
- Fix ObsContext's raw observation fields (obs, float_metadata, obs_type,
  channel, platform) to be required, matching the original healda
  UnifiedObservation contract, instead of incorrectly Optional; rename
  values -> obs to match upstream. Update HealDAv2.forward's obs param to
  obs_ctx to avoid an obs.obs collision.
- Add full named_parameters() gradient-coverage checks to the HealDAv2 tests
  (previously only spot-checked a single parameter), matching the coverage
  already used for PixelCrossAttention/ObsTokenizerFiLM.
- Simplify test ObsContext builders: drop unnecessary dataclasses.replace
  indirection (tokens can be passed directly, per ObsContext's own contract)
  and inline single-use cross-attention factories.

Commit message authored by AI

Guard HealDAv2, VideoDiT, ObsTokenizerFiLM, PixelCrossAttention,
VideoDiTBlock, and TemporalAttention forwards with torch.compiler
checks. Rename fused Triton helper to _fused_film_tokenizer_triton.

Commit message authored by AI

Add CI-tested Examples to HealDAv2 and VideoDiT (correct obs packing
order and group_map), fix healda cross-refs, export VideoDiT, store
meta_dim on ObsTokenizerFiLM, and restore the fused Triton helper doc.

Commit message authored by AI

…d DiT naming

- Rename forward's noise_labels -> t and add the missing class_labels
  parameter (threaded through to VideoDiT/DiT as condition), matching
  v1 HealDA.forward and physicsnemo.models.dit.DiT.forward naming.
- Rename HealDAv2's emb_channels -> condition_embed_dim to match v1's
  naming and stop colliding in meaning with the detokenizer's own
  condition_dim.
- Add prepare_obs_context() to build a packed ObsContext directly from
  raw per-observation arrays, and export it (replacing the internal
  PixelGroupMap export) from the healda package.
- Tidy HealDAv2/VideoDiT docstrings for conciseness and accuracy.

Commit message authored by AI

Consolidate cross_attention, temporal_attention, and pixel_cross_attention
into one module and update imports/doc refs. Tidy obs_context module docs.

Commit message authored by AI

Reorder so TemporalAttention leads, then a section separator, then the
cross-attention contract and PixelCrossAttention.

Commit message authored by AI

…ts tests

- Move VideoDiTBlock into video_dit.py and delete video_dit_block.py, per
  review feedback to reduce file count.
- Merge test_video_dit_block.py's tests into test_video_dit.py and delete
  the standalone file.
- Replace the real HEALPixPatchTokenizer/Detokenizer in test_video_dit.py
  with minimal linear-projection stubs, dropping the earth2grid test
  dependency entirely so VideoDiT's attention/temporal/cross-attention/
  conditioning wiring can be tested without it.

Commit message authored by AI

…2-arch

Persist the winning Triton @autotune tile config only when
HEALDA_PIXEL_ATTN_AUTOTUNE_CACHE_DIR points at a directory (unset = off, no
disk I/O). Drop the PHYSICSNEMO_CACHE_DIR default and HEALDA_PIXEL_ATTN_PREWARM
opt-out; the cache is orthogonal to Triton's compiled-kernel cache. Document the
behavior at the model/layer level so users need not read the kernel module.

Commit message authored by AI