Add the api_version keyword to get_namespace

Author: asmeurer

array_api_compat/common/_helpers.py

@@ -49,7 +49,11 @@ def is_array_api_obj(x):
         or _is_torch_array(x) \
         or hasattr(x, '__array_namespace__')
 
-def get_namespace(*xs, _use_compat=True):
+def _check_api_version(api_version):
+    if api_version is not None and api_version != '2021.12':
+        raise ValueError("Only the 2021.12 version of the array API specification is currently supported")
+
+def get_namespace(*xs, api_version=None, _use_compat=True):
     """
     Get the array API compatible namespace for the arrays `xs`.
 
@@ -61,28 +65,34 @@ def your_function(x, y):
             xp = array_api_compat.get_namespace(x, y)
             # Now use xp as the array library namespace
             return xp.mean(x, axis=0) + 2*xp.std(y, axis=0)
+
+    api_version should be the newest version of the spec that you need support
+    for (currently the compat library wrapped APIs only support v2021.12).
     """
     namespaces = set()
     for x in xs:
         if isinstance(x, (tuple, list)):
             namespaces.add(get_namespace(*x, _use_compat=_use_compat))
         elif hasattr(x, '__array_namespace__'):
-            namespaces.add(x.__array_namespace__())
+            namespaces.add(x.__array_namespace__(api_version=api_version))
         elif _is_numpy_array(x):
+            _check_api_version(api_version)
             if _use_compat:
                 from .. import numpy as numpy_namespace
                 namespaces.add(numpy_namespace)
             else:
                 import numpy as np
                 namespaces.add(np)
         elif _is_cupy_array(x):
+            _check_api_version(api_version)
             if _use_compat:
                 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:
                 from .. import torch as torch_namespace
                 namespaces.add(torch_namespace)

tests/test_get_namespace.py

@@ -4,22 +4,17 @@
 
 
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch"])
-def test_get_namespace(library):
+@pytest.mark.parametrize("api_version", [None, '2021.12'])
+def test_get_namespace(library, api_version):
     lib = pytest.importorskip(library)
 
     array = lib.asarray([1.0, 2.0, 3.0])
-    namespace = array_api_compat.get_namespace(array)
+    namespace = array_api_compat.get_namespace(array, api_version=api_version)
 
-    expected_namespace = getattr(array_api_compat, library)
-    assert namespace is expected_namespace
-
-
-@pytest.mark.parametrize("array_namespace", ["cupy.array_api", "numpy.array_api"])
-def test_get_namespace_returns_actual_namespace(array_namespace):
-    xp = pytest.importorskip(array_namespace)
-    X = xp.asarray([1, 2, 3])
-    xp_ = get_namespace(X)
-    assert xp_ is xp
+    if 'array_api' in library:
+        assert namespace == lib
+    else:
+        assert namespace == getattr(array_api_compat, library)
 
 def test_get_namespace_multiple():
     import numpy as np
@@ -38,3 +33,5 @@ def test_get_namespace_errors():
     y = torch.asarray([1, 2])
 
     pytest.raises(TypeError, lambda: get_namespace(x, y))
+
+    pytest.raises(ValueError, lambda: get_namespace(x, api_version='2022.12'))