Decompose the v2 SFNO at encoder/processor/decoder cut-points (translate 5/6)

[mcgibbon]

The latent-splice transfer arm needs a donor's SFNO processor frozen between
learned transforms that replace the donor's physical-variable interface, and
the multi-scale warm-start needs all three of a composite's stages initialized
from a plain ACE checkpoint. Both need the monolithic noise-conditioned SFNO
opened up into separately-buildable, separately-freezable pool components whose
parameters keep the donor's state_dict names, so any subset of the three can
be name-matched onto a donor checkpoint.

sfno_cut_point is one TransformSelector entry with a part field selecting
the stage. All three parts are configured with the same sfno: block — the
donor's own NoiseConditionedSFNOBuilder config — and each builds only the
parameters belonging to its stage; together they partition the donor's
parameters exactly. Channel counts come from the domains and are validated
against embed_dim/big_skip, and against the donor's own input width when a
donor_checkpoint is configured.

Both cut-points carry a single tensor, because that is what a pool component
consumes and produces: the embed_dim latent followed by the big-skip residual
(absent when big_skip is False), so a latent domain at a cut-point declares
embed_dim + in_chans channels. Stage boundaries follow the monolithic
forward with one deliberate regrouping — the context-conditioned big-skip
normalization moves from the encoder (where the monolith computes it) into the
processor. That puts every context-conditioned operation in the one component
that draws the noise, so composing the three parts reproduces the monolith
bit-for-bit rather than approximately; the alternative leaves the skip
normalization conditioned on a second, independent noise draw. A parametrized
test asserts that equivalence across twenty configurations — big skip on/off,
normalized and affine skip norms, output and residual filtering, position
embeddings, isotropic noise, global layer norm, context position embeddings,
both gradient-checkpointing levels, a leading ensemble dimension, and six
block-internal options (no MLP, LoRA, grouped and reduced-ratio spectral
filters, a local DISCO block, the makani filter). It is what pins this module
against changes to the net it mirrors, and forward now carries a comment
pointing back at it.

The parts hold references to the submodules of a net built by the existing
builder, so the stages they keep are built exactly as the donor builds them,
and every shape and flag they need is read off that net rather than re-read
from the config. There are no changes under fme/ace, and the only change
under fme/core is that pointing comment.

Changes:

• fme.translate.cutpoint.SFNOCutPointConfig: the sfno_cut_point registry
  entry — part, the donor sfno block, donor_checkpoint (+
  donor_module_index), and conditional for the processor.

• Donor initialization filters the checkpoint to this part's names, and rejects
  a donor that is missing any of the part's parameters, that shapes any of them
  differently, or whose own input width disagrees with the declared cut-point.
  Without those checks overwrite_weights copies leading slices, so a
  mis-declared cut-point would build a part that is only partly the donor's
  with no error — the width check is what covers a lone frozen processor, which
  keeps no parameter sized by the donor's input channels. Initialization runs
  before TransformConfig.parameter_init, so an explicit weights_path still
  wins; buffers are deliberately not transferred, so a donor-warm-started
  clip_latent_global_means encoder relearns its envelope.

• fme.translate.modules.TransformModuleConfig.build_for_load: opt-out hook,
  defaulting to build, for builders that read a checkpoint of their own at
  build time. TransformConfig.build_for_load routes through it, so a saved
  component reloads without its donor checkpoint still existing.

• fme.translate.ComponentPool.set_epoch: fans the latent global-mean envelope
  reset out to transforms as Stepper.set_epoch does to its modules, which a
  cut-point encoder configured with clip_latent_global_means needs.

• SphericalFourierNeuralOperatorNet.forward: comment naming the module that
  mirrors it.

• fme/translate/examples/README.md: the PR 5 note now names the config shape.

• Tests added

• If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated

[mcgibbon]

Pre-review agent: applied fix — this shape check is new (commit beb030b).

Before it, _apply_donor_weights only checked names. overwrite_weights -> overwrite_weight_initial_slice copies the leading slice when a destination axis is longer than the donor's, so a cut-point domain declaring more channels than embed_dim + in_chans built fine and loaded the donor partially, with no error. Reproduced against the tiny donor in test_cutpoint.py (donor in_chans=2, embed_dim=4, latent declared as 7 instead of 6):

dest conditional_model.decoder.0.weight (4, 7, 1, 1)   donor (4, 6, 1, 1)
-> built, donor's 6 input channels copied, the 7th left at random init

Under-declaring already raised, but with a confusing pre-existing message ("Dest parameter has size 5 ... which needs to be greater than loaded parameter size 6").

The PR description's claim that a part configured differently from its donor "is not silently half-initialized" was true for a wrong num_layers but not for a wrong channel count; it is true now. test_donor_shape_mismatch_raises pins it.

