cuquantum.densitymat.GPUCallback

class cuquantum.densitymat.GPUCallback(callback, is_inplace=False, gradient_callback=None, gradient_dir=None)[source]

Wrapper for GPU callback functions.

Note

  • Out-of-place callback functions have a signature of (t: float, args: cp.ndarray) -> cp.ndarray.

  • Inplace callback functions have a signature of (t: float, args: cp.ndarray, arr: cp.ndarray) -> None.

  • Callback functions are invoked in a stream and device context and are assumed to return (for out-of-place callbacks) a cp.ndarray on the current device.

  • Callback functions receive args argument as cp.ndarray of shape (num_params, batch_size).

  • In-place gradient callbacks are required to accumulate the gradient contributions of the callback instead of overwriting the output array.

  • Backward gradient callback functions have a signature of (t: float, args: cp.ndarray, adj: cp.ndarray) -> cp.ndarray. adj is the adjoint computed through backward pass and of the same shape as the output of the regular callback.

  • In-place gradient callbacks have a signature of (t: float, args: cp.ndarray, adj: cp.ndarray, args_grad: cp.ndarray) -> None. args_grad is the array to accumulate the gradient contributions into, and of the same shape as args.

Parameters

  • callback – The callback function to be wrapped, which can be either an inplace callback modifying its third argument, or an out-of-place callback returning a cp.ndarray.

  • is_inplace – Specifies whether the callback is inplace.

  • gradient_callback – The gradient callback function to be wrapped. If not specified, the callback is not considered differentiable and the contribution of this callback to the operator action gradient is not taken into account.

  • gradient_dir – Whether or not the gradient is evaluated in the forward or backward direction. All differentiable callbacks in a Operator must have the same gradient direction.

Methods

__init__(callback: Union[Callable[[float, cupy.ndarray, cupy.ndarray], None], Callable[[float, cupy.ndarray], cupy.ndarray]], is_inplace: bool = False, gradient_callback: Optional[Union[Callable[[float, cupy.ndarray, cupy.ndarray, cupy.ndarray], None], Callable[[float, cupy.ndarray, cupy.ndarray], cupy.ndarray]]] = None, gradient_dir: str | None = None)[source]

Initialize a Callback object.

Parameters

  • callback – The callback function to be wrapped.

  • is_inplace – Specifies whether the callback is inplace.

  • gradient_callback – The gradient callback function to be wrapped.

  • gradient_dir – The gradient direction.

Attributes

callback

Return the callback function wrapped by this Callback.

gradient_callback

Return the gradient callback function wrapped by this Callback.

gradient_dir

Return the gradient direction of the gradient callback function wrapped by this Callback.

has_gradient

Returns whether a gradient callback function is available for this Callback instance.

inplace_callback

Return the inplace version of the callback function wrapped by this Callback.

inplace_gradient_callback

Return the inplace version of the gradient callback function wrapped by this Callback.

is_gpu_callback

Return whether the callback function operates on GPU arrays.

is_inplace

Returns True if the callback function wrapped by this Callback is acting in-place instead of returning its result.