Add sharding support and fix the world axis

[amacati]

Rationale

We want to support sharding of the simulation so that we can run on multiple devices. Multi-device support unlocks performance improvements on the CPU and might be interesting for future parallel GPU simulations. This PR adds the tools for sharding.

What was previously broken

While working on this, I noticed that some of our functions are fundamentally broken. This mostly affects resets and sharding. When we reset our data, we previously made a shape check if the array matches the number of worlds in its leading dimension. This leads to wrong positives, i.e. the gravity_vector with shape=(3,) would get picked up for n_worlds=3. This is obviously wrong. There is no way to tell which array just happens to coincide with the simulation n_worlds, and which ones are truly parallel across worlds.

Design

The world axis

We have two realistic options: either we keep an explicit list of all arrays that are batched over worlds somewhere, or we somehow mark each array.

The first option feels clunky, since we would store information about the data in two places, and users would have to use a separate registration mechanism for plugins.

The second option needs some way to attach metadata to arrays. Luckily, flax.struct.field already has a metadata field. We now use this for every array that has a world axis. Fields are marked using field(metadata={WORLD_INDEXED_KEY: True}).

While adding this, I noticed that many arrays document false shape information. This is now also fixed.

Checking which arrays are world-indexed

We have a new utils function, world_mask, that maps out (possibly nested) structs. It replicates their structure exactly, but replaces leaves with a bool flag that indicates if they have a per-world axis. For each array with the WORLD_INDEXED_KEY flag set, it inserts a True.

Proper reset policies

The world axis allows us to formulate proper policies around resets. Obviously, masked resets can only reset arrays that are replicated per-world. Shared arrays cannot be decided. The new reset function uses a revised version of pytree_replace() that takes the world_mask PyTree and only resets the proper fields. Importantly, this also applies for all-True masks.

If None is passed, we instead reset everything, including the shared data, because the case is trivially decidable, and the behavior is what users would expect.

Support for sharding across devices

Having introduced all this machinery, sharding, luckily, becomes largely trivial. We shard all world-axis arrays evenly across devices and replicate shared arrays. Selecting world-batched arrays is trivial with the new world_mask function.

One caveat is that we currently cannot support sharing in explicit mode. SciPy introduces some scatter operations that cannot be resolved in explicit mode, so we pre-select jax.sharding.AxisType.Auto (see https://docs.jax.dev/en/latest/parallel.html#auto-sharding-mode-decides-shardings-automatically-during-compilation). Fixing this will require work upstream in SciPy and is out of scope for now.

Related to efficient sharding: #98

[amacati]

This new merge reworks the mechanism to use core_ndim instead. Previously, we would have to redefine the metadata field if an array got batched, e.g. for domain randomization. Now, we automatically detect that an array no longer has its core dimensions, and hence must be batched. Everything else stays the same.

This is important, because the metadata is a class variable. We cannot change the metadata for a single simulation, because that would break any other existing simulation that has not batched a particular field. With core_ndim, metadata can stay static, batched parameters get automatically promoted to have a world axis, and sizing down parameters again automatically demotes them to shared parameters.

@ratheron