[mcgibbon]

Pre-review agent: remaining hole in the same class — a processor-only splice's cut-point width is unpinned.

The shape check above closes the encoder and decoder cases, and the processor case when normalize_big_skip is set (norm_big_skip is sized by in_chans). But the latent-splice arm's central configuration — a lone frozen cut-point processor between two learned transforms, with the default normalize_big_skip: false — has no parameter sized by in_chans at all: blocks.* are sized by embed_dim only, and norm_big_skip is a NoLayerNorm. So an over-declared latent domain builds, name-matches the donor exactly, and passes silently:

# donor in_chans=2, embed_dim=4  =>  correct cut-point width is 6
_part_config("processor", sfno, donor_checkpoint=ckpt).build(
    n_in_channels=7, n_out_channels=7, ...   # builds, loads, no error
)

The processor then treats x[:, 4:] as a 3-channel residual and hands 3 residual channels to the decoder. With a learned decoder on the far side (which is the splice arm) nothing downstream complains — the decoder just learns around a latent domain that does not match the donor's interface.

There is no way to catch this from the sfno block alone: in_chans is a property of the donor's data, not of its module config. Two options, your call:

1. Read the donor's input channel count out of the checkpoint's stepper config (len(in_names) on the step config) and validate _donor_in_channels against it. Closes the hole completely and needs no new config surface, at the cost of the cut-point knowing about the donor's stepper config, not just its weights.
2. Add an optional donor_in_channels: int | None to this config, cross-checked against the cut-point width in __post_init__. No checkpoint reading, but it is one more number a config has to get right.

I'd take (1). I did not implement either — it is config-surface/design, not a mechanical fix.

[mcgibbon]

Claude: fixed in 529a238 — thanks, this was the right catch. Rather than reading in_names off the donor stepper config, _validate_donor_width reads the donor’s input width straight out of conditional_model.encoder.0.weight's input axis, which is present in every part's donor checkpoint whether or not that part keeps the encoder, and needs no second config load. It runs before the per-parameter checks and applies to all three parts, so a mis-declared cut-point now gets one clear error naming both widths instead of a shape-mismatch dump.

Your shape check stays as the backstop for mismatches that leave the width self-consistent, so I re-pointed test_donor_shape_mismatch_raises at a differing embed_dim (where the width check passes and only the shapes disagree) and added test_donor_cut_point_width_mismatch_raises over all three parts × normalize_big_skip on/off — the processor-with-normalization-off case is the one that was silent.

[mcgibbon]

Pre-review agent: builder-pattern rule violation, and it collides with the wrapper-only constraint.

SFNOCutPointConfig reads five of NoiseConditionedSFNOBuilder's fields — embed_dim and big_skip here and in _donor_in_channels, plus checkpointing, clip_latent_global_means and filter_output in _build_part. AGENTS.md's builder-pattern rules make a config reading another config's fields a violation "even one it directly contains", with the fix being a @property on the child. There is no existing precedent for reading these fields from outside — grep -rn '\.embed_dim' fme/ outside self/params/context_config returns only this file.

The fix would be two-to-five @propertys on NoiseConditionedSFNOBuilder, which is under fme/ace and so out of the wrapper-only scope this PR set for itself. So it needs your call rather than a unilateral fix. Three ways out:

1. Add the properties to NoiseConditionedSFNOBuilder and accept a ~6-line fme/ace diff.
2. Shrink the violation without touching fme/ace: checkpointing, filter_output and img_shape are already public attributes of the built SphericalFourierNeuralOperatorNet (net.checkpointing, net.filter_output, net.img_shape), as are net.embed_dim/net.big_skip. Reading them off net in _build_part would leave only embed_dim/big_skip read from the config, and only because _validate_channels has to run before the net exists. (clip_latent_global_means is private on the net, so it would still come from the config.)
3. Grandfather it, noted in the description as a deliberate exception.

(2) also removes drift points: _SFNOEncoder currently re-derives Distributed.get_instance().get_local_slices(img_shape) from dataset_info.img_shape rather than from the net that already computed exactly that.

[mcgibbon]

Claude: took your option 2, in ab7f789. _build_part now reads embed_dim, big_skip, checkpointing, filter_output and img_shape off the built net instead of the config — all public attributes — which also kills the get_local_slices re-derivation you flagged, since the slices now come from the same img_shape the net used.

Two config reads are left, both structural rather than stylistic: _validate_channels needs embed_dim/big_skip before the net exists (validating in __post_init__-adjacent code is the point of validating early), and clip_latent_global_means is the one flag the net keeps private (net._clip_latent_global_means), so reading it from the config beats reaching into a private attribute or inferring it from the presence of the _gm_min buffer. Commented at the call site.

