### Initialize Scale Module

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Provides an example of initializing the Scale module, which is used for learnable scaling parameters within nGPT components, following the parametrization from section 2.5 of the paper.

```python
import torch
from nGPT_pytorch.nGPT import Scale
```

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

### Install nGPT-pytorch

Source: https://github.com/lucidrains/ngpt-pytorch/blob/main/README.md

Install the nGPT-pytorch library using pip. This command is used to add the library to your project's dependencies.

```bash
pip install nGPT-pytorch
```

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

### Initialize and Use NormLinear Module

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Shows how to create and use the NormLinear module, a linear layer with L2-normalized weights. It supports automatic weight normalization during training via PyTorch parametrization. Includes an example of manually normalizing weights if parametrization is disabled.

```python
import torch
from nGPT_pytorch.nGPT import NormLinear

# Create normalized linear layer
linear = NormLinear(
    dim=512,                 # Input dimension
    dim_out=1024,            # Output dimension
    norm_dim_in=True,        # Normalize along input dimension (rows)
    parametrize=True,        # Use PyTorch parametrization for automatic normalization
    norm_eps=0.0,            # Normalization epsilon
    groups=1                 # Number of groups for group normalization
)

# Forward pass
x = torch.randn(2, 100, 512)
output = linear(x)  # Shape: (2, 100, 1024)

# Manually re-normalize weights (if parametrize=False)
linear.norm_weights_()

# Access underlying weight matrix
weight = linear.weight  # Already normalized via parametrization
print(weight.norm(dim=-1))  # Should be ~1.0 for each row
```

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

### Initialize and Use FeedForward Module

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Demonstrates the creation and forward pass of the FeedForward module, which uses normalized SwiGLU feedforward layers. Note that the output is not automatically normalized and requires a manual L2 normalization step.

```python
import torch
from nGPT_pytorch.nGPT import FeedForward

# Create feedforward module
ff = FeedForward(
    dim=512,                 # Input/output dimension
    expand_factor=4,         # Hidden layer expansion (actual is 2/3 * expand_factor)
    s_hidden_init=1.0,       # Initial hidden scale
    s_hidden_scale=1.0,      # Hidden scale multiplier
    s_gate_init=1.0,         # Initial gate scale
    s_gate_scale=1.0,        # Gate scale multiplier
    norm_eps=0.0,            # Normalization epsilon
    num_hyperspheres=1       # Number of hyperspheres
)

# Forward pass on normalized input
x = torch.randn(2, 1024, 512)
x = torch.nn.functional.normalize(x, dim=-1)

output = ff(x)  # Output shape: (2, 1024, 512)
# Output is NOT automatically normalized - caller should apply l2norm
```

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

### Initialize and Use Attention Module

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Demonstrates the creation and forward pass of the Attention module, including configuration for normalization, causal masking, rotary embeddings, and flash attention. Shows how to use value residual connections for subsequent layers.

```python
import torch
from nGPT_pytorch.nGPT import Attention

# Create attention module
attention = Attention(
    dim=512,                 # Input dimension
    dim_head=64,             # Dimension per head
    heads=8,                 # Number of heads
    norm_qk=True,            # Normalize queries and keys
    causal=True,             # Causal masking for autoregressive
    s_qk_init=1.0,           # Initial scale for query-key
    s_qk_scale=None,         # Auto-calculated based on dim
    flash_kwargs=dict(       # Flash attention configuration
        enable_flash=True,
        enable_math=True,
        enable_mem_efficient=True,
        enable_cudnn=True
    ),
    norm_eps=0.0,            # Normalization epsilon
    num_hyperspheres=1       # Number of hyperspheres
)

# Forward pass on normalized input
x = torch.randn(2, 1024, 512)
x = torch.nn.functional.normalize(x, dim=-1)

# Create rotary embeddings
from rotary_embedding_torch import RotaryEmbedding
rotary = RotaryEmbedding(dim=64)

# Attention with rotary embeddings and value residual
output, values = attention(
    x,
    rotary_embed=rotary,
    return_values=True,
    value_residual=None  # First layer has no residual
)

# Subsequent layer with value residual connection
output2, values2 = attention(
    output,
    rotary_embed=rotary,
    return_values=True,
    value_residual=values  # Pass first layer values
)
```

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

### Instantiate and Use nGPT Model

Source: https://github.com/lucidrains/ngpt-pytorch/blob/main/README.md

Instantiate the nGPT model with specified parameters and demonstrate its usage for calculating loss and generating logits. Ensure torch is imported and the model is initialized before use.

```python
import torch
from nGPT_pytorch import nGPT

model = nGPT(
    num_tokens = 256,
    dim = 512,
    depth = 4,
    attn_norm_qk = True
)

x = torch.randint(0, 256, (2, 2048))

loss = model(x, return_loss = True)
loss.backward()

logits = model(x) # (2, 2048, 256)
```

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

### Configure Training Loop with Mixed Precision

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Sets up an nGPT model for training with Adam optimizer and gradient scaling. Includes hooks for manual weight normalization if enabled.

