Question Regarding Model Sharding

[ttdelgadott]

Hi, I'm working on reimplementing FGN within the publicly available GenCast codebase. I was wondering if model sharding across ensembles is implemented similarly to AIFS-CRPS? And if so, could you provide some sample functions (eg all-gather) of how to handle the sharding of gradients across TPUs in JAX? Thank you!

[andrewlkd]

Hey,

We indeed make use of sample sharding. Unfortunately, we don't have any snippets we can share that are compatible with this repo because we have switched to using https://docs.jax.dev/en/latest/jax.experimental.pjit.html internally.

Best,

Andrew