jax

a1224782 - Add a test to make sure axis_index behaves as expected under vmap when axis_data.name is not in axes.