### Basic NNX-Torch Installation

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Installs the core NNX layers with PyTorch. Use uv for faster installations.

```bash
# Using uv (recommended)
uv add nnx-torch

# Using pip
pip install nnx-torch
```

--------------------------------

### Install Development Dependencies

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Use uv to install the project with all development dependencies.

```bash
uv sync --extra dev --extra all
```

--------------------------------

### Development Installation for NNX-Torch

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Sets up the project for development, including cloning the repository and installing development tools.

```bash
# Clone the repository
git clone https://github.com/yourname/nnx-torch.git
cd nnx-torch
```

--------------------------------

### NNX-Torch Installation with Optional Dependencies

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Installs NNX-Torch with specific optional features. Replace 'uv add' with 'pip install' if not using uv.

```bash
# For FlexAttention support (requires PyTorch 2.3+)
uv add nnx-torch[flex]

# For Linear Attention backends (GLA, DeltaNet, etc.)
uv add nnx-torch[linear]

# For Triton kernel support
uv add nnx-torch[triton]

# Full CUDA-accelerated stack
uv add nnx-torch[cuda]

# Development tools (testing, linting, type checking)
uv add nnx-torch[dev]

# Install everything (all optional dependencies)
uv add nnx-torch[all]
```

--------------------------------

### CUDA PyTorch Installation for NNX-Torch

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Installs NNX-Torch with CUDA support by specifying the PyTorch index URL. Choose the CUDA version that matches your system.

```bash
# Install with CUDA 12.1 support
uv add nnx-torch[cuda] --extra-index-url https://download.pytorch.org/whl/cu121

# Install with CUDA 11.8 support
uv add nnx-torch[cuda] --extra-index-url https://download.pytorch.org/whl/cu118
```

--------------------------------

### Create LLaMA-style and GPT-style TransformerLayers

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Provides examples for creating TransformerLayer blocks, configurable for different architectures like LLaMA and GPT. Users can specify attention and FFN classes, normalization types (Pre-LN or Post-LN), and causal masking. The LLaMA example uses RMSNorm, GatedFFN with silu activation, and SDPAttention, while the GPT example uses LayerNorm, FFN with GELU activation, and SDPAttention.

```python
import torch
from nnx.layers import TransformerLayer
from nnx.layers.feedforward import GatedFFN, FFN, MoEFFN
from nnx.attention import SDPAttention, RoPEAttention

# LLaMA-style block (RMSNorm, SwiGLU, causal)
lama_layer = TransformerLayer(
    embed_dim=4096,
    num_heads=32,
    ffn_dim=None,  # Defaults to 8/3 * embed_dim for GatedFFN
    dropout=0.0,
    attention_cls=SDPAttention,
    ffn_cls=GatedFFN,
    ffn_kwargs={"activation": "silu"},  # SwiGLU
    norm_type="rmsnorm",
    pre_norm=True,
    causal=True
)

# GPT-style block (LayerNorm, GELU FFN)
gpt_layer = TransformerLayer(
    embed_dim=768,
    num_heads=12,
    ffn_dim=3072,
    attention_cls=SDPAttention,
    ffn_cls=FFN,
    ffn_kwargs={"activation": "gelu"},
    norm_type="layernorm",
    pre_norm=True,
    causal=True
)
```

--------------------------------

### Initialize Transformer Block

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Create a transformer block using SDPA attention and RMSNorm.

```python
import torch
from nnx.layers import TransformerBlock, RMSNorm
from nnx.attention import SDPA

# Create a transformer block with SDPA attention
block = TransformerBlock(
    d_model=512,
    num_heads=8,
    d_ff=2048,
    attention=SDPA(),
    norm=RMSNorm(512),
)

# Forward pass with HuggingFace-style attention mask
x = torch.randn(2, 128, 512)  # (batch, seq_len, d_model)
mask = torch.ones(2, 128)  # 1 = real token, 0 = padding

output = block(x, attention_mask=mask)
```

--------------------------------

### Initialize MoE Transformer Layer

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Demonstrates the initialization of a Transformer layer using a Mixture-of-Experts (MoE) feed-forward network. It configures the number of experts and routing strategy for sparse computation.