```python
import torch
from torch.optim import Adam
from torch.amp import GradScaler
from nGPT_pytorch import nGPT

# Model configuration
model = nGPT(
    num_tokens=256,
    dim=512,
    depth=8,
    tied_embedding=True,
    add_value_residual=True,
    attn_norm_qk=False,  # Optional QK norm
    manual_norm_weights=False  # Use parametrization (recommended)
).cuda()

# Optimizer and gradient scaler for mixed precision
optimizer = Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()

# If using manual weight normalization (manual_norm_weights=True):
# model.register_step_post_hook(optimizer)  # Auto-normalize after each step

# Training loop
for batch in dataloader:
    model.train()

    # Mixed precision forward pass
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        loss = model(batch, return_loss=True)

    # Gradient accumulation
    scaler.scale(loss / accumulation_steps).backward()

    if (step + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

    # Validation
    if step % validate_every == 0:
        model.eval()
        with torch.no_grad():
            val_loss = model(val_batch, return_loss=True)
            print(f"Validation loss: {val_loss.item():.3f}")
```

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

### Import Evolutionary Crossover Utilities

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Imports functions for breeding nGPT models by combining weights from parent models.

```python
import torch
from nGPT_pytorch import nGPT
from nGPT_pytorch.evo import (
    cross_over_ngpt,
    cross_over_attention,
    cross_over_feedforward,
    cross_over_linear,
    cross_over_scale
)
```

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

### Initialize and Use Learnable Scale

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Configures a learnable scale parameter for residual interpolation. The scale is stored as a fraction of the initial value.

```python
scale = Scale(
    dim=512,                 # Dimension of scale parameter
    init=0.1,                # Initial effective scale value
    scale=0.01               # Internal storage scale (init/scale stored)
)

# Get current scale values
current_scale = scale()  # Returns tensor of shape (512,)
print(current_scale.mean())  # ~0.1

# Used in residual interpolation: out = x.lerp(branch, scale())
x = torch.randn(2, 100, 512)
branch = torch.randn(2, 100, 512)
interpolated = x.lerp(branch, scale())
```

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

### Run Enwik8 Test

Source: https://github.com/lucidrains/ngpt-pytorch/blob/main/README.md

Execute the training script for the Enwik8 dataset. This command initiates the testing process for the nGPT model on the specified dataset.

```bash
python train.py
```

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

### Initialize and use nGPT model

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Configures an autoregressive language model with token embeddings and performs forward passes for training or inference.

```python
import torch
from nGPT_pytorch import nGPT

# Create a normalized GPT model for character-level language modeling
model = nGPT(
    num_tokens=256,          # Vocabulary size
    dim=512,                 # Model dimension
    depth=8,                 # Number of transformer layers
    dim_head=64,             # Dimension per attention head
    heads=8,                 # Number of attention heads
    attn_norm_qk=True,       # Enable query-key normalization
    ff_expand_factor=4.0,    # Feedforward expansion factor
    tied_embedding=True,     # Tie input/output embeddings
    add_value_residual=True, # Enable value residual connections (ResFormer)
    causal=True,             # Causal attention for autoregressive modeling
    num_hyperspheres=1,      # Number of hyperspheres for multi-head grouping
    norm_eps=0.0             # Allow exact unit norm (>0 allows slight deviation)
)

# Generate random input tokens
x = torch.randint(0, 256, (2, 2048))

# Forward pass - get logits
logits = model(x)  # Output shape: (2, 2048, 256)

# Training pass - compute autoregressive loss internally
loss = model(x, return_loss=True)
loss.backward()

# Get logits with hidden state outputs for analysis
logits, hiddens = model(x, return_hiddens=True)
# hiddens shape: (num_layers * 2 + 1, batch, seq_len, dim)

# Inference with attention mask for variable-length sequences
mask = torch.ones(2, 2048, dtype=torch.bool)
mask[:, 1024:] = False  # Mask out second half
logits = model(x, mask=mask)
```

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

### Initialize and use nTransformer backbone

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Configures a standalone transformer encoder/decoder for continuous embeddings and demonstrates processing with or without pre-normalization.

```python
import torch
from nGPT_pytorch import nTransformer

# Create a normalized transformer backbone
transformer = nTransformer(
    dim=256,                 # Input/output dimension
    depth=4,                 # Number of transformer layers
    dim_head=64,             # Dimension per attention head
    heads=8,                 # Number of attention heads
    attn_norm_qk=True,       # Enable query-key normalization
    ff_expand_factor=4.0,    # Feedforward expansion factor
    causal=False,            # Non-causal attention for bidirectional processing
    rotary_embed=True,       # Enable rotary positional embeddings
    num_hyperspheres=1,      # Hypersphere grouping
    alpha_init=None          # Auto-calculated as 1/depth
)

# Process continuous embeddings (e.g., from image patches or RL observations)
x = torch.randn(2, 1024, 256)

# Forward pass with automatic input normalization
embeddings = transformer(x, norm_input=True)
assert embeddings.shape == x.shape  # Shape preserved

# Forward pass assuming pre-normalized input
x_normed = torch.nn.functional.normalize(x, dim=-1)
embeddings = transformer(x_normed, norm_input=False)

# With attention mask for variable-length sequences
mask = torch.ones(2, 1024, dtype=torch.bool)
embeddings = transformer(x, mask=mask, norm_input=True)
```

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

