Merge pull request #22734 from jakevdp:array-api-methods

PiperOrigin-RevId: 657299896

Author: jax authors

jax/experimental/array_api/_utility_functions.py jax/_src/numpy/array_api_metadata.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The JAX Authors.
+# 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.
@@ -12,23 +12,50 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""
+This module contains metadata related to the `Python array API`_.
+
+.. _Python array API: https://data-apis.org/array-api/
+"""
 from __future__ import annotations
 
 import jax
 from jax._src.sharding import Sharding
 from jax._src.lib import xla_client as xc
 from jax._src import dtypes as _dtypes, config
 
-# TODO(micky774): Add to jax.numpy.util when finalizing jax.experimental.array_api
-# deprecation
-class __array_namespace_info__:
 
-  def __init__(self):
-    self._capabilities = {
-      "boolean indexing": True,
-      "data-dependent shapes": False,
-    }
+# TODO(jakevdp, vfdev-5): export this in jax.numpy once migration is complete.
+__array_api_version__ = '2023.12'
+
+
+# TODO(jakevdp, vfdev-5): export this in jax.numpy once migration is complete.
+def __array_namespace_info__() -> ArrayNamespaceInfo:
+  return ArrayNamespaceInfo()
+
+
+def __array_namespace__(self, /, *, api_version: None | str = None):
+  """Return the `Python array API`_ namespace for JAX.
+
+  .. _Python array API: https://data-apis.org/array-api/
+  """
+  if api_version is not None and api_version != __array_api_version__:
+    raise ValueError(f"{api_version=!r} is not available; "
+                     f"available versions are: {[__array_api_version__]}")
+  # TODO(jakevdp, vfdev-5): change this to jax.numpy once migration is complete.
+  import jax.experimental.array_api
+  return jax.experimental.array_api  # pytype: disable=module-attr
+
+
+class ArrayNamespaceInfo:
+  """Metadata for the `Python array API`_
 
+  .. _Python array API: https://data-apis.org/array-api/
+  """
+  _capabilities = {
+    "boolean indexing": True,
+    "data-dependent shapes": False,
+  }
 
   def _build_dtype_dict(self):
     array_api_types = {

jax/_src/numpy/array_methods.py

@@ -39,6 +39,7 @@
 from jax._src.array import ArrayImpl
 from jax._src.lax import lax as lax_internal
 from jax._src.lib import xla_client as xc
+from jax._src.numpy import array_api_metadata
 from jax._src.numpy import lax_numpy
 from jax._src.numpy import reductions
 from jax._src.numpy import ufuncs
@@ -665,6 +666,7 @@ def max(self, values, *, indices_are_sorted=False, unique_indices=False,
 }
 
 _array_methods = {
+  "__array_namespace__": array_api_metadata.__array_namespace__,
   "all": reductions.all,
   "any": reductions.any,
   "argmax": lax_numpy.argmax,

jax/experimental/array_api/__init__.py

@@ -36,11 +36,14 @@
 
 from __future__ import annotations
 
-from jax.experimental.array_api._version import __array_api_version__ as __array_api_version__
-
 from jax.experimental.array_api import fft as fft
 from jax.experimental.array_api import linalg as linalg
 
+from jax._src.numpy.array_api_metadata import (
+    __array_api_version__ as __array_api_version__,
+    __array_namespace_info__ as __array_namespace_info__,
+)
+
 from jax.numpy import (
     abs as abs,
     acos as acos,
@@ -197,11 +200,3 @@
     clip as clip,
     hypot as hypot,
 )
-
-from jax.experimental.array_api._utility_functions import (
-    __array_namespace_info__ as __array_namespace_info__,
-)
-
-from jax.experimental.array_api import _array_methods
-_array_methods.add_array_object_methods()
-del _array_methods

jax/experimental/array_api/_array_methods.py

@@ -274,6 +274,10 @@
 except ImportError:
   pass
 
+from jax._src.numpy.array_api_metadata import (
+  __array_api_version__ as __array_api_version__
+)
+
 from jax._src.numpy.index_tricks import (
   c_ as c_,
   index_exp as index_exp,