I did not add properties to NoiseConditionedSFNOBuilder: with the net-sourced reads in place they would only serve the two pre-build reads, which is not worth an fme/ace diff in a PR scoped to avoid one.

[mcgibbon]

Pre-review agent: a donor-warm-started encoder silently drops the donor's latent global-mean envelope. Needs a decision, not obviously a bug.

overwrite_weights copies parameters only — overwrite_weight_initial_slice calls module.get_parameter(name) and the AttributeError for a buffer is swallowed. The _gm_min/_gm_max buffers registered a few lines up therefore stay at ±inf after a donor load, even though they are present in the donor's state_dict under exactly the matching names. Verified:

donor    conditional_model._gm_min: tensor([ 0.0147, -0.2090, -0.0826,  0.2722])
cutpoint conditional_model._gm_min: tensor([inf, inf, inf, inf])
cutpoint conditional_model._gm_max: tensor([-inf, -inf, -inf, -inf])

And _apply_global_mean_clip's eval branch is guarded by torch.isfinite(self._gm_max).all(), so the consequence is that the clip is a silent no-op — the donor's inference-time bound on the latent's global mean is gone with no warning.

Whether that is wrong depends on the use:

• multi-scale warm-start: right as-is. The encoder trains on, ComponentPool.set_epoch resets the envelope each epoch anyway, and the donor's envelope was measured on a different latent distribution.
• latent-splice with a frozen cut-point encoder, or a composite exported for inference (PR 4) without a training pass: wrong. The envelope never gets established and a behavior the donor was trained with disappears. (A frozen encoder inside a training run does still accumulate it — the update is under no_grad but not gated on requires_grad — so only the never-trained path loses it.)

I documented the current behavior on donor_checkpoint so it reads as deliberate rather than accidental, and left the behavior alone. If you want it inherited, note that the general "copy matching buffers" version is not safe (SHT buffers are rank-local under spatial parallelism, and a same-shape-different-shard copy would be wrong), so it would have to be scoped to the envelope specifically.

[mcgibbon]

Claude: agreed, and leaving it as documented behaviour. Not inheriting the envelope is right for the warm-start (it retrains, and the envelope is reset per epoch anyway), and overwrite_weights skipping buffers is what makes it so. The frozen-encoder inference export you describe is a real concern, but it belongs to the PR-4 composite export, which does not exist yet — and clip_latent_global_means is off in every committed config today, so nothing in flight is affected. Worth re-checking when the export lands.

[mcgibbon]

Pre-review agent: build_for_load is a concrete method on an ABC with three in-repo subclasses, so it is neither @abstractmethod nor @final — the abstract-or-final rule. It cannot be @final (this PR overrides it in SFNOCutPointConfig), so satisfying the rule literally means making it abstract and implementing return self.build(...) in SameGridTransformConfig and InterpolateTransformConfig.

Against that: the whole point of the hook is that the default is the safe behavior, and making it abstract puts the burden on every future builder to remember to opt out of something it does not do. I checked both existing siblings and neither reads a checkpoint at build time, so the delegate-to-build default is correct for them today, and TransformConfig._build_raw's for_load flag routes cleanly.

Flagging it because it is a rule violation, not because I think abstract is better. Your call whether the exception is worth stating in the description.

One alternative that sidesteps the rule entirely: give build a for_load: bool = False parameter instead of a second method, so there is one abstract method and builders that care branch on it. That trades the safe-default property for one fewer method — probably not worth it, but it is the other shape.

[mcgibbon]

Claude: keeping it concrete with a default. Making it abstract would force same_grid and interpolate — neither of which reads a checkpoint at build time — to write an identical passthrough, which is the boilerplate the default exists to avoid, and the repo has no @final convention to lean on here. The seam is narrow on purpose: it only exists so TransformConfig.build_for_load can skip build-time external initialization, and a builder that has none correctly does nothing different.

[mcgibbon]

Pre-review agent: the main thing I'd want you to decide: mirror the forward, or split sfnonet.py's forward into three methods and call them?

I walked the mirrored logic against SphericalFourierNeuralOperatorNet.forward line by line and it is faithful — checkpointing levels, the pos_embed distributed slice, pos_drop, the whole clip_latent_global_means envelope block, concat order, filter_output (the monolith applies it unconditionally, but as Identity(Identity(x)) when off, so gating it in the decoder is equivalent), the _UnconditionalCutPoint leading-dim flatten matching NoiseConditionedModel.forward, and the deliberate norm_big_skip regrouping. I found no divergence, and I extended the equivalence test to the two checkpointing levels and a leading ensemble dimension, which the parametrization had not reached (both pass).

So the code is right. The question is whether ~120 lines of duplicated forward logic in fme/translate is the right shape, given that the net it mirrors is under fme/core, not fme/ace. The scope constraint this PR set for itself is "no changes under fme/ace" — which the alternative would honour:

