> This page is for version 0.4.0 · 26.04.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.

# nemo_automodel.components.models.praxis.masks

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

| Name                                                                    | Description                                                                            |
| ----------------------------------------------------------------------- | -------------------------------------------------------------------------------------- |
| [`MaskCache`](#nemo_automodel-components-models-praxis-masks-MaskCache) | Small LRU cache of built masks keyed by the caller's layout key.                       |
| [`Slot`](#nemo_automodel-components-models-praxis-masks-Slot)           | A run of tokens of one kind: `frames` latent frames of `tokens_per_frame` tokens each. |
| [`TokenTags`](#nemo_automodel-components-models-praxis-masks-TokenTags) | Per-token tags of one sequence, all int32 tensors of shape \[tokens] on one device.    |

### Functions

| Name                                                                                    | Description                                                                       |
| --------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- |
| [`block_mask`](#nemo_automodel-components-models-praxis-masks-block_mask)               | Build a flex `BlockMask` over \[q_len, kv_len] tiles of `block_size` from `rule`. |
| [`chunk_causal_rule`](#nemo_automodel-components-models-praxis-masks-chunk_causal_rule) | The teacher-forcing rule of the trunk towers over one sequence.                   |
| [`dense_mask`](#nemo_automodel-components-models-praxis-masks-dense_mask)               | Evaluate `rule` on the full grid.                                                 |
| [`expert_cross_rule`](#nemo_automodel-components-models-praxis-masks-expert_cross_rule) | Rule of the action expert's cross attention.                                      |
| [`expert_self_rule`](#nemo_automodel-components-models-praxis-masks-expert_self_rule)   | Rule of the action expert's self attention: one block per chunk.                  |
| [`window_tags`](#nemo_automodel-components-models-praxis-masks-window_tags)             | Tags of a window sequence laid out as `slots` in order.                           |

### Data

[`KIND_CLEAN`](#nemo_automodel-components-models-praxis-masks-KIND_CLEAN)

[`KIND_GLOBAL`](#nemo_automodel-components-models-praxis-masks-KIND_GLOBAL)

[`KIND_NOISED`](#nemo_automodel-components-models-praxis-masks-KIND_NOISED)

[`MaskKey`](#nemo_automodel-components-models-praxis-masks-MaskKey)

[`MaskRule`](#nemo_automodel-components-models-praxis-masks-MaskRule)

### API

```python
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()`

---

```python
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
```

```python
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`

---

```python
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`

---

```python
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`.

```python
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.

```python
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.

```python
nemo_automodel.components.models.praxis.masks.expert_cross_rule(
    q_tags: nemo_automodel.components.models.praxis.masks.TokenTags,
    kv_tags: nemo_automodel.components.models.praxis.masks.TokenTags,
    window: int
) -> nemo_automodel.components.models.praxis.masks.MaskRule
```

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.

```python
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.

```python
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.

```python
nemo_automodel.components.models.praxis.masks.KIND_CLEAN = 2
```

```python
nemo_automodel.components.models.praxis.masks.KIND_GLOBAL = 0
```

```python
nemo_automodel.components.models.praxis.masks.KIND_NOISED = 1
```

```python
nemo_automodel.components.models.praxis.masks.MaskKey = tuple[int | str | bool | tuple[Slot, ...], ...]
```

```python
nemo_automodel.components.models.praxis.masks.MaskRule = Callable[[Tensor, Tensor], Tensor]
```