MaskIndices

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

Abstract Signature:

MaskIndices(n: int, mask_func: Callable, k: int = 0)

PyTorch

API: —
Strategy: Plugin (numpy_fallback)

JAX (Core)

API: jax.numpy.mask_indices
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: numpy.mask_indices
Strategy: Direct Mapping
Official Docs ↗

Keras

API: —
Strategy: Plugin (numpy_fallback)

TensorFlow

API: —
Strategy: Plugin (numpy_fallback)

Apple MLX

API: —
Strategy: Plugin (numpy_fallback)

Flax NNX

API: jax.numpy.mask_indices
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.numpy.mask_indices
Strategy: Direct Mapping
Official Docs ↗