cupynumeric.mask_indices#

cupynumeric.mask_indices(
n: int,
mask_func: Callable[[ndarray, int], ndarray],
k: int = 0,
) → tuple[ndarray, ...]#

Return the indices to access (n, n) arrays, given a masking function.

Assume mask_func is a function that, for a square array a of size (n, n) with a possible offset argument k, when called as mask_func(a, k) returns a new array with zeros in certain locations (functions like cupynumeric.triu() or cupynumeric.tril() do precisely this). Then this function returns the indices where the non-zero values would be located.

Parameters:

  • n (int) – The returned indices will be valid to access arrays of shape (n, n).

  • mask_func (callable) – A function whose call signature is similar to that of cupynumeric.triu(), cupynumeric.tril(). That is, mask_func(x, k) returns a boolean array, shaped like x. k is an optional argument to the function.

  • k (scalar) – An optional argument which is passed through to mask_func. Functions like cupynumeric.triu(), cupynumeric,tril() take a second argument that is interpreted as an offset.

Returns:

indices – The n arrays of indices corresponding to the locations where mask_func(np.ones((n, n)), k) is True.

Return type:

tuple of arrays.

See also

numpy.mask_indices

Notes

WARNING: mask_indices expects mask_function to call cuPyNumeric functions for good performance. In case non-cuPyNumeric functions are called by mask_function, cuPyNumeric will have to materialize all data on the host which might result in running out of system memory.

Availability:

Multiple GPUs, Multiple CPUs

On this page