Add a get_xp decorator to support multiple array namespaces

Author: asmeurer

numpy_array_api_compat/_helpers.py

@@ -19,7 +19,7 @@ def is_array_api_obj(x):
     """
     return _is_numpy_array(x) or hasattr(x, '__array_namespace__')
 
-def get_namespace(*xs):
+def get_namespace(*xs, _use_compat=True):
     """
     Get the array API compatible namespace for the arrays `xs`.
 
@@ -30,7 +30,10 @@ def get_namespace(*xs):
         if hasattr(x, '__array_namespace__'):
             namespaces.add(x.__array_namespace__)
         elif _is_numpy_array(x):
-            namespaces.add(compat_namespace)
+            if _use_compat:
+                namespaces.add(compat_namespace)
+            else:
+                namespaces.add(np)
         else:
             # TODO: Support Python scalars?
             raise ValueError("The input is not a supported array type")

numpy_array_api_compat/_internal.py

@@ -0,0 +1,32 @@
+"""
+Internal helpers
+"""
+
+from functools import wraps
+from inspect import signature
+
+from ._helpers import get_namespace
+
+def get_xp(f):
+    """
+    Decorator to automatically replace xp with the corresponding array module
+
+    Use like
+
+    @get_xp
+    def func(x, /, xp, kwarg=None):
+        return xp.func(x, kwarg=kwarg)
+
+    Note that xp must be able to be passed as a keyword argument.
+    """
+    @wraps(f)
+    def inner(*args, **kwargs):
+        xp = get_namespace(*args, _use_compat=False)
+        return f(*args, xp=xp, **kwargs)
+
+    sig = signature(f)
+    new_sig = sig.replace(parameters=[sig.parameters[i] for i in sig.parameters if i != 'xp'])
+
+    inner.__signature__ = new_sig
+
+    return inner