[JAX] Replace users of deprecated jax.lib.xla_bridge functions with more public names.

google / trax

Trax — Deep Learning with Clear Code and Speed

Apache License 2.0

8.07k stars 813 forks source link

Closed copybara-service[bot] closed 2 years ago

copybara-service[bot] commented 2 years ago

[JAX] Replace users of deprecated jax.lib.xla_bridge functions with more public names.

Change in preparation for removing some function exports from jax.lib.xla_bridge.