cunumeric.mask_indices#
- cunumeric.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 asmask_func(a, k)
returns a new array with zeros in certain locations (functions likecunumeric.triu()
orcunumeric.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
cunumeric.triu()
,cunumeric.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
cunumeric.triu()
,cunumeric,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
Notes
WARNING: mask_indices expects mask_function to call cuNumeric functions for good performance. In case non-cuNumeric functions are called by mask_function, cuNumeric will have to materialize all data on the host which might result in running out of system memory.
- Availability:
Multiple GPUs, Multiple CPUs