> This page is for version 0.5.0 · 26.06.
> 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.backbone

Two-tower video backbone of the chunk-causal window.

The encoder tower reads `[registers | anchor | text | noised chunks | clean chunks]` and
returns the conditioning of the noised tokens plus the tap features for the action
expert. The decoder tower reads `[registers | noised chunks]` at the unpatched latent
resolution, modulated per token by the encoder output, and returns the video velocity.

Rotary positions: each latent frame owns a block of `FRAME_STRIDE` positions; inside a
frame a token sits at `1 + row * ROW_STRIDE + column`. Registers are at position 0. In
the encoder the anchor takes frame block 0, the text frame block 1 and the noised frames
start at frame block 2; a clean frame shares the positions of its noised twin.

## Module Contents

### Classes

| Name                                                                                                   | Description                                                                                                       |
| ------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------- |
| [`BackboneOutput`](#nemo_automodel-components-models-praxis-backbone-BackboneOutput)                   | `velocity` of shape \[batch, channels, frames, height, width] (`None` without decoding) and the encoder features. |
| [`DecoderOutputProjection`](#nemo_automodel-components-models-praxis-backbone-DecoderOutputProjection) | Zero-centred RMSNorm followed by a linear map to latent channels.                                                 |
| [`DecoderTower`](#nemo_automodel-components-models-praxis-backbone-DecoderTower)                       | The decoder tower (see the module docstring).                                                                     |
| [`EncoderFeatures`](#nemo_automodel-components-models-praxis-backbone-EncoderFeatures)                 | Encoder outputs that the action expert reads.                                                                     |
| [`EncoderTower`](#nemo_automodel-components-models-praxis-backbone-EncoderTower)                       | The encoder tower (see the module docstring).                                                                     |
| [`FourierTimestep`](#nemo_automodel-components-models-praxis-backbone-FourierTimestep)                 | Fixed random Fourier features of the flow time: `0.5 * cos(2 pi (t f + phi))`.                                    |
| [`MoeTrainingSettings`](#nemo_automodel-components-models-praxis-backbone-MoeTrainingSettings)         | Load-balancing terms of the routers during training (both zero at inference).                                     |
| [`TextInputProjection`](#nemo_automodel-components-models-praxis-backbone-TextInputProjection)         | MLP from text-encoder states to trunk tokens, ending in a zero-centred RMSNorm.                                   |
| [`VideoBackbone`](#nemo_automodel-components-models-praxis-backbone-VideoBackbone)                     | Encoder tower and decoder tower of one denoising step.                                                            |

### Functions

| Name                                                                                       | Description                                                                                  |
| ------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------------------------- |
| [`_expand_rows`](#nemo_automodel-components-models-praxis-backbone-_expand_rows)           | Repeat row `k` of `rows` (\[batch, chunks, hidden]) over `chunk_tokens[k]` tokens.           |
| [`_make_stack`](#nemo_automodel-components-models-praxis-backbone-_make_stack)             | -                                                                                            |
| [`encoder_positions`](#nemo_automodel-components-models-praxis-backbone-encoder_positions) | Rotary positions of the encoder sequence `[registers \| anchor \| text \| noised \| clean]`. |
| [`patchify`](#nemo_automodel-components-models-praxis-backbone-patchify)                   | Inverse of `unpatch`: fold `(ph, pw)` spatial patches into the channels.                     |
| [`unpatch`](#nemo_automodel-components-models-praxis-backbone-unpatch)                     | Move the channel patch of `[batch, channels * ph * pw, frames, h, w]` into space.            |
| [`video_positions`](#nemo_automodel-components-models-praxis-backbone-video_positions)     | Rotary positions of a video run in frame-major order.                                        |

### Data

[`FRAME_STRIDE`](#nemo_automodel-components-models-praxis-backbone-FRAME_STRIDE)

[`ROW_STRIDE`](#nemo_automodel-components-models-praxis-backbone-ROW_STRIDE)

### API

```python
class nemo_automodel.components.models.praxis.backbone.BackboneOutput(
    velocity: torch.Tensor | None,
    features: nemo_automodel.components.models.praxis.backbone.EncoderFeatures
)
```

Dataclass

`velocity` of shape \[batch, channels, frames, height, width] (`None` without decoding) and the encoder features.

**`features`** `EncoderFeatures`

---

**`velocity`** `Tensor | None`

---

```python
class nemo_automodel.components.models.praxis.backbone.DecoderOutputProjection(
    dim: int,
    c_out: int
)
```

**Bases:** `Module`

Zero-centred RMSNorm followed by a linear map to latent channels.

**`linear`** `= nn.Linear(dim, c_out)`

---

**`norm`** `= ZeroCenteredRMSNorm(dim, cast_dtype=None)`

---

```python
nemo_automodel.components.models.praxis.backbone.DecoderOutputProjection.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Project decoder tokens.

**Parameters:**

**`x`** `Tensor`

Tensor of shape \[batch, tokens, hidden].

---

**Returns:** `Tensor`

Tensor of shape \[batch, tokens, c_out].

```python
class nemo_automodel.components.models.praxis.backbone.DecoderTower(
    cfg: nemo_automodel.components.models.praxis.config.BackboneConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings
)
```

**Bases:** `Module`

The decoder tower (see the module docstring).

**`_masks`** `= MaskCache()`

---

**`block_size`** `= cfg.attention.flex_block_size`

---

**`core`**

---

**`modulation_dtype`** `= cfg.modulation_dtype`

---

**`patch_size`** `= (tower.unpatch_factor[1], tower.unpatch_factor[2])`

---

**`registers`**

---

**`t_embedding`** `= FourierTimestep(cfg.timestep_dim)`

---

**`time_proj`** `= nn.Linear(cfg.timestep_dim, 2 * cfg.dim)`

---

**`use_flex`** `= backend.attn == 'flex'`

---

**`video_in`** `= nn.Linear(tower.video_in.c_in, cfg.dim)`

---

**`video_out`** `= DecoderOutputProjection(cfg.dim, c_out)`

---

```python
nemo_automodel.components.models.praxis.backbone.DecoderTower._mask(
    slots: tuple[nemo_automodel.components.models.praxis.masks.Slot, ...],
    chunk: int,
    device: torch.device
) -> torch.nn.attention.flex_attention.BlockMask | torch.Tensor
```

```python
nemo_automodel.components.models.praxis.backbone.DecoderTower.forward(
    noised: torch.Tensor,
    cond: torch.Tensor,
    timesteps: torch.Tensor,
    rope_shift: int = 0
) -> torch.Tensor
```

Predict the velocity of the noised chunks.

**Parameters:**

**`noised`** `Tensor`

Noised latents of shape \[batch, channels, K \* F, h, w].

---

**`cond`** `Tensor`

Encoder conditioning of shape \[batch, K \* F, h, w, hidden].

---

**`timesteps`** `Tensor`

Flow time of each chunk, shape \[batch, K].

---

**`rope_shift`** `int` — default: 0

Rotary offset of the video positions in training.

---

**Returns:** `Tensor`

Velocity of shape \[batch, channels, K \* F, h, w].

```python
class nemo_automodel.components.models.praxis.backbone.EncoderFeatures(
    taps: list[torch.Tensor],
    embeddings: torch.Tensor,
    positions: torch.Tensor,
    global_spans: dict[str, tuple[int, int]],
    noised_span: tuple[int, int],
    clean_span: tuple[int, int] | None,
    frames_per_chunk: int,
    tokens_per_frame: int,
    n_chunks: int
)
```

Dataclass

Encoder outputs that the action expert reads.

**`clean_span`** `tuple[int, int] | None`

`(start, length)` of the clean chunk tokens, `None` for one chunk.

---

**`embeddings`** `Tensor`

The encoder input sequence (block 0 input) of shape \[batch, tokens, hidden].

---

**`frames_per_chunk`** `int`

Latent frames per chunk.

---

**`global_spans`** `dict[str, tuple[int, int]]`

`(start, length)` of the `register`, `anchor` and `text` runs.

---

**`n_chunks`** `int`

Chunks in the window.

---

**`noised_span`** `tuple[int, int]`

`(start, length)` of the noised chunk tokens.

---

**`positions`** `Tensor`

Rotary positions of the sequence, shape \[batch, tokens].

---

**`taps`** `list[Tensor]`

Block outputs at the tap blocks, each of shape \[batch, tokens, hidden].

---

**`tokens_per_frame`** `int`

Tokens per latent frame.

---

```python
class nemo_automodel.components.models.praxis.backbone.EncoderTower(
    cfg: nemo_automodel.components.models.praxis.config.BackboneConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings
)
```

**Bases:** `Module`

The encoder tower (see the module docstring).

**`_masks`** `= MaskCache()`

---

**`block_size`** `= cfg.attention.flex_block_size`

---

**`core`**

---

**`modulation_dtype`** `= cfg.modulation_dtype`

---

**`registers`**

---

**`t_embedding`** `= FourierTimestep(cfg.timestep_dim)`

---

**`tap_blocks`** `= tower.tap_blocks`

---

**`text_in`** `= TextInputProjection(tower.text_in.c_in, cfg.dim)`

---

**`time_proj`** `= nn.Linear(cfg.timestep_dim, 2 * cfg.dim)`

---

**`use_flex`** `= backend.attn == 'flex'`

---

**`video_in`** `= nn.Linear(tower.video_in.c_in, cfg.dim)`

---

**`video_out`** `= nn.Linear(cfg.dim, cfg.dim)`

---

```python
nemo_automodel.components.models.praxis.backbone.EncoderTower._mask(
    slots: tuple[nemo_automodel.components.models.praxis.masks.Slot, ...],
    chunk: int,
    device: torch.device
) -> torch.nn.attention.flex_attention.BlockMask | torch.Tensor
```

```python
nemo_automodel.components.models.praxis.backbone.EncoderTower._modulate_chunks(
    tokens: torch.Tensor,
    timesteps: torch.Tensor,
    chunk_tokens: list[int]
) -> torch.Tensor
```

```python
nemo_automodel.components.models.praxis.backbone.EncoderTower._video_tokens(
    latents: torch.Tensor
) -> torch.Tensor
```

`[batch, channels, frames, h, w]` latents to `[batch, frames * h * w, hidden]` tokens.

```python
nemo_automodel.components.models.praxis.backbone.EncoderTower.forward(
    text: torch.Tensor,
    anchor: torch.Tensor,
    clean: torch.Tensor | None,
    noised: torch.Tensor,
    timesteps: torch.Tensor,
    rope_shifts: tuple[int, int, int] = (0, 0, 0)
) -> tuple[torch.Tensor, nemo_automodel.components.models.praxis.backbone.EncoderFeatures]
```

Encode a window of `K` chunks.

**Parameters:**

**`text`** `Tensor`

Text-encoder states of shape \[batch, text_tokens, text_dim].

---

**`anchor`** `Tensor`

Anchor latent of shape \[batch, channels, 1, h, w].

---

**`clean`** `Tensor | None`

Clean copies of chunks `0 .. K-2` of shape \[batch, channels, (K-1) \* F, h, w], or `None` for `K == 1`.

---

**`noised`** `Tensor`

Noised chunks of shape \[batch, channels, K \* F, h, w].

---

**`timesteps`** `Tensor`

Flow time of each noised chunk, shape \[batch, K].

---

**`rope_shifts`** `tuple[int, int, int]` — default: (0, 0, 0)

`(anchor, text, video)` rotary offsets of training.

---

**Returns:** `Tensor`

The conditioning of the noised tokens, shape \[batch, K \* F, h, w, hidden], and the

```python
class nemo_automodel.components.models.praxis.backbone.FourierTimestep(
    dim: int
)
```

**Bases:** `Module`

Fixed random Fourier features of the flow time: `0.5 * cos(2 pi (t f + phi))`.

**`freqs`** `Tensor`

---

**`phases`** `Tensor`

---

```python
nemo_automodel.components.models.praxis.backbone.FourierTimestep.forward(
    t: torch.Tensor
) -> torch.Tensor
```

Embed timesteps.

**Parameters:**

**`t`** `Tensor`

Tensor of shape \[..., 1] or \[...]; the last axis is broadcast against the features.

---

**Returns:** `Tensor`

Tensor of shape \[..., dim] in the dtype of `t`.

```python
class nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings(
    gate_bias_update_factor: float = 0.0,
    aux_loss_coeff: float = 0.0
)
```

Dataclass

Load-balancing terms of the routers during training (both zero at inference).

**`aux_loss_coeff`** `float = 0.0`

Weight of the sequence-wise balance loss of each router
(`experts * sum_i f_i p_i`), injected through the routing weights.

---

**`gate_bias_update_factor`** `float = 0.0`

Step of the per-expert selection bias toward balanced load,
applied by `PraxisPolicy.update_moe_gate_bias` after each optimizer step.

---

```python
class nemo_automodel.components.models.praxis.backbone.TextInputProjection(
    c_in: int,
    dim: int
)
```

**Bases:** `Module`

MLP from text-encoder states to trunk tokens, ending in a zero-centred RMSNorm.

**`fc1`** `= nn.Linear(c_in, 4 * dim)`

---

**`fc2`** `= nn.Linear(4 * dim, dim)`

---

**`norm`** `= ZeroCenteredRMSNorm(dim, cast_dtype=None)`

---

```python
nemo_automodel.components.models.praxis.backbone.TextInputProjection.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Project text states.

**Parameters:**

**`x`** `Tensor`

Tensor of shape \[batch, text_tokens, c_in].

---

**Returns:** `Tensor`

Tensor of shape \[batch, text_tokens, hidden].

```python
class nemo_automodel.components.models.praxis.backbone.VideoBackbone(
    cfg: nemo_automodel.components.models.praxis.config.BackboneConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings | None = None
)
```

**Bases:** `Module`

Encoder tower and decoder tower of one denoising step.

**`decoder`** `= DecoderTower(cfg, backend, dtype, moe_training)`

---

**`encoder`** `= EncoderTower(cfg, backend, dtype, moe_training)`

---

```python
nemo_automodel.components.models.praxis.backbone.VideoBackbone.forward(
    text: torch.Tensor,
    anchor: torch.Tensor,
    clean: torch.Tensor | None,
    noised: torch.Tensor,
    timesteps: torch.Tensor,
    decode: bool = True,
    rope_shifts: tuple[int, int, int, int] = (0, 0, 0, 0)
) -> nemo_automodel.components.models.praxis.backbone.BackboneOutput
```

Run one denoising step of a window.

**Parameters:**

**`text`** `Tensor`

Text-encoder states of shape \[batch, text_tokens, text_dim].

---

**`anchor`** `Tensor`

Anchor latent of shape \[batch, channels, 1, h, w].

---

**`clean`** `Tensor | None`

Clean copies of shape \[batch, channels, (K-1) \* F, h, w], or `None`.

---

**`noised`** `Tensor`

Noised chunks of shape \[batch, channels, K \* F, h, w].

---

**`timesteps`** `Tensor`

Flow time of each chunk, shape \[batch, K].

---

**`decode`** `bool` — default: True

Run the decoder tower; otherwise `velocity` is `None`.

---

**`rope_shifts`** `tuple[int, int, int, int]` — default: (0, 0, 0, 0)

`(encoder anchor, encoder text, encoder video, decoder video)` offsets.

---

**Returns:** `BackboneOutput`

class:`BackboneOutput` with the velocity of shape \[batch, channels, K \* F, h, w].

```python
nemo_automodel.components.models.praxis.backbone._expand_rows(
    rows: torch.Tensor,
    chunk_tokens: list[int]
) -> torch.Tensor
```

Repeat row `k` of `rows` (\[batch, chunks, hidden]) over `chunk_tokens[k]` tokens.

```python
nemo_automodel.components.models.praxis.backbone._make_stack(
    cfg: nemo_automodel.components.models.praxis.config.BackboneConfig,
    tower: nemo_automodel.components.models.praxis.config.TowerConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings
) -> nemo_automodel.components.models.praxis.layers.TransformerStack
```

```python
nemo_automodel.components.models.praxis.backbone.encoder_positions(
    num_registers: int,
    anchor_hw: tuple[int, int],
    num_text: int,
    noised_thw: tuple[int, int, int],
    num_clean_frames: int,
    shifts: tuple[int, int, int]
) -> torch.Tensor
```

Rotary positions of the encoder sequence `[registers | anchor | text | noised | clean]`.

**Parameters:**

**`num_registers`** `int`

Register tokens (position 0).

---

**`anchor_hw`** `tuple[int, int]`

Latent grid of the anchor frame.

---

**`num_text`** `int`

Text tokens.

---

**`noised_thw`** `tuple[int, int, int]`

Frames and latent grid of the noised chunks.

---

**`num_clean_frames`** `int`

Frames of the clean copies.

---

**`shifts`** `tuple[int, int, int]`

Random `(anchor, text, video)` offsets of training; zeros at inference.

---

**Returns:** `Tensor`

Int64 tensor of shape \[tokens].

```python
nemo_automodel.components.models.praxis.backbone.patchify(
    x: torch.Tensor,
    patch: tuple[int, int]
) -> torch.Tensor
```

Inverse of `unpatch`: fold `(ph, pw)` spatial patches into the channels.

**Returns:** `Tensor`

Tensor of shape \[batch, channels \* ph \* pw, frames, h // ph, w // pw].

```python
nemo_automodel.components.models.praxis.backbone.unpatch(
    x: torch.Tensor,
    patch: tuple[int, int]
) -> torch.Tensor
```

Move the channel patch of `[batch, channels * ph * pw, frames, h, w]` into space.

**Returns:** `Tensor`

Tensor of shape \[batch, channels, frames, h \* ph, w \* pw].

```python
nemo_automodel.components.models.praxis.backbone.video_positions(
    frames: int,
    height: int,
    width: int,
    start: int
) -> torch.Tensor
```

Rotary positions of a video run in frame-major order.

**Returns:** `Tensor`

Int64 tensor of shape \[frames \* height \* width].

```python
nemo_automodel.components.models.praxis.backbone.FRAME_STRIDE = 129600
```

```python
nemo_automodel.components.models.praxis.backbone.ROW_STRIDE = 360
```