Clean up usage of jax.clear_backends

PiperOrigin-RevId: 616207814

Author: yueshengys

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)