Only check api_version in __array_namespace__ when a wrapper is used

Author: asmeurer

array_api_compat/common/_helpers.py

@@ -248,10 +248,10 @@ def your_function(x, y):
     namespaces = set()
     for x in xs:
         if is_numpy_array(x):
-            _check_api_version(api_version)
             from .. import numpy as numpy_namespace
             import numpy as np
             if use_compat is True:
+                _check_api_version(api_version)
                 namespaces.add(numpy_namespace)
             elif use_compat is False:
                 namespaces.add(np)
@@ -263,32 +263,32 @@ def your_function(x, y):
                 else:
                     namespaces.add(numpy_namespace)
         elif is_cupy_array(x):
-            _check_api_version(api_version)
             if _use_compat:
+                _check_api_version(api_version)
                 from .. import cupy as cupy_namespace
                 namespaces.add(cupy_namespace)
             else:
                 import cupy as cp
                 namespaces.add(cp)
         elif is_torch_array(x):
-            _check_api_version(api_version)
             if _use_compat:
+                _check_api_version(api_version)
                 from .. import torch as torch_namespace
                 namespaces.add(torch_namespace)
             else:
                 import torch
                 namespaces.add(torch)
         elif is_dask_array(x):
-            _check_api_version(api_version)
             if _use_compat:
+                _check_api_version(api_version)
                 from ..dask import array as dask_namespace
                 namespaces.add(dask_namespace)
             else:
                 import dask.array as da
                 namespaces.add(da)
         elif is_jax_array(x):
-            _check_api_version(api_version)
             if use_compat is True:
+                _check_api_version(api_version)
                 raise ValueError("JAX does not have an array-api-compat wrapper")
             elif use_compat is False:
                 import jax.numpy as jnp

tests/test_array_namespace.py

@@ -12,7 +12,7 @@
 from ._helpers import import_, all_libraries, wrapped_libraries
 
 @pytest.mark.parametrize("use_compat", [True, False, None])
-@pytest.mark.parametrize("api_version", [None, "2021.12"])
+@pytest.mark.parametrize("api_version", [None, "2021.12", "2022.12"])
 @pytest.mark.parametrize("library", all_libraries + ['array_api_strict'])
 def test_array_namespace(library, api_version, use_compat):
     xp = import_(library)
@@ -69,14 +69,14 @@ def test_array_namespace_errors_torch():
     pytest.raises(TypeError, lambda: array_namespace(x, y))
 
 def test_api_version():
-    x = np.asarray([1, 2])
-    np_ = import_("numpy", wrapper=True)
-    assert array_namespace(x, api_version="2022.12") == np_
-    assert array_namespace(x, api_version=None) == np_
-    assert array_namespace(x) == np_
+    x = torch.asarray([1, 2])
+    torch_ = import_("torch", wrapper=True)
+    assert array_namespace(x, api_version="2022.12") == torch_
+    assert array_namespace(x, api_version=None) == torch_
+    assert array_namespace(x) == torch_
     # Should issue a warning
     with warnings.catch_warnings(record=True) as w:
-        assert array_namespace(x, api_version="2021.12") == np_
+        assert array_namespace(x, api_version="2021.12") == torch_
         assert len(w) == 1
         assert "2021.12" in str(w[0].message)