Repository: google/jax
Pull request #22825 (2024-08-01 21:20:24)

[array api] fix deprecation to support old import pattern
Followup to #22818

The previous approach made it so that a common import pattern didn't work:
```python
>>> import jax.experimental.array_api as nx
ModuleNotFoundError: No module named 'jax.experimental.array_api'
```
We need to keep a stub submodule here, but patch `sys.modules` to return the new path. With this change, it looks like this:
```python
In [1]: import jax.experimental.array_api as nx
<ipython-input-1-9d2abc00e9fa>:1: DeprecationWarning: jax.experimental.array_api import is no longer required as of JAX v0.4.32; jax.numpy supports the array API by default.
  import jax.experimental.array_api as nx

In [2]: import jax.numpy as jnp

In [3]: nx is jnp
Out[3]: True
```


Patch:
diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml
index 9869d0256..743aae9b4 100644
--- a/.github/workflows/ci-build.yaml
+++ b/.github/workflows/ci-build.yaml
@@ -141,7 +141,7 @@ jobs:
         PY_COLORS: 1
       run: |
         pytest -n auto --tb=short --doctest-glob='*.md' --doctest-glob='*.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md --ignore=docs/jax.experimental.array_api.rst
-        pytest -n auto --tb=short --doctest-modules jax --ignore=jax/config.py --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/_src/lib/triton.py --ignore=jax/_src/lib/mosaic_gpu.py --ignore=jax/interpreters/mlir.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py --ignore=jax/_src/tpu_custom_call.py --ignore=jax/experimental/mosaic --ignore=jax/experimental/pallas --ignore=jax/_src/pallas --ignore=jax/experimental/maps.py
+        pytest -n auto --tb=short --doctest-modules jax --ignore=jax/config.py --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/_src/lib/triton.py --ignore=jax/_src/lib/mosaic_gpu.py --ignore=jax/interpreters/mlir.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py --ignore=jax/_src/tpu_custom_call.py --ignore=jax/experimental/mosaic --ignore=jax/experimental/pallas --ignore=jax/_src/pallas --ignore=jax/experimental/array_api
 
 
   documentation_render:
diff --git a/jax/BUILD b/jax/BUILD
index 69be704ba..66df0d2f7 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -987,9 +987,13 @@ pytype_library(
 
 pytype_library(
     name = "experimental_array_api",
+    srcs = glob(
+        [
+            "experimental/array_api/*.py",
+        ],
+    ),
     visibility = [":internal"] + jax_visibility("array_api"),
     deps = [
-        ":experimental",
         ":jax",
     ],
 )
diff --git a/jax/experimental/__init__.py b/jax/experimental/__init__.py
index dcca44773..caf27ec7a 100644
--- a/jax/experimental/__init__.py
+++ b/jax/experimental/__init__.py
@@ -25,20 +25,3 @@ from jax._src.callback import (
 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
diff --git a/jax/experimental/array_api/__init__.py b/jax/experimental/array_api/__init__.py
new file mode 100644
index 000000000..69b25d0b6
--- /dev/null
+++ b/jax/experimental/array_api/__init__.py
@@ -0,0 +1,32 @@
+# 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
+
+# Added 2024-08-01
+_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
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index 16ee5e431..0ca47a8b9 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -21,8 +21,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "7f8cc3357fb6837ffcf1f4f05dab3d95dbd5a88c"
-XLA_SHA256 = "c760e02611d89c2c57372cfec906c25ff6b4eb92884cfbddf84dfd85c0dfa50a"
+XLA_COMMIT = "a3bc38ac798daf0e5999dc98982d97dd1a84379f"
+XLA_SHA256 = "f4fea8c7d5219505afd8d72380ae06aecbefc4efe09c0ca6bcc59f359035169a"
 
 def repo():
     tf_http_archive(


Test patch:
diff --git a/tests/array_api_test.py b/tests/array_api_test.py
index 41951fea3..2f3042fa2 100644
--- a/tests/array_api_test.py
+++ b/tests/array_api_test.py
@@ -249,8 +249,8 @@ class ArrayAPISmokeTest(absltest.TestCase):
   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):