> 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.laguna.model

## Module Contents

### Classes

| Name                                                                                            | Description                                                                |
| ----------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`LagunaAttention`](#nemo_automodel-components-models-laguna-model-LagunaAttention)             | Laguna attention: explicit per-layer heads, QK RMSNorm, and output gating. |
| [`LagunaBlock`](#nemo_automodel-components-models-laguna-model-LagunaBlock)                     | Decoder block that uses dense MLP for configured layers and MoE otherwise. |
| [`LagunaForCausalLM`](#nemo_automodel-components-models-laguna-model-LagunaForCausalLM)         | Causal LM wrapper for Laguna with Automodel checkpoint adapters.           |
| [`LagunaModel`](#nemo_automodel-components-models-laguna-model-LagunaModel)                     | Backbone model for Laguna SFT.                                             |
| [`LagunaRMSNorm`](#nemo_automodel-components-models-laguna-model-LagunaRMSNorm)                 | RMSNorm with fp32 variance, matching the Laguna reference implementation.  |
| [`LagunaRotaryEmbedding`](#nemo_automodel-components-models-laguna-model-LagunaRotaryEmbedding) | Rotary embedding with Laguna's nested full-attention/SWA RoPE support.     |

### Functions

| Name                                                                                                                    | Description                                                         |
| ----------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- |
| [`_apply_rotary_pos_emb`](#nemo_automodel-components-models-laguna-model-_apply_rotary_pos_emb)                         | Apply RoPE to query and key states.                                 |
| [`_config_with_rope`](#nemo_automodel-components-models-laguna-model-_config_with_rope)                                 | -                                                                   |
| [`_convert_bool_4d_mask_to_additive`](#nemo_automodel-components-models-laguna-model-_convert_bool_4d_mask_to_additive) | Convert a bool 4D attention mask to additive attention-mask layout. |
| [`_derive_padding_mask`](#nemo_automodel-components-models-laguna-model-_derive_padding_mask)                           | Derive an MoE padding mask from a sequence or attention mask.       |
| [`_eager_attention_forward`](#nemo_automodel-components-models-laguna-model-_eager_attention_forward)                   | Run eager attention over projected query/key/value states.          |
| [`_ensure_additive_mask`](#nemo_automodel-components-models-laguna-model-_ensure_additive_mask)                         | Normalize an optional precomputed mask into additive layout.        |
| [`_fallback_additive_mask`](#nemo_automodel-components-models-laguna-model-_fallback_additive_mask)                     | Build a causal additive attention mask.                             |
| [`_normalize_gating_mode`](#nemo_automodel-components-models-laguna-model-_normalize_gating_mode)                       | -                                                                   |
| [`_repeat_kv`](#nemo_automodel-components-models-laguna-model-_repeat_kv)                                               | Repeat key/value heads for grouped-query attention.                 |
| [`_rotate_half`](#nemo_automodel-components-models-laguna-model-_rotate_half)                                           | Rotate the two halves of the RoPE feature axis.                     |
| [`_sdpa_attention_forward`](#nemo_automodel-components-models-laguna-model-_sdpa_attention_forward)                     | Run SDPA attention over projected query/key/value states.           |

### Data

[`ModelClass`](#nemo_automodel-components-models-laguna-model-ModelClass)

[`__all__`](#nemo_automodel-components-models-laguna-model-__all__)

### API

```python
class nemo_automodel.components.models.laguna.model.LagunaAttention(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    layer_idx: int
)
```

**Bases:** `Module`

Laguna attention: explicit per-layer heads, QK RMSNorm, and output gating.

**`attention_dropout`** `= float(config.attention_dropout or 0.0)`

---

**`attention_type`**

---

**`g_proj`**

---

**`gating_mode`** `= _normalize_gating_mode(gating)`

---

**`head_dim`** `= config.head_dim`

---

**`is_sliding`** `= self.attention_type == 'sliding_attention'`

---

**`k_norm`**

---

**`k_proj`**

---

**`num_heads`**

---

**`num_key_value_groups`** `= self.num_heads // self.num_key_value_heads`

---

**`num_key_value_heads`** `= config.num_key_value_heads`

---

**`o_proj`**

---

**`q_norm`**

---

**`q_proj`**

---

**`scaling`** `= self.head_dim ** -0.5`

---

**`sliding_window`** `= config.sliding_window if self.is_sliding else None`

---

**`v_proj`**

---

```python
nemo_automodel.components.models.laguna.model.LagunaAttention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]
```

Run Laguna attention for one decoder layer.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

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

Tuple of cosine and sine RoPE tensors, each of shape
\[batch, sequence, rotary\_dim].

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional additive mask broadcastable to
\[batch, heads, sequence, key\_sequence].

---

**`**kwargs`** `Any` — default: \{}

Additional attention backend arguments.

---

**Returns:** `torch.Tensor`

Tuple of attention output \[batch, sequence, hidden] and optional attention weights

```python
nemo_automodel.components.models.laguna.model.LagunaAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.models.laguna.model.LagunaBlock(
    layer_idx: int,
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Decoder block that uses dense MLP for configured layers and MoE otherwise.

**`attention_type`** `= config.layer_types[layer_idx]`

---

**`input_layernorm`**

---

**`mlp`** `= MoE(moe_config, backend)`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```python
nemo_automodel.components.models.laguna.model.LagunaBlock.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run a decoder block.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional additive attention mask broadcastable to
\[batch, heads, sequence, key\_sequence].

---

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

Tuple of cosine and sine RoPE tensors, each of shape
\[batch, sequence, rotary\_dim].

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional bool tensor of shape \[batch, sequence], where True marks padding.

---

**`**kwargs`** `Any` — default: \{}

Additional attention backend arguments.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.models.laguna.model.LagunaBlock.init_weights(
    buffer_device: torch.device
) -> None
```

```python
class nemo_automodel.components.models.laguna.model.LagunaForCausalLM(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
```

**Bases:** [HFCheckpointingMixin](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin), `Module`, [MoEFSDPSyncMixin](/nemo-automodel/nemo_automodel/components/moe/fsdp_mixin#nemo_automodel-components-moe-fsdp_mixin-MoEFSDPSyncMixin)

Causal LM wrapper for Laguna with Automodel checkpoint adapters.

**`_keep_in_fp32_modules_strict`** `= ['mlp.gate.e_score_correction_bias', 'rotary_emb']`

---

**`backend`** `= backend or BackendConfig()`

---

**`lm_head`**

---

**`model`**

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.UNTIED_ONLY`

---

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    past_key_values: typing.Any = None,
    use_cache: bool | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.CausalLMOutputWithPast
```

Run the Laguna causal language model.

**Parameters:**

**`input_ids`** `torch.Tensor | None` — default: None

Optional token IDs of shape \[batch, sequence].

---

**`inputs_embeds`** `torch.Tensor | None` — default: None

Optional embeddings of shape \[batch, sequence, hidden].

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional position IDs of shape \[batch, sequence].

---

**`attention_mask`** `torch.Tensor | dict[str, torch.Tensor] | None` — default: None

Optional 2D bool/int mask of shape \[batch, sequence], a 4D additive
mask, or a mapping with per-attention-type masks keyed by "full\_attention" and
"sliding\_attention".

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional bool tensor of shape \[batch, sequence], where True marks tokens
excluded from MoE routing.

---

**`past_key_values`** `Any` — default: None

Unsupported KV-cache state.

---

**`use_cache`** `bool | None` — default: None

Unsupported cache flag.

---

**`logits_to_keep`** `Union[int, torch.Tensor]` — default: 0

If 0, compute logits for all sequence positions. If an int or tensor,
compute logits only for the selected trailing positions.

---

**`output_hidden_states`** `bool | None` — default: None

When true, include final hidden states in the output.

---

**`**kwargs`** `Any` — default: \{}

Additional attention backend arguments.

---

**Returns:** `CausalLMOutputWithPast`

Causal LM output with logits and optional hidden states.

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.from_config(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)
```

classmethod

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.get_input_embeddings()
```

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.get_output_embeddings()
```

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None
```

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.set_input_embeddings(
    value
)
```

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.set_output_embeddings(
    new_embeddings
)
```

```python
nemo_automodel.components.models.laguna.model.LagunaForCausalLM.update_moe_gate_bias() -> None
```

```python
class nemo_automodel.components.models.laguna.model.LagunaModel(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)
```

**Bases:** `Module`

Backbone model for Laguna SFT.

**`embed_tokens`**

---

**`has_sliding_layers`** `= 'sliding_attention' in config.layer_types`

---

**`layers`**

---

**`moe_config`** `= moe_config or MoEConfig(**moe_defaults)`

---

**`norm`**

---

**`rotary_emb`**

---

**`swa_rotary_emb`**

---

```python
nemo_automodel.components.models.laguna.model.LagunaModel._build_causal_mask_mapping(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None,
    position_ids: torch.Tensor
) -> dict[str, torch.Tensor]
```

Build additive attention masks for full and sliding Laguna layers.

**Parameters:**

**`inputs_embeds`** `torch.Tensor`

Input embedding tensor of shape \[batch, sequence, hidden].

---

**`attention_mask`** `torch.Tensor | dict[str, torch.Tensor] | None`

Optional sequence mask of shape \[batch, sequence], 4D bool/additive
mask, or mapping with masks keyed by "full\_attention" and "sliding\_attention".

---

**`position_ids`** `torch.Tensor`

Position tensor of shape \[batch, sequence].

---

**Returns:** `dict[str, torch.Tensor]`

Mapping from attention type to additive masks of shape

```python
nemo_automodel.components.models.laguna.model.LagunaModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run the Laguna decoder stack.

**Parameters:**

**`input_ids`** `torch.Tensor | None` — default: None

Optional token IDs of shape \[batch, sequence].

---

**`inputs_embeds`** `torch.Tensor | None` — default: None

Optional embeddings of shape \[batch, sequence, hidden].

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional position IDs of shape \[batch, sequence].

---

**`attention_mask`** `torch.Tensor | dict[str, torch.Tensor] | None` — default: None

Optional 2D bool/int mask of shape \[batch, sequence], a 4D additive
mask, or a mapping with per-attention-type masks keyed by "full\_attention" and
"sliding\_attention".

---

**`padding_mask`** `torch.Tensor | None` — default: None

Optional bool tensor of shape \[batch, sequence], where True marks tokens
excluded from MoE routing.

---

**`**kwargs`** `Any` — default: \{}

Additional attention backend arguments.

---

**Returns:** `torch.Tensor`

Final hidden states of shape \[batch, sequence, hidden].

```python
nemo_automodel.components.models.laguna.model.LagunaModel.init_weights(
    buffer_device: torch.device | None = None
) -> None
```

```python
class nemo_automodel.components.models.laguna.model.LagunaRMSNorm(
    hidden_size: int,
    eps: float = 1e-06,
    dtype: torch.dtype = torch.bfloat16
)
```

**Bases:** `Module`

RMSNorm with fp32 variance, matching the Laguna reference implementation.

**`weight`** `= nn.Parameter(torch.ones(hidden_size, dtype=dtype))`

---

```python
nemo_automodel.components.models.laguna.model.LagunaRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Apply RMS normalization while computing variance in fp32.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.models.laguna.model.LagunaRMSNorm.reset_parameters() -> None
```

```python
class nemo_automodel.components.models.laguna.model.LagunaRotaryEmbedding(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    device: torch.device | None = None
)
```

**Bases:** `Module`

Rotary embedding with Laguna's nested full-attention/SWA RoPE support.

**`inv_freq`** `Tensor`

---

**`max_seq_len_cached`** `= config.max_position_embeddings`

---

**`original_max_seq_len`** `= config.max_position_embeddings`

---

**`rope_type`** `= config.rope_parameters['rope_type']`

---

```python
nemo_automodel.components.models.laguna.model.LagunaRotaryEmbedding._compute_default_rope_parameters(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    device: torch.device | None = None,
    seq_len: int | None = None
) -> tuple[torch.Tensor, float]
```

staticmethod

```python
nemo_automodel.components.models.laguna.model.LagunaRotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Compute RoPE cosine and sine tables.

**Parameters:**

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

Activation tensor of shape \[batch, sequence, hidden], used for device and dtype.

---

**`position_ids`** `torch.Tensor`

Position tensor of shape \[batch, sequence].

---

**Returns:** `tuple[torch.Tensor, torch.Tensor]`

Tuple of cosine and sine tensors, each of shape \[batch, sequence, rotary\_dim].

```python
nemo_automodel.components.models.laguna.model._apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Apply RoPE to query and key states.

**Parameters:**

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

Query tensor of shape \[batch, heads, sequence, head\_dim].

---

**`k`** `torch.Tensor`

Key tensor of shape \[batch, key\_value\_heads, sequence, head\_dim].

---

**`cos`** `torch.Tensor`

Cosine tensor of shape \[batch, sequence, rotary\_dim].

---

**`sin`** `torch.Tensor`

Sine tensor of shape \[batch, sequence, rotary\_dim].

---

**Returns:** `tuple[torch.Tensor, torch.Tensor]`

Tuple of rotated query and key tensors with the same shapes as `q` and `k`.

```python
nemo_automodel.components.models.laguna.model._config_with_rope(
    config: nemo_automodel.components.models.laguna.config.LagunaConfig,
    rope_parameters: dict[str, typing.Any]
) -> nemo_automodel.components.models.laguna.config.LagunaConfig
```

```python
nemo_automodel.components.models.laguna.model._convert_bool_4d_mask_to_additive(
    attention_mask: torch.Tensor,
    dtype: torch.dtype
) -> torch.Tensor
```

Convert a bool 4D attention mask to additive attention-mask layout.

**Parameters:**

**`attention_mask`** `torch.Tensor`

Tensor of shape \[batch, heads\_or\_one, sequence, key\_sequence].
Bool masks use True for visible entries; non-bool masks are already additive.

---

**`dtype`** `torch.dtype`

Floating-point dtype for the additive mask.

---

**Returns:** `torch.Tensor`

Additive mask tensor with the same shape, or the original non-bool mask.

```python
nemo_automodel.components.models.laguna.model._derive_padding_mask(
    attention_mask: torch.Tensor
) -> torch.Tensor
```

Derive an MoE padding mask from a sequence or attention mask.

**Parameters:**

**`attention_mask`** `torch.Tensor`

Tensor of shape \[batch, sequence] with nonzero valid tokens, or
\[batch, heads\_or\_one, sequence, key\_sequence] in bool or additive mask layout.

---

**Returns:** `torch.Tensor`

Bool tensor of shape \[batch, sequence], where True marks padding.

```python
nemo_automodel.components.models.laguna.model._eager_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor]
```

Run eager attention over projected query/key/value states.

**Parameters:**

**`module`** `nn.Module`

Attention module carrying head-repeat metadata.

---

**`query`** `torch.Tensor`

Query tensor of shape \[batch, heads, sequence, head\_dim].

---

**`key`** `torch.Tensor`

Key tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

---

**`value`** `torch.Tensor`

Value tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

---

**`attention_mask`** `torch.Tensor | None`

Optional additive mask broadcastable to
\[batch, heads, sequence, key\_sequence].

---

**`scaling`** `float`

Multiplicative scale for attention logits.

---

**`dropout`** `float` — default: 0.0

Attention dropout probability.

---

**`**kwargs`** `Any` — default: \{}

Ignored attention backend arguments.

---

**Returns:** `torch.Tensor`

Tuple of context tensor \[batch, sequence, heads, head\_dim] and attention weights

```python
nemo_automodel.components.models.laguna.model._ensure_additive_mask(
    mask: torch.Tensor | None,
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None,
    sliding_window: int | None
) -> torch.Tensor
```

Normalize an optional precomputed mask into additive layout.

**Parameters:**

**`mask`** `torch.Tensor | None`

Optional attention mask of shape \[batch, heads\_or\_one, sequence, key\_sequence].
Bool masks use True for visible entries; additive masks are passed through.

---

**`batch_size`** `int`

Batch size used when `mask` is absent.

---

**`seq_len`** `int`

Query and key sequence length used when `mask` is absent.

---

**`dtype`** `torch.dtype`

Floating-point dtype for the additive mask.

---

**`device`** `torch.device`

Device for the additive mask.

---

**`attention_mask`** `torch.Tensor | None`

Optional sequence mask of shape \[batch, sequence] with nonzero valid tokens.

---

**`sliding_window`** `int | None`

Optional sliding-window size for local attention.

---

**Returns:** `torch.Tensor`

Additive mask tensor of shape \[batch, heads\_or\_one, sequence, key\_sequence].

```python
nemo_automodel.components.models.laguna.model._fallback_additive_mask(
    batch_size: int,
    seq_len: int,
    dtype: torch.dtype,
    device: torch.device,
    attention_mask: torch.Tensor | None = None,
    sliding_window: int | None = None
) -> torch.Tensor
```

Build a causal additive attention mask.

**Parameters:**

**`batch_size`** `int`

Batch size for the returned mask.

---

**`seq_len`** `int`

Query and key sequence length.

---

**`dtype`** `torch.dtype`

Floating-point dtype for the additive mask.

---

**`device`** `torch.device`

Device for the returned mask.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional sequence mask of shape \[batch, sequence] with nonzero valid tokens.

---

**`sliding_window`** `int | None` — default: None

Optional sliding-window size for local attention.

---

**Returns:** `torch.Tensor`

Additive mask tensor of shape \[batch, 1, sequence, sequence].

```python
nemo_automodel.components.models.laguna.model._normalize_gating_mode(
    gating: bool | str
) -> str | None
```

```python
nemo_automodel.components.models.laguna.model._repeat_kv(
    hidden_states: torch.Tensor,
    n_rep: int
) -> torch.Tensor
```

Repeat key/value heads for grouped-query attention.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, key\_value\_heads, sequence, head\_dim].

---

**`n_rep`** `int`

Number of repeats per key/value head.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, key\_value\_heads \* n\_rep, sequence, head\_dim].

```python
nemo_automodel.components.models.laguna.model._rotate_half(
    x: torch.Tensor
) -> torch.Tensor
```

Rotate the two halves of the RoPE feature axis.

**Parameters:**

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

Tensor of shape \[..., rotary\_dim], where rotary\_dim is even.

---

**Returns:** `torch.Tensor`

Tensor of shape \[..., rotary\_dim].

```python
nemo_automodel.components.models.laguna.model._sdpa_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, None]
```

Run SDPA attention over projected query/key/value states.

**Parameters:**

**`module`** `nn.Module`

Attention module carrying head-repeat metadata.

---

**`query`** `torch.Tensor`

Query tensor of shape \[batch, heads, sequence, head\_dim].

---

**`key`** `torch.Tensor`

Key tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

---

**`value`** `torch.Tensor`

Value tensor of shape \[batch, key\_value\_heads, key\_sequence, head\_dim].

---

**`attention_mask`** `torch.Tensor | None`

Optional additive mask broadcastable to
\[batch, heads, sequence, key\_sequence].

---

**`scaling`** `float`

Ignored because SDPA applies its own scale.

---

**`dropout`** `float` — default: 0.0

Attention dropout probability.

---

**`**kwargs`** `Any` — default: \{}

Ignored attention backend arguments.

---

**Returns:** `tuple[torch.Tensor, None]`

Tuple of context tensor \[batch, sequence, heads, head\_dim] and `None` for weights.

```python
nemo_automodel.components.models.laguna.model.ModelClass = LagunaForCausalLM
```

```python
nemo_automodel.components.models.laguna.model.__all__ = ['LagunaForCausalLM', 'LagunaModel', 'ModelClass']
```