jax

cac2b8d5 - [Pallas/MGPU] Undo transforms before giving refs back to users

[Pallas/MGPU] Undo transforms before giving refs back to users This is a second attempt at this change. The first one was rolled back because of reported failures. Reverts 411928b9668570bbc3795522aba94cece6894881 PiperOrigin-RevId: 680943744