> This page is for version Latest.
> 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.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.muse_spark.layers

Text layers for the native Muse Spark Open implementation.

## Module Contents

### Classes

| Name                                                                                                           | Description                                                              |
| -------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| [`MuseSparkAttention`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkAttention)                 | Muse GQA with Q/K and output RMSNorm, output gate, RoPE, SDPA, and TE.   |
| [`MuseSparkDecoderLayer`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkDecoderLayer)           | One FP32 residual-stream Muse Spark Open decoder layer.                  |
| [`MuseSparkGate`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkGate)                           | FP32 MetaP-scaled sigmoid router with bias used for selection only.      |
| [`MuseSparkInputScaleRMSNorm`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkInputScaleRMSNorm) | Apply the learned gain before parameter-free RMS normalization.          |
| [`MuseSparkMoE`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkMoE)                             | Compressed Muse experts with per-expert norms and grouped-GEMM dispatch. |
| [`MuseSparkRMSNorm`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkRMSNorm)                     | FP32-accumulating RMSNorm with optional zero-centered checkpoint gain.   |
| [`MuseSparkResidual`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkResidual)                   | Learned norm-preserving residual connection from the Muse checkpoint.    |
| [`MuseSparkRotaryEmbedding`](#nemo_automodel-components-models-muse_spark-layers-MuseSparkRotaryEmbedding)     | FP32 adjacent-pair rotary frequencies shared by all text layers.         |

### Functions

| Name                                                                                                             | Description                                                                  |
| ---------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
| [`_checkpoint_recomputed_math`](#nemo_automodel-components-models-muse_spark-layers-_checkpoint_recomputed_math) | Keep pure compiled math and autotuner scratch tensors outside SAC caching.   |
| [`_compiled_norm_core`](#nemo_automodel-components-models-muse_spark-layers-_compiled_norm_core)                 | Lazily share each of the five pure compiled cores within this process.       |
| [`_input_scale_norm_core`](#nemo_automodel-components-models-muse_spark-layers-_input_scale_norm_core)           | Apply the effective gain before FP32 RMS normalization.                      |
| [`_norm_residual_core`](#nemo_automodel-components-models-muse_spark-layers-_norm_residual_core)                 | Fuse a branch norm and residual blend without changing the branch cast.      |
| [`_output_norm_gate_core`](#nemo_automodel-components-models-muse_spark-layers-_output_norm_gate_core)           | Fuse parameter-free output RMSNorm and the checkpoint's sigmoid gate.        |
| [`_qk_norm_rope_core`](#nemo_automodel-components-models-muse_spark-layers-_qk_norm_rope_core)                   | Normalize and rotate Q/K while preserving the checkpoint's dtype boundaries. |
| [`_repeat_kv`](#nemo_automodel-components-models-muse_spark-layers-_repeat_kv)                                   | Repeat BSHD key/value heads from `[B,Nkv,S,D]` to query-head count.          |
| [`_residual_core`](#nemo_automodel-components-models-muse_spark-layers-_residual_core)                           | Blend the residual and branch using the checkpoint's learned gate.           |
| [`_rms_norm_core`](#nemo_automodel-components-models-muse_spark-layers-_rms_norm_core)                           | Normalize in FP32, then apply the effective gain before the output cast.     |
| [`apply_rotary_embedding`](#nemo_automodel-components-models-muse_spark-layers-apply_rotary_embedding)           | Apply adjacent-pair RoPE with the checkpoint's arithmetic precision.         |
| [`build_muse_spark_moe_config`](#nemo_automodel-components-models-muse_spark-layers-build_muse_spark_moe_config) | Translate the released Muse router and compressed-expert contract.           |
| [`calculate_gain_center`](#nemo_automodel-components-models-muse_spark-layers-calculate_gain_center)             | Return the checkpoint-defined zero-centered RMSNorm gain offset.             |
| [`next_dtype`](#nemo_automodel-components-models-muse_spark-layers-next_dtype)                                   | Resolve the config's model dtype without importing recipe utilities.         |

### Data

[`_NormArgs`](#nemo_automodel-components-models-muse_spark-layers-_NormArgs)

[`_NormOutput`](#nemo_automodel-components-models-muse_spark-layers-_NormOutput)

### API

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Muse GQA with Q/K and output RMSNorm, output gate, RoPE, SDPA, and TE.

**`_fused_output_norm_gate`**

---

**`_fused_qk_norm_rope`**

---

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

---

**`k_proj`**

---

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

---

**`num_heads`** `= config.num_attention_heads`

---

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

---

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

---

**`o_gate_proj`**

---

**`o_norm`**

---

**`o_proj`**

---

**`q_proj`**

---

**`qk_norm`**

---

**`rope_in_fp32`** `= config.model_type != 'muse_spark'`

---

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

---

**`sliding_window`** `= config.layer_sliding_windows[layer_idx] or None`

---

**`use_rope`** `= config.layer_uses_rope(layer_idx)`

---

**`v_proj`**

---

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention._shape_qkv(
    hidden_states: torch.Tensor,
    is_thd: bool
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Project `[B,S,H]` or `[T,H]` states into matching BSHD/THD QKV.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkAttention._te_window_size(
    is_thd: bool,
    max_real_seqlen: int | None
) -> tuple[int, int]
```

Return TE's left/right window, treating short packed documents as full causal.

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

Attend over ordinary or packed states using this layer's CP transport.

**Parameters:**

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

CP-local states \[batch, local_sequence, hidden] or
packed \[local_tokens, hidden].

---

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

Rotary cosine/sine tensors \[batch, local_sequence,
head_dim] or packed \[local_tokens, head_dim].

---

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

Optional mask broadcastable to \[batch, heads,
sequence, sequence]; unused for packed THD.

---

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

Attention metadata, including qkv_format and global packed
document boundaries cu_seqlens \[documents + 1].

---

**Returns:** `torch.Tensor`

Attention output with the same local layout as hidden_states.

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkDecoderLayer(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `GradientCheckpointingLayer`

One FP32 residual-stream Muse Spark Open decoder layer.

**`_fused_norm_residual`**

---

**`add_attn`**

---

**`add_ffn`**

---

**`feed_forward`**

---

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

---

**`mlp`** `MuseSparkMoE`

Expose the checkpoint-named feed-forward through the generic MoE protocol.

---

**`post_attn_norm`**

---

**`post_ffn_norm`**

---

**`pre_ffn_norm`**

---

**`self_attn`** `= MuseSparkAttention(config, layer_idx, backend)`

---

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

Transform batched or packed residual states through attention and MoE.

**Parameters:**

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

Residual activations shaped \[B, S, H] or \[T, H].

---

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

Cosine and sine tensors shaped \[B, S, head_dim / 2].

---

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

Optional mask in the selected attention backend's layout.

---

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

Optional token padding mask shaped \[B, S] or \[T].

---

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

Additional arguments for the selected attention backend.

---

**Returns:** `torch.Tensor`

Independent residual states with the input shape, dtype and device.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkDecoderLayer.nemo_checkpoint_context_fn(
    context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
) -> collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
```

Preserve norm compilation limits on the checkpoint recomputation thread.

**Parameters:**

**`context_fn`** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None`

Existing forward/recompute context factory, if present.

---

**Returns:** `Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None`

A factory composing the existing contexts with the forward thread's

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkGate(
    config: nemo_automodel.components.moe.config.MoEConfig,
    router_output_multiplier: float,
    division_eps: float
)
```

**Bases:** [Gate](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-Gate)

FP32 MetaP-scaled sigmoid router with bias used for selection only.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkGate.forward(
    hidden_states: torch.Tensor,
    token_mask: torch.Tensor,
    cp_mesh: typing.Any = None
) -> tuple[torch.Tensor, torch.Tensor, None]
```

Route local `[tokens, hidden]` states and return FP32 top-k weights.

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkInputScaleRMSNorm()
```

**Bases:** [MuseSparkRMSNorm](#nemo_automodel-components-models-muse_spark-layers-MuseSparkRMSNorm)

Apply the learned gain before parameter-free RMS normalization.

**`_norm_core`** `= staticmethod(_input_scale_norm_core)`

---

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkMoE(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** [MoE](/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-MoE)

Compressed Muse experts with per-expert norms and grouped-GEMM dispatch.

**`gate`**

---

**`post_expert_norm`**

---

**`pre_expert_norm`**

---

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkMoE.forward(
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
    cp_mesh: typing.Any = None
) -> torch.Tensor
```

Route `[..., hidden]` states through compressed, normalized experts.

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm(
    dim: int,
    eps: float,
    dtype: torch.dtype,
    has_weight: bool = True,
    zero_centered_gamma: bool = False,
    gain_center: float = 1.0,
    compile_norm: bool = False
)
```

**Bases:** `Module`

FP32-accumulating RMSNorm with optional zero-centered checkpoint gain.

**`_forward_core`**

---

**`_norm_core`** `= staticmethod(_rms_norm_core)`

---

**`weight`**

---

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm.effective_weight() -> torch.Tensor | None
```

Return the actual multiplicative gain with shape `[hidden]`.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Normalize activations with the selected eager or compiled core.

**Parameters:**

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

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

---

**Returns:** `torch.Tensor`

Independent normalized tensor with the input shape, dtype and device.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkRMSNorm.reset_parameters() -> None
```

Initialize the serialized gain without changing its effective value.

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkResidual(
    dim: int,
    dtype: torch.dtype,
    temperature: float = 0.3,
    eps: float = 0.001
)
```

**Bases:** `Module`

Learned norm-preserving residual connection from the Muse checkpoint.

**`gate`** `= nn.Parameter(torch.empty(dim, dtype=dtype))`

---

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkResidual.forward(
    stream: torch.Tensor,
    branch: torch.Tensor
) -> torch.Tensor
```

Blend activations while preserving the residual stream dtype.

**Parameters:**

**`stream`** `torch.Tensor`

Residual tensor of shape \[..., hidden], ordinarily FP32.

---

**`branch`** `torch.Tensor`

Tensor with the same shape and device as stream.

---

**Returns:** `torch.Tensor`

Independent tensor with the stream shape, dtype and device.

```python
class nemo_automodel.components.models.muse_spark.layers.MuseSparkRotaryEmbedding(
    dim: int,
    theta: float
)
```

**Bases:** `Module`

FP32 adjacent-pair rotary frequencies shared by all text layers.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkRotaryEmbedding.forward(
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]
```

Return cosine/sine tensors shaped `[batch, sequence, head_dim / 2]`.

```python
nemo_automodel.components.models.muse_spark.layers.MuseSparkRotaryEmbedding.reset_parameters() -> None
```

Recreate FP32 frequencies after meta-device materialization.

```python
nemo_automodel.components.models.muse_spark.layers._checkpoint_recomputed_math(
    fn: collections.abc.Callable[nemo_automodel.components.models.muse_spark.layers._NormArgs, nemo_automodel.components.models.muse_spark.layers._NormOutput]
) -> collections.abc.Callable[nemo_automodel.components.models.muse_spark.layers._NormArgs, nemo_automodel.components.models.muse_spark.layers._NormOutput]
```

Keep pure compiled math and autotuner scratch tensors outside SAC caching.

Muse's selective checkpoint policy saves routing results, not norm or
activation math. Autotuning a new shape during recomputation can execute
scratch operations absent from the original forward. Exclude these pure
islands from dispatch-mode accounting in both passes; autograd and the
enclosing activation checkpoint still run.

```python
nemo_automodel.components.models.muse_spark.layers._compiled_norm_core(
    core: collections.abc.Callable[nemo_automodel.components.models.muse_spark.layers._NormArgs, nemo_automodel.components.models.muse_spark.layers._NormOutput]
) -> collections.abc.Callable[nemo_automodel.components.models.muse_spark.layers._NormArgs, nemo_automodel.components.models.muse_spark.layers._NormOutput]
```

Lazily share each of the five pure compiled cores within this process.

The cache retains callables, never model parameters or activations. Eager
instances keep their original callables; enabling one model does not change
another model's dispatch. Compilation starts on the first real tensor call.

```python
nemo_automodel.components.models.muse_spark.layers._input_scale_norm_core(
    hidden_states: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float
) -> torch.Tensor
```

Apply the effective gain before FP32 RMS normalization.

**Parameters:**

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

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

---

**`weight`** `torch.Tensor | None`

Optional effective FP32 gain of shape \[hidden], on the input device.

---

**`eps`** `float`

Stabilizer added to the mean squared scaled activation.

---

**Returns:** `torch.Tensor`

Independent tensor of shape \[..., hidden], in the input dtype and device.

```python
nemo_automodel.components.models.muse_spark.layers._norm_residual_core(
    stream: torch.Tensor,
    branch: torch.Tensor,
    weight: torch.Tensor | None,
    gate: torch.Tensor,
    input_scale: bool,
    norm_eps: float,
    temperature: float,
    residual_eps: float
) -> torch.Tensor
```

Fuse a branch norm and residual blend without changing the branch cast.

**Parameters:**

**`stream`** `torch.Tensor`

Residual tensor of shape \[..., hidden], ordinarily FP32.

---

**`branch`** `torch.Tensor`

Tensor of shape \[..., hidden], matching the stream shape and device.
Normalization casts back to this dtype before residual blending.

---

**`weight`** `torch.Tensor | None`

Optional effective FP32 norm gain of shape \[hidden].

---

**`gate`** `torch.Tensor`

Trainable residual gate of shape \[hidden], on the stream device.

---

**`input_scale`** `bool`

Whether the gain precedes rather than follows normalization.

---

**`norm_eps`** `float`

RMS normalization stabilizer.

---

**`temperature`** `float`

Residual gate temperature.

---

**`residual_eps`** `float`

Lower bound on the squared residual coefficient.

---

**Returns:** `torch.Tensor`

Independent tensor of shape \[..., hidden], in the stream dtype and device.

```python
nemo_automodel.components.models.muse_spark.layers._output_norm_gate_core(
    output: torch.Tensor,
    gate: torch.Tensor,
    eps: float
) -> torch.Tensor
```

Fuse parameter-free output RMSNorm and the checkpoint's sigmoid gate.

**Parameters:**

**`output`** `torch.Tensor`

Attention output \[B,S,Hq,D] or packed \[T,Hq,D].

---

**`gate`** `torch.Tensor`

Gate projection with the same shape, dtype and device as output.

---

**`eps`** `float`

FP32 RMS normalization stabilizer.

---

**Returns:** `torch.Tensor`

Gated output in the original dtype. Both normalization and sigmoid

```python
nemo_automodel.components.models.muse_spark.layers._qk_norm_rope_core(
    query: torch.Tensor,
    key: torch.Tensor,
    cosine: torch.Tensor,
    sine: torch.Tensor,
    eps: float,
    query_scale: float,
    use_rope: bool,
    upcast_rope: bool
) -> tuple[torch.Tensor, torch.Tensor]
```

Normalize and rotate Q/K while preserving the checkpoint's dtype boundaries.

**Parameters:**

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

Query tensor \[B,S,Hq,D] or packed \[T,Hq,D].

---

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

Key tensor with matching leading dimensions and Hkv heads.

---

**`cosine`** `torch.Tensor`

Rotary cosines \[B,S,D/2], or \[1,T,D/2] for packed inputs.

---

**`sine`** `torch.Tensor`

Rotary sines with the same shape and dtype as cosine.

---

**`eps`** `float`

FP32 RMS normalization stabilizer.

---

**`query_scale`** `float`

Multiplier applied after the normalized query's dtype cast.

---

**`use_rope`** `bool`

Whether this layer uses adjacent-pair rotary embeddings.

---

**`upcast_rope`** `bool`

Whether rotation uses FP32 (legacy) or the Q/K dtype (v2).

---

**Returns:** `torch.Tensor`

Independent normalized query and key tensors in their original layouts

```python
nemo_automodel.components.models.muse_spark.layers._repeat_kv(
    hidden_states: torch.Tensor,
    repeats: int
) -> torch.Tensor
```

Repeat BSHD key/value heads from `[B,Nkv,S,D]` to query-head count.

```python
nemo_automodel.components.models.muse_spark.layers._residual_core(
    stream: torch.Tensor,
    branch: torch.Tensor,
    gate: torch.Tensor,
    temperature: float,
    eps: float
) -> torch.Tensor
```

Blend the residual and branch using the checkpoint's learned gate.

**Parameters:**

**`stream`** `torch.Tensor`

Residual tensor of shape \[..., hidden], ordinarily FP32.

---

**`branch`** `torch.Tensor`

Tensor of shape \[..., hidden], matching the stream shape and device.

---

**`gate`** `torch.Tensor`

Trainable gate of shape \[hidden], on the stream device.

---

**`temperature`** `float`

Gate temperature.

---

**`eps`** `float`

Lower bound on the squared residual coefficient.

---

**Returns:** `torch.Tensor`

Independent tensor of shape \[..., hidden], in the stream dtype and device.

```python
nemo_automodel.components.models.muse_spark.layers._rms_norm_core(
    hidden_states: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float
) -> torch.Tensor
```

Normalize in FP32, then apply the effective gain before the output cast.

**Parameters:**

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

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

---

**`weight`** `torch.Tensor | None`

Optional effective FP32 gain of shape \[hidden], on the input device.

---

**`eps`** `float`

Stabilizer added to the mean squared activation.

---

**Returns:** `torch.Tensor`

Independent tensor of shape \[..., hidden], in the input dtype and device.

```python
nemo_automodel.components.models.muse_spark.layers.apply_rotary_embedding(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    upcast: bool = True
) -> torch.Tensor
```

Apply adjacent-pair RoPE with the checkpoint's arithmetic precision.

The original checkpoint computes the rotation in FP32. The HF release
rounds both products in the Q/K dtype before adding them.

```python
nemo_automodel.components.models.muse_spark.layers.build_muse_spark_moe_config(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig,
    dtype: torch.dtype,
    overrides: dict[str, typing.Any] | None = None
) -> nemo_automodel.components.moe.config.MoEConfig
```

Translate the released Muse router and compressed-expert contract.

```python
nemo_automodel.components.models.muse_spark.layers.calculate_gain_center(
    gain_center_type: str,
    residual_id: int
) -> float
```

Return the checkpoint-defined zero-centered RMSNorm gain offset.

```python
nemo_automodel.components.models.muse_spark.layers.next_dtype(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig
) -> torch.dtype
```

Resolve the config's model dtype without importing recipe utilities.

```python
nemo_automodel.components.models.muse_spark.layers._NormArgs = ParamSpec('_NormArgs')
```

```python
nemo_automodel.components.models.muse_spark.layers._NormOutput = TypeVar('_NormOutput')
```