# sfnonet.py
def forward(self, x, context):
    latent, residual = self._encode(x)
    latent = self._process(latent, residual, context)
    return self._decode(latent)

The norm_big_skip regrouping is expressible in exactly that split, so the three cut-point parts become thin wrappers holding a net and calling one method each — no mirrored control flow, and no drift surface at all. The cost is a fme/core diff and three new (private) methods on a widely-used class.

The case for the mirror as written: fme/core stays untouched, and the equivalence test is real coverage. The case against: the test only covers the fourteen configurations it enumerates, and the coupling is invisible from sfnonet.py — nothing there tells someone editing forward that a file in another package mirrors it. local_blocks, spectral_ratio < 1, lora_rank, use_mlp=False and filter_num_groups > 1 are all unparametrized (all of them only affect the blocks, which live wholly inside the processor, so the risk is low — but that argument has to be re-made every time forward changes).

If you keep the mirror, a pointing comment on SphericalFourierNeuralOperatorNet.forward naming fme/translate/cutpoint.py would at least make the coupling discoverable from the side that gets edited.

[mcgibbon]

Claude: keeping the mirror, and I added the pointing comment you suggested (ab7f789) — SphericalFourierNeuralOperatorNet.forward now names fme/translate/cutpoint.py, which turns your main objection (the coupling is invisible from the side that gets edited) into a signpost.

The reason the split-forward alternative does not actually win: "thin wrappers holding a net and calling one method each" means each part holds the whole net. That breaks the property the tests assert — that the three parts partition the donor's parameters — and costs a 3× checkpoint, dead parameters in the optimizer and EMA, and DDP unused-parameter errors on every part. To avoid that you would have to strip the unused stages out of each net instance after building it, which encodes exactly the same structural knowledge the mirror does (which submodules each stage needs), leaves dead SHT buffers behind, and yields a stranger object than a small module holding the submodules it uses. So the choice is not "mirror vs. no duplication" but "mirror vs. mutate a net into a half-net", and I think the mirror is the better of those two.

On the coverage argument: I widened the parametrization to twenty configurations, adding the six block-internal options you named — use_mlp=False, lora_rank, filter_num_groups>1, spectral_ratio<1, local_blocks, and makani-linear — and checked that each one actually changes the parameter set, so none is passing vacuously. All the block-internal ones live wholly inside the processor as you say, so this is belt-and-braces, but it means the argument no longer has to be re-made by hand.

[mcgibbon]

Pre-review agent: pre-review of b82f3052, diffed against origin/feature/translate-skeleton (the base is the long-lived series branch, not main). Fixes applied in beb030b8.

Audit passes

WARN 1. Silent-failure pass. One found and fixed, two found and left for you.

• Fixed (beb030b8): a cut-point domain declaring more channels than embed_dim + in_chans passed _validate_channels and then loaded the donor a leading slice at a time, leaving the rest randomly initialized with no error - overwrite_weights copies the initial slice when a destination axis is longer. _apply_donor_weights now compares shapes as well as names. Details and the reproduction are in-thread.
• Left for you: the same hole stays open for a processor-only splice with the default normalize_big_skip: false, where no parameter is sized by in_chans - thread, with two options and a recommendation.
• Left for you: a donor-warm-started encoder silently drops the donor's clip_latent_global_means envelope, so the clip becomes a no-op at inference - thread. Correct for the warm-start use, wrong for a frozen-encoder inference export; needs your call, so I documented the behavior rather than changing it.

Checked and clear: build_for_load's delegate-to-build default against both existing builders (neither reads a checkpoint at build time); conditional=False against a labeled donor raises labels must be provided from ConditionalLayerNorm rather than silently zeroing the label conditioning (at forward time rather than __post_init__, matching ModuleSelector's existing behavior); donor_checkpoint=None is a clean no-op; a donor with weights is None and an out-of-range donor_module_index both raise (I also made the index guard reject negatives).

OK 2. Behavior-change pass. The only modified existing code paths are TransformConfig._build_raw (new for_load=False, default preserves behavior), TransformSelector/TransformModuleConfig (additive), and ComponentPool.set_epoch, which went from a no-op comment to a fan-out that exactly mirrors Stepper.set_epoch's duck-typed request_latent_global_mean_envelope_reset loop (fme/ace/stepper/single_module.py:1441) - same traversal, same predicate, and epoch is unused in both. Nothing else observable changed.

WARN 3. Differential-test pass. The equivalence test would fail if the decomposition were wrong - the "without the donor the part is randomly initialized" counter-check in test_donor_checkpoint_initializes_the_part and the without_clip arm in test_latent_global_mean_clip_matches_the_donor are both genuine, and the clip test's +50.0 offset does put the input outside the envelope. Two gaps closed in beb030b8:

