asc2.mask

asc2.mask(*, count: RuntimeInt, other: RuntimeNumeric | None = None) → Generator[None, Any, None]
asc2.mask(*, bits: Iterable[RuntimeInt], other: RuntimeNumeric | None = None) → Generator[None, Any, None]

[Experimental] A context manager for masked operations on tensors.

Within the mask context, operations are applied only to the specified elements, with other elements optionally set to a different value.

Two masking modes are supported:

  1. Count-based masking: Apply operations to the first count elements along the innermost dimension. Elements beyond count are set to other (default 0).

  2. Bit-based masking: Apply operations to elements where the bit index (computed from position) falls within the range specified by bits. bits must contain exactly two integers defining the range.

Parameters:

  • count – The number of elements to apply operations to (from the start). Mutually exclusive with bits.

  • bits – A tuple of two integers defining a bit-based range for masking. Mutually exclusive with count.

  • other – The value to use for elements outside the mask. Default is 0.

Raises:

  • TypeError – If other is not a scalar, count is not an integer, or bits does not contain integers

  • RuntimeError – If bits does not contain exactly two integers

  • ValueError – If neither or both of count and bits are provided

Warning

This is an experimental API which is not guaranteed to work for every relevant vector operation. Its interface, availability, and functional coverage may change in the future.

Note

Exactly one of count or bits must be provided. This context manager can only be applied to vector operations (e.g., add, exp), not to load/store operations.

Examples

Apply addition only to first 8 elements, others set to 0:

with asc2.mask(count=8, other=0):
    result = tensor_a + tensor_b

Apply exp only to elements within bit range, others set to -1:

with asc2.mask(bits=[0, 64], other=-1):
    result = asc2.exp(tensor)