Port the native-grid SamudraMulti model layer onto main

[alxmrs]

Why

Jesse's SamudraMulti rebuild lives on codex/decoder-root-cause-report — 217 commits ahead of the merge-base, ~10k lines of src/ changes, no PR. main still ships the architecture those reports show is broken: a Perceiver encoder compressing fixed 3°×5° patches into one token, feeding a Perceiver IO decoder that has to learn spatial routing and channel rendering jointly.

	1° normalized MSE
main SamudraMulti, 1-step, full data	0.29469
Samudra 2 reference	0.0236
Native-grid rebuild, single-res 1°, lead 1	0.0205

We want that architecture on main so we can reproduce it, then take it to long rollouts and to the ¼° goal.

Why a port and not a merge

I tried the whole-branch merge first. It bottoms out on the data layer, where three designs are in flight at once:

Where	Design
main	DataSource + DatasetSpec + OceanData
codex/decoder-root-cause-report	CanonicalDataset + CanonicalReader
#823 (draft)	CanonicalSource + DataLayout + BatchPreprocessor

#823's own summary says it replaces "the overlapping DatasetSpec / CanonicalDataset / OceanData abstractions." Resolving ~900 lines of conflict onto CanonicalDataset would be deleted work the moment #823 lands.

The model layer turns out to be completely decoupled from that argument — zero references to CanonicalDataset, rust_data, or data_backend; it reaches the data pipeline only through GridContext, masks, and resolutions. So it ports independently, and the data-layer rewrite stays where it belongs, in #800 and #823.

What's here

• samudra_multi.py / base.py — encode → process → decode, plus reconstruct_once (zero-depth inverse) and latent_forecast (encode once, advance N steps in latent space, decode at the requested lead).
• modules/encoder.py — DirectPatchEncoder with native_projection: one learned vector per native input cell rather than per fixed patch. This is the fix for the encoder-bandwidth failure, where ½°→½° reconstruction retained only 49.6% of target high-wavenumber power even though input and output shared a grid. Also PatchMomentEncoder and CanonicalResampleEncoder for the coarse-latent line.
• modules/decoder.py — ResampleProjectionDecoder with project_before_resample: decode physical channels on the native grid, then transport each with its own wet mask. On an exact checkpoint swap that ordering removed 78% of the excess ½°→1° error while leaving both same-grid routes bit-identical.
• modules/augment_input.py — BoundaryEncoder (one time-aligned forcing state per processor call) and ProcessorGeometryConditioner (zero-initialized geometry sidecar).
• aggregator/validate/spatial.py — high-wavenumber power and patch-seam jump diagnostics. For the ¼° goal these matter more than MSE; aggregate MSE hid the velocity-spectrum failures in several of these experiments.
• utils/ctx.py — optional input_mask so interpolation weights renormalize per channel.
• stepper.py / utils/output.py — initialize_rollout plus a rollout_state carried across inference chunks, so a latent model never decode/re-encodes at a chunk boundary.
• config.py — PerceiverConfig / EncoderConfig / DecoderConfig / SamudraMultiConfig spliced onto main's, plus the Checkpointing / LayerCheckpointing split making "selective" a SamudraMulti-only mode.

Please look at