• checkpointing was unparametrized, so neither >= 1 (encoder/decoder) nor >= 3 (blocks) - the mirrored levels - was exercised. Added both; both pass. (They surface the pre-existing use_reentrant/requires_grad warnings from torch.utils.checkpoint, which is new noise in the suite but genuine information.)
• All inputs were 4-D, so _UnconditionalCutPoint's leading-dimension flatten was never exercised. Added a leading ensemble dimension; passes.

Still unparametrized, all judged low-risk because they only affect the blocks, which live wholly inside the processor: local_blocks, spectral_ratio < 1, lora_rank, filter_num_groups > 1, use_mlp=False. conditional=True has no test at all - I verified by hand that a conditional processor builds and runs against a labeled DatasetInfo, but there is no committed coverage of it.

OK 4. Sibling / mirror coverage pass. build_for_load is the new hook; the two sibling TransformSelector entries (same_grid, interpolate) both correctly take the default, and I confirmed neither reads an external checkpoint at build time. The sibling conditional builders (SwinTransformer, NoiseConditionedSwinTransformer, LocalNet) get no cut-point, which is right for this PR's scope - worth one line in the description saying so, since sfno_cut_point is deliberately specific to NoiseConditionedSFNO.

OK 5. Config back-compat pass. No prior config or checkpoint is touched: sfno_cut_point is a new registry entry, build_for_load is additive with a behavior-preserving default, and fme/translate is not reachable from any existing entrypoint. test_pool_state_round_trips_without_the_donor covers the new state path, and build_for_load is what makes a saved component reload after the donor checkpoint is gone.

OK 6. Split heuristic. 6 files, one new feature, no refactor and no back-compat shim. Don't split.

WARN 7. Abstract-or-final / inheritance-depth pass. Inheritance is one level (TransformModuleConfig -> SFNOCutPointConfig; the nn.Module parts are direct subclasses). TransformModuleConfig.build_for_load is a concrete method on a class with three in-repo subclasses and is neither @abstractmethod nor @final - a literal rule violation, but it cannot be @final because this PR overrides it, and the safe default is the point of the hook. Thread with the alternatives; I don't think abstract is better.

WARN 8. AI-agent artifact pass. Clean on the usual tells - no Any returns, no isinstance/type: ignore pairs, no in-function imports, no abbreviated identifiers, no stale comments, nothing removed. Two doc inaccuracies fixed in beb030b8: _donor_in_channels claimed "nothing they build is sized by it", but norm_big_skip is built at that size when normalize_big_skip is set (it is just never kept without a big skip); and the module docstring's "exactly ... for any configuration" silently depended on the encoder stage consuming no RNG (the composed path runs the encoder before the processor's noise draw, the monolith draws noise first), which holds only because NoiseConditionedSFNOBuilder does not expose drop_rate - now stated, and the equivalence test runs in training mode so it would catch a regression.

FAIL 9. Builder-pattern pass. SFNOCutPointConfig reads five of NoiseConditionedSFNOBuilder's fields (embed_dim, big_skip, checkpointing, clip_latent_global_means, filter_output) - a config reading another config's fields, which the rules call a violation even for a directly-contained config, and there is no existing precedent in the repo for reading these from outside. The canonical fix is a @property on the child, which lands under fme/ace and so collides with this PR's wrapper-only constraint. Thread with three options, including one that shrinks the violation to two fields without touching fme/ace by reading the already-public net.checkpointing/net.filter_output/net.embed_dim/net.big_skip/net.img_shape. I did not apply it - it is a scope decision.

State of the PR

The central design claim holds. Composing encoder -> processor -> decoder does reproduce SphericalFourierNeuralOperatorNet.forward bit-for-bit, and moving the context-conditioned norm_big_skip into the processor so all noise-consuming ops share one draw is the right call - I checked the mirrored forward against the original line by line and found no divergence, and the two coverage gaps I closed (checkpointing levels, leading ensemble dim) both pass unchanged. The structure is good: one registry entry with a part field, the parts partitioning the donor's parameters exactly, build_for_load as the seam that lets a saved component reload without its donor. Tests: 68 pass in ~6s; pre-commit clean on all changed files.

Two things I would want settled before merge, neither of which is a defect in what is written:

1. Mirror or split sfnonet.py? The duplicated forward logic is the PR's main cost, and the net lives under fme/core, not fme/ace - so splitting forward into _encode/_process/_decode and calling them would satisfy the stated scope constraint and delete the drift surface entirely. Thread argues both sides. If the mirror stays, a pointing comment on forward would at least make the coupling discoverable from the side that gets edited.
2. The processor-only width hole and the envelope-inheritance question above - both are decisions about what the cut-point should know about its donor, and both change config or loading semantics, so I left them to you.