```python
moe_layer = TransformerLayer(
    embed_dim=512,
    num_heads=8,
    ffn_cls=MoEFFN,
    ffn_kwargs={
        "num_experts": 8,
        "top_k": 2,
        "ffn_dim": 2048
    },
    causal=True
)
```

--------------------------------

### Initialize MoEFFN with Auxiliary Loss

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Creates a sparse Top-K Mixture-of-Experts FFN. Includes logic for accessing the auxiliary load-balancing loss required for training.

```python
moe = MoEFFN(
    embed_dim=512,
    ffn_dim=2048,
    num_experts=8,
    top_k=2,
    expert_cls=GatedFFN,
    expert_kwargs={"activation": "silu"},
    router_bias=False
)
```

--------------------------------

### Implement Packed Projections

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Combine linear projections to reduce kernel launch overhead.

```python
from nnx.attention import SDPAttention
from nnx.layers import GatedFFN

# Packed QKV projections in attention (single linear for Q, K, V)
attn = SDPAttention(512, num_heads=8)

# Packed gate+up projections in SwiGLU/GeGLU FFN
ffn = GatedFFN(512)
```

--------------------------------

### Project Directory Structure

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Visual representation of the repository layout.

```text
nnx-torch/
├── nnx/
│   ├── attention/       # Attention backends (SDPA, Flex, Linear, RWKV)
│   ├── layers/          # Building blocks (embeddings, FFN, norms, transformers)
│   └── utils/           # Utilities (masking, etc.)
├── tests/               # Test suite
├── pyproject.toml       # Project configuration
└── README.md
```

--------------------------------

### Run Test Suite

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Commands to execute tests using pytest with optional coverage reporting.

```bash
# Run all tests
uv run pytest

# Run with coverage
uv run pytest --cov=nnx --cov-report=html
```

--------------------------------

### Implement Scaled Dot Product Attention

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

SDPAttention provides a wrapper around PyTorch's scaled_dot_product_attention. It automatically selects the most efficient kernel based on hardware and supports both self-attention and cross-attention patterns.

```python
import torch
from nnx.attention import SDPAttention

# Create SDPA attention module
attn = SDPAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    bias=True,
    head_dim=None  # Defaults to embed_dim // num_heads = 64
)

# Self-attention with causal mask
x = torch.randn(4, 256, 512)
mask = torch.ones(4, 256)  # All real tokens

# Causal self-attention (autoregressive)
output = attn(x, attention_mask=mask, causal=True)

# Bidirectional self-attention (encoder-style)
output = attn(x, attention_mask=mask, causal=False)

# Cross-attention between query and key/value sequences
query = torch.randn(4, 64, 512)   # Decoder queries
key = torch.randn(4, 256, 512)    # Encoder keys
value = torch.randn(4, 256, 512)  # Encoder values
encoder_mask = torch.ones(4, 256)

output = attn(query, key=key, value=value, attention_mask=encoder_mask)
# output shape: (4, 64, 512)
```

--------------------------------

### Configure GatedFFN Variants

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Initializes Gated Linear Unit variants like SwiGLU or GeGLU. These are used in modern LLMs to provide improved gradient flow compared to standard activation functions.

```python
swiglu = GatedFFN(
    embed_dim=512,
    ffn_dim=None,
    activation="silu",
    dropout=0.0,
    bias=False
)

geglu = GatedFFN(
    embed_dim=512,
    ffn_dim=2048,
    activation="gelu",
    dropout=0.1,
    bias=True
)
```

--------------------------------

### Apply Gradient Checkpointing

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Enable gradient checkpointing at the layer, stack, or FFN level to reduce memory usage.

```python
from nnx.layers import TransformerLayer, TransformerStack, GatedFFN

# Per-layer checkpointing
layer = TransformerLayer(512, num_heads=8, gradient_checkpointing=True)

# Stack-level checkpointing
stack = TransformerStack(
    n_layers=12, embed_dim=512, num_heads=8,
    gradient_checkpointing=True,
)

# FFN-level checkpointing
ffn = GatedFFN(512, gradient_checkpointing=True)
```

--------------------------------

### Build LLaMA-style decoder model

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Demonstrates the construction of a complete LLaMA-style decoder using modular nnx-torch components. It integrates token embeddings, a transformer stack with RoPE attention, and gated feed-forward networks.

