vortex_torch.cache.mask

Classes

MaskSlice(start, end, dim[, alpha, beta])

Position-dependent slice mask over one inner axis (cache side).

class MaskSlice(start, end, dim, alpha=1.0, beta=0.0)[source]

Bases: vOp

Position-dependent slice mask over one inner axis (cache side).

Math:

For target axis dim (1\(D_0\), 2\(D_1\)) and index \(i\) along it (other axes broadcast unchanged):

\[\begin{split}Y_{\dots,i,\dots} = \begin{cases} \alpha, & \text{start} \le i < \text{end}, \\ \beta, & \text{otherwise}. \end{cases}\end{split}\]
__init__:

MaskSlice(start, end, dim, alpha=1.0, beta=0.0) — write \(\alpha\) on [start, end) of axis dim (1 or 2) and \(\beta\) elsewhere.

__call__:

y = op(x, loc=loc, ctx=ctx)x [B, D_0, D_1] → same shape. A pure position writer (x values are unused).

Note:

only dim {1, 2}.

Parameters: