> This page is for version Nightly (default).
> 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.layers

Transformer layers shared by the two towers of the Praxis video backbone.

Every block is pre-norm with a zero-centred RMSNorm (scale `1 + weight`), gated
grouped-query attention over integer rotary positions, and either a dense SwiGLU or an
Automodel `MoE` feed-forward. The decoder
tower modulates each block input with a per-token shift and scale from the encoder.

Norms, rotary rotation, modulation and the router run in float32; everything else in
the parameter dtype (bf16 in the released checkpoint).

## Module Contents

### Classes

| Name                                                                                               | Description                                                                                      |
| -------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------ |
| [`GatedAttention`](#nemo_automodel-components-models-praxis-layers-GatedAttention)                 | Grouped-query attention with a sigmoid output gate per head.                                     |
| [`IntegerRotaryEmbedding`](#nemo_automodel-components-models-praxis-layers-IntegerRotaryEmbedding) | Rotary (cos, sin) table for large integer positions.                                             |
| [`TransformerBlock`](#nemo_automodel-components-models-praxis-layers-TransformerBlock)             | Pre-norm block `x + attn(norm(x) * s)`, `x + ffn(norm(x) * s)` with `s = 1 / sqrt(layer_index)`. |
| [`TransformerStack`](#nemo_automodel-components-models-praxis-layers-TransformerStack)             | A stack of `TransformerBlock` sharing one rotary table.                                          |
| [`ZeroCenteredRMSNorm`](#nemo_automodel-components-models-praxis-layers-ZeroCenteredRMSNorm)       | RMSNorm whose scale is `1 + weight`; `cast_dtype` selects the compute dtype.                     |

### Functions

| Name                                                                                   | Description                                                                                     |
| -------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- |
| [`apply_rotary`](#nemo_automodel-components-models-praxis-layers-apply_rotary)         | Rotate `x` with a rotary table in float32 (rotate-half layout).                                 |
| [`build_moe_config`](#nemo_automodel-components-models-praxis-layers-build_moe_config) | MoE settings of the Praxis trunk: fp32 sigmoid router with a selection bias, one shared expert. |
| [`modulate`](#nemo_automodel-components-models-praxis-layers-modulate)                 | Return `x * scale + shift` in the dtype of `x`.                                                 |

### Data

[`ModulationDtype`](#nemo_automodel-components-models-praxis-layers-ModulationDtype)

[`_compiled_flex_attention`](#nemo_automodel-components-models-praxis-layers-_compiled_flex_attention)

### API

```python
class nemo_automodel.components.models.praxis.layers.GatedAttention(
    dim: int,
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    qk_norm_eps: float,
    attn_backend: str
)
```

**Bases:** `Module`

Grouped-query attention with a sigmoid output gate per head.

Order: project, rotate q and k, RMS-normalize q and k per head, attend, gate, project out.

**`gate_proj`** `= nn.Linear(dim, num_heads, bias=False)`

---

**`k_norm`** `= ZeroCenteredRMSNorm(head_dim, eps=qk_norm_eps)`

---

**`k_proj`**

---

**`out_proj`** `= nn.Linear(num_heads * head_dim, dim, bias=True)`

---

**`q_norm`** `= ZeroCenteredRMSNorm(head_dim, eps=qk_norm_eps)`

---

**`q_proj`** `= nn.Linear(dim, num_heads * head_dim, bias=False)`

---

**`v_proj`**

---

```python
nemo_automodel.components.models.praxis.layers.GatedAttention.forward(
    x: torch.Tensor,
    rope: torch.Tensor,
    mask: torch.nn.attention.flex_attention.BlockMask | torch.Tensor | None
) -> torch.Tensor
```

Attend within one sequence.

**Parameters:**

**`x`** `Tensor`

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

---

**`rope`** `Tensor`

Float32 rotary table of shape \[batch, tokens, head_dim // 2, 2].

---

**`mask`** `BlockMask | Tensor | None`

A flex `BlockMask` over \[tokens, tokens], a boolean tensor of shape
\[tokens, tokens] (true where attention is allowed), or `None` for full attention.

---

**Returns:** `Tensor`

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

```python
class nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding(
    dim: int,
    effective_dim: int,
    base: int,
    table_size: int
)
```

**Bases:** `Module`

Rotary (cos, sin) table for large integer positions.

Only the first `effective_dim // 2` frequency pairs rotate; the remaining pairs have
zero frequency. A position `p` is split as `p = high * table_size + low` and the
angle of each part is formed in float64 so that positions in the millions keep their
precision; the two parts are combined with the angle-sum identity.

**`theta`** `= self.make_theta(torch.device('cpu'))`

---

```python
nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding.forward(
    positions: torch.Tensor
) -> torch.Tensor
```

Build the rotary table.

**Parameters:**

**`positions`** `Tensor`

Integer tensor of shape \[batch, tokens].

---

**Returns:** `Tensor`

Float32 tensor of shape \[batch, tokens, dim // 2, 2] holding `(cos, sin)` per pair.

```python
nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding.make_theta(
    device: torch.device
) -> torch.Tensor
```

Per-pair frequencies of shape \[dim // 2] in float32 on `device` (zero beyond `effective_dim`).

```python
class nemo_automodel.components.models.praxis.layers.TransformerBlock(
    dim: int,
    attention: nemo_automodel.components.models.praxis.layers.GatedAttention,
    ffn: torch.nn.Module,
    layer_index: int,
    norm_eps: float,
    modulation_dim: int | None,
    modulation_dtype: nemo_automodel.components.models.praxis.layers.ModulationDtype
)
```

**Bases:** `Module`

Pre-norm block `x + attn(norm(x) * s)`, `x + ffn(norm(x) * s)` with `s = 1 / sqrt(layer_index)`.

`s` is a persistent buffer so the checkpoint value (rounded to the parameter dtype) is used.

**`ln_in`** `= ZeroCenteredRMSNorm(dim, eps=norm_eps)`

---

**`ln_out`** `= ZeroCenteredRMSNorm(dim, eps=norm_eps)`

---

**`mlp`** `MLP | MoE`

Expose the feed-forward sub-block (`ffn`) under the name the shared MoE parallelizer looks up.

---

**`modulator`**

---

**`s`** `Tensor`

---

```python
nemo_automodel.components.models.praxis.layers.TransformerBlock.forward(
    x: torch.Tensor,
    rope: torch.Tensor,
    mask: torch.nn.attention.flex_attention.BlockMask | torch.Tensor | None,
    x_mod: torch.Tensor | None
) -> torch.Tensor
```

Run the block.

**Parameters:**

**`x`** `Tensor`

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

---

**`rope`** `Tensor`

Float32 rotary table of shape \[batch, tokens, head_dim // 2, 2].

---

**`mask`** `BlockMask | Tensor | None`

Attention mask (see `GatedAttention.forward`).

---

**`x_mod`** `Tensor | None`

Per-token conditioning of shape \[batch, tokens, modulation_dim] when the block
has a modulator, else `None`.

---

**Returns:** `Tensor`

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

```python
class nemo_automodel.components.models.praxis.layers.TransformerStack(
    depth: int,
    dense_blocks: tuple[int, ...],
    dim: int,
    num_heads: int,
    num_kv_heads: int,
    head_dim: int,
    qk_norm_eps: float,
    norm_eps: float,
    dense_hidden: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    rope: nemo_automodel.components.models.praxis.layers.IntegerRotaryEmbedding,
    modulation_dim: int | None,
    modulation_dtype: nemo_automodel.components.models.praxis.layers.ModulationDtype,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype
)
```

**Bases:** `Module`

A stack of `TransformerBlock` sharing one rotary table.

Blocks listed in `dense_blocks` use a dense SwiGLU; the others an Automodel `MoE`.

**`blocks`** `= nn.ModuleList(blocks)`

---

```python
nemo_automodel.components.models.praxis.layers.TransformerStack.forward(
    x: torch.Tensor,
    positions: torch.Tensor,
    mask: torch.nn.attention.flex_attention.BlockMask | torch.Tensor | None,
    x_mod: torch.Tensor | None,
    tap_blocks: tuple[int, ...]
) -> tuple[torch.Tensor, list[torch.Tensor]]
```

Run every block and collect the tap outputs.

**Parameters:**

**`x`** `Tensor`

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

---

**`positions`** `Tensor`

Integer tensor of shape \[batch, tokens] with the rotary positions.

---

**`mask`** `BlockMask | Tensor | None`

Attention mask (see `GatedAttention.forward`).

---

**`x_mod`** `Tensor | None`

Per-token conditioning of shape \[batch, tokens, modulation_dim], or `None`.

---

**`tap_blocks`** `tuple[int, ...]`

Indices of the blocks whose outputs are returned.

---

**Returns:** `Tensor`

The output of shape \[batch, tokens, hidden] and the tap outputs, each of shape

```python
class nemo_automodel.components.models.praxis.layers.ZeroCenteredRMSNorm(
    dim: int,
    eps: float = 1e-06,
    cast_dtype: torch.dtype | None = torch.float32
)
```

**Bases:** `Module`

RMSNorm whose scale is `1 + weight`; `cast_dtype` selects the compute dtype.

**`weight`** `= nn.Parameter(torch.zeros(dim))`

---

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

Normalize over the last axis.

**Parameters:**

**`x`** `Tensor`

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

---

**Returns:** `Tensor`

Tensor of shape \[..., hidden] in the dtype of `x`.

```python
nemo_automodel.components.models.praxis.layers.apply_rotary(
    x: torch.Tensor,
    table: torch.Tensor
) -> torch.Tensor
```

Rotate `x` with a rotary table in float32 (rotate-half layout).

**Parameters:**

**`x`** `Tensor`

Tensor of shape \[batch, heads, tokens, head_dim]; channel `j` pairs with `j + head_dim // 2`.

---

**`table`** `Tensor`

Float32 tensor of shape \[batch, tokens, head_dim // 2, 2] from `IntegerRotaryEmbedding`.

---

**Returns:** `Tensor`

Tensor of shape \[batch, heads, tokens, head_dim] in the dtype of `x`.

```python
nemo_automodel.components.models.praxis.layers.build_moe_config(
    dim: int,
    expert_hidden: int,
    num_routed_experts: int,
    top_k_routed: int,
    routed_scaling_factor: float,
    dtype: torch.dtype,
    gate_bias_update_factor: float,
    aux_loss_coeff: float = 0.0
) -> nemo_automodel.components.moe.config.MoEConfig
```

MoE settings of the Praxis trunk: fp32 sigmoid router with a selection bias, one shared expert.

The routed sum is scaled by `routed_scaling_factor` and the top-k scores are
renormalized to sum to one, as in DeepSeek-V3.

```python
nemo_automodel.components.models.praxis.layers.modulate(
    x: torch.Tensor,
    shift: torch.Tensor,
    scale: torch.Tensor,
    dtype: nemo_automodel.components.models.praxis.layers.ModulationDtype
) -> torch.Tensor
```

Return `x * scale + shift` in the dtype of `x`.

**Parameters:**

**`x`** `Tensor`

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

---

**`shift`** `Tensor`

Tensor broadcastable to `x`.

---

**`scale`** `Tensor`

Tensor broadcastable to `x`.

---

**`dtype`** `ModulationDtype`

`"mixed"` and `"float32"` compute in float32, `"bfloat16"` in bf16.

---

**Returns:** `Tensor`

Tensor of shape \[batch, tokens, hidden] in the dtype of `x`.

```python
nemo_automodel.components.models.praxis.layers.ModulationDtype = Literal['bfloat16', 'float32', 'mixed']
```

```python
nemo_automodel.components.models.praxis.layers._compiled_flex_attention = torch.compile(flex_attention)
```