jax

6ca3f74a - [mgpu] Set attributes on temporary module used for `cond` lowering.

[mgpu] Set attributes on temporary module used for `cond` lowering. The dummy module use for Pallas->MGPU `cond` lowering was missing the attributes used by `mgpu.utils.get_arch` to establish the current architecture. PiperOrigin-RevId: 912441721