fix: Sharding for learnable residual connections.

[jakob-schloer]

Description

Add model-sharding support for SpectralOrnsteinConnection so it works correctly when num_gpus_per_model > 1.

What problem does this change solve?

SpectralOrnsteinConnection was incompatible with model sharding (grid split across multiple GPUs). Two issues existed:

1. The inverse SHT produces a spatial weight field with the full grid shape, but when sharded, x only contains the local grid shard. Broadcasting fails on the mismatched grid dimension.
2. When truncate=True, the forward SHT requires all latitude rings to produce correct spectral coefficients. Operating on a grid shard yields incorrect results silently.

What issue or task does this change relate to?

Additional notes

The gather_tensor in _apply_truncation introduces an all-gather collective per forward pass which would become a problem for very large grids.

Todos

• Waiting for merging the shard_shapes refactor refactor(models, training): shard_shapes #964
• Testing with deterministic model
• Testing with ensemble model
• Integration test

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.