def batch(x, num_devices):
    """Batch `x` so that so that leading axis is num_devices."""
    n = len(x)
    b = int(n / num_devices)
    return jnp.stack([jnp.stack(x[idx : idx + b]) for idx in range(0, n, b)])
Batch `x` so that so that leading axis is num_devices.
tools/train/scalable_shampoo/distributed_shampoo.py