refactor(models): graph partitioning

[japols]

Description

Refactor graph partitioning logic for sharding and chunking into a unified GraphPartition abstraction that exploits dst-sorted edges for fast subgraph operations via slicing.

Previously, edge sharding across GPUs and chunking within a GPU used separate code paths with expensive bipartite_subgraph calls. This PR consolidates both into a single GraphPartition dataclass that precomputes dst/edge splits and materialises subgraphs via simple slice operations.

This lays the groundwork for extending edge sharding to processors (non-bipartite graphs) and for more efficient halo exchange communication schemes.

What problem does this change solve?

• faster subgraph operations via slicing instead of index_select
• clearer graph sharding/chunking code in mapper.py
• no sorting overhead when using Triton GT

What issue or task does this change relate to?

Additional notes

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.

[japols]

Benchmark tests pass with ~10-20% throughput improvements and ~10% reduced peak memory usage.