jax

81f63ed1 - Fix a bug in `device_put` lowering introduced by https://github.com/google/jax/pull/21754

Fix a bug in `device_put` lowering introduced by https://github.com/google/jax/pull/21754 Also adds a test that triggers the issue. Confirmed that the test fails without the fix. PiperOrigin-RevId: 644117038