jax

b1d37d1d - Merge pull request #21071 from mattjj:vmap-spmd-axis-name-errors