[JAX] Replace uses of jax.xla.Device with jax.Device.

google-research / jax3d

Apache License 2.0

737 stars 96 forks source link

Closed copybara-service[bot] closed 1 year ago

copybara-service[bot] commented 1 year ago

[JAX] Replace uses of jax.xla.Device with jax.Device.

jax.xla was an accidental export from the jax namespace. jax.Device is the public name for JAX devices, as of JAX 0.4.3.

This is a trivial and safe change: jax.xla.Device and jax.Device are aliases.