> 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.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.mistral_large4.rope_utils

Mistral Large 4 differentiable rotary and SSMax floating-point numerics.

## Module Contents

### Classes

| Name                                                                                                     | Description                                                                  |
| -------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`FP32LayerNorm`](#nemo_automodel-components-models-mistral_large4-rope_utils-FP32LayerNorm)             | Normalize indexer keys with an FP32 parameter-owning module call.            |
| [`MistralLarge4Rotary`](#nemo_automodel-components-models-mistral_large4-rope_utils-MistralLarge4Rotary) | Shared learnable LeRoPE frequencies; build angles in FP32 after dtype casts. |
| [`SSMax`](#nemo_automodel-components-models-mistral_large4-rope_utils-SSMax)                             | Per-head s \* log(position + 1) + b in the query's compute dtype.            |

### Functions

| Name                                                                                                           | Description                                                               |
| -------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`apply_interleaved_rope`](#nemo_automodel-components-models-mistral_large4-rope_utils-apply_interleaved_rope) | Rotate adjacent channel pairs in FP32, then restore the activation dtype. |

### API

```python
class nemo_automodel.components.models.mistral_large4.rope_utils.FP32LayerNorm()
```

**Bases:** `LayerNorm`

Normalize indexer keys with an FP32 parameter-owning module call.

```python
nemo_automodel.components.models.mistral_large4.rope_utils.FP32LayerNorm.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Normalize in FP32 and restore the input activation dtype.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[..., index_dim], with arbitrary leading dimensions.

---

**Returns:** `torch.Tensor`

Independent tensor with the input shape/dtype. Parameters remain FP32.

```python
class nemo_automodel.components.models.mistral_large4.rope_utils.MistralLarge4Rotary(
    config: nemo_automodel.components.models.mistral_large4.config.MistralLarge4Config
)
```

**Bases:** `Module`

Shared learnable LeRoPE frequencies; build angles in FP32 after dtype casts.

**`alpha`**

---

```python
nemo_automodel.components.models.mistral_large4.rope_utils.MistralLarge4Rotary.forward(
    positions: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Build differentiable interleaved YaRN tables.

**Parameters:**

**`positions`** `torch.Tensor`

Integer tensor of shape \[tokens] containing document-relative positions.

---

**Returns:** `torch.Tensor`

FP32 cosine and sine tensors, each of shape \[tokens, rotary_dim / 2].

```python
class nemo_automodel.components.models.mistral_large4.rope_utils.SSMax(
    heads: int,
    dtype: torch.dtype
)
```

**Bases:** `Module`

Per-head s \* log(position + 1) + b in the query's compute dtype.

**`b`** `= nn.Parameter(torch.ones(heads, dtype=dtype))`

---

**`s`** `= nn.Parameter(torch.zeros(heads, dtype=dtype))`

---

```python
nemo_automodel.components.models.mistral_large4.rope_utils.SSMax.forward(
    q: torch.Tensor,
    log_positions: torch.Tensor
) -> torch.Tensor
```

Scale a query without modifying storage shared with another consumer.

**Parameters:**

**`q`** `torch.Tensor`

Tensor of shape \[tokens, heads, head_dim].

---

**`log_positions`** `torch.Tensor`

FP32 tensor of shape \[tokens].

---

**Returns:** `torch.Tensor`

Independent tensor of shape \[tokens, heads, head_dim] in q.dtype.

```python
nemo_automodel.components.models.mistral_large4.rope_utils.apply_interleaved_rope(
    x: torch.Tensor,
    tables: tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor
```

Rotate adjacent channel pairs in FP32, then restore the activation dtype.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[tokens, heads, rotary_dim] or \[tokens, rotary_dim].

---

**`tables`** `tuple[torch.Tensor, torch.Tensor]`

FP32 cosine and sine tensors of shape \[tokens, rotary_dim / 2].

---

**Returns:** `torch.Tensor`

Independent tensor with the same shape and dtype as x.