Provides an interface to import numerical libraries using the GPU (if available).
The main goal of this is to smooth over the differences between numpy and cupy so that
the rest of the code can use them interchangably. We also need to monkey patch scipy’s ivp solver
to work on cupy arrays.
Linters HATE this module because it’s really abusing the import system (by design).
Perform imports for libraries with APIs matching numpy, scipy.integrate.ivp, scipy.sparse.
Reimport xp, xp_sparse, xp_ivp from the global context (in case they’ve been update to cupy).
Decorator that pulls xp, xp_sparse, xp_ivp from the global bucky context into the wrapped function.
to_cpu_noop(x, stream=None, order='C', out=None)
Simple class to mock the optuna warning if we don’t have optuna.
These imports will use a monkey-patched version of these modules
that has had all it’s numpy references replaced with CuPy.
if optimize is True, place the kernel optimization context in xp.optimize_kernels,
otherwise make it a nullcontext (noop)
returns nothing but imports a version of ‘xp’, ‘ivp’, and ‘sparse’ to the global scope of this module
optimize (bool) – Enable kernel optimization in cupy >=v8.0.0. This will slow down initial
function call (mostly reduction operations) but will offer better
performance for repeated calls (e.g. in the RHS call of an integrator).
exit_code – Non-zero value indicates error code, or zero on success.
NotImplementedError – If the user calls a monkeypatched function of the libs that isn’t