Make use_compat a public flag to array_namespace()

If None, it does the default behavior. If True or False it forces the use or
non-use of the compat wrapper.

This commit also changes it so that NumPy 2.0 does not return the wrapper
library at all by default (when use_compat=None), since it is completely array
API compatible on its own.

Author: asmeurer

array_api_compat/common/_helpers.py

@@ -178,7 +178,7 @@ def _check_api_version(api_version):
     elif api_version is not None and api_version != '2022.12':
         raise ValueError("Only the 2022.12 version of the array API specification is currently supported")
 
-def array_namespace(*xs, api_version=None, _use_compat=True):
+def array_namespace(*xs, api_version=None, use_compat=None):
     """
     Get the array API compatible namespace for the arrays `xs`.
 
@@ -191,6 +191,12 @@ def array_namespace(*xs, api_version=None, _use_compat=True):
         The newest version of the spec that you need support for (currently
         the compat library wrapped APIs support v2022.12).
 
+    use_compat: bool or None
+        If None (the default), the native namespace will be returned if it is
+        already array API compatible, otherwise a compat wrapper is used. If
+        True, the compat library wrapped library will be returned. If False,
+        the native library namespace is returned.
+
     Returns
     -------
 
@@ -234,16 +240,28 @@ def your_function(x, y):
     is_jax_array
 
     """
+    if use_compat not in [None, True, False]:
+        raise ValueError("use_compat must be None, True, or False")
+
+    _use_compat = use_compat in [None, True]
+
     namespaces = set()
     for x in xs:
         if is_numpy_array(x):
             _check_api_version(api_version)
-            if _use_compat:
-                from .. import numpy as numpy_namespace
+            from .. import numpy as numpy_namespace
+            import numpy as np
+            if use_compat is True:
                 namespaces.add(numpy_namespace)
-            else:
-                import numpy as np
+            elif use_compat is False:
                 namespaces.add(np)
+            else:
+                # numpy 2.0 has __array_namespace__ and is fully array API
+                # compatible.
+                if hasattr(x, '__array_namespace__'):
+                    namespaces.add(x.__array_namespace__(api_version=api_version))
+                else:
+                    namespaces.add(numpy_namespace)
         elif is_cupy_array(x):
             _check_api_version(api_version)
             if _use_compat:
@@ -266,14 +284,22 @@ def your_function(x, y):
                 from ..dask import array as dask_namespace
                 namespaces.add(dask_namespace)
             else:
-                raise TypeError("_use_compat cannot be False if input array is a dask array!")
+                import dask.array as da
+                namespaces.add(da)
         elif is_jax_array(x):
             _check_api_version(api_version)
-            # jax.experimental.array_api is already an array namespace. We do
-            # not have a wrapper submodule for it.
-            import jax.experimental.array_api as jnp
+            if use_compat is True:
+                raise ValueError("JAX does not have an array-api-compat wrapper")
+            elif use_compat is False:
+                import jax.numpy as jnp
+            else:
+                # jax.experimental.array_api is already an array namespace. We do
+                # not have a wrapper submodule for it.
+                import jax.experimental.array_api as jnp
             namespaces.add(jnp)
         elif hasattr(x, '__array_namespace__'):
+            if use_compat is True:
+                raise ValueError("The given array does not have an array-api-compat wrapper")
             namespaces.add(x.__array_namespace__(api_version=api_version))
         else:
             # TODO: Support Python scalars?

tests/_helpers.py

@@ -2,8 +2,11 @@
 
 import pytest
 
-wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"]
-all_libraries = wrapped_libraries + ["jax.numpy"]
+wrapped_libraries = ["cupy", "torch", "dask.array"]
+all_libraries = wrapped_libraries + ["numpy", "jax.numpy"]
+import numpy as np
+if np.__version__[0] == '1':
+    wrapped_libraries.append("numpy")
 
 def import_(library, wrapper=False):
     if library == 'cupy':

tests/test_array_namespace.py

@@ -9,24 +9,29 @@
 import array_api_compat
 from array_api_compat import array_namespace
 
-from ._helpers import import_, all_libraries
+from ._helpers import import_, all_libraries, wrapped_libraries
 
-@pytest.mark.parametrize("library", all_libraries)
+@pytest.mark.parametrize("use_compat", [True, False, None])
 @pytest.mark.parametrize("api_version", [None, "2021.12"])
-def test_array_namespace(library, api_version):
+@pytest.mark.parametrize("library", all_libraries + ['array_api_strict'])
+def test_array_namespace(library, api_version, use_compat):
     xp = import_(library)
 
     array = xp.asarray([1.0, 2.0, 3.0])
-    namespace = array_api_compat.array_namespace(array, api_version=api_version)
+    if use_compat is True and library in ['array_api_strict', 'jax.numpy']:
+        pytest.raises(ValueError, lambda: array_namespace(array, use_compat=use_compat))
+        return
+    namespace = array_api_compat.array_namespace(array, api_version=api_version, use_compat=use_compat)
 
-    if "array_api" in library:
-        assert namespace == xp
+    if use_compat is False or use_compat is None and library not in wrapped_libraries:
+        if library == "jax.numpy" and use_compat is None:
+            import jax.experimental.array_api
+            assert namespace == jax.experimental.array_api
+        else:
+            assert namespace == xp
     else:
         if library == "dask.array":
             assert namespace == array_api_compat.dask.array
-        elif library == "jax.numpy":
-            import jax.experimental.array_api
-            assert namespace == jax.experimental.array_api
         else:
             assert namespace == getattr(array_api_compat, library)