revert #22734

Reverts 5ce66dc

PiperOrigin-RevId: 657638187

Author: Jake VanderPlas

jax/_src/numpy/array_methods.py

@@ -39,7 +39,6 @@
 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
@@ -666,7 +665,6 @@ 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,14 +36,11 @@
 
 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,
@@ -200,3 +197,11 @@
     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

@@ -0,0 +1,31 @@
+# Copyright 2023 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.
+
+from __future__ import annotations
+
+import jax
+from jax._src.array import ArrayImpl
+from jax.experimental.array_api._version import __array_api_version__
+
+
+def _array_namespace(self, /, *, api_version: None | str = None):
+  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__]}")
+  return jax.experimental.array_api
+
+
+def add_array_object_methods():
+  # TODO(jakevdp): set on tracers as well?
+  setattr(ArrayImpl, "__array_namespace__", _array_namespace)

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

@@ -1,4 +1,4 @@
-# Copyright 2024 The JAX Authors.
+# Copyright 2023 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,50 +12,23 @@
 # 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__:
 
-# 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`_
+  def __init__(self):
+    self._capabilities = {
+      "boolean indexing": True,
+      "data-dependent shapes": False,
+    }
 
-  .. _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/numpy/__init__.py

@@ -274,10 +274,6 @@
 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,