### Implement Autoregressive Text Generation

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Provides sampling utilities including min-p filtering and Gumbel-softmax sampling for autoregressive generation.

```python
import torch
from nGPT_pytorch import nGPT

def min_p_filter(logits, min_p=0.1):
    """Filter low-probability tokens based on min_p threshold."""
    probs = logits.softmax(dim=-1)
    max_probs = probs.amax(dim=-1, keepdim=True)
    limit = min_p * max_probs
    return torch.where(probs < limit, float('-inf'), logits)

def gumbel_sample(logits, temperature=1.0, dim=-1):
    """Sample from logits using Gumbel-softmax trick."""
    noise = torch.zeros_like(logits).uniform_(0, 1)
    gumbel = -torch.log(-torch.log(noise.clamp(min=1e-20)))
    return ((logits / max(temperature, 1e-10)) + gumbel).argmax(dim=dim, keepdim=True)

def generate(model, prompt, max_length=512, temperature=1.5, min_p=0.1):
    """Autoregressive generation with min-p filtering."""
    model.eval()
    output = prompt.clone()

    for _ in range(max_length - prompt.shape[-1]):
        with torch.no_grad():
            logits = model(output)[:, -1]  # Get last position logits
            logits = min_p_filter(logits, min_p=min_p)
            next_token = gumbel_sample(logits, temperature=temperature)
            output = torch.cat([output, next_token], dim=-1)

    return output

# Usage
model = nGPT(num_tokens=256, dim=512, depth=8).cuda()
prompt = torch.randint(0, 256, (1, 128)).cuda()
generated = generate(model, prompt, max_length=512, temperature=1.5, min_p=0.1)
```

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

### Wrap Modules with Residual Connection

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Wraps a module with a normalized residual connection using spherical linear interpolation. Input tensors should be normalized before passing through the residual module.

```python
import torch
from nGPT_pytorch.nGPT import Residual, FeedForward, l2norm

# Create feedforward with residual wrapper
ff = FeedForward(dim=512, expand_factor=4)
residual_ff = Residual(
    fn=ff,                   # Module to wrap
    dim=512,                 # Dimension for scale
    init=0.25,               # Initial interpolation factor (1/depth typically)
    scale=None,              # Auto-calculated as dim^-0.5
    groups=1,                # Hypersphere groups
    norm_eps=0.0             # Normalization epsilon
)

# Forward pass - input should be normalized
x = torch.randn(2, 100, 512)
x = l2norm(x, dim=-1)

output = residual_ff(x)  # Normalized interpolation of x and ff(x)
print(output.norm(dim=-1).mean())  # ~1.0
```

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

### Perform Component-Level Crossover

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Iterates through model layers to perform fine-grained crossover on individual attention and feedforward modules.

```python
for (attn1, ff1), (attn2, ff2) in zip(parent1.layers, parent2.layers):
    # Cross over attention module
    child_attn = cross_over_attention(attn1.fn, attn2.fn)

    # Cross over feedforward module
    child_ff = cross_over_feedforward(ff1.fn, ff2.fn)
```

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

### Perform L2 Normalization

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Demonstrates the usage of the l2norm function for normalizing tensors. Covers basic L2 normalization, soft normalization with an epsilon range, and grouped normalization for multi-head attention scenarios.

```python
import torch
from nGPT_pytorch.nGPT import l2norm

# Basic L2 normalization
x = torch.randn(2, 1024, 512)
x_normed = l2norm(x, dim=-1)  # Normalize along last dimension
print(x_normed.norm(dim=-1))  # All values ~1.0

# Soft normalization (allow norms in range [1-eps, 1+eps])
x_soft = l2norm(x, dim=-1, norm_eps=0.01)  # Norms in [0.99, 1.01]

# Grouped normalization for multi-head attention
# Splits tensor into groups, normalizes each, then concatenates
x_heads = torch.randn(2, 8, 1024, 64)  # (batch, heads, seq, dim)
x_grouped = l2norm(x_heads.flatten(-2, -1), dim=-1, groups=8)
```

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

### Perform nGPT Model Crossover

Source: https://context7.com/lucidrains/ngpt-pytorch/llms.txt

Crosses over weights between two parent models to create a child model, with an optional pre-allocated child for memory efficiency.

```python
parent1 = nGPT(num_tokens=256, dim=512, depth=4, heads=8)
parent2 = nGPT(num_tokens=256, dim=512, depth=4, heads=8)

# Train parents on different data or with different random seeds...

# Breed a child model by crossing over weights
child = cross_over_ngpt(parent1, parent2)
# Child inherits ~50% of weights from each parent (random selection)

# Can also provide pre-allocated child for efficiency
child_preallocated = nGPT(num_tokens=256, dim=512, depth=4, heads=8)
child = cross_over_ngpt(parent1, parent2, child=child_preallocated)
```

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