jax

77944e22 - Improve error message for sharding rules when the dimension does not divide mesh axes it's sharded on