Allow Python scalars in array_namespace

They are just ignored. This makes array_namespace easier to use for functions
that accept either arrays or scalars.

I'm not sure if I should have this behavior by default, or if it should be
enabled by a flag.

Author: asmeurer

array_api_compat/common/_helpers.py

@@ -235,7 +235,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,
+        or complex), which are ignored.
 
     api_version: str
         The newest version of the spec that you need support for (currently
@@ -298,7 +299,9 @@ def your_function(x, y):
 
     namespaces = set()
     for x in xs:
-        if is_numpy_array(x):
+        if isinstance(x, (bool, int, float, complex)):
+            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

@@ -92,3 +92,17 @@ 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))
+
+    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