Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.9/dist-packages/scipy/_lib/array_api_compat/array_api_compat/_internal.py: 93%

14 statements  

« prev     ^ index     » next       coverage.py v7.3.1, created at 2023-09-23 06:43 +0000

1""" 

2Internal helpers 

3""" 

4 

5from functools import wraps 

6from inspect import signature 

7 

8def get_xp(xp): 

9 """ 

10 Decorator to automatically replace xp with the corresponding array module. 

11 

12 Use like 

13 

14 import numpy as np 

15 

16 @get_xp(np) 

17 def func(x, /, xp, kwarg=None): 

18 return xp.func(x, kwarg=kwarg) 

19 

20 Note that xp must be a keyword argument and come after all non-keyword 

21 arguments. 

22 

23 """ 

24 def inner(f): 

25 @wraps(f) 

26 def wrapped_f(*args, **kwargs): 

27 return f(*args, xp=xp, **kwargs) 

28 

29 sig = signature(f) 

30 new_sig = sig.replace(parameters=[sig.parameters[i] for i in sig.parameters if i != 'xp']) 

31 

32 if wrapped_f.__doc__ is None: 

33 wrapped_f.__doc__ = f"""\ 

34Array API compatibility wrapper for {f.__name__}. 

35 

36See the corresponding documentation in NumPy/CuPy and/or the array API 

37specification for more details. 

38 

39""" 

40 wrapped_f.__signature__ = new_sig 

41 return wrapped_f 

42 

43 return inner