jax

83c8ce4a - [Pallas/Mosaic GPU] Fix lowering of `_device_id_to_logical`.

[Pallas/Mosaic GPU] Fix lowering of `_device_id_to_logical`. This fixes `test_remote_dma_2d` tests. * Before: [`gpu_pallas_distributed_test` fails](https://github.com/jax-ml/jax/actions/runs/24652765874/job/72079022408) * After: [`gpu_pallas_distributed_test` passes](https://github.com/jax-ml/jax/actions/runs/24656255289/job/72090412291?pr=36973) PiperOrigin-RevId: 902546176