Don't bypass numpy compat just because it has __array_namespace__

asmeurer · asmeurer · commit 2d309019b5e9 · 2024-01-17T14:28:14.000-07:00
data-apis#77 (comment)
diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py
@@ -71,9 +71,7 @@ def your_function(x, y):
     """
     namespaces = set()
     for x in xs:
-        if hasattr(x, '__array_namespace__'):
-            namespaces.add(x.__array_namespace__(api_version=api_version))
-        elif _is_numpy_array(x):
+        if _is_numpy_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import numpy as numpy_namespace
@@ -97,6 +95,8 @@ def your_function(x, y):
             else:
                 import torch
                 namespaces.add(torch)
+        elif hasattr(x, '__array_namespace__'):
+            namespaces.add(x.__array_namespace__(api_version=api_version))
         else:
             # TODO: Support Python scalars?
             raise TypeError(f"{type(x).__name__} is not a supported array type")
