Avoid use of deprecated xla_bridge.get_backend().live_buffers()

xla_bridge.get_backend is deprecated, and the public API for this is jax.live_arrays(). This is a drop-in replacement with no change of behavior.

PiperOrigin-RevId: 723227225

Author: Jake VanderPlas

trax/optimizers/trainer.py

@@ -440,8 +440,7 @@ def _collect_weights(self, layer):
 
   def _free_accelerators(self, exceptions=(), keep_constants=True):
     """Deletes all live buffers from accelerator with no safety guarantees."""
-    backend = jax.lib.xla_bridge.get_backend()
-    live_buffers = backend.live_buffers()
+    live_buffers = jax.live_arrays()
     logging.info('Deleting %d live buffers.', len(live_buffers))
     exceptions_buffers = []
     for x in fastmath.tree_flatten(exceptions):