Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/opt_einsum/backends/cupy.py: 100%

1""" 

2Required functions for optimized contractions of numpy arrays using cupy. 

3""" 

4 

5import numpy as np 

6 

7from ..sharing import to_backend_cache_wrap 

8 

9__all__ = ["to_cupy", "build_expression", "evaluate_constants"] 

10 

11 

12@to_backend_cache_wrap 

13def to_cupy(array): # pragma: no cover 

14 import cupy 

15 

16 if isinstance(array, np.ndarray): 

17 return cupy.asarray(array) 

18 

19 return array 

20 

21 

22def build_expression(_, expr): # pragma: no cover 

23 """Build a cupy function based on ``arrays`` and ``expr``. 

24 """ 

25 def cupy_contract(*arrays): 

26 return expr._contract([to_cupy(x) for x in arrays], backend='cupy').get() 

27 

28 return cupy_contract 

29 

30 

31def evaluate_constants(const_arrays, expr): # pragma: no cover 

32 """Convert constant arguments to cupy arrays, and perform any possible 

33 constant contractions. 

34 """ 

35 return expr(*[to_cupy(x) for x in const_arrays], backend='cupy', evaluate_constants=True)