```python
import torch
import torch.nn as nn
from nnx.layers import TransformerStack, TokenEmbedding, GatedFFN
from nnx.attention import RoPEAttention

class LLaMADecoder(nn.Module):
    def __init__(self, vocab_size=32000, embed_dim=4096, num_heads=32, n_layers=32):
        super().__init__()
        self.embed = TokenEmbedding(vocab_size, embed_dim, padding_idx=0)
        self.transformer = TransformerStack(
            n_layers=n_layers, embed_dim=embed_dim, num_heads=num_heads,
            attention_cls=RoPEAttention, ffn_cls=GatedFFN, norm_type="rmsnorm", causal=True
        )
        self.lm_head = nn.Linear(embed_dim, vocab_size, bias=False)
        self.lm_head.weight = self.embed.embed.weight

    def forward(self, input_ids, attention_mask=None):
        x = self.embed(input_ids)
        x = self.transformer(x, attention_mask=attention_mask)
        return self.lm_head(x)
```

--------------------------------

### Create and Use ALiBiAttention

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Illustrates the initialization and forward pass of ALiBiAttention, which incorporates Attention with Linear Biases. This method adds fixed head-specific biases to attention logits based on query-key distance, improving generalization to longer sequences without explicit position embeddings. Key parameters include embed_dim, num_heads, and max_len.

```python
import torch
from nnx.attention import ALiBiAttention

# Create ALiBi attention
attn = ALiBiAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    bias=True,
    max_len=4096,
    head_dim=64
)

# Forward pass - ALiBi bias is automatically added
x = torch.randn(2, 512, 512)
mask = torch.ones(2, 512)

# ALiBi provides position information via attention bias
output = attn(x, attention_mask=mask, causal=True)
# output shape: (2, 512, 512)
```

--------------------------------

### Create and Use RWKV6TimeMixing

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Shows how to initialize and use RWKV6TimeMixing, implementing RWKV-6 style time-mixing with data-dependent interpolation (DDlerp). This layer dynamically computes interpolation ratios from the input, allowing per-token variation in blending short and long contexts. Configuration includes embed_dim, layer_id, n_layers, and n_heads for multi-head WKV.

```python
import torch
from nnx.attention import RWKV6TimeMixing

# Create RWKV-6 time mixing layer
rwkv6 = RWKV6TimeMixing(
    embed_dim=512,
    layer_id=0,
    n_layers=12,
    n_heads=8  # Multi-head WKV
)

# Forward pass
x = torch.randn(2, 256, 512)
mask = torch.ones(2, 256)

output = rwkv6(x, attention_mask=mask)
# output shape: (2, 256, 512)
```

--------------------------------

### Enable Causal Attention Dispatch

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Set causal=True to leverage optimized causal attention kernels in SDPA.

```python
from nnx.layers import TransformerLayer

layer = TransformerLayer(512, num_heads=8, causal=True)
# Automatically uses native causal attention in SDPA
output = layer(x)
```

--------------------------------

### Standard FFN Implementation

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Configures a standard two-layer feed-forward network with support for various activation functions like GELU or ReLU.

```python
ffn = FFN(
    embed_dim=512,
    ffn_dim=2048,
    activation="gelu",
    dropout=0.1,
    bias=True
)
```

--------------------------------

### Create and Use RWKVTimeMixing

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Demonstrates the creation and usage of RWKVTimeMixing, a time-mixing layer inspired by RWKV-4. This layer replaces self-attention with a recurrent WKV operator, achieving O(T) time and memory complexity. It features learnable time-shift interpolation ratios and requires parameters like embed_dim, layer_id, and n_layers for decay initialization.

```python
import torch
from nnx.attention import RWKVTimeMixing

# Create RWKV-4 time mixing layer
rwkv = RWKVTimeMixing(
    embed_dim=512,
    layer_id=0,      # Layer index for decay initialization
    n_layers=12,     # Total layers for decay initialization
    head_size=64     # Head dimension for grouped WKV
)

# Forward pass (inherently causal)
x = torch.randn(2, 256, 512)
mask = torch.ones(2, 256)  # 1 = real, 0 = padding

output = rwkv(x, attention_mask=mask)
# output shape: (2, 256, 512)
```

