Set __array_api_version__ with the api_version flag

Author: asmeurer

array_api_strict/__init__.py

@@ -16,7 +16,11 @@
 
 """
 
-__array_api_version__ = "2022.12"
+# Warning: __array_api_version__ could change globally with
+# set_array_api_strict_flags(). This should always be accessed as an
+# attribute, like xp.__array_api_version__, or using
+# array_api_strict.get_array_api_strict_flags()['api_version'].
+from ._flags import API_VERSION as __array_api_version__
 
 __all__ = ["__array_api_version__"]
 

array_api_strict/_array_object.py

@@ -497,7 +497,7 @@ def __array_namespace__(
         the API version for the array_api_strict module globally. This can
         also be achieved with the
         {func}`array_api_strict.set_array_api_strict_flags` function. If you
-        want some way to only set the version locally, use the
+        want to only set the version locally, use the
         {class}`array_api_strict.ArrayApiStrictFlags` context manager.
 
         """

array_api_strict/_flags.py

@@ -14,6 +14,8 @@
 import functools
 import os
 
+import array_api_strict
+
 supported_versions = (
     "2021.12",
     "2022.12",
@@ -37,7 +39,6 @@
     "linalg",
     "fft",
 )
-
 # Public functions
 
 def set_array_api_strict_flags(
@@ -118,6 +119,7 @@ def set_array_api_strict_flags(
         if api_version not in supported_versions:
             raise ValueError(f"Unsupported standard version {api_version!r}")
         API_VERSION = api_version
+        array_api_strict.__array_api_version__ = API_VERSION
 
     if data_dependent_shapes is not None:
         DATA_DEPENDENT_SHAPES = data_dependent_shapes
@@ -206,6 +208,7 @@ def reset_array_api_strict_flags():
     """
     global API_VERSION, DATA_DEPENDENT_SHAPES, ENABLED_EXTENSIONS
     API_VERSION = default_version
+    array_api_strict.__array_api_version__ = API_VERSION
     DATA_DEPENDENT_SHAPES = True
     ENABLED_EXTENSIONS = default_extensions
 

array_api_strict/tests/test_array_object.py

@@ -402,9 +402,17 @@ def test_array_keys_use_private_array():
 def test_array_namespace():
     a = ones((3, 3))
     assert a.__array_namespace__() == array_api_strict
+    assert array_api_strict.__array_api_version__ == "2022.12"
+
     assert a.__array_namespace__(api_version=None) is array_api_strict
+    assert array_api_strict.__array_api_version__ == "2022.12"
+
     assert a.__array_namespace__(api_version="2022.12") is array_api_strict
+    assert array_api_strict.__array_api_version__ == "2022.12"
+
     with pytest.warns(UserWarning):
         assert a.__array_namespace__(api_version="2021.12") is array_api_strict
+    assert array_api_strict.__array_api_version__ == "2021.12"
+
     pytest.raises(ValueError, lambda: a.__array_namespace__(api_version="2021.11"))
     pytest.raises(ValueError, lambda: a.__array_namespace__(api_version="2023.12"))

array_api_strict/tests/test_flags.py

@@ -63,6 +63,14 @@ def test_flags():
         'enabled_extensions': ('linalg', 'fft'),
     }
 
+def test_api_version():
+    # Test defaults
+    assert xp.__array_api_version__ == '2022.12'
+
+    # Test setting the version
+    set_array_api_strict_flags(api_version='2021.12')
+    assert xp.__array_api_version__ == '2021.12'
+
 def test_data_dependent_shapes():
     a = asarray([0, 0, 1, 2, 2])
     mask = asarray([True, False, True, False, True])