jax

3919f0bd - [Mosaic:GPU] Replace jtu.test_device_matches(["cuda"]) with jtu.is_device_cuda().