Serialize StepperState across segmented-inference restarts

[mcgibbon]

Segmented inference threads the full PrognosticState through the rollout in memory, but restart.nc only ever persisted a lossy projection of it — the prognostic data and time. Everything else on the terminal BatchData (the StepperState = corrector + random state, plus labels and data_mask) was silently dropped, so on resume the noise generator restarted at each segment boundary (the cross-restart reproducibility Spencer raised on #1330) and the corrector re-captured its pinned global_dry_air_mass from the segment boundary instead of the true IC.

This PR makes BatchData own its own serialization, so restart.nc becomes a faithful, self-contained, still-inspectable netCDF of the whole state.

Changes:

• BatchData.to_xarray_dataset() / from_xarray_dataset() (fme/ace/data_loading/batch_data.py): the full BatchData↔Dataset round-trip. to_xarray_dataset writes prognostic variables under their plain names + time (single-timestep time dimension squeezed, matching the prior restart shape), plus the round-trippable extras (stepper_state, labels, data_mask) under reserved _fme_state__ variables with a _fme_schema_version marker attribute — but only when present, so a batch carrying none serializes to a plain data+time dataset byte-identical to what was written before this feature. from_xarray_dataset is the strict inverse (recovers data, time, horizontal_dims, and the extras with exact dtypes). StepperState.to_state_dict/from_state_dict remain the tensor intermediate (the stepper stays opaque); the generator state (uint8 ByteTensor) round-trips through a netCDF ubyte variable; the corrector's global_dry_air_mass carries the shared sample dim so start_indices subselection subsets it in step with the prognostics. n_ensemble/horizontal_dims are derived on load, not serialized (avoids the double-broadcast bug).
• Writer: restart _write is now ds = data.to_xarray_dataset() plus the writer's presentation (per-variable metadata attrs, coords, dataset metadata). write_stepper_state and the .pt write are gone.
• Reader (get_initial_condition): takes an InitialConditionRequirements (a new container in fme/ace/requirements.py bundling prognostic_names/labels/n_ensemble) and dispatches. A dataset with the schema marker (BatchData.dataset_has_embedded_state) is rebuilt via from_xarray_dataset and validated against the requirements (BatchData.validate_initial_condition) rather than re-derived. A plain netCDF without the marker (external ICs and legacy restart files) takes the existing lenient path unchanged: stepper_state=None, labels from config. Both paths return CPU state; InferenceGriddedData.__init__ moves it to the compute device (BatchData.to_device moves the stepper_state; the generator stays CPU — keeps the GPU-corrector fix).
• PrognosticState.with_stepper_state removed (the reader reconstructs the whole BatchData); with_random_state keeps the seed path.
• Segmented driver: the next segment's IC is simply InitialConditionConfig(path=restart.nc) — the full state rides the restart.
• config.seed precedence (inference.py + evaluator.py, via data.apply_config_seed, which delegates to PrognosticState.apply_config_seed): applied only when the IC carries no restored random_state; a restored state that supersedes it logs an info line.

Scope: inference segmented path only; training-restart is untouched (matches #1330's scope).

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

Stacked originally on #1330 (now merged to main).

[mcgibbon]

Claude: Design note on config.seed precedence. A restored random_state (from a segment's restart sidecar) takes precedence over config.seed: the seed is applied only when the initial condition carries no random state, so segment 0 seeds and later segments continue the restored generator. I deliberately did NOT add a check that config.seed matches the seed that originally created the restored state. The serialized form is the generator's advanced Mersenne-Twister state (get_state()), not the seed, so such a check would require threading the config seed into the opaque RandomState/StepperState — which those types deliberately exclude (they wrap only an advancing generator). Semantically the generator continues on resume and config.seed is inert there; ignoring a changed seed is what 'resume' means, not a corruption — output stays bitwise-identical to the single-run seeded rollout. To avoid silent confusion, a restored random state that supersedes config.seed logs an info line. The restart mechanism does no other config↔restart cross-checking (e.g. of prognostic contents), so a seed-only check would also be inconsistent.

[mcgibbon]

Claude (pre-review agent): Applied fix pushed to this branch (d2ae7ff).

[High] Restored stepper state was left on CPU → GPU device mismatch. get_stepper_state() deserializes the sidecar on CPU (from_state_dict builds CPU tensors), and run_inference_from_config attached it to the IC without moving it to the compute device. On a GPU segmented run with a dry-air-conserving corrector, the restored corrector_state.global_dry_air_mass (CPU) is subtracted from the on-device prediction in _adjust_gen_dry_air_to_target (corrector/atmosphere.py:426, where .to(precision) changes dtype only, not device) → RuntimeError: expected all tensors on the same device. This defeats one of the PR's stated goals (restoring the corrector's pinned mass) on GPU. CPU tests can't surface it, and the acceptance test uses a corrector-free NoiseConditionedSFNO, so the corrector-restore path had no rollout-level coverage.

Fix: .to_device() on the restored StepperState at the attach site. RandomState.to_device is a no-op, so the generator correctly stays on CPU; only corrector_state moves. Re-ran the changed suites + test_inference/test_evaluator/test_batch_data (all pass) and pre-commit (ruff/mypy clean).

[mcgibbon]

This ignore is a little scary, but I couldn't think of a good way to avoid it.

[mcgibbon]

Claude: The skip is intentional and is exactly what makes a resumed rollout continue the same generator — dropping it would reseed at each segment boundary and break the bitwise-identical-to-single-run property this PR establishes. It isn't silent: when a restored random state supersedes config.seed we log an INFO line. I considered raising when both a seed and a restored state are present, but in segmented runs config.seed is legitimately set once and carried into every segment via the deep-copied config, so raising would break the normal flow (segment 0 seeds, segments 1+ inherit the same config with the seed still set). The INFO log is the mitigation. Happy to elevate it to a warning if you'd prefer it more visible — let me know.

[spencerkclark]

Noting this with a log statement seems reasonable to me. If we find a need for it down the road we could introduce a configuration flag to ignore the random state in the restart file, but I think this current default behavior is more often than not what people would want.

This has a flavor to it like the behavior of GFDL models around setting the current model time at initialization. If reading from restart files, it will be set from those unless a flag is set to override it from the namelist.

[mcgibbon]

Claude (pre-review agent): Verdict — ready for review (after the one fix I applied above).

Reviewed the delta on top of #1330 only. The design is clean and the round-trip is correct.

Verified

• Round-trip: RandomState serializes generator.get_state() (CPU ByteTensor), not the seed, so restore continues the exact draw stream (unit-tested). The <name>.present marker scheme reconstructs None vs present-but-empty sub-states exactly (parametrized over all combinations). torch.load(..., weights_only=True) is safe — the payload is only tensors.
• Acceptance test is genuinely non-vacuous: I confirmed it fails when the restore is disabled (temporarily patched out with_stepper_state, then reverted). Segment-1 output is bitwise-equal to the single 6-step run, and the seed-sensitivity assert rules out a deterministic artifact.
• Back-compat: a restart dir with no sidecar → stepper_state_path=None → starts fresh; set-but-missing raises FileNotFoundError, unset raises ValueError — both tested.
• All four WriterABC implementers (DataWriter, PairedDataWriter, NullDataWriter, CoupledPairedDataWriter) implement write_stepper_state; none left abstract.
• Seed precedence: segment 0 seeds; segments 1+ with a restored random_state skip the seed (info-logged); the non-segmented seeded and unseeded paths are unchanged. Evaluator shares apply_config_seed with no behavior change. Training-restart path untouched.
• pre-commit (ruff/ruff-format/mypy) clean on all changed files; the Seedable random state threaded through StepperState for reproducible stochastic inference #1330 suites (test_random_state, test_segmented, test_inference, test_evaluator, test_batch_data) pass — 118 tests.

Non-blocking suggestions (for the author, optional)

• Consider a CPU end-to-end segmented test with a dry-air-conserving corrector asserting segmented == single. It wouldn't catch the device issue on CPU, but the corrector-restore path currently has only writer/config round-trip coverage, not rollout coverage.
• CoupledPairedDataWriter.write_stepper_state writes a sidecar per (ocean/atmos) dir, but there's no coupled restore wiring yet — this is forward-compatible plumbing (coupled inference doesn't use run_segmented_inference), just noting it's write-only for now.

[mcgibbon]

This is quite a bit of a comment, and could easily go out of date. Is there any way to make this information obvious with code structure instead? Or is this comment even necessary?

Another question, why is it that restart_path would always be passed but sidecar_path is only sometimes passed? And would it be clearer to call this stepper_state_path instead of "sidecar" (which is jargon our team doesn't use)?

[mcgibbon]

Claude: Done in 7392e9d. (1) Purged "sidecar" everywhere — comments, docstrings, the evaluator's "restart sidecar", the local var (sidecar_path → stepper_state_path), and the test names (test_stepper_state_sidecar_* → test_stepper_state_file_*); prose now says "stepper state file" / "restart stepper state" (the public field was already stepper_state_path). (2) Cut the multi-line comment in run_segmented_inference to a single intent line. (3) On the asymmetry (restart_path always passed, stepper_state_path only sometimes): kept the conditional deliberately. restart.nc is the IC and always exists, but the stepper state file exists only when a segment had state to persist (seeded and/or corrector) — a deterministic segment or a pre-feature restart dir has none. A set-but-missing stepper_state_path is a genuine user error and trips the FileNotFoundError guard, so the driver only sets it when the file is actually present.

[mcgibbon]

Claude: This asymmetry is now moot — resolved by the redesign. The .pt sidecar is gone; the full PrognosticState (including StepperState) is embedded directly in restart.nc, so the next segment's IC is just InitialConditionConfig(path=restart_path) with no stepper_state_path and no conditional. See the design-change comment on the PR.

[mcgibbon]

Can/should we assert the random_state is the random_state we passed in?

[mcgibbon]

Claude: Added in 7392e9d. The test now captures the RandomState instance placed on the IC and asserts terminal.random_state is random_state (identity, not just presence) — the generator advances in place and the device/ensemble helpers return self, so the exact instance is threaded through the step that rebuilds StepperState.

[mcgibbon]

Claude: Design change — replaced the .pt sidecar with embedding the full state directly in the restart netCDF (Jeremy's request). Rationale: restart.nc was a lossy projection of PrognosticState (only data+time), and StepperState was just the first dropped field. Rather than add a second file for one field, the writer now embeds the whole round-trippable state — stepper_state, labels, data_mask — into the same restart.nc.

• Format: netCDF (not torch.save), so the full-state restart stays inspectable with standard tools; prognostic variables keep their plain names and the file opens as a normal xarray Dataset. Extras live under a reserved _fme_state__ prefix with a _fme_restart_schema_version attr as the presence marker.
• Single file, no sidecar. The reader auto-detects the marker: present → reconstruct + attach (incl. StepperState); absent → behaves exactly as today (stepper_state=None, labels from config), so any legacy/plain restart.nc or user-supplied IC resumes unchanged. Only present fields are written, so a deterministic run's restart.nc is byte-identical to before.
• Generator round-trip verified: the uint8 ByteTensor from get_state() round-trips through a netCDF ubyte variable and set_state reproduces the exact stream. The corrector's per-sample global_dry_air_mass carries the sample dim so start_indices subselection subsets it in step with the prognostics; the generator has no sample dim and is untouched.
• n_ensemble/horizontal_dims are intentionally NOT serialized (derived on load) to avoid reintroducing the double-broadcast bug.
• stepper_state_path/get_stepper_state are gone; the segmented driver's next segment is just InitialConditionConfig(path=restart.nc). This dissolves the restart_path/stepper_state_path asymmetry.

The @pytest.mark.slow segmented seed-reproducibility acceptance test still passes, now resuming through the embedded state, plus new tests for the netCDF round-trip, byte-clean deterministic restart, plain-netCDF back-compat, and start_indices on a full-state restart.

[mcgibbon]

Halting a full review, these comments have wide-reaching implications.

[mcgibbon]

Instead of having a BatchData.with_stepper_state, can we have a method that serializes/deserializes BatchData to/from an xarray.Dataset (which we could then store to disk when we save restarts)? This would make the functionality easier to unit test.

[mcgibbon]

Claude: Done in b95e7e1. Added BatchData.to_xarray_dataset() / from_xarray_dataset() as the full BatchData↔Dataset round-trip: to_xarray_dataset produces the complete restart dataset (prognostic vars under plain names + time, single-timestep squeeze preserved, plus the reserved _fme_state__ extras when present with the _fme_schema_version marker), and from_xarray_dataset is its strict inverse (recovers data, time, horizontal_dims, and embedded stepper_state/labels/data_mask with exact dtypes). The prognostic-data + time serialization that used to live in the writer's _write is folded in. Removed PrognosticState.with_stepper_state and inlined its one use in with_random_state (a dataclasses.replace of the stepper_state) — the reader now reconstructs the whole BatchData (incl. stepper_state) via from_xarray_dataset, so the attach helper is unneeded. Added pure BatchData↔Dataset round-trip unit tests (no writer/inference machinery) parametrized over every field combination.

[mcgibbon]

This should live with BatchData, not with the inference data writers.

[mcgibbon]

Claude: Done — restart.py is deleted and all of its logic moved onto BatchData: the _fme_state__ prefix, the schema-version marker, _fme_dtype dtype preservation, the generator ubyte encoding, and the corrector/labels/data_mask encoding now live in to_xarray_dataset/from_xarray_dataset and the private _encode_reserved_state/_decode_reserved_state helpers (the old RestartExtras/encode_restart_extras/decode_restart_extras/has_restart_extras are gone). The writer's _write is now ds = data.to_xarray_dataset() + metadata/coords/dataset_metadata; the reader gates on BatchData.dataset_has_embedded_state and calls from_xarray_dataset. Contract as implemented: from_xarray_dataset is a strict inverse of to_xarray_dataset only — it is not asked to interpret arbitrary external IC datasets or legacy plain restart files (those keep the lenient get_initial_condition path). Generator uint8 ↔ netCDF ubyte round-trip and set_state stream reproduction are verified (through-netCDF test in test_batch_data.py).

[mcgibbon]

Claude: Independent pre-review (Opus, worktree-isolated, no authoring context). I reviewed the full current diff vs main, ran every affected suite, and probed the acceptance test for vacuousness. Verdict up front: ready for review.

Audit passes

• Silent-failure ✅ — routing is marker-based and loud (dataset_has_embedded_state on the _fme_schema_version attr); from_state_dict (state["generator_state"]) and _restore_tensor (_DTYPE_BY_NAME[...]) raise KeyError on malformed input rather than defaulting; the plain/legacy path is explicit (stepper_state=None, labels from config).
• Round-trip fidelity ✅ — to_xarray_dataset/from_xarray_dataset verified an exact inverse over every combination (stepper none/empty/corrector/random/both × labels × data_mask), multi-timestep, and through real netCDF I/O; generator uint8 state continues the identical draw stream; bool→uint8→bool dtype restored; cftime time round-trips (squeezed single-timestep and unsqueezed).
• Reserved-name safety ✅ — _fme_state__ excluded in the reconstruct data loop and the plain reader only enumerates prognostic_names; reserved vars get variable-private dims; dot-bearing var/dim names validated by the real-netCDF tests.
• start_indices ✅ — corrector global_dry_air_mass (shared sample dim) subsets in step; generator state (own dim) untouched (test_full_state_restart_start_indices).
• Device ✅ — restored stepper_state.to_device() moves the corrector to the compute device; RandomState.to_device/to_cpu are no-ops so the generator stays CPU (GPU-corrector fix preserved).
• Back-compat ✅ — test_plain_restart_netcdf_is_backcompat confirms a no-state run writes a marker-free, reserved-var-free netCDF and reloads exactly as before.
• Removed-symbol coherence ✅ — no remaining reference to with_stepper_state, write_stepper_state, stepper_state_path, or a restart.py; with_random_state inlines the old behavior and preserves an existing corrector sub-state; both DataWriter.write and PairedDataWriter.write route through the changed _write.
• Differential test ✅ — proved the @pytest.mark.slow acceptance test non-vacuous by neutering RandomState.from_state_dict to reseed instead of restore: test_segmented_stochastic_inference_matches_single_run then fails (two-segment ≠ single run), and passes once reverted. The state round-trip tests are parametrized across all sub-state combinations.

Tests — all green: test_random_state.py + test_batch_data.py (91 passed, 3 skipped), segmented + writer non-slow (8), data_writer/ + test_inference.py + test_evaluator.py (147), the slow acceptance test, and test_single_module.py::test_predict_threads_random_state_alongside_corrector_state. pre-commit (ruff / ruff-format / mypy) clean on all changed source files.

Non-blocking notes (no fix applied — both are defensiveness judgment calls, not bugs; flagging for the reviewer):

1. _reserved_var_dims infers the per-sample axis by axis == 0 and size == n_samples. Correct for the corrector (n_samples,1,1), and the flat generator vector is left on its own dim as intended — but only because its length (~5000+ bytes) never coincides with n_samples. If a flat reserved var ever had a leading axis of length exactly n_samples, it would be mislabeled sample and wrongly subselected by start_indices. A size heuristic rather than an explicit per-sample tag; safe today, worth knowing.
2. The writer metadata loop now iterates ds.data_vars (which includes time), where the old code added time after the per-variable-metadata loop. Harmless in practice: variable_metadata is built only from requested physical variable names (_get_variable_metadata over self._names), so "time" is never a key — but it is a latent behavior change if that ever stops holding. Related, initial_condition.nc/restart.nc now embed full state when the terminal state carries corrector/random state (previously plain); this is the intended faithful serialization and nothing reads those back as a plain IC.

Does it accomplish its purpose / is it easy to review? Yes. The serialization is cleanly owned by BatchData, the marker-gated reader keeps external/legacy ICs on the untouched path, and the PR description matches the diff (title and body accurate). The test suite makes the reproducibility guarantee legible and the acceptance test is genuinely non-vacuous. Ready for review.

[mcgibbon]

I don't understand why this conditional is here, perhaps it would be more obvious if the indented code and the code below were each contained in helper functions?

[mcgibbon]

Claude: Done. get_initial_condition is now a thin dispatcher:

if BatchData.dataset_has_embedded_state(ds):
    return _initial_condition_from_state(ds, prognostic_names, labels, n_ensemble)
return _initial_condition_from_variables(ds, prognostic_names, labels, n_ensemble)

_initial_condition_from_variables is the unchanged lenient plain-tensor path (external ICs and legacy restarts: builds prognostic tensors from the named variables, sets labels from config, broadcasts the ensemble). _initial_condition_from_state is the embedded-full-state path (rebuild via from_xarray_dataset, validate, wrap, move stepper_state to device — see the comment on line 156). The docstring now describes the validate-not-rederive split.

[mcgibbon]

Unnecessary comment.

[mcgibbon]

Claude: Deleted the comment above config_copy.initial_condition = InitialConditionConfig(...) in run_segmented_inference.

[mcgibbon]

When we're loading BatchData from disk, we shouldn't need to set labels, broadcast the ensemble, get the correct prognostic names, etc. as this was already done before we saved it. You just need to ensure that if values are passed for these, they get validated against the loaded value. For example, we should error out if we attempted to load with labels that doesn't match batch_data.labels. This should be contained in a validation helper, perhaps a BatchData.validate_initial_condition that takes in the necessary arguments from this scope.

[mcgibbon]

Claude: Done — the embedded path now validates rather than re-derives. Added BatchData.validate_initial_condition(self, prognostic_names, labels, n_ensemble) which raises on inconsistency: (a) any name in prognostic_names absent from self.data (error lists the missing names and the present ones); (b) labels given but not matching — self.labels None while config provides labels, or self.labels.names != list(labels); (c) n_ensemble != self.n_ensemble (for a loaded full-state IC that's != 1), with a message that a full-state restart can't be re-broadcast because its sample dim already carries the ensemble.

_initial_condition_from_state is now: batch_data = BatchData.from_xarray_dataset(ds) → batch_data.validate_initial_condition(prognostic_names, labels, n_ensemble) → get_start(prognostic_names, n_ic_timesteps=1) (only the wrap-into-IC step; no broadcast, no label-setting) → .to_device() the restored stepper_state (RandomState.to_device is a no-op; GPU-corrector fix kept). Unit tests added for validate_initial_condition: passes when consistent; raises on a labels mismatch, a labels-provided-but-none-saved case, a missing prognostic name, and an n_ensemble mismatch.

[mcgibbon]

Claude: Independent pre-review of the latest round (c29d46f), focused on the four deltas. Verdict: ready for review — no findings, no fixes applied.

Ran on CPU: test_batch_data.py, test_segmented.py (incl. the @pytest.mark.slow test_segmented_stochastic_inference_matches_single_run), test_inference.py, data_writer/, corrector/, core/test_random_state.py, stepper/test_random_state.py — all pass. pre-commit (ruff/ruff-format/mypy) clean on all changed files.

1. get_initial_condition split — dispatch is correct (marker → _initial_condition_from_state; no marker → _initial_condition_from_variables). The plain path body is byte-for-byte identical to main's get_initial_condition body; behavior unchanged.

2. validate_initial_condition — each check is correct: missing-name, label mismatch incl. loaded-None-but-config-provided, and n_ensemble. On the n_ensemble check: n_ensemble is not serialized so a loaded BatchData is always n_ensemble=1; both call sites pass n_ensemble=1 to get_initial_condition (inference.py:375 has no arg; the ensemble broadcast is applied downstream at inference.py:399 via n_ensemble_per_ic), and run_segmented_inference hard-rejects n_ensemble_per_ic>1. So the check evaluates 1!=1 (passes) for every supported case and correctly refuses a re-broadcast otherwise. Dropping the old broadcast_ensemble call changes nothing: broadcast_ensemble(1) was an identity on data/time/labels/stepper_state and the loaded n_ensemble is already 1. get_start is still needed as the single-timestep wrap and is unchanged. Label check compares names only, which is right (embedded values are the source of truth). Missing-name now fails early with a clear message instead of get_start silently dropping the var.

3. Explicit per-sample encoding — namespacing is consistent: to_state_dict emits corrector_state.global_dry_air_mass and StepperState.per_sample_state_keys() emits the same, so the corrector var gets the sample dim (subselected by start_indices in step with prognostics) and the generator var (RandomState → empty set) never does. The .present markers are scalars (ndim 0 → no dims). labels/data_mask are sample-dim'd on their own explicit paths (correct — both are per-sample). The new regression test builds a batch whose sample count equals the generator-state length and asserts the generator var keeps a private dim — it genuinely exercises the coincidence and would fail under the old size==n_samples heuristic.

4. Writer metadata loop — skipping time and _fme_state__* is safe: those vars are never in variable_metadata (keyed by physical variables), and _write is snapshot-only (restart.nc / initial_condition.nc), so time-series/non-restart writes are unaffected.

Round-trip inverse, generator uint8/set_state continuation across segments (the slow acceptance test), and plain/legacy → stepper_state=None back-compat all hold.

Minor, non-blocking, pre-existing and out of scope: n_ensemble is not persisted across restart, so a manual non-segmented resume from a restart written by an n_ensemble_per_ic>1 run would lose the ensemble-metric grouping hint (sample dim is preserved). Identical to pre-PR behavior; the supported segmented path enforces n_ensemble=1, and validate_initial_condition guards the dangerous re-broadcast case.

[mcgibbon]

Instead of putting the low-level code here and digging several levels deep, please call data.apply_config_seed, a new method you will write, which will in turn call self._initial_condition = self._initial_condition.apply_config_seed. You can retain the same behavior you have now, logging if it's already been seeded.

[mcgibbon]

Claude: Done in e8427ed. run_evaluator_from_config (and run_inference_from_config) now call data.apply_config_seed(config.seed). InferenceGriddedData.apply_config_seed does self._initial_condition = self._initial_condition.apply_config_seed(seed), and the seed-or-defer logic (including the info log when a restored random state supersedes the config seed) lives on PrognosticState.apply_config_seed. The module-level helper in evaluator.py — and its reach into data._initial_condition from outside — is gone. Unit tests added for the PrognosticState method (seeds when unseeded, returns self on None, defers to a restored state); the segmented acceptance test still covers the whole chain.

[mcgibbon]

This feels wrong somehow, I can't put my finger on it. Like it feels like there should be a method on a class somewhere, maybe? Or we should put the last three arguments into a container class?

[mcgibbon]

Claude: Done in e8427ed — went with the container class. InitialConditionRequirements (fme/ace/requirements.py, beside PrognosticStateDataRequirements) bundles prognostic_names/labels/n_ensemble — "what the run requires of its initial condition" — and flows through get_initial_condition, both build helpers, and BatchData.validate_initial_condition(requirements), so the validate call is now the container handed straight down. Callers (standalone inference, coupled) construct it at the top.

Chasing the "feels wrong" further turned up a real wart in _initial_condition_from_state: its explicit stepper_state-to-device re-wrap was (a) redundant — InferenceGriddedData.__init__ already calls initial_condition.to_device(), and BatchData.to_device moves stepper_state (corrector to device, generator a no-op on CPU) — and (b) asymmetric, since the plain-variables path returns everything on CPU. Both paths now return CPU state and the device move happens in one place; the two GPU tests that asserted the early move now compare on CPU.

[mcgibbon]

Claude: Pre-review of round-4 commit e8427ed (review-response), read in the context of the full PR diff against main. Verdict: ready for review — no blocking findings.

What I verified:

• Seed-precedence behavior unchanged. PrognosticState.apply_config_seed (batch_data.py) is logically identical to the old evaluator.apply_config_seed module function it replaces: seed is None → no-op; a restored stepper_state.random_state present → log + defer to the restored generator; otherwise with_random_state(RandomState.from_seed(seed)). self._data.stepper_state is the same object the old code reached via data.initial_condition.as_batch_data().stepper_state. Both call sites (evaluator.py, inference.py) call data.apply_config_seed(config.seed) at the same point in the flow. Covered by the new test_apply_config_seed_* units and, end-to-end, by test_segmented_stochastic_inference_matches_single_run (passes).

• Device-move removal is safe. Removing the explicit stepper_state.to_device() in _initial_condition_from_state does not change where tensors live by the time the stepper uses them: every consumer of get_initial_condition's output routes the IC through an InferenceGriddedData.__init__, which calls initial_condition.to_device() → BatchData.to_device moves stepper_state (batch_data.py:367). Confirmed for the ace path (run_inference_from_config → get_forcing_data → get_inference_data → InferenceGriddedData) and the coupled path (CoupledPrognosticState.to_device recurses into each PrognosticState.to_device). The evaluator path builds its IC from the loader via requirements and never calls get_initial_condition, so it is unaffected. Between get_initial_condition and the to_device, only .time is read (for start-index lookup), never stepper_state. The two GPU tests that call get_initial_condition directly (no InferenceGriddedData in between) are correctly switched to compare on CPU.

• API-change fallout covered. get_initial_condition's loose args are bundled into InitialConditionRequirements; both fme/coupled call sites and all tests are updated, and no other callers exist. Both get_initial_condition and the new InitialConditionRequirements are exported from fme/ace/__init__.py.

• Conventions / hygiene. New dataclass follows the requirements.py "Parameters:" docstring convention and sits with its kin. Orphaned imports from the moved helper (RandomState, InferenceGriddedData) are removed from evaluator.py and are unused elsewhere there. ruff, ruff format --check, and mypy (pre-commit) all pass on the 10 changed files; test_batch_data, test_inference, test_evaluator, test_segmented, and fme/coupled/inference/test_inference all pass under FME_FORCE_CPU=1 (140 passed, 2 skipped).

[spencerkclark]

Thanks—it's going to be super nice to have this for formal and informal segmented runs. It seems like you have covered all the bases and it's great that the full BatchData state is now possible to roundtrip through a single netCDF.

I had a couple questions regarding whether it would be possible to simplify the schema some by adopting existing conventions rather than defining our own, but otherwise just minor nits.

[spencerkclark]

nit: this is slightly tidier and more closely mirrors the later tensor logic:

Suggested change

	time_da = ds[_TIME_DIM]
	squeezed = list(time_da.dims) == [_SAMPLE_DIM]
	if squeezed:
	time = xr.DataArray(time_da.values[:, None], dims=[_SAMPLE_DIM, _TIME_DIM])
	else:
	time = xr.DataArray(time_da.values, dims=[_SAMPLE_DIM, _TIME_DIM])
	time = ds[_TIME_DIM]
	squeezed = list(time.dims) == [_SAMPLE_DIM]
	if squeezed:
	time = time.expand_dims(dim=_TIME_DIM, axis=1)

[spencerkclark]

Question: are the self.labels.names interpretable as a coordinate along the _LABEL_INDEX_DIM? In other words could they equivalently be encoded as a dimension coordinate rather than in json as a Dataset-level attribute? It could make things slightly more intuitive to human readers of the restart files.

[spencerkclark]

Question: xarray already has a mechanism for automatically roundtripping bool dtype variables to and from netCDF. Could leveraging that remove the need to define our own dtype encoding mechanism (modulo conversion between NumPy arrays and torch tensors)?

[spencerkclark]

nit: not new, but this works for HEALPix data too, correct?

[spencerkclark]

nit: I realize much of this is inherited from earlier code, but is there a reason we cannot simplify this to the following? It should be possible to handle the squeeze all at once at the Dataset level.

Suggested change

	if self.time.sizes[_TIME_DIM] == 1:
	time_dim = self.dims.index(_TIME_DIM)
	dims_to_write = self.dims[:time_dim] + self.dims[time_dim + 1 :]
	def present(x: torch.Tensor) -> torch.Tensor:
	return x.squeeze(dim=time_dim)
	time_array = self.time.isel(time=0)
	else:
	dims_to_write = self.dims
	def present(x: torch.Tensor) -> torch.Tensor:
	return x
	time_array = self.time
	data_arrays: dict[str, xr.DataArray] = {}
	for name, tensor in self.data.items():
	data_arrays[name] = xr.DataArray(
	present(tensor).detach().cpu().numpy(), dims=dims_to_write
	)
	data_arrays[_TIME_DIM] = time_array
	extra_arrays, attrs = self._encode_reserved_state()
	data_arrays.update(extra_arrays)
	ds = xr.Dataset(data_arrays)
	ds.attrs.update(attrs)
	return ds
	data_arrays: dict[str, xr.DataArray] = {}
	for name, tensor in self.data.items():
	data_arrays[name] = xr.DataArray(
	tensor.detach().cpu().numpy(), dims=self.dims
	)
	data_arrays[_TIME_DIM] = self.time
	extra_arrays, attrs = self._encode_reserved_state()
	data_arrays.update(extra_arrays)
	ds = xr.Dataset(data_arrays)
	ds.attrs.update(attrs)
	if ds.sizes[_TIME_DIM] == 1:
	ds = ds.squeeze(_TIME_DIM)
	return ds

[spencerkclark]

Noting this with a log statement seems reasonable to me. If we find a need for it down the road we could introduce a configuration flag to ignore the random state in the restart file, but I think this current default behavior is more often than not what people would want.

This has a flavor to it like the behavior of GFDL models around setting the current model time at initialization. If reading from restart files, it will be set from those unless a flag is set to override it from the namelist.