Merge pull request #147 from asmeurer/array_namespace-scalars

Allow Python scalars in array_namespace

Author: asmeurer

array_api_compat/common/_helpers.py

@@ -428,7 +428,8 @@ def array_namespace(*xs, api_version=None, use_compat=None):
     Parameters
     ----------
     xs: arrays
-        one or more arrays.
+        one or more arrays. xs can also be Python scalars (bool, int, float,
+        complex, or None), which are ignored.
 
     api_version: str
         The newest version of the spec that you need support for (currently
@@ -491,7 +492,9 @@ def your_function(x, y):
 
     namespaces = set()
     for x in xs:
-        if is_numpy_array(x):
+        if isinstance(x, (bool, int, float, complex, type(None))):
+            continue
+        elif is_numpy_array(x):
             from .. import numpy as numpy_namespace
             import numpy as np
             if use_compat is True:

tests/test_array_namespace.py

@@ -114,3 +114,19 @@ def test_api_version():
 def test_get_namespace():
     # Backwards compatible wrapper
     assert array_api_compat.get_namespace is array_api_compat.array_namespace
+
+def test_python_scalars():
+    a = torch.asarray([1, 2])
+    xp = import_("torch", wrapper=True)
+
+    pytest.raises(TypeError, lambda: array_namespace(1))
+    pytest.raises(TypeError, lambda: array_namespace(1.0))
+    pytest.raises(TypeError, lambda: array_namespace(1j))
+    pytest.raises(TypeError, lambda: array_namespace(True))
+    pytest.raises(TypeError, lambda: array_namespace(None))
+
+    assert array_namespace(a, 1) == xp
+    assert array_namespace(a, 1.0) == xp
+    assert array_namespace(a, 1j) == xp
+    assert array_namespace(a, True) == xp
+    assert array_namespace(a, None) == xp