--------------------------------

### Build Attention Modules with Factory Function

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

The build_attention factory function allows for the dynamic creation of various attention backends such as SDPA, RoPE, and ALiBi. It simplifies model architecture switching by abstracting the instantiation of different attention mechanisms.

```python
import torch
from nnx.attention import build_attention

# Create SDPA attention (default, uses PyTorch's optimized kernels)
attn_sdpa = build_attention(
    attn_type="sdpa",
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    bias=True
)

# Create RoPE attention (rotary position embeddings)
attn_rope = build_attention(
    attn_type="rope",
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    head_dim=64,
    base=10000.0,  # RoPE base frequency
    max_len=4096
)

# Create ALiBi attention (attention with linear biases)
attn_alibi = build_attention(
    attn_type="alibi",
    embed_dim=512,
    num_heads=8,
    dropout=0.1
)

# Forward pass with HuggingFace-style mask
x = torch.randn(2, 128, 512)  # (batch, seq_len, embed_dim)
mask = torch.ones(2, 128)     # 1 = real token, 0 = padding
mask[:, 100:] = 0             # Last 28 tokens are padding

output = attn_sdpa(x, attention_mask=mask, causal=True)
# output shape: (2, 128, 512)
```

--------------------------------

### Configure Attention Mechanisms

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Configure MHA, GQA, or MQA attention patterns by adjusting head counts.

```python
from nnx.attention import SDPAttention

# Standard Multi-Head Attention (MHA)
attn = SDPAttention(512, num_heads=8)

# Grouped Query Attention (GQA) — 4 query heads share 2 KV heads
attn = SDPAttention(512, num_heads=4, num_key_value_heads=2)

# Multi-Query Attention (MQA) — all query heads share 1 KV head
attn = SDPAttention(512, num_heads=4, num_key_value_heads=1)
```

--------------------------------

### Configure TransformerStack

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Creates a multi-layer Transformer stack suitable for causal decoding. It supports custom attention classes, gated FFNs, and normalization types.

```python
decoder = TransformerStack(
    n_layers=12,
    embed_dim=768,
    num_heads=12,
    ffn_dim=3072,
    dropout=0.1,
    attention_cls=SDPAttention,
    ffn_cls=GatedFFN,
    ffn_kwargs={"activation": "silu"},
    norm_type="rmsnorm",
    pre_norm=True,
    causal=True,
    final_norm=True
)
```

--------------------------------

### Create and Use RoPEAttention

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Shows how to initialize and use RoPEAttention for multi-head attention with Rotary Position Embeddings. This module applies RoPE to queries and keys, encoding absolute positions as rotations. It supports self-attention and cross-attention, with parameters like embed_dim, num_heads, and base for angle calculation.

```python
import torch
from nnx.attention import RoPEAttention

# Create RoPE attention (LLaMA-style)
attn = RoPEAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.1,
    bias=True,
    base=10000.0,   # Angle base (LLaMA-3 uses 500000)
    max_len=4096,   # Maximum sequence length for cache
    head_dim=64     # Must be even for RoPE
)

# Self-attention with RoPE
x = torch.randn(2, 256, 512)
mask = torch.ones(2, 256)

output = attn(x, attention_mask=mask, causal=True)
# output shape: (2, 256, 512)

# Cross-attention also applies RoPE separately to Q and K
query = torch.randn(2, 64, 512)
key_value = torch.randn(2, 256, 512)
output = attn(query, key=key_value, value=key_value)
```

--------------------------------

### Create Gated Linear Attention (GLA) and Retention Attention

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Demonstrates the creation of LinearAttention modules with 'gla' and 'retention' variants. Requires CUDA and flash-linear-attention. The GLA variant is suitable for gated linear attention, while the retention variant is for RetNet-style retention. Both accept configuration parameters like embed_dim, num_heads, and head_dim.