Description nits: it says the equivalence test covers "twelve configurations" (now fourteen), and its claim that a part configured differently from its donor "is not silently half-initialized" was true for a wrong num_layers but not for a wrong channel count until beb030b8. Both want a touch-up before merge.

[mcgibbon]

Claude: one consequence for infra 3 that this PR settles, flagged here rather than in the description since it is about follow-on work.

The design note (§1) describes the latent-splice arm as the pool "admitting a bare processor sub-component as a backbone". It cannot be one as the pool stands: backbones are ace Steppers (normalization, ocean, corrector, derived variables), and a bare processor is by construction none of those — the note itself says it "steps without the donor's corrector/normalization shell". So sfno_cut_point registers all three parts, processor included, as TransformSelector entries.

That is the right shape — a processor is a module, not a stepper — but it means ForwardPredictionObjective(encoder, backbone, decoder, ...) in infra 3 has to accept a transform in its backbone slot, not only a Backbone. Worth deciding there whether that slot becomes "any component that maps a domain to itself" or whether the latent-splice arm gets its own objective shape.

Also for the record: no CI runs on this PR (or any PR in the series) — all four ace workflows gate pull_request on branches: [main], so a PR targeting feature/translate-skeleton gets GitGuardian only. Local python -m pytest fme/translate/ fme/core/models/conditional_sfno/ fme/ace/registry/test_stochastic_sfno.py is 198 passed / 3 skipped, and pre-commit is clean on the changed files. Tracked as a research task for a decision on whether to widen the workflow triggers.

[mcgibbon]

Pre-review agent (delta re-review): verified, no false positives found. conditional_model.encoder.0.weight is a safe universal key, and I confirmed it from both ends.

From the source: the encoder loop starts at current_dim = self.in_chans and appends LoRAConv2d(current_dim, encoder_hidden_dim, 1, ...), so encoder[0] is always the conv built from in_chans. encoder_layers=0 is the case worth naming — the loop body never runs, so the unconditional trailing LoRAConv2d(current_dim, embed_dim, 1, bias=False) becomes encoder[0], still with current_dim untouched at in_chans. And LoRAConv2d subclasses nn.Conv2d with groups=1 here, so weight is the plain Conv2d weight and shape[1] is exactly in_channels, not in_channels // groups.

Empirically: the key is present with shape[1] == in_chans across 19 configurations — encoder_layers 0/1/4, lora_rank=2, big_skip=False, makani-linear, local_blocks, spectral_ratio=0.5, filter_num_groups=2, use_mlp=False, pos_embed=False, context_pos_embed_dim, isotropic noise, normalize_big_skip+affine_norms, clip_latent_global_means, global_layer_norm, checkpointing=3, filter_output+filter_residual — and recovers 1, 5 and 17 input channels correctly.

No legitimate configuration now raises that previously built:

• big_skip=False makes expected collapse to embed_dim, which _validate_channels already required for every non-encoder part, so the new check is redundant-but-consistent there and can never contradict it.
• For the encoder with big_skip, _validate_channels already pinned latent == embed_dim + n_in_channels, so the new check reduces to n_in_channels == donor_in_channels — previously caught one step later by the per-parameter shape check, now caught earlier with a much better message.
• embed_dim disagreeing with the donor leaves the width self-consistent and falls through to the per-parameter check, which is what your re-pointed test_donor_shape_mismatch_raises asserts. I confirmed all three parts reach it that way.

The graceful return on an absent key is the right call — it defers to the loud missing branch rather than inventing a verdict. One residual worth a sentence in the docstring: it means the width check is best-effort, so a donor that somehow carries blocks.* without encoder.0.weight would skip it silently. Not reachable today (donor_checkpoint takes an ACE stepper checkpoint, and any SFNO in one has an encoder), so this is a note, not a request.

I'd take this over reading in_names off the stepper config, for a reason I missed the first time: it works uniformly for all three parts without the cut-point needing to know anything about the donor's stepper, only its module.

[mcgibbon]

Claude: thanks for checking both ends — the encoder_layers=0 case going through the trailing unconditional conv, and LoRAConv2d keeping groups=1 so shape[1] is the true input width, are the two things that would have made this fragile.

[mcgibbon]

Pre-review agent (delta re-review): verified a pure refactor, not a behavior change. Every field now read off the net is a verbatim pass-through, checked by source and by construction:

