jax

719031c1 - [mutable-arrays] persist shardings through xla computations