```python
from nnx.attention import LinearAttention
import torch

# Create GLA (Gated Linear Attention) - requires CUDA and flash-linear-attention
attn_gla = LinearAttention(
    embed_dim=512,
    num_heads=8,
    variant="gla",      # Options: "gla", "delta", "based", "retention"
    head_dim=64,
    bias=False,
    expand_k=1.0,       # Key expansion ratio
    expand_v=1.0,       # Value expansion ratio
    chunk_size=64
)

# Create RetNet-style retention
attn_retention = LinearAttention(
    embed_dim=512,
    num_heads=8,
    variant="retention",
    head_dim=64
)

# Forward pass (always causal due to recurrence)
x = torch.randn(2, 1024, 512).cuda().bfloat16()
mask = torch.ones(2, 1024).cuda()

output = attn_gla(x, attention_mask=mask)
# output shape: (2, 1024, 512)
```

--------------------------------

### Enable KV Caching

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Cache key/value tensors to optimize autoregressive generation.

```python
from nnx.attention import SDPAttention

attn = SDPAttention(512, num_heads=8, use_cache=True)

# First pass
out1 = attn(x, use_cache=True)
past_kv = out1.past_key_value

# Subsequent passes with cached KV
out2 = attn(new_token, past_key_value=past_kv, use_cache=True)
```

--------------------------------

### Apply ALiBiEmbedding in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

ALiBi (Attention with Linear Biases) applies head-specific linear biases to attention scores based on query-key distance, enabling extrapolation to longer sequence lengths.

```python
import torch
from nnx.layers import ALiBiEmbedding

alibi = ALiBiEmbedding(num_heads=8, max_len=4096)

seq_len = 512
device = torch.device("cpu")
bias = alibi(seq_len, device)
```

--------------------------------

### Implement TokenEmbedding in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

TokenEmbedding provides a learnable lookup table for tokens with optional scaling by the square root of the embedding dimension, as seen in standard Transformer architectures.

```python
import torch
from nnx.layers import TokenEmbedding

embed = TokenEmbedding(vocab_size=32000, embed_dim=512, padding_idx=0, scale=True)

tokens = torch.randint(0, 32000, (2, 128))
embeddings = embed(tokens)
```

--------------------------------

### Implement CrossAttentionLayer

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Sets up a cross-attention block for sequence-to-sequence modeling. It handles both self-attention for the decoder and cross-attention against encoder outputs.

```python
cross_layer = CrossAttentionLayer(
    embed_dim=512,
    num_heads=8,
    ffn_dim=2048,
    dropout=0.1,
    attention_cls=SDPAttention,
    norm_type="rmsnorm",
    causal=True
)
```

--------------------------------

### Configure FlexAttention with Custom Kernels

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

FlexAttention enables the use of custom score modification functions compiled into CUDA kernels via PyTorch 2.3+. This is useful for implementing non-standard attention biases like ALiBi or complex masking patterns.

```python
import torch
from nnx.attention import FlexAttention

# Custom ALiBi-style score modification function
def alibi_score_mod(score, b, h, q_idx, kv_idx):
    # Add linear bias based on distance
    bias = -torch.abs(q_idx - kv_idx).float() * (h + 1) * 0.5
    return score + bias

# Create FlexAttention with custom score modifier
attn = FlexAttention(
    embed_dim=512,
    num_heads=8,
    dropout=0.0,
    bias=True,
    score_mod=alibi_score_mod,  # Custom score modification
    block_mask=None  # Optional BlockMask for sparse patterns
)

# Forward pass
x = torch.randn(2, 128, 512).cuda()  # FlexAttention requires CUDA
mask = torch.ones(2, 128).cuda()

output = attn(x, attention_mask=mask, causal=False)
# output shape: (2, 128, 512)
```

--------------------------------

### Convert HuggingFace masks to additive bias

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Converts boolean or 0-1 attention masks into additive bias tensors. It maps real tokens to 0.0 and padding tokens to -inf, making them compatible with softmax-based attention mechanisms.

```python
import torch
from nnx.utils import hf_to_additive

hf_mask = torch.ones(2, 128)
hf_mask[:, 100:] = 0

additive_bias = hf_to_additive(hf_mask, dtype=torch.float32)
```

--------------------------------

### Combine multiple attention masks

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Aggregates multiple additive attention masks via addition. It safely ignores None values and returns None if no valid masks are provided.

```python
from nnx.utils import combine_masks, hf_to_additive, make_causal_mask

padding = hf_to_additive(torch.tensor([[1, 1, 1, 0, 0]]))
causal = make_causal_mask(5, torch.device("cpu"))

combined = combine_masks(padding, causal)
```

