def shard_batch(batch, sharding):
    """Shards a batch across devices along its first dimension.

    Args:
        batch: A pytree of arrays.
        sharding: A jax Sharding object with shape (num_devices,).
    """
    return jax.tree_map(
        lambda x: jax.device_put(
            x, sharding.reshape(sharding.shape[0], *((1,) * (x.ndim - 1)))
        ),
        batch,
    )
Shards a batch across devices along its first dimension.

Args:
batch: A pytree of arrays.
sharding: A jax Sharding object with shape (num_devices,).
serl_launcher/serl_launcher/common/common.py