field	path
embed_dim	builder.embed_dim -> SFNONetConfig.embed_dim -> self.embed_dim = params.embed_dim
big_skip	builder.big_skip -> SFNONetConfig.big_skip -> self.big_skip = params.big_skip
checkpointing	builder.checkpointing -> SFNONetConfig.checkpointing -> self.checkpointing = params.checkpointing
filter_output	builder.filter_output -> SFNONetConfig.filter_output -> self.filter_output = params.filter_output
img_shape	dataset_info.img_shape -> get_lat_lon_sfnonet(img_shape=...) -> self.img_shape = img_shape, forwarded verbatim (only modes_lat/modes_lon are derived from it, and nothing writes back)

No transformation, no defaulting, no clamping anywhere on those five paths. Confirmed empirically too: net.<field> == builder.<field> for all five across {}, big_skip=False, checkpointing=3, filter_output=True, encoder_layers 0 and 3, embed_dim=8, and clip_latent_global_means=True.

I also checked the one you left on the config: net._clip_latent_global_means == builder.clip_latent_global_means holds identically, so reading it from the config is equivalent and not a latent divergence — it is only there because the net keeps it private.

And img_shape is strictly better sourced from the net than from dataset_info: _SFNOEncoder re-derives get_local_slices(img_shape) for the pos_embed slice, and it now provably derives it from the same value the net used, rather than from a second reference to the same DatasetInfo that happened to agree.

The two remaining config reads (_validate_channels pre-build, clip_latent_global_means) are both forced, and the comment says why. Good as-is.

[mcgibbon]

Claude: good — verbatim pass-through was the assumption the refactor rests on, and it needed checking rather than assuming.

[mcgibbon]

Pre-review agent (delta re-review): new, non-blocking — the one legitimate reason a part can have parameters its donor lacks is now refused, and the escape hatch I pointed at did not work. Fixed the message in 7cb3032b; the behavior question is yours.

Two things I checked by running them:

1. lora_rank: 2 over a non-LoRA donor is rejected: "has no weights for 12 parameter(s) ... e.g. conditional_model.blocks.0.inner_skip.lora_down.weight". But LoRAConv2d's own docstring advertises exactly this ("Can load a checkpoint saved from nn.Conv2d even when lora_rank > 0"), and lora.py:123 does nn.init.zeros_(self.lora_up.weight) — so an absent LoRA adapter is a genuine identity, not a half-initialized layer. A LoRA fine-tune of a frozen donor processor is a plausible near-term use of this registry entry, and donor_checkpoint refuses it.

2. My previous error message told the user to fall back to parameter_init.weights_path. That does not work: overwrite_weights requires from_names.issubset(to_names), so pointing weights_path at a full donor state_dict raises "Dest module is missing parameters {'conditional_model.decoder.0.bias', 'conditional_model.pos_embed', ...}" — the part legitimately lacks the other stages' keys. The hatch only opens if you first save a state_dict hand-filtered to the part's names.

7cb3032b corrects both messages to say that, and names LoRA adapters as the case on the missing branch, so a user who hits either has a route rather than a dead end. It also drops the now-duplicated "cut-point domain whose channel count is..." wording from the shape branch, since _validate_donor_width owns that diagnosis as of 529a238a.

If you want the LoRA warm-start to work through donor_checkpoint directly, the principled version is to exempt zero-initialized LoRA adapters from the missing check. I did not do it — that is a feature, and the current refusal is at least loud and now correctly signposted.

[mcgibbon]

Claude: your message fix is right and I am keeping it. Leaving the exemption itself unimplemented, deliberately.

Exempting zero-initialized LoRA adapters from the missing check is a real feature — LoRA over a frozen non-LoRA donor is a sensible transfer-learning recipe, and lora_up's zero init does make an absent adapter an exact identity, so it would be well-defined. But nothing needs it today (lora_rank defaults to 0 and no config in the repo sets it), and implementing it means exempting parameters by name pattern — which is precisely the shape of thing that would let a genuine partial load back through the check you added this pass. Adding an unexercised feature that weakens a guard is the wrong trade, so: error message names the case and the working route, and the exemption waits for a config that wants it.

[mcgibbon]

Pre-review agent (delta re-review): beb030b82...ab7f78968, plus one message fix from me in 7cb3032b. Verdict: ready for review. Nothing blocking.

The five items

1. Processor width hole (529a238a) — verified, no false positives. Your key choice is better than my recommendation and I was wrong to prefer reading in_names off the stepper config: conditional_model.encoder.0.weight works uniformly for all three parts without the cut-point needing to know anything about the donor's stepper, only its module. I confirmed encoder[0] always carries in_chans — including the encoder_layers=0 case you flagged, where the loop body never runs and the trailing unconditional conv becomes encoder[0] with current_dim still at in_chans — and that LoRAConv2d subclasses nn.Conv2d with groups=1, so shape[1] is exactly in_channels. Empirically present with the right width across 19 configurations and 1/5/17 input channels. No legitimate configuration now raises that previously built; details and the reasoning per part are in-thread. The graceful return on an absent key is the right call — it defers to the loud missing branch instead of inventing a verdict.

