jax

92fcd21c - [Pallas MGPU] Remove unnecessary synchronization if we’re not copying out refs.