Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsPraxisnemo_automodel.components.models.praxis.masks

nemo_automodel.components.models.praxis.masks

View as MarkdownOpen in Claude

Attention masks of the chunk-causal window.

Every token of a sequence carries three integer tags: its kind (global, noised video, clean video), its chunk index and its latent frame index. A mask rule is a pure function of the tags of the query and the key, so one rule serves both the flex BlockMask (CUDA) and a dense boolean mask (CPU and SDPA).

Module Contents

Classes

NameDescription
MaskCacheSmall LRU cache of built masks keyed by the caller’s layout key.
SlotA run of tokens of one kind: frames latent frames of tokens_per_frame tokens each.
TokenTagsPer-token tags of one sequence, all int32 tensors of shape [tokens] on one device.

Functions

NameDescription
block_maskBuild a flex BlockMask over [q_len, kv_len] tiles of block_size from rule.
chunk_causal_ruleThe teacher-forcing rule of the trunk towers over one sequence.
dense_maskEvaluate rule on the full grid.
expert_cross_ruleRule of the action expert’s cross attention.
expert_self_ruleRule of the action expert’s self attention: one block per chunk.
window_tagsTags of a window sequence laid out as slots in order.

Data

KIND_CLEAN

KIND_GLOBAL

KIND_NOISED

MaskKey

MaskRule

API

class nemo_automodel.components.models.praxis.masks.MaskCache(
    max_entries: int = 64
)

Small LRU cache of built masks keyed by the caller’s layout key.

_entries
OrderedDict[MaskKey, BlockMask | Tensor] = OrderedDict()
nemo_automodel.components.models.praxis.masks.MaskCache.get(
    key: nemo_automodel.components.models.praxis.masks.MaskKey,
    build: typing.Callable[[], torch.nn.attention.flex_attention.BlockMask | torch.Tensor]
) -> torch.nn.attention.flex_attention.BlockMask | torch.Tensor
class nemo_automodel.components.models.praxis.masks.Slot(
    kind: int,
    frames: int,
    tokens_per_frame: int
)
Dataclass

A run of tokens of one kind: frames latent frames of tokens_per_frame tokens each.

A global slot has frames == 1 and tokens_per_frame tokens.

frames
int
kind
int
tokens_per_frame
int
class nemo_automodel.components.models.praxis.masks.TokenTags(
    kind: torch.Tensor,
    chunk: torch.Tensor,
    frame: torch.Tensor
)
Dataclass

Per-token tags of one sequence, all int32 tensors of shape [tokens] on one device.

chunk
Tensor
frame
Tensor
kind
Tensor
length
int
nemo_automodel.components.models.praxis.masks.block_mask(
    rule: nemo_automodel.components.models.praxis.masks.MaskRule,
    q_len: int,
    kv_len: int,
    device: torch.device,
    block_size: int
) -> torch.nn.attention.flex_attention.BlockMask

Build a flex BlockMask over [q_len, kv_len] tiles of block_size from rule.

nemo_automodel.components.models.praxis.masks.chunk_causal_rule(
    tags: nemo_automodel.components.models.praxis.masks.TokenTags
) -> nemo_automodel.components.models.praxis.masks.MaskRule

The teacher-forcing rule of the trunk towers over one sequence.

  • every token attends every global token; a global token attends only global tokens;
  • a clean token attends the clean tokens of its own and earlier chunks;
  • a noised token attends the noised tokens of its own chunk and the clean tokens of earlier chunks;
  • a clean token never attends a noised token.
nemo_automodel.components.models.praxis.masks.dense_mask(
    rule: nemo_automodel.components.models.praxis.masks.MaskRule,
    q_len: int,
    kv_len: int,
    device: torch.device
) -> torch.Tensor

Evaluate rule on the full grid.

Returns: Tensor

Boolean tensor of shape [q_len, kv_len], true where attention is allowed.

Rule of the action expert’s cross attention.

A query of chunk c (first frame c * chunk_size, stored in q_tags.frame) attends every global key, the clean frames strictly before its chunk and at most window frames before it, and the noised frames of its own chunk.

nemo_automodel.components.models.praxis.masks.expert_self_rule(
    q_tags: nemo_automodel.components.models.praxis.masks.TokenTags
) -> nemo_automodel.components.models.praxis.masks.MaskRule

Rule of the action expert’s self attention: one block per chunk.

nemo_automodel.components.models.praxis.masks.window_tags(
    slots: tuple[nemo_automodel.components.models.praxis.masks.Slot, ...],
    chunk_size: int,
    device: torch.device
) -> nemo_automodel.components.models.praxis.masks.TokenTags

Tags of a window sequence laid out as slots in order.

Frame f of a video slot belongs to chunk f // chunk_size. Global tokens get chunk and frame -1.

Parameters:

slots
tuple[Slot, ...]

The runs of the sequence in order.

chunk_size
int

Latent frames per chunk.

device
torch.device

Device of the tag tensors.

Returns: TokenTags

class:TokenTags of the whole sequence.

nemo_automodel.components.models.praxis.masks.KIND_CLEAN = 2
 
nemo_automodel.components.models.praxis.masks.KIND_GLOBAL = 0
 
nemo_automodel.components.models.praxis.masks.KIND_NOISED = 1
 
nemo_automodel.components.models.praxis.masks.MaskKey = tuple[int | str | bool | tuple[Slot, ...], ...]
 
nemo_automodel.components.models.praxis.masks.MaskRule = Callable[[Tensor, Tensor], Tensor]