jax

6f55c374 - Make sure vary_unreduced_cast and reduced_vary_cast order the axes wrt the context mesh. Also code-golf some code in shard_map