1. SamudraMultiConfig.build signature. Dropped static_data_for_corrector (correctors deleted in Delete correctors and static-data plumbing #822) and the unused tensor_map / normalize / dataset_spec params; CanonicalDataset → DataSource.
2. reject_selective_checkpointing added to SamudraConfig and SamudraMiniConfig so the widened Checkpointing literal can't reach a UNet that doesn't understand it.
3. @jder — a question for later, not a change here. Remove data "schedules" #783 removed dst and the "mix" schedule because "it was always a bit unclear what 'mix' schedules would do for autoregression." The four-route training that produced these numbers needs differently-resolved input/label pairs, so dst has to come back eventually. I think the latent-AR contract answers the original objection — mix was ambiguous under decode/re-encode, but under encode-once → roll latent → render at the requested grid it's well defined. Not in this PR; single-scale doesn't need it. Want your read before the four-route work.

Not here yet

train.py contract (freeze_model_parameters / frozen_model_prefixes, train_processor_depths cycling, validation_boundary_ablations), identity.py, model and train configs, and dst. None of it is needed to review this layer, and all of it is needed before any run reproduces a published number.

Testing

• uv run pytest -m "not manual and not cuda" — 358 passed, 2 skipped, 10 xfailed
• uvx pre-commit run --all-files — all hooks pass (ruff, ruff-format, mypy, schemas, reuse)

No training or eval run has been done against this port yet.

🤖 Generated with Claude Code

https://claude.ai/code/session_018b4Bk5YegprhJojXj6wbff

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0aeac9ec6e

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Reject native projection with the default decoder

When encoder.native_projection is set, this branch makes the encoder emit one latent cell for every native input cell, but SamudraMultiConfig.build can still leave the decoder as the default PerceiverDecoder. With the shipped/default patch extents (for example 3°×5° or 6°×10°), PerceiverDecoder.forward builds position encodings for patch cells (H/patch_h * W/patch_w) and adds them to native-cell tokens (H * W), so the first forward in this advertised native-grid mode hits a shape mismatch unless the user also knows to select a resampling/direct decoder. Please reject or auto-pair incompatible decoder modes at build time.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Raise instead of dropping mismatched correction masks

When continuous_resample_attention_residual receives a valid mask whose spatial shape does not match the processor grid (for example, a physical-grid ctx.input_mask passed while the latent grid is coarser), this fallback silently disables masking for the attention correction. That makes the residual attend to source cells the caller marked invalid instead of surfacing the bad mask, so the diagnostic/training result can be corrupted without an error; raise like the other correction path rather than replacing it with None.

AGENTS.md reference: AGENTS.md:L310-L310

Useful? React with 👍 / 👎.

[alxmrs]

Code review pass — two real defects fixed in 7f4963c

I ran a review over the port. Two findings were genuine defects in what landed here (both caused by the model layer's defaults assuming a training contract this PR doesn't bring), and both were failing silently, which the repo guidance explicitly forbids. Fixed:

1. Inference computed a different function than training. initialize_rollout/inference took the latent-autoregressive path whenever a BoundaryEncoder existed — which SamudraMultiConfig.build creates for every non-bypassed config. So samudra eval encoded once and rolled the latent, while training still decoded and re-encoded every step through forward_once. Confirmed numerically: ~10% max-abs divergence over three steps on a freshly built model, diverging structurally rather than just numerically. Train samudra_multi_om4 then eval, and the zarr output came from a computation the weights were never fit for.

Now gated behind an explicit latent_autoregression flag, default false, so inference reproduces the step function training actually used. It flips on together with the latent-depth training contract in the follow-up, and that's the natural place to assert the two agree.

2. project_before_resample silently did unmasked interpolation. decode() reads GridContext.input_mask, but nothing populates it — both dataset constructors build a context without one. So setting decoder.project_before_resample: true (documented as preserving per-channel wet masks across resolutions, and worth 78% of the excess ½°→1° error in the original experiments) got plain bilinear with land bleeding into ocean cells. No error, no warning. ResampleProjectionDecoder.__init__ even guards that project_before_resample requires coordinate_resampling "so the source validity mask can renormalize interpolation" — a guard for a mask that never arrived.

Now raises, but only when the decode actually transports between grids. Same-grid decoding reduces exactly to the learned 1×1 channel map, so the single-scale route stays usable while input_mask is unwired — which is what Phase 1 needs.

Three tests added covering both contracts and the same-grid exemption. 362 passed, pre-commit clean.

Known-and-deliberate, listed so review doesn't re-find them

• train.py never passes processor_depth, so physical_forecast_loss_weight, latent_teacher_loss_weight, processor_iterations != 1, and processor_residual are inert config today. They land with the training contract.
• NormalizedSpatialDiagnosticsAggregator is not registered in ValidateAggregator. The module is here; the wiring comes with the training contract.

@jder — three pre-existing issues on your branch, unchanged here

Not touched, since they behave differently in the context of your full branch and you should decide:

1. ResampleAttentionResidualDecoder.forward (decoder.py:830) lacks the mismatched-mask guard its sibling ContinuousResampleAttentionResidualDecoder.forward has. Latent today only because input_mask is unset; it raises the moment that's wired up with a patch encoder.
2. ContinuousCoordinateAttentionCorrection._routing (decoder.py:650) hardcodes 180.0 / len(output_lat) and 360.0 / len(output_lon) for output spacing, then divides by source spacing derived from real coordinates. Within ~0.3% for global OM4 grids, wrong for any regional or non-uniform output grid.
3. ablate_boundary_forcing("batch_shuffle") (stepper.py:32) rolls the batch axis, so it's a no-op at batch_size: 1 — the shipped value in train_samudra_multi.yaml. It would report a shuffled-forcing ablation while feeding correct forcing.

Also patch_seam_jump_ratio returns NaN for patch_size == (1, 1) (every row and column is a seam, interior is empty), and NormalizedSpatialDiagnosticsAggregator.get_logs raises on a rank that saw zero batches while other ranks enter all_reduce_mean — a DDP deadlock rather than a clean failure. Both matter once the aggregator is registered.

[alxmrs]

Second review pass — two more fixed in 1b48790, one finding doesn't reproduce

Same theme as the last round: unported configuration accepted and quietly ignored instead of refused.

Fixed: forward(processor_depth=...) dropped pred_residuals. The depth path returns the decode directly, but BaseModel.forward adds prog_tensor back when pred_residuals is set — so depth training would fit an absolute-state function while inference rolls out a residual one. initialize_rollout already refuses that pairing; the training side now does too. Latent today (nothing passes a depth), but it's the identical failure mode to the one fixed last round, and cheaper to guard than to debug later.

Fixed: physical_forecast_loss_weight / latent_teacher_loss_weight silently inert. Only read on the depth branch. Setting physical_forecast_loss_weight: 0.0 to train latent-only ran ordinary full-weight physical-loss training and reported nothing. Now rejected at config validation until train_processor_depths lands; the validator comes out with that change.

Does not reproduce: processor_iterations: 0 as a DDP hazard. Reported as ungraded parameters killing the first backward on multi-GPU. But forward_once raises whenever a BoundaryEncoder is present and the count isn't one, and one is present for every non-bypassed config — so it dies loudly on the first forward and never reaches a backward. With bypass_processor: true the processor is a parameterless nn.Identity. No hazard on either path, so no change.

Also correcting the record from round one: processor_residual was listed there as inert config. It isn't — process() applies the residual scale and forward_once calls process(), so it's live on the default training path. I didn't act on that one at the time.

Still deliberately unchanged (all previously flagged to @jder above): the ResampleAttentionResidualDecoder mask-shape asymmetry, patch_seam_jump_ratio returning NaN for 1×1 patches, and batch_shuffle being a no-op at batch size 1. On the first of those — the sibling silently drops a mismatched mask while this one raises. Given the repo's "don't swallow errors" guidance, raising is arguably the correct half of that asymmetry and the silent drop is the side worth revisiting. That's the author's call, not a drive-by change.

363 passed, pre-commit clean.

[alxmrs]

Review 1/N

[alxmrs]

Maybe we can wait to include these spectral metrics until we land #834, so we can re-use some of the spectral code used at eval time? WDYT?

[alxmrs]

It's possible that these val aggregators are complements to the new spectral metrics and thus should be structured as they are.

[alxmrs]

I think the aurora sources which we depend on may have something that calculates this. Worth a review to see if we can import something instead of reimplement.

[alxmrs]

Happy this is extracted.

[alxmrs]

Nice

[alxmrs]

Thanks — all addressed in 55cf656. Replies below, including two where I pushed back.

spatial.py — agreed, and I went further: removed it from this PR

You asked whether to wait for #834 and reuse the eval-time spectral code, then noted these might be complements that belong as they are. I think both readings are right, and the way to avoid guessing is to not ship it yet: NormalizedSpatialDiagnosticsAggregator is never constructed by ValidateAggregator, so it's dead code here. It comes back when it's actually wired into validation, by which point #834 has landed and the reuse question answers itself with the code in front of us.

My guess is they are complements — these run per-batch on normalized GPU tensors inside the training loop and DDP-reduce; #834's operate on physical-unit xarray over a whole rollout against observations. Different data, different lifecycle. But the FFT itself is genuinely shared and worth extracting once, and I'd rather make that call against real code.

Bonus: it moots two latent defects review found in that file — patch_seam_jump_ratio returns NaN for one-cell patches (exactly the native_projection config this PR adds), and get_logs raises on a rank that saw zero batches while its peers enter all_reduce_mean, deadlocking DDP rather than failing cleanly. Both need fixing before it's registered.

Aurora already has this — you were right, and I still don't want to swap it here

Checked it numerically. Our normalized log cell area equals aurora.model.posencoding.patch_root_area to 2.5e-5 (float32 rounding), correlation 1.0. Same formula, R²π(sin φ₁ − sin φ₂)(θ₁ − θ₂).

But the normalization is zero-mean/unit-RMS, so every constant factor — Earth radius included — cancels. Swapping buys code reuse, not correctness. And patch_root_area wants explicit corner arrays, which is most of our function's body, so it'd replace about two lines.

The version actually worth adopting is aurora.area.compute_patch_areas, and the reason is curvilinear grids: it takes 2-D lat/lon and computes real spherical-polygon areas, where ours assumes separable 1-D coordinates. That's the tripolar-grid direction. Doing it now needs a flip-and-meshgrid adapter (aurora requires latitude decreasing along rows; ours increases) and changes a numeric feature inside a port whose whole job is reproducing a published number. So: later, tied to the curvilinear work. The docstring now names compute_patch_areas as the replacement so nobody re-derives this.

Splitting the log-area function — done, though your YAGNI counter was correct

One caller, so on its own merits YAGNI wins. What tipped it: the seam falls exactly where both known future changes land — the aurora swap above, and the curvilinear rewrite. Isolating it makes each a one-function change instead of surgery inside a 40-line function. Verified bit-identical to the previous concat.

ProcessorGeometryConditioner — yes to both, docstring expanded

You have it exactly right. It owns a Conv2d because it holds learnable state, and that's what separates it from the concatenating helpers around it. Two details behind the design:

It adds rather than concatenates — geometry is projected to the processor's own width and summed into the latent. So processor input width doesn't depend on whether geometry is on, which is what lets the same weights be applied zero or N times in a row. Concatenation would change the width and break the repeated-application contract.

And the decoupling is the whole point, with evidence: adding position/scale directly to encoder output measurably hurt reconstruction, while supplying it once per processor call did not. Keeping geometry out of the representation the decoder must invert is why the frozen-inverse result holds. The zero-init makes turning it on an exact no-op, so it can't perturb an already-trained inverse.

decoder.py chunked loop — comment added

Explains that each output cell attends to a bounded neighborhood, that at quarter degree materializing all the logits at once is what exhausts memory, that chunking the query axis bounds peak memory without changing the result (each chunk's softmax is over its own neighborhood), and what the all-invalid branch is for — land-only neighborhoods would otherwise softmax over all -inf and produce NaN.

blocks.py / decoder.py:454

No action. Worth noting _normalization fixed a real pre-existing hole: NormType allowed "layer" but ConvNeXtBlock raised NotImplementedError for it.

363 passed, pre-commit clean.

[alxmrs]

Self review, mostly understanding.

[alxmrs]

This type of reflection is not helpful in documentation.

[alxmrs]

This class is very similar to the one below it. Are both needed? Could they be abstracted or combined into one?

[alxmrs]

Why do we do this?

[alxmrs]

Maybe this should go in augment_input? Where did the old function that preceded it live?

[alxmrs]

What is the shape of the latent_resolution tensor? is it a tensor of integers?

[alxmrs]

Is "mix" schedule even still in this codebase? I thought it was removed. If so, we should delete this comment.