jax

0e30a3ac - [mutable-arrays] read values should have the same explicit sharding as ref