jax

82f4db2a - Use `ct_zero_from_primal` in shard_map_transpose to create zeros from residuals