cupy.get_array_module — CuPy 13.4.1 documentation (original) (raw)

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, the cupy module is returned.

Parameters:

args – Values to determine whether NumPy or CuPy should be used.

Returns:

cupy or numpy 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)))