Refactor how get_xp works

Now instead of guessing the array library from the input with get_namespace,
the array library is hard-coded into the wrapped function based on which
subnamespace it is imported from.

Author: asmeurer

array_api_compat/_internal.py

@@ -5,36 +5,40 @@
 from functools import wraps
 from inspect import signature
 
-from .common._helpers import get_namespace
-
-def get_xp(f):
+def get_xp(xp):
     """
-    Decorator to automatically replace xp with the corresponding array module
+    Decorator to automatically replace xp with the corresponding array module.
 
     Use like
 
-    @get_xp
+    import numpy as np
+
+    @get_xp(np)
     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.
+    Note that xp must be a keyword argument and come after all non-keyword
+    arguments.
+
     """
-    @wraps(f)
-    def inner(*args, **kwargs):
-        xp = get_namespace(*args, _use_compat=False)
-        return f(*args, xp=xp, **kwargs)
+    def inner(f):
+        sig = signature(f)
+
+        @wraps(f)
+        def wrapped_f(*args, **kwargs):
+            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'])
+        new_sig = sig.replace(parameters=[sig.parameters[i] for i in sig.parameters if i != 'xp'])
 
-    if inner.__doc__ is None:
-        inner.__doc__ = f"""\
+        if wrapped_f.__doc__ is None:
+            wrapped_f.__doc__ = f"""\
 Array API compatibility wrapper for {f.__name__}.
 
 See the corresponding documentation in NumPy/CuPy and/or the array API
 specification for more details.
 
 """
-    inner.__signature__ = new_sig
+        # wrapped_f.__signature__ = new_sig
+        return wrapped_f
 
     return inner