jax

8e346ddb - Replace `with_partitions` and `with_unreduced` with `.update` on Partitions