2. Builder-pattern (ab7f78968) — verified a pure refactor. All five net-sourced fields are verbatim pass-throughs; I traced embed_dim, big_skip, checkpointing, filter_output through SFNONetConfig and img_shape through get_lat_lon_sfnonet (which derives modes_lat/modes_lon from it but never writes it back), and confirmed net.<field> == builder.<field> empirically across eight configurations. No transformation, defaulting or clamping on any path, so this is a refactor and not a behavior change. clip_latent_global_means is identical too, so leaving it on the config is equivalent rather than a latent divergence. Thread. Sourcing img_shape from the net is a genuine improvement, not just rule compliance: _SFNOEncoder's pos_embed slice now provably derives from the same value the net used.

3. Mirror vs split — you're right, I was wrong, and I'd make the argument more strongly than you did. I had not thought through that a part holding a net holds all of its parameters. That breaks the partition property, triples the checkpoint, and hands DDP unused parameters in every part — and the fix (strip the unused stages out of each net) encodes the same structural knowledge the mirror does while additionally leaving a net whose own forward is a lie. So it is strictly worse than the mirror, not a trade.

Worth recording that the genuinely clean third option is also closed, since it is the one someone will propose next: make each stage its own nn.Module inside sfnonet.py and have forward compose them, so the parts hold one stage module each with no mirror and no half-nets. That changes state_dict key names — it adds a level of nesting no matter what the sub-modules are called — which breaks inference loading for every existing SFNO checkpoint. AGENTS.md calls that the one hard back-compat guarantee in the repo, so it is off the table without a migration shim that would cost more than the mirror does.

With those two ruled out, the mirror is the right answer rather than the tolerable one, and the pointing comment on forward closes the discoverability gap I raised. I'd suggest folding the reasoning into the module docstring or the PR description — it is the kind of thing that gets re-litigated in six months otherwise.

4. Coverage — verified non-vacuous, independently. I did not take the assertion on trust; I diffed the parameter name-and-shape set of each new configuration against the default build:

override	added	removed	reshaped	example
use_mlp: False	0	8	0	blocks.0.mlp.fwd.0.weight
lora_rank: 2	20	0	0	blocks.0.inner_skip.lora_down.weight
filter_num_groups: 2	0	0	2	blocks.0.filter.filter.weight
spectral_ratio: 0.5	4	0	2	blocks.0.filter.filter.pre_proj.weight
local_blocks: [0]	1	2	0	blocks.0.filter.filter.conv.weight
filter_type: makani-linear	0	2	2	blocks.0.filter.filter.bias

Every one changes the parameter set, so none passes vacuously. Twenty configurations now, and I agree with the framing in the comment — the point is that the "these are block-internal so they cannot reach the stage boundaries" argument no longer has to be re-made by hand.

5. Envelope and build_for_load abstractness — both answers hold. The envelope reasoning is right that the warm-start case wants it relearned; documenting the behavior is the correct resting place given it cuts both ways. On build_for_load, the safe-default argument is the one I'd have made myself.

One new finding (non-blocking, fixed in 7cb3032b)

The missing check refuses the one configuration where a part legitimately has parameters its donor lacks: lora_rank over a non-LoRA donor. lora.py:123 zero-initializes lora_up, so an absent adapter is a genuine identity, and LoRAConv2d's docstring advertises loading a plain-Conv2d checkpoint. Worse, the escape hatch my previous message pointed at does not open — overwrite_weights requires from_names.issubset(to_names), so parameter_init.weights_path with a full donor state_dict raises "Dest module is missing parameters {...decoder..., pos_embed, ...}". I ran both. 7cb3032b corrects both messages to say the file has to be filtered to the part's names, and names LoRA adapters as the case on the missing branch. Whether donor_checkpoint should exempt zero-initialized adapters outright is a feature call, in-thread.

Before merge

Two description edits, and they matter because the description becomes the squash commit message:

• It still says "No changes under fme/ace or fme/core", which ab7f78968 invalidated — the pointing comment lands in fme/core/models/conditional_sfno/sfnonet.py. The claim is worth keeping in a truer form ("no behavior change outside fme/translate; the only fme/core edit is a comment pointing at the mirror"), since wrapper-only was the task's constraint.
• "twelve configurations" is now twenty, and the list of what they cover has grown by six.

Tests: 81 in fme/translate and 198 passed / 3 skipped over fme/translate + fme/core/models/conditional_sfno + fme/ace/registry/test_stochastic_sfno.py, matching your numbers; pre-commit clean on all three changed files. Reminder that none of this is CI-verified on this branch — all four workflows gate pull_request on branches: [main], so the series runs on local verification only until it merges.