Add support for StereoRoPE positional embedding and Strata features

[zyhu-hu]

PhysicsNeMo Pull Request

Description

Contributes the Strata weather-emulator architecture in two layers:

1. physicsnemo.nn — a stereographic 2D rotary position embedding primitive
   (the continuous-coordinate sibling of physicsnemo.nn's RotaryPositionEmbedding2D).
2. physicsnemo.experimental.models.strata — the DiT3D and PixelDiT
   regression models that consume it.

On the "DiT" name. "DiT" denotes the architecture family, not a training
objective. DiT3D is a conceptual 3D analog of physicsnemo.models.dit.DiT:
it follows the same DiT-style template (patch-embed → pre-norm transformer →
linear decode) but is an independent reimplementation, not a reuse of that
class's components — it defines its own 3D blocks and shares only the na3d
functional, Mlp, and the RoPE primitives. PixelDiT additionally keeps DiT's
defining adaptive-layer-norm (adaLN) conditioning — only the conditioning
signal differs (the semantic stage's features vs. a diffusion timestep) — and
its name also credits the PixelDiT paper (arXiv:2511.20645) it adapts.

Both are deterministic regression models — not generative diffusion
models: no diffusion process and no noise / timestep / class-label / text
conditioning (DiT3D carries no adaLN at all; PixelDiT's adaLN is driven by
the semantic stage, not a timestep).

Part 1 — StereoRoPE (physicsnemo/nn/module/rope.py, re-exported from physicsnemo.nn)

Token latitude/longitude are mapped to a local tangent plane via a stereographic
projection; the resulting continuous (x, y) coordinates drive a 2D RoPE that
reuses physicsnemo.nn's apply_rotary_pos_emb verbatim (the per-pair
rotation math is identical — no new rotation code). New public API:

• StereographicRotaryPositionEmbedding2D — project(lat, lon, length_scale, …)
  maps lat/lon to continuous tangent-plane coords; build_tables / forward(q, k, x_pos, y_pos)
  rotate the query/key tensors. length_scale is a required argument (no sensible
  default can be inferred from data).
• build_axial_rope_cos_sin_2d_continuous — continuous-coordinate generalization of
  build_axial_rope_cos_sin_2d (the integer grid is its special case, asserted by a
  test), built on a reusable build_rope_cos_sin_1d_continuous core.
• stereographic_projection — antipode-guarded tangent-plane projection.
• spherical_centroid — pole- and seam-robust tile center (3D unit-vector mean),
  used for project's default centering (correct at the poles, where a per-axis
  lat/lon mean is not).

Placement: a small, parameter-free primitive sharing the rotation convention
and channel layout of physicsnemo.nn's existing RoPE, so per MOD-002a's maturity
exception it lives in production
nn/module/rope.py (MOD-000a reusable layer), not experimental. It is a plain
torch.nn.Module, not a physicsnemo.Module model.

Part 2 — DiT3D / PixelDiT (physicsnemo/experimental/models/strata/)

