mdhaber / scipy

Scipy library main repository
http://scipy.org/scipylib/
BSD 3-Clause "New" or "Revised" License
1 stars 5 forks source link

ENH: stats: allow bootstrap to use CuPy #63

Closed mdhaber closed 2 years ago

mdhaber commented 3 years ago

Experimental use of CuPy in scipy.special.bootstrap

rgommers commented 3 years ago

This is an interesting experiment. I believe we will not want to use an explicit import of CuPy within SciPy, rather we will want to use the array API standard and __array_namespace__. But that's not hard to change in the near future (once support in CuPy and NumPy is complete). The interesting part is that there is a large benefit.

mdhaber commented 2 years ago

I'm seeing ~10x speedup on this. (Update: 50x using rand and casting instead of randint)

import numpy as np
import cupy as cp
from scipy import stats

data_np = np.random.rand(10000)
rng_np = np.random.RandomState(0)
res_np = stats.bootstrap((data_np,), np.std, batch=1000,
                         random_state=rng_np, xp=np)  # 2.8 s ± 9.26 ms per loop 

data_cp = cp.array(data_np)
rng_cp = cp.random.RandomState(0)
res_cp = stats.bootstrap((data_cp,), cp.std, batch=1000, 
                         random_state=rng_cp, xp=cp)  # 240 ms ± 518 µs per loop  

No need to leave this PR open. Concept has been demonstrated. Actual implementation will depend on how SciPy decides to handle other array backends in general.