jax

41220ca5 - Fix reshape sharding rule to handle merging and splitting in presence of singleton dimensions