Update __array_namespace__ to accept api_version. (#12)

Quantco · Jul 1, 2024 · e56a4ad · e56a4ad
1 parent 0563003
commit e56a4ad
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -13,6 +13,10 @@ Changelog
 
 - ndonnx now exports type annotations.
 
+**Bug fixes**
+
+- ``__array_namespace__`` now accepts the optional ``api_version`` argument to specify the version of the Array API to use.
+
 0.4.0 (2024-05-16)
 ------------------
 

diff --git a/ndonnx/_array.py b/ndonnx/_array.py
@@ -329,8 +329,14 @@ def spox_var(self) -> Var:
         return self._core().var
 
     @staticmethod
-    def __array_namespace__():
-        return ndx
+    def __array_namespace__(*, api_version: str | None = None):
+        if api_version is None:
+            api_version = "2023.12"
+
+        if api_version == "2023.12":
+            return ndx
+        else:
+            raise ValueError(f"Unsupported API version {api_version}")
 
     def __float__(self) -> float:
         eager_value = self.to_numpy()