Closed mdhaber closed 2 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.
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
xp.rand
, multiplying is much faster, bringing the total execution time from 240ms to 53ms.cupy.scipy.special.ndtr
. The slow part is again generating the resamples - specifically a reshape
operation in _jackknife_resample
. I could probably find a more efficient way to generate the jacknife resamples on the GPU.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.
Experimental use of CuPy in scipy.special.bootstrap