jax

311fc965 - Keep the axis_env `axes` in mt.varying when remove_one_sized_mesh_axes is True