cupy.get_array_module#
- cupy.get_array_module(*args)[source]#
Returns the array module for arguments.
This function is used to implement CPU/GPU generic code. If at least one of the arguments is a
cupy.ndarray
object, thecupy
module is returned.- Parameters:
args – Values to determine whether NumPy or CuPy should be used.
- Returns:
cupy
ornumpy
is returned based on the types of the arguments.- Return type:
module
Example
A NumPy/CuPy generic function can be written as follows
>>> def softplus(x): ... xp = cupy.get_array_module(x) ... return xp.maximum(0, x) + xp.log1p(xp.exp(-abs(x)))