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 as mask_func(a, k) returns a new array with zeros in certain locations (functions like cunumeric.triu() or cunumeric.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.