jax.clear_backends() is not doing what it is intended to do, users should try to avoid using it.

We decide to move it into `jax.extend`. This CL is the first step which adds a new module `jax.extend.backend`.

PiperOrigin-RevId: 615934218

Author: yueshengys

jax/BUILD

@@ -1024,6 +1024,7 @@ pytype_strict_library(
     visibility = [":jax_extend_users"],
     deps = [
         "//jax/extend",
+        "//jax/extend:backend",
         "//jax/extend:core",
         "//jax/extend:linear_util",
         "//jax/extend:random",

jax/extend/BUILD

@@ -26,6 +26,7 @@ pytype_strict_library(
     name = "extend",
     srcs = ["__init__.py"],
     deps = [
+        ":backend",
         ":core",
         ":linear_util",
         ":random",
@@ -45,6 +46,12 @@ pytype_strict_library(
     deps = ["//jax:core"],
 )
 
+pytype_strict_library(
+    name = "backend",
+    srcs = ["backend.py"],
+    deps = ["//jax"],
+)
+
 pytype_strict_library(
     name = "random",
     srcs = ["random.py"],

jax/extend/__init__.py

@@ -29,6 +29,7 @@
 """
 
 from jax.extend import (
+    backend as backend,
     core as core,
     linear_util as linear_util,
     random as random,

jax/extend/backend.py

@@ -0,0 +1,20 @@
+# Copyright 2024 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
+from jax._src.api import (
+  clear_backends as clear_backends,
+)

tests/extend_test.py

@@ -18,6 +18,7 @@
 import jax.extend as jex
 import jax.numpy as jnp
 
+from jax._src import api
 from jax._src import abstract_arrays
 from jax._src import linear_util
 from jax._src import prng
@@ -39,6 +40,7 @@ def test_symbols(self):
     self.assertIs(jex.random.unsafe_rbg_prng_impl, prng.unsafe_rbg_prng_impl)
 
     # Assume these are tested elsewhere, only check equivalence
+    self.assertIs(jex.backend.clear_backends, api.clear_backends)
     self.assertIs(jex.core.array_types, abstract_arrays.array_types)
     self.assertIs(jex.linear_util.StoreException, linear_util.StoreException)
     self.assertIs(jex.linear_util.WrappedFun, linear_util.WrappedFun)