bucky.numerical_libs
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.
Notes
Linters HATE this module because it’s really abusing the import system (by design).
enable_cupy(optimize=False)
enable_cupy
Perform imports for libraries with APIs matching numpy, scipy.integrate.ivp, scipy.sparse.
reimport_numerical_libs(context=None)
reimport_numerical_libs
Reimport xp, xp_sparse, xp_ivp from the global context (in case they’ve been update to cupy).
sync_numerical_libs(func)
sync_numerical_libs
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)
to_cpu_noop
bucky.numerical_libs.
ExperimentalWarning
_qr
empty_like_pinned
empty_pinned
optimize_kernels
reimport_cache
scatter_add
special
to_cpu
xp
xp_ivp
xp_sparse
zeros_like_pinned
zeros_pinned
MockExperimentalWarning
Bases: Warning
Warning
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).
bool
exit_code – Non-zero value indicates error code, or zero on success.
int
NotImplementedError – If the user calls a monkeypatched function of the libs that isn’t fully implemented.
bucky.make_input_graph
bucky.postprocess