jax

8649132d - [pallas] Support DMA start partial discharge and run_scoped() does its own partial discharge.

[pallas] Support DMA start partial discharge and run_scoped() does its own partial discharge. This CL lays the ground for a future CL that makes run_scoped discharge to not request the discharge of the temporary buffers it creates. This causes issues becausa a) dma_start can't discharge some but not all its references b) run_scoped() lowering depends on run_scoped discharge to remove the run_scoped operation (or it goes in an infinite loop). PiperOrigin-RevId: 722126566