jax

c99f50d6 - Merge pull request #32533 from jakevdp:dep-shard-map