[array api] fix deprecation to support old import pattern

Author: jakevdp

jax/experimental/__init__.py

@@ -25,20 +25,3 @@
 from jax._src.earray import (
     EArray as EArray
 )
-
-from jax import numpy as _array_api
-
-
-_deprecations = {
-  # Deprecated 01 Aug 2024
-  "array_api": (
-    "jax.experimental.array_api import is no longer required as of JAX v0.4.32; "
-    "jax.numpy supports the array API by default.",
-    _array_api
-  ),
-}
-
-from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
-__getattr__ = _deprecation_getattr(__name__, _deprecations)
-del _deprecation_getattr
-del _array_api

jax/experimental/array_api/__init__.py

@@ -0,0 +1,31 @@
+# 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
+
+import sys as _sys
+import warnings as _warnings
+
+import jax.numpy as _array_api
+
+_warnings.warn(
+    "jax.experimental.array_api import is no longer required as of JAX v0.4.32; "
+    "jax.numpy supports the array API by default.",
+    DeprecationWarning, stacklevel=2
+)
+
+_sys.modules['jax.experimental.array_api'] = _array_api
+
+del _array_api, _sys, _warnings

tests/array_api_test.py

@@ -249,8 +249,8 @@ def test_array_namespace_method(self):
   def test_deprecated_import(self):
     msg = "jax.experimental.array_api import is no longer required"
     with self.assertWarnsRegex(DeprecationWarning, msg):
-      from jax.experimental import array_api
-    self.assertIs(array_api, ARRAY_API_NAMESPACE)
+      import jax.experimental.array_api as nx
+    self.assertIs(nx, ARRAY_API_NAMESPACE)
 
 
 class ArrayAPIInspectionUtilsTest(jtu.JaxTestCase):