jax

ca4d7884 - Make in_avals and out_avals match as well as we can with in_shardings and out_shardings respectively during lowering.