refactor(models, training): shard_shapes

[japols]

Description

This PR simplifies the sharding metadata used to track tensor shapes across ranks.

Previously, we stored the full tensor shape for each rank (e.g. list[list[int]]). This required layers to manually track reshapes and dimension changes whenever tensors were transformed, which made the sharding logic fragile and tightly coupled to tensor layouts.

This refactor introduces ShardSizes = Union[list[int], None], representing the per-rank shard sizes along only the sharded dimension. Layers now propagate this information through a bundled GraphShardInfo / BipartiteGraphInfo which tracks shard metadata for both nodes and edges. The full shape expansion only happens at the level of the communication primitive where shapes are assumed equal for non-sharded dimensions across ranks.

Also refactor the all-to-all primitives for head/channel <-> grid sharding into a single common all-to-all primitive.

Additional notes

For now I've tested sharding for a global model across combinations of:

• ensemble/single
• keep_batch_sharded True/False
• transformer/graphtransformer
• head/edge sharding

please feel free to also test your favourite use case.

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.

[cathalobrien]

nice, the benchmark tests pass

[japols]

waiting for #931 to be merged