[JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu").#1791

Merged

copybara-service[bot] merged 1 commit intomasterfrom test_582427745

Nov 15, 2023

+2-2

Commits on Nov 15, 2023

[JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu").

hawkinsp
authored and
copybara-github
committed