Clean up usage of jax.clear_backends

PiperOrigin-RevId: 615956199
jax-ml · Mar 15, 2024 · 5c187f6 · 5c187f6
1 parent f569031
commit 5c187f6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/clear_backends_test.py b/tests/clear_backends_test.py
@@ -14,9 +14,9 @@
 """Tests for release_backend_clients."""
 
 from absl.testing import absltest
-
 import jax
 from jax import config
+from jax._src import api
 from jax._src import test_util as jtu
 from jax._src import xla_bridge as xb
 
@@ -29,7 +29,7 @@ def test_clear_backends(self):
     g = jax.jit(lambda x, y: x * y)
     self.assertEqual(g(1, 2), 2)
     self.assertNotEmpty(xb.get_backend().live_executables())
-    jax.clear_backends()
+    api.clear_backends()
     self.assertEmpty(xb.get_backend().live_executables())
     self.assertEqual(g(1, 2), 2)