• DiT3D — a 3D transformer backbone (its own 3D reimplementation of the
  DiT-style template, not a reuse of physicsnemo.models.dit.DiT's components):
  3D patch embedding, 3D neighborhood attention (reusing physicsnemo.nn.functional.na3d)
  with optional depth-axis and gated attention, optional axial or stereographic
  RoPE, activation checkpointing, and bf16 autocast (CUDA). set_tile_size re-tiles
  a trained model for inference at a new resolution (rebuilding the cached axial tables;
  the stereographic / none modes build per forward, so only the expected shape updates) —
  supporting regional / tiled-global rollout.
• PixelDiT — two-stage: a DiT3D semantic stage (coarse patches) conditions a
  pixel-resolution stage whose blocks inject the semantic tokens via pixel-wise
  adaptive layer norm (pixel_proj / bilinear_dw). bilinear_dw supports any
  semantic vertical patch size (trilinear depth upsample, with an exact 2D-bilinear
  fallback when pixel depth == semantic depth, so the common case is numerically
  unchanged). Pixel-stage and semantic-stage RoPE are independent; pixel blocks also
  support activation checkpointing. Adapts the pixel-wise-AdaLN idea from PixelDiT
  (arXiv:2511.20645) — an adaptation, not a faithful port: deterministic regression
  (no diffusion/timestep/label conditioning), an independent AdaLN reimplementation,
  with the original bilinear_dw conditioning path added beyond the paper.

Decoupled geometry: lat/lon is an optional forward input used only for
stereographic RoPE — no earth2grid / grid-library dependency, no hard spherical-grid
assumption. The token-coordinate builders are pure free functions in coords.py
shared by both stages (no cross-stage instance coupling). Per MOD-002a, new models
are introduced in experimental.

Tests

• test/nn/module/test_rope.py (17 new tests): continuous-builder ≡ integer-builder
  consistency, rotation handedness (a flipped rotate_half is caught), projection
  geometry + closed-form 2·tan(Δ/2), antipode finiteness, spherical_centroid
  pole/seam robustness, project centering + required length_scale, module
  shapes/validation/norm preservation, relative-position invariance, batched-coord
  broadcast, end-to-end lat/lon → project → forward, exact forward-vs-build_tables,
  physicsnemo.Module checkpoint round-trip, bf16 dtype preservation, theta wiring.
• test/experimental/models/strata/: constructor/attribute, validation, forward-shape
  across RoPE / AdaLN modes, non-regression goldens ({args, state_dict, y} reload —
  version-robust, CPU-reproducible on the attn_kernel=-1 SDPA path), physicsnemo.Module
  checkpoint round-trip, the NATTEN NA3D path on CUDA, gradient-flow (no dead params;
  PixelDiT semantic stage reached through the pixel stage), DepthwiseConv (chunked ≡
  plain, deepcopy uses its own weights, bias=False device move), pure-bfloat16/half
  cast of both models (SDPA path on CPU/CUDA, plus the real NA3D + depth-axis + gated
  attention path on CUDA), set_tile_size re-tiling (axial buffer rebuild + non-axial shape
  update, with the new tile accepted and the old one rejected), shape variation (incl.
  vertical patch > 1), torch.compile, bf16 autocast, activation checkpointing (DiT3D +
  PixelDiT, train-mode), PixelDiT RoPE in both stages independently, and the coords.py
  builders (with value-pinned stereographic geometry).
• About the CPU skips. On a GPU runner the strata suite reports 74 passed, 3 skipped. The three skips are test_dit3d_natten_forward[cpu],
  test_pixeldit_natten_forward[cpu], and test_dit3d_natten_bf16_cast[cpu] — they
  pytest.skip by design because NATTEN's neighborhood-attention kernels are
  CUDA-only (no CPU kernel exists), not because of any missing coverage or
  environment gap. Their [cuda] counterparts run and pass on a GPU, and the CPU side
  of these models is exercised by the dense SDPA fallback (attn_kernel=-1), which
  the goldens use for deterministic CPU reproduction. CI without a GPU will skip the
  NA3D cases identically.
• ruff, interrogate (100%), license headers, and import-linter (10/0) all clean.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• [na] An issue is linked to this pull request.
  author note: no existing issue is relevant to this PR
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards.

Dependencies

None new. 3D neighborhood attention uses physicsnemo.nn.functional.na3d (NATTEN —
already an optional physicsnemo extra), with a CPU SDPA fallback (attn_kernel=-1).
Geometry is an optional forward input, so there is no earth2grid / grid-library
dependency. import-linter's "Prevent Non-listed external imports" contract passes
unchanged.

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR contributes the Strata weather-emulator architecture: a StereographicRotaryPositionEmbedding2D primitive in physicsnemo.nn (continuous-coordinate sphere-aware RoPE using stereographic projection) and two deterministic regression models — DiT3D (a 3D patch-embedding + neighborhood-attention transformer) and PixelDiT (DiT3D semantic stage conditioning a pixel-resolution stage via pixel-wise AdaLN) — placed under physicsnemo.experimental.models.strata.

• rope.py: adds StereographicRotaryPositionEmbedding2D, build_axial_rope_cos_sin_2d_continuous, stereographic_projection, and spherical_centroid; all reuse the existing apply_rotary_pos_emb rotation kernel and are gated behind the existing module's head_dim % 4 validation.
• dit3d.py / pixel.py: full physicsnemo.Module models with ModelMetaData, DiT-style zero-init output heads, bf16 autocast, activation checkpointing, and optional axial or stereographic RoPE; DiT3D.set_tile_size rebuilds axial buffers for inference re-tiling.
• Tests: 17 new RoPE tests and a comprehensive strata suite (non-regression goldens, gradient flow, torch.compile, NATTEN, set_tile_size, DepthwiseConv correctness); all CI-skipped cases are by-design CUDA-only NATTEN paths.

Important Files Changed

Filename	Overview
physicsnemo/nn/module/rope.py	Adds StereographicRotaryPositionEmbedding2D, build_rope_cos_sin_1d_continuous, build_axial_rope_cos_sin_2d_continuous, stereographic_projection, and spherical_centroid. Well-documented, antipode-guarded, correctly composes with existing apply_rotary_pos_emb; no issues found.
physicsnemo/experimental/models/strata/dit3d.py	New DiT3D model; construction, shape validation, activation checkpointing, and set_tile_size are well-implemented. Minor: set_tile_size doesn't validate divisibility of new height/width by patch size when rope_mode=axial.
physicsnemo/experimental/models/strata/pixel.py	Two-stage PixelDiT model with pixel-wise AdaLN conditioning; activation checkpointing uses closed-over tensors (s_cond, s_cond_bilinear) rather than explicit checkpoint arguments, which works with use_reentrant=False but diverges from DiT3D's pattern.
physicsnemo/experimental/models/strata/layers.py	Natten3DSelfAttention, DiT3DBlock, PatchEmbed3D, FinalLayer3D building blocks; fp32 output heads, correct RoPE application, proper dtype handling. No issues.
physicsnemo/experimental/models/strata/depthwise_conv.py	DepthwiseConv with vmap-based chunked path correctly passes weight/bias as explicit arguments for deepcopy safety. Warning threshold uses byte count compared to 2^32 bytes, which is overly conservative for float32 (warns at 2^30 elements, actual limit ~2^31).
physicsnemo/experimental/models/strata/coords.py	Pure helpers for axial and stereographic token coordinate building; spherical_centroid used correctly for patch pooling and tile centring; no issues.
test/experimental/models/strata/test_dit3d.py	Comprehensive test suite covering construction, forward shapes, non-regression goldens, checkpoint round-trips, gradient flow, set_tile_size, bf16 casts, and activation checkpointing. No issues.
test/nn/module/test_rope.py	17 new tests for the stereographic RoPE additions; verifies consistency with the integer-grid builder, handedness, projection geometry, norm preservation, and dtype fidelity.

_{Reviews (1): Last reviewed commit: "Test pure bf16 cast on the real NA3D att..." | Re-trigger Greptile}

[greptile-apps]

Tensors captured by closure in activation_checkpoint calls

s_cond, s_cond_bilinear, pixel_dhw, semantic_dhw, and rope_tables are all closed over by _run rather than passed as explicit arguments to activation_checkpoint. With use_reentrant=False this is safe — PyTorch's non-reentrant implementation correctly handles closed-over tensors and the PR's gradient-flow tests validate it — but it differs from the DiT3D.forward_tokens pattern (where rope_tables is bound via a default arg r=block_rope). Passing at minimum s_cond and s_cond_bilinear as explicit arguments would align with the non-reentrant best practice and make the gradient path easier to audit.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[pzharrington]

@zyhu-hu I think this is reasonable ask, for code clarity

[pzharrington]

Still overly verbose, these details can be put in documentation rather than changelog. Just listing the added public-level modules and functions is fine

[pzharrington]

No need for a separate file, might as well put this in layers.py

[pzharrington]

Stale references to stereographic module should be removed, and there are only two RoPE modules now without it.

[pzharrington]

This happens before input validation and could lead to "too many values to unpack" error. Since forward_tokens guarantees the shape equals the configured one, you could just use self.depth/height/width for unpatchify, and drop the early unpack here.

[pzharrington]

This only checks for pos is None but not shape, and the backbone validates pos.shape only when the backbone's rope_mode="stereographic". So the backbone="none"/"axial", pixel="stereographic" config never shape-checks pos; a wrong (B,2,H,W) falls into a cryptic rearrange error in build_stereographic_token_coords. Worth a one-line shape check in Strata.forward.

[pzharrington]

Use jaxtyping or make private to the class

[pzharrington]

Use jaxtyping or make private to the class

[pzharrington]

Use jaxtyping annotations

[pzharrington]

Nothing here ever initializes the pixel blocks' attention and MLP Linears (they keep PyTorch's nn.Linear default), intended or no? It sorta contradicts the docstring claim of matching StrataTransformer3D.initialize_weights

[pzharrington]

No need to export the MetaData classes

[pzharrington]

Same here and in transformer.py, no need to export the MetaData, it's just a module-local helper

[pzharrington]

Since it's a block, maybe move it to layers.py?

[pzharrington]

/ok to test 78e69d3