chunking the members and latent cells for longer rollout

[kctezcan]

Description

We cannot run long inference rollouts with multiple members due to OOM issues.

These fixes allow that.

You can run inference with the following command for 1 member and 40 forecast steps (should work at least for 4 members):

uv run --offline inference --from-run-id ni41n7gy --options \
  test_config.samples_per_mini_epoch=1 test_config.output.num_samples=1 \
  data_loading.num_workers=0 training_config.forecast.num_steps=40 \
  validation_config.validation_noise_levels=[] diffusion_rollout=True \
  fe_diffusion_num_ensemble_members=1 \
  decode_members_per_pass=1 decode_cells_per_pass=1536 \
  test_config.forecast.chunk_size=4 test_config.compute_full_validation_loss=False \
  training_config.losses='{"physical": {"type": "LossPhysical", "weight": 0.0, "loss_fcts": {"mse": {}}, "target_and_aux_calc": "Physical"}}' \
  load_decoder_chkpt.run_id=mapw62ot load_decoder_chkpt.mini_epoch=-1 \
  test_config.start_date=2023-10-11T00:00 test_config.end_date=2023-11-11T00:00 test_config.shuffle=False \
  test_config.losses='{"physical": {"type": "LossPhysical", "weight": 0.0, "loss_fcts": {"mse": {}}, "target_and_aux_calc": "Physical"}}'

Issue Number

No issue

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

FastEvaluation

• I have updated the public documentation if necessary

[clessig]

This should not be merged. It covers functionality also in #2765 but goes further. But we need to think how to implement it properly.