--------------------------------

### Process Variable-Length Sequences

Source: https://github.com/darioarena87/nnx-torch/blob/master/README.md

Use NestedTensor to handle variable-length sequences without padding overhead.

```python
from nnx.layers import TransformerStack

stack = TransformerStack(
    n_layers=6, embed_dim=512, num_heads=8,
    use_nested_tensor=True,  # Enable NestedTensor processing
)

# Variable-length sequences (mask: 1=real, 0=padding)
x = torch.randn(4, 128, 512)
mask = torch.tensor(
    [
        [1, 1, 1, 1, 0, 0],
        [1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 1, 0],
        [1, 1, 0, 0, 0, 0],
    ], dtype=torch.long
)

output = stack(x, attention_mask=mask)
```

--------------------------------

### Apply LearnedPositional embeddings in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

LearnedPositional provides absolute positional embeddings that are learned during training. It supports offsets for incremental decoding scenarios like KV-caching.

```python
import torch
from nnx.layers import LearnedPositional

pos_emb = LearnedPositional(max_len=2048, embed_dim=512, dropout=0.1)

x = torch.randn(2, 256, 512)
output = pos_emb(x, offset=0)

x_new = torch.randn(2, 1, 512)
output_new = pos_emb(x_new, offset=256)
```

--------------------------------

### Apply RotaryEmbedding (RoPE) in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

RotaryEmbedding (RoPE) encodes positions as rotations in complex space, allowing the model to capture relative positional relationships effectively in attention mechanisms.

```python
import torch
from nnx.layers import RotaryEmbedding

rope = RotaryEmbedding(head_dim=64, max_len=4096, base=10000.0)

q = torch.randn(2, 8, 256, 64)
k = torch.randn(2, 8, 256, 64)
q_rot, k_rot = rope.rotate_queries_keys(q, k, offset=0)
```

--------------------------------

### Apply SinusoidalPositional encoding in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

SinusoidalPositional adds fixed frequency-based positional information to input embeddings. It helps models distinguish token order without requiring learned parameters.

```python
import torch
from nnx.layers import SinusoidalPositional

pos_enc = SinusoidalPositional(embed_dim=512, max_len=4096, dropout=0.1)

x = torch.randn(2, 256, 512)
output = pos_enc(x)
```

--------------------------------

### Normalize features with ScaleNorm in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

ScaleNorm performs L2 normalization followed by scaling with a single learned parameter. It serves as a parameter-efficient alternative to LayerNorm.

```python
import torch
from nnx.layers import ScaleNorm

norm = ScaleNorm(dim=512, eps=1e-5)

x = torch.randn(2, 128, 512)
output = norm(x)
print(f"Scale: {norm.scale.item():.4f}")
```

--------------------------------

### Apply AdaptiveRMSNorm in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

AdaptiveRMSNorm provides RMS normalization conditioned on external inputs, such as timestep or class embeddings. It is commonly used in diffusion models to modulate activations.

```python
import torch
from nnx.layers import AdaptiveRMSNorm

norm = AdaptiveRMSNorm(dim=512, cond_dim=256, eps=1e-6, use_shift=True)

x = torch.randn(2, 128, 512)
cond = torch.randn(2, 256)
output = norm(x, cond)
```

--------------------------------

### Generate causal attention masks

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

Creates an upper-triangular causal mask to prevent attention to future tokens. It returns a tensor with 0.0 on and below the diagonal and -inf above it.

```python
import torch
from nnx.utils import make_causal_mask

seq_len = 128
device = torch.device("cpu")
causal_mask = make_causal_mask(seq_len, device, dtype=torch.float32)
```

--------------------------------

### Normalize features with CosineNorm in PyTorch

Source: https://context7.com/darioarena87/nnx-torch/llms.txt

CosineNorm normalizes feature vectors to the unit hypersphere using L2 normalization. This is particularly useful for operations sensitive to input magnitude, such as dot-product attention.

```python
import torch
from nnx.layers import CosineNorm

norm = CosineNorm(eps=1e-5)

x = torch.randn(2, 128, 512)
output = norm(x)
print(f"Output norm: {output[0, 0].norm().item():.4f}")
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.