jax

ad30807b - Make auto_axes a no-op when the mesh is empty. This is required because library code can be written for no mesh context or a mesh context and auto-axes should work in both places. The other option is for user to do this kind of switching but I feel it's not too bad for JAX to do this since auto_axes already is a no-op if called under auto mode.