### Setup experiments with uv

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

Install uv and run the training script.

```bash
$ pip install uv
```

```bash
$ uv run train_mac.py
```

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

### MemorySwiGluMLP Example

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Example of initializing MemorySwiGluMLP and integrating it with NeuralMemory. This setup uses a standard chunk size suitable for the SwiGLU architecture.

```python
from titans_pytorch.memory_models import MemorySwiGluMLP
from titans_pytorch import NeuralMemory

model = MemorySwiGluMLP(dim=384, depth=2, expansion_factor=4.0)

mem = NeuralMemory(
    dim=384,
    model=model,
    chunk_size=64
)
```

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

### Minimal MAC Transformer Example

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

Shows how to initialize and use the MemoryAsContextTransformer for training and sampling. Requires torch.

```python
from titans_pytorch import MemoryAsContextTransformer

model = MemoryAsContextTransformer(
    num_tokens=256,
    dim=256,
    depth=4,
    segment_len=128,
    num_longterm_mem_tokens=16,
)

# Training
token_ids = torch.randint(0, 256, (2, 512))
loss = model(token_ids, return_loss=True)
loss.backward()

# Sampling
sampled = model.sample(
    prompt=token_ids[:, :32],
    seq_len=256,
    temperature=1.5
)
```

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

### Setting up NeuralMemory with Custom MemorySwiGluMLP

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

This example demonstrates how to initialize NeuralMemory with a custom MemorySwiGluMLP model. Ensure the model and NeuralMemory are imported.

```python
from titans_pytorch import NeuralMemory
from titans_pytorch.memory_models import (
    MemoryMLP,
    MemorySwiGluMLP,
    FactorizedMemoryMLP
)

# Custom model
model = MemorySwiGluMLP(dim=384, depth=2)
mem = NeuralMemory(
    dim=384,
    model=model,
    chunk_size=64
)
```

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

### Install titans-pytorch

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

Install the package via pip.

```bash
$ pip install titans-pytorch
```

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

### Minimal Neural Memory Example

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

Demonstrates the basic usage of the NeuralMemory module. Ensure torch is imported and tensors are created.

```python
import torch
from titans_pytorch import NeuralMemory

mem = NeuralMemory(
    dim=384,
    chunk_size=64,
    heads=4
)

seq = torch.randn(2, 1024, 384)
retrieved, state = mem(seq)

assert retrieved.shape == seq.shape
```

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

### Basic Training of MemoryAsContextTransformer

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Demonstrates the basic setup and forward pass for training the MemoryAsContextTransformer model. Ensure optimizer is defined before calling `optimizer.step()`.

```python
from titans_pytorch import MemoryAsContextTransformer

model = MemoryAsContextTransformer(
    num_tokens=1024,
    dim=512,
    depth=8,
    segment_len=256,
    num_persist_mem_tokens=8,
    num_longterm_mem_tokens=32,
)

# Forward pass with loss
token_ids = torch.randint(0, 1024, (4, 2048))
loss = model(token_ids, return_loss=True)
loss.backward()
optimizer.step()
```

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

### FactorizedMemoryMLP Example

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Example of initializing FactorizedMemoryMLP and integrating it with NeuralMemory. This configuration is suitable for scenarios requiring smaller chunk sizes due to the factorized parameters.

```python
from titans_pytorch.memory_models import FactorizedMemoryMLP
from titans_pytorch import NeuralMemory

# Low-rank memory model for smaller chunks
model = FactorizedMemoryMLP(dim=384, depth=2, k=16)

mem = NeuralMemory(
    dim=384,
    model=model,
    chunk_size=16  # can use smaller chunks with factorized params
)
```

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

### Setting up NeuralMemory with Default MemoryMLP

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

This snippet shows the basic setup of NeuralMemory using the default MemoryMLP model. Ensure NeuralMemory is imported.

```python
from titans_pytorch import NeuralMemory
from titans_pytorch.memory_models import (
    MemoryMLP,
    MemorySwiGluMLP,
    FactorizedMemoryMLP
)

# Simple setup
mem = NeuralMemory(dim=384, chunk_size=64)  # uses default MemoryMLP
```

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

### Full MAC Transformer Usage and Sampling

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Illustrates the setup and usage of the MemoryAsContextTransformer, including loss calculation and sequence sampling. Requires importing MemoryAsContextTransformer and torch.

```python
from titans_pytorch import MemoryAsContextTransformer

model = MemoryAsContextTransformer(
    num_tokens=50000,
    dim=512,
    depth=12,
    segment_len=256,
    num_longterm_mem_tokens=16,
)

token_ids = torch.randint(0, 50000, (4, 2048))
loss = model(token_ids, return_loss=True)
loss.backward()

# Sampling
sampled = model.sample(prompt=token_ids[:, :256], seq_len=512)
```

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

### NestedAttention Usage Example

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/attention-variants.md

Demonstrates how to instantiate and use the NestedAttention module, including basic forward passes and caching for sequential inference.

```python
import torch
from titans_pytorch.nested_attention import NestedAttention

# Nested attention with 3 branches
attn = NestedAttention(
    dim=512,
    dim_head=64,
    heads=8
)

tokens = torch.randn(2, 1024, 512)

# Forward pass
output = attn(tokens)
assert output.shape == tokens.shape

# With caching
output, cache = attn(tokens, return_kv_cache=True)

# Inference
next_token = torch.randn(2, 1, 512)
next_output, cache = attn(next_token, cache=cache, return_kv_cache=True)
assert next_output.shape == (2, 1, 512)
```

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

### Example Cache Structure During Inference

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/types.md

Illustrative example of the cache tuple's content during inference, showing sequence position, KV cache shape, and neural memory states per layer.

```python
cache = (
    seq_index=100,                    # Position in sequence
    kv_caches=Tensor(
        shape=(depth, 2, batch, heads, attn_window_size, dim_head)
        # [layer, (k|v), batch, heads, pos, dim_head]
    ),
    neural_mem_caches=[
        NeuralMemState(...),          # For layer 1 (if has neural memory)
        None,                         # For layer 2 (no neural memory)
        NeuralMemState(...),          # For layer 3 (if has neural memory)
        ...
    ]
)
```

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

### Autoregressive Sampling with Prompt

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Demonstrates autoregressive sampling starting from a given prompt, generating new tokens up to a specified sequence length.

```python
prompt = torch.randint(0, 256, (1, 4))

# Generate 100 new tokens
generated = transformer.sample(
    prompt=prompt,
    seq_len=104,
    temperature=1.5,
    show_progress=True
)

full_sequence = torch.cat([prompt, generated], dim=1)
assert full_sequence.shape == (1, 104)
```

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

### Default Loss Function Implementation

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/types.md

An example implementation of a default loss function using Mean Squared Error (MSE).

```python
def default_loss_fn(pred: Tensor, target: Tensor) -> Tensor:
    return (pred - target).pow(2).mean(dim=-1)  # MSE
```

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

### Using ImplicitMLPAttention as a Memory Module

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/attention-variants.md

Demonstrates how to instantiate ImplicitMLPAttention and integrate it as a neural memory model within the MemoryAsContextTransformer. This setup is useful for custom attention mechanisms where a memory component is required.

```python
from titans_pytorch import MemoryAsContextTransformer
from titans_pytorch.memory_models import MemoryAttention
from titans_pytorch.implicit_mlp_attention import ImplicitMLPAttention

# Note: Not directly integrated into MAC, but works as memory module
memory_attn = ImplicitMLPAttention(
    dim=256,
    mlp_hiddens=(64, 128, 64),
    heads=4
)

transformer = MemoryAsContextTransformer(
    num_tokens=256,
    dim=256,
    depth=8,
    segment_len=128,
    neural_memory_model=memory_attn,  # Use implicit MLP as memory
)
```

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

### NeuralMemory Forward Pass Example

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/neural-memory.md

Demonstrates the initial forward pass to retrieve information from memory and a subsequent pass to continue a sequence using the previous state. Ensure the NeuralMemory class is imported.

```python
import torch
from titans_pytorch import NeuralMemory

mem = NeuralMemory(
    dim=384,
    chunk_size=64,
    heads=4
)

seq = torch.randn(2, 1024, 384)

# Initial forward pass
retrieved, state = mem(seq)
assert retrieved.shape == seq.shape

# Continue sequence with state
next_seq = torch.randn(2, 512, 384)
next_retrieved, next_state = mem(next_seq, state=state)
assert next_retrieved.shape == next_seq.shape
```

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

### Default Adaptive Step Transform Implementation

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/types.md

An example implementation of a default adaptive step transform that uses sigmoid to scale the output within a maximum learning rate.

```python
def default_adaptive_step_transform(
    adaptive_step: Tensor,
    max_lr: float = 1e-2
) -> Tensor:
    return adaptive_step.sigmoid() * max_lr
```

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

### Inference with MemoryAsContextTransformer

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Performs inference using the MemoryAsContextTransformer to get logits for given token IDs.

```python
# Inference
logits = transformer(token_ids)
```

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

### Initialize and use MemoryAsContextTransformer

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

Shows how to configure and use the MemoryAsContextTransformer for training and sampling.

```python
import torch
from titans_pytorch import MemoryAsContextTransformer

transformer = MemoryAsContextTransformer(
    num_tokens = 256,
    dim = 256,
    depth = 2,
    segment_len = 128,              # local attention window size
    num_persist_mem_tokens = 4,
    num_longterm_mem_tokens = 16,
)

token_ids = torch.randint(0, 256, (1, 1023))

loss = transformer(token_ids, return_loss = True) # (1, 1023, 256)
loss.backward()

# after much training

sampled = transformer.sample(token_ids[:, :4], 512)
```

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

### Initialize and use NeuralMemory

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

Demonstrates initializing the NeuralMemory module and performing a forward pass with a sequence.

```python
import torch
from titans_pytorch import NeuralMemory

mem = NeuralMemory(
    dim = 384,
    chunk_size = 64 # set to smaller chunk size for better perf on smaller sequence lengths (but more memory usage)
).cuda()

seq = torch.randn(2, 1024, 384).cuda()
retrieved, mem_state = mem(seq)

assert seq.shape == retrieved.shape
```

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

### Setting up NeuralMemory with FactorizedMemoryMLP for Constrained Scenarios

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

This snippet illustrates setting up NeuralMemory with FactorizedMemoryMLP for scenarios with memory constraints. Ensure imports are present.

```python
from titans_pytorch import NeuralMemory
from titans_pytorch.memory_models import (
    MemoryMLP,
    MemorySwiGluMLP,
    FactorizedMemoryMLP
)

# For constrained scenarios
small_model = FactorizedMemoryMLP(dim=128, depth=2, k=8)
mem = NeuralMemory(
    dim=128,
    model=small_model,
    chunk_size=8
)
```

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

### Small Model Configuration (Testing)

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/configuration.md

Use this configuration for testing purposes. It sets up a smaller model with fewer parameters.

```python
from titans_pytorch import MemoryAsContextTransformer

model = MemoryAsContextTransformer(
    num_tokens=256,
    dim=64,
    depth=2,
    segment_len=32,
    heads=2,
    dim_head=32,
    num_persist_mem_tokens=0,
    num_longterm_mem_tokens=0,
)
```

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

### Custom Memory Configuration with MemorySwiGluMLP

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Demonstrates configuring the MemoryAsContextTransformer with a custom neural memory model (MemorySwiGluMLP) and specific keyword arguments for its behavior. It also shows how to apply the neural memory to specific layers.

```python
from titans_pytorch.memory_models import MemorySwiGluMLP

model = MemoryAsContextTransformer(
    num_tokens=256,
    dim=512,
    depth=12,
    segment_len=512,
    neural_memory_model=MemorySwiGluMLP(dim=512, depth=2),
    neural_memory_kwargs={
        'chunk_size': 64,
        'momentum': True,
        'momentum_order': 2,
        'max_grad_norm': 2.0,
    },
    neural_memory_layers=(2, 4, 6, 8, 10, 12),  # Apply to specific layers
    neural_mem_weight_residual=True,  # Pass weights between layers
)
```

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

### Initialize and Train MemoryAsContextTransformer

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Initializes the MemoryAsContextTransformer and demonstrates its use for training by computing and backpropagating the loss.

```python
import torch
from titans_pytorch import MemoryAsContextTransformer

transformer = MemoryAsContextTransformer(
    num_tokens=256,
    dim=256,
    depth=4,
    segment_len=128,
    num_persist_mem_tokens=4,
    num_longterm_mem_tokens=16,
)

# Training
token_ids = torch.randint(0, 256, (2, 512))
loss = transformer(token_ids, return_loss=True)
loss.backward()
```

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

### MemorySwiGluMLP Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Initializes a SwiGLU-style feedforward network with gated linear units. Suitable for modern architectures requiring efficient hidden dimension expansion.

```python
class MemorySwiGluMLP(Module):
    def __init__(
        self,
        dim: int,
        depth: int = 1,
        expansion_factor: float = 4.
    )
```

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

### Attention Backend Configuration

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/types.md

Dictionary for the Attend backend. Primarily used to enable flash attention if available.

```python
attend_kwargs: dict = dict(
    # Typical:
    flash: bool = True,  # Use flash attention if available
    # Others depend on Attend implementation
)
```

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

### Import NeuralMemory and MemoryAsContextTransformer

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Recommended import for most users to access core components.

```python
from titans_pytorch import (
    NeuralMemory,
    MemoryAsContextTransformer,
)
```

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

### Sampling with Flex Attention on CUDA

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Illustrates how to enable and use flex attention for sampling, requiring CUDA availability and the latest PyTorch version. The model is moved to the GPU.

```python
# Must run on CUDA with latest PyTorch
if torch.cuda.is_available():
    model = MemoryAsContextTransformer(
        num_tokens=256,
        dim=256,
        depth=6,
        segment_len=128,
        use_flex_attn=True,  # Enable flex_attention
    ).cuda()
    
    generated = model.sample(prompt.cuda(), seq_len=512)
```

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

### Sample from MemoryAsContextTransformer

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Generates a sequence of tokens autoregressively using the sample method with specified parameters.

```python
# Sampling
sampled = transformer.sample(
    prompt=token_ids[:, :20],
    seq_len=256,
    temperature=1.5
)
```

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

### NeuralMemory with Custom Memory Models

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Shows how to integrate custom memory models with NeuralMemory. Two options are presented: a pre-built MemorySwiGluMLP and an efficient FactorizedMemoryMLP.

```python
from titans_pytorch import NeuralMemory, MemorySwiGluMLP
from titans_pytorch.memory_models import FactorizedMemoryMLP

# Option 1: Pre-built model
model = MemorySwiGluMLP(dim=384, depth=2)
mem = NeuralMemory(dim=384, model=model, chunk_size=64)

# Option 2: Factorized for efficiency
model = FactorizedMemoryMLP(dim=256, depth=2, k=16)
mem = NeuralMemory(dim=256, model=model, chunk_size=16)
```

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

### Import for Custom Models

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Import necessary components when building custom memory models.

```python
from titans_pytorch import (
    NeuralMemory,
    MemoryMLP,
    MemorySwiGluMLP,
    FactorizedMemoryMLP,
)
```

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

### Neural Memory Configuration Options

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/types.md

Dictionary for NeuralMemory constructor. Includes common options like momentum and attention pooling. Excludes 'dim' and 'chunk_size'.

```python
neural_memory_kwargs: dict = dict(
    # Common options:
    momentum: bool = True,
    momentum_order: int = 1,
    max_grad_norm: float | None = None,
    attn_pool_chunks: bool = False,
    per_head_learned_parameters: bool = True,
    # ... any other NeuralMemory.__init__ parameter except dim, chunk_size
)
```

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

### sample

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Performs autoregressive sampling to generate new token sequences, with options for caching, temperature, and logit filtering.

```APIDOC
## sample

### Description
Autoregressive sampling with optional caching.

### Method
`sample`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
* **prompt** (Tensor) - Required - Shape (batch, prompt_len) with initial tokens.
* **seq_len** (int) - Required - Total sequence length (including prompt).
* **temperature** (float) - Optional - Default: 1.5 - Sampling temperature; 0=argmax, >1=more random.
* **filter_fn** (Callable) - Optional - Default: min_p_filter - Logit filter function (min_p or top_k).
* **filter_kwargs** (dict) - Optional - Default: dict(min_p=0.1) - Kwargs for filter_fn.
* **show_progress** (bool) - Optional - Default: True - Display progress bar.
* **use_cache** (bool) - Optional - Default: False - Use KV cache for faster generation (requires extra memory).

### Request Example
```python
import torch

# Assuming transformer is an instance of MemoryAsContextTransformer
# prompt = torch.randint(0, 256, (1, 4))
# generated = transformer.sample(
#     prompt=prompt,
#     seq_len=104,
#     temperature=1.5,
#     show_progress=True
# )
```

### Response
#### Success Response
- **sampled** (Tensor) - Shape (batch, seq_len - prompt_len), sampled token IDs.

#### Response Example
```python
# generated_tokens = transformer.sample(prompt, seq_len=104)
```
```

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

### Initialize MemoryMLP and NeuralMemory

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Instantiates a MemoryMLP model and a NeuralMemory object. Use MemoryMLP for a simple feedforward MLP memory, configuring its dimension, depth, and expansion factor.

```python
from titans_pytorch.memory_models import MemoryMLP
from titans_pytorch import NeuralMemory

model = MemoryMLP(dim=384, depth=2, expansion_factor=4.0)

mem = NeuralMemory(
    dim=384,
    model=model,
    chunk_size=64
)
```

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

### Medium Model Configuration (Production)

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/configuration.md

This configuration is suitable for production environments, offering a balance between performance and resource usage. It includes settings for neural memory.

```python
model = MemoryAsContextTransformer(
    num_tokens=50000,
    dim=512,
    depth=12,
    segment_len=256,
    heads=8,
    dim_head=64,
    num_persist_mem_tokens=4,
    num_longterm_mem_tokens=16,
    neural_memory_kwargs=dict(
        momentum=True,
        momentum_order=2,
        max_grad_norm=2.0,
    ),
)
```

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

### MemorySwiGluMLP

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

A memory model utilizing the SwiGLU activation function.

```APIDOC
## MemorySwiGluMLP

### Description
SwiGLU-style memory model.

### Parameters
- **dim** (int) - Dimension of the model.
- **depth** (int) - Depth of the MLP.
- **expansion_factor** (int) - Expansion factor for the MLP layers.

### Location
titans_pytorch/memory_models.py:167
```

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

### Training Loop with MemoryAsContextTransformer

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

Standard PyTorch training loop using MemoryAsContextTransformer. Ensure to initialize the model and optimizer correctly.

```python
model = MemoryAsContextTransformer(...)
optimizer = torch.optim.AdamW(model.parameters())

for batch in dataloader:
    loss = model(batch, return_loss=True)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
```

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

### Initialize GatedResidualMemoryMLP and NeuralMemory

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Instantiates a GatedResidualMemoryMLP model and a NeuralMemory object. This model features gated residual connections and learned interpolation, suitable for more complex memory interactions.

```python
from titans_pytorch.memory_models import GatedResidualMemoryMLP
from titans_pytorch import NeuralMemory

model = GatedResidualMemoryMLP(dim=256, depth=3, expansion_factor=4.0)

mem = NeuralMemory(
    dim=256,
    model=model,
    chunk_size=32
)
```

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

### Inference with Caching for Single-Token Generation

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Shows how to perform inference with caching to generate tokens sequentially. The cache is updated and passed in subsequent forward calls.

```python
prompt = torch.randint(0, 1024, (1, 256))

# Initial forward
logits, cache = model(prompt, return_cache=True)

# Single-token generation
next_token_id = logits[:, -1].argmax(dim=-1, keepdim=True)

for _ in range(512):
    logits, cache = model(
        next_token_id,
        cache=cache,
        return_cache=True
    )
    next_token_id = logits[:, -1].argmax(dim=-1, keepdim=True)
```

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

### Core PyTorch and Utility Imports

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Lists essential external dependencies for Titans PyTorch, including PyTorch, einops, and tensordict. Some imports are optional or require specific PyTorch versions.

```python
import torch
from torch import nn, Tensor
from torch.nn import Module, Linear, Parameter, ParameterList

from einops import rearrange, repeat, reduce, einsum, pack, unpack
from einops.layers.torch import Rearrange, Reduce

from tensordict import TensorDict
from assoc_scan import AssocScan
from rotary_embedding_torch import RotaryEmbedding
from axial_positional_embedding import ContinuousAxialPositionalEmbedding
from x_transformers.attend import Attend
from hyper_connections import mc_get_init_and_expand_reduce_stream_functions

# Optional flex_attention (torch >= 2.5)
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
```

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

### Initialize MemoryAttention Module

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Instantiates the MemoryAttention module with specified dimensions and scaling factors. This module is often used within a NeuralMemory wrapper.

```python
from titans_pytorch.memory_models import MemoryAttention
from titans_pytorch import NeuralMemory

model = MemoryAttention(dim=384, scale=8.0, expansion_factor=2.0)

mem = NeuralMemory(
    dim=384,
    model=model,
    chunk_size=64
)
```

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

### ImplicitMLPAttention Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/attention-variants.md

Initializes the ImplicitMLPAttention module. Configure dimensions, MLP hidden layers, activation function, attention heads, and normalization settings.

```python
class ImplicitMLPAttention(Module):
    def __init__(
        self,
        dim: int,
        mlp_hiddens: tuple[int, ...],
        *, 
        activation: Module = nn.SiLU(),
        heads: int = 8,
        talking_heads: bool = True,
        prenorm: bool = True,
        keys_rmsnorm: bool = True
    )
```

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

### NeuralMemory with Implicit MLP Attention

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Demonstrates using ImplicitMLPAttention as a custom memory model within NeuralMemory. Requires importing ImplicitMLPAttention and NeuralMemory.

```python
from titans_pytorch.implicit_mlp_attention import ImplicitMLPAttention

# Use as custom memory model
attn_memory = ImplicitMLPAttention(
    dim=384,
    mlp_hiddens=(64, 128, 64),
    heads=4
)

mem = NeuralMemory(
    dim=384,
    model=attn_memory,
    chunk_size=64
)
```

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

### NeuralMemory Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/neural-memory.md

Initializes the NeuralMemory module. Configure dimensions, attention heads, memory model, and various training parameters like loss functions and learning rate modulation.

```python
NeuralMemory(
    dim: int,
    chunk_size: int | tuple[int, int] = 1,
    batch_size: int | None = None,
    dim_head: int | None = None,
    heads: int = 1,
    model: Module | None = None,
    store_memory_loss_fn: Callable = default_loss_fn,
    adaptive_step_transform: Callable | None = None,
    default_step_transform_max_lr: float = 1.,
    per_parameter_lr_modulation: bool = False,
    max_mem_layer_modulation: float = 1.,
    per_head_learned_parameters: bool = True,
    attn_pool_chunks: bool = False,
    momentum: bool = True,
    momentum_order: int = 1,
    learned_momentum_combine: bool = False,
    learned_combine_include_zeroth: bool = False,
    num_kv_per_token: int = 1,
    qkv_receives_diff_views: bool = False,
    pre_rmsnorm: bool = True,
    post_rmsnorm: bool = False,
    qk_rmsnorm: bool = False,
    max_grad_norm: float | None = None,
    use_accelerated_scan: bool = False,
    activation: Module | None = None,
    init_adaptive_step_bias: float | None = None,
    init_momentum_bias: float | None = None,
    init_decay_bias: float | None = None,
    accept_weight_residual: bool = False,
    spectral_norm_surprises: bool = False,
    gated_transition: bool = False,
    mem_model_norm_add_residual: bool = True,
    store_with_lookahead_value: bool = False,
    default_model_kwargs: dict = dict(depth=2, expansion_factor=4.)
)
```

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

### Exported from memory_models

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Imports various memory model architectures. Choose the appropriate model based on your needs for dimensionality, depth, and efficiency.

```python
from titans_pytorch import MemoryMLP
from titans_pytorch import MemoryAttention
from titans_pytorch import FactorizedMemoryMLP
from titans_pytorch import MemorySwiGluMLP
from titans_pytorch import GatedResidualMemoryMLP
```

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

### Import Internal Components

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Import for accessing internal utilities, rarely needed for typical usage.

```python
from titans_pytorch.neural_memory import MultiheadRMSNorm, AveragePool
from titans_pytorch.mac_transformer import SegmentedAttention, min_p_filter
```

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

### Low Learning Rate Configuration (Stable)

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/configuration.md

Sets a low learning rate for stable training. This configuration is conservative and suitable for fine-tuning.

```python
neural_memory_kwargs=dict(
    default_step_transform_max_lr=0.1,   # Decrease from default 1.0
    init_adaptive_step_bias=-2.0,        # Start very conservative
)
```

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

### Using ImplicitMLPAttention as a NeuralMemory Module

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/attention-variants.md

Instantiates ImplicitMLPAttention and integrates it as a custom memory module into NeuralMemory. Ensure the attention module adheres to specific input/output shapes and gradient computation requirements.

```python
from titans_pytorch import NeuralMemory
from titans_pytorch.implicit_mlp_attention import ImplicitMLPAttention

# Use implicit MLP as memory
memory_module = ImplicitMLPAttention(
    dim=384,
    mlp_hiddens=(64, 128, 64),
    heads=4
)

mem = NeuralMemory(
    dim=384,
    model=memory_module,
    chunk_size=64
)
```

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

### FactorizedMemoryMLP Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/memory-models.md

Initializes a low-rank factorized MLP. Use this for memory-efficient storage with smaller chunk sizes. The rank 'k' controls parameter count.

```python
class FactorizedMemoryMLP(Module):
    def __init__(
        self,
        dim: int,
        depth: int,
        k: int = 32
    )
```

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

### Importing Neural Memory Components

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Import various normalization, pooling, and attention-related components directly from the `neural_memory` submodule. Includes helper functions for internal use.

```python
from titans_pytorch.neural_memory import (
    MultiheadRMSNorm,
    AveragePool,
    AttentionPool,
    default_adaptive_step_transform,
    default_loss_fn,
    # Helper functions (private):
    exists,
    default,
    softclamp_max,
    softclamp_grad_norm,
    newtonschulz5,
    pack_one_with_inverse,
)
```

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

### Importing MAC Transformer Components

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Import components for the MAC Transformer, including attention mechanisms, mask creation, sampling utilities, and feedforward layers from the `mac_transformer` submodule.

```python
from titans_pytorch.mac_transformer import (
    SegmentedAttention,
    create_mac_block_mask,
    min_p_filter,
    gumbel_sample,
    FeedForward,
    GEGLU,
    # Utilities:
    AttnIntermediates,
)
```

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

### Exported from neural_memory

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Imports core components for test-time training and memory state management. Use these for direct interaction with the neural memory system.

```python
from titans_pytorch import NeuralMemory
from titans_pytorch import NeuralMemState
from titans_pytorch import mem_state_detach
```

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

### Create MAC Block Mask

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Generates a block mask for the MAC architecture's flex attention. Use this to define attention patterns based on sequence length, window size, and persistent memory.

```python
def create_mac_block_mask(
    seq_len: int,
    window_size: int,
    persist_mem_len: int,
    sliding: bool = False
) -> BlockMask:
    pass
```

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

### MemoryAsContextTransformer Parameters

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

Key parameters for the MemoryAsContextTransformer module, including vocabulary size, model dimension, layer depth, and attention configurations.

```markdown
| Parameter | Default | Purpose |
|-----------|---------|---------|
| `num_tokens` | — | Vocabulary size |
| `dim` | — | Model dimension |
| `depth` | — | Layer count |
| `segment_len` | — | Local attention window |
| `num_longterm_mem_tokens` | 0 | Distributed memory tokens |
| `num_persist_mem_tokens` | 0 | Global memory tokens |
| `heads` | 8 | Attention heads |
| `use_flex_attn` | False | Efficient attention backend |
| `neural_memory_layers` | All | Where to apply memory |
| `neural_memory_kwargs` | {} | Memory configuration |
```

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

### MemoryAsContextTransformer Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/mac-transformer.md

Defines the MemoryAsContextTransformer class with various parameters for configuring memory, attention, and model dimensions. Use this to instantiate the transformer model.

```python
class MemoryAsContextTransformer(Module):
    def __init__(
        self,
        *, 
        num_tokens: int,
        dim: int,
        depth: int,
        segment_len: int,
        neural_memory_segment_len: int | None = None,
        neural_mem_gate_attn_output: bool = False,
        neural_memory_add_value_residual: bool = False,
        num_longterm_mem_tokens: int = 0,
        num_persist_mem_tokens: int = 0,
        neural_memory_batch_size: int | None = None,
        neural_memory_qkv_receives_diff_views: bool = False,
        dim_head: int = 64,
        heads: int = 8,
        ff_mult: int = 4,
        num_residual_streams: int = 4,
        neural_memory_model: Module | None = None,
        neural_memory_kwargs: dict = dict(),
        neural_memory_layers: tuple[int, ...] | None = None,
        use_flex_attn: bool = False,
        sliding_window_attn: bool = False,
        neural_mem_weight_residual: bool = False,
        token_emb: Module | None = None,
    )
```

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

### Large Model with Efficient Attention

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/configuration.md

Configure a large model with efficient attention mechanisms. This requires CUDA and the latest PyTorch version. It includes advanced neural memory settings.

```python
model = MemoryAsContextTransformer(
    num_tokens=100000,
    dim=1024,
    depth=24,
    segment_len=512,
    heads=16,
    dim_head=64,
    num_persist_mem_tokens=8,
    num_longterm_mem_tokens=32,
    use_flex_attn=True,  # Requires CUDA + latest PyTorch
    neural_memory_kwargs=dict(
        momentum=True,
        momentum_order=1,
        max_grad_norm=1.0,
        per_head_learned_parameters=True,
    ),
)
```

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

### NeuralMemory Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/neural-memory.md

Initializes a NeuralMemory module with specified dimensions, chunk sizes, and various configuration options for memory storage and retrieval.

```APIDOC
## NeuralMemory Constructor

### Description
Initializes a NeuralMemory module with specified dimensions, chunk sizes, and various configuration options for memory storage and retrieval.

### Parameters

#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters

- **dim** (int) - Required - Feature dimension of input sequences
- **chunk_size** (int | tuple[int, int]) - Optional - Default: 1 - Chunk size for storing to memory. If tuple, (retrieve_chunk_size, store_chunk_size)
- **batch_size** (int | None) - Optional - Default: None - Mini-batch size for weight updates during sequence processing
- **dim_head** (int | None) - Optional - Default: None - Feature dimension per head; defaults to `dim` if heads=1
- **heads** (int) - Optional - Default: 1 - Number of attention heads
- **model** (Module | None) - Optional - Default: None - Memory model for storing/retrieving; defaults to MemoryMLP
- **store_memory_loss_fn** (Callable) - Optional - Default: default_loss_fn - Loss function for computing gradients: `(pred, target) -> loss`
- **adaptive_step_transform** (Callable | None) - Optional - Default: None - Transform for learning rate; applied to output of `to_adaptive_step`
- **default_step_transform_max_lr** (float) - Optional - Default: 1. - Maximum learning rate when using default transform
- **per_parameter_lr_modulation** (bool) - Optional - Default: False - Allow per-weight modulation of learning rate
- **max_mem_layer_modulation** (float) - Optional - Default: 1. - Scale factor for per-parameter learning rate
- **per_head_learned_parameters** (bool) - Optional - Default: True - If True, memory model parameters are per-head; else shared across heads
- **attn_pool_chunks** (bool) - Optional - Default: False - Use attention pooling instead of average pooling for chunks
- **momentum** (bool) - Optional - Default: True - Apply momentum/velocity accumulation via associative scan
- **momentum_order** (int) - Optional - Default: 1 - Order of momentum (1=first-order, 2=second-order)
- **learned_momentum_combine** (bool) - Optional - Default: False - Learn combination weights for multiple momentum orders
- **learned_combine_include_zeroth** (bool) - Optional - Default: False - Include zeroth-order (original surprise) in momentum combination
- **num_kv_per_token** (int) - Optional - Default: 1 - Number of key-value pairs per token for multiple updates
- **qkv_receives_diff_views** (bool) - Optional - Default: False - Allow queries and key-values from different input views
- **pre_rmsnorm** (bool) - Optional - Default: True - Apply RMSNorm before projecting to keys/values/queries
- **post_rmsnorm** (bool) - Optional - Default: False - Apply multi-head RMSNorm after value projection
- **qk_rmsnorm** (bool) - Optional - Default: False - Apply RMSNorm to queries and keys separately
- **max_grad_norm** (float | None) - Optional - Default: None - Soft-clamp maximum gradient norm
- **use_accelerated_scan** (bool) - Optional - Default: False - Use CUDA-accelerated associative scan
- **activation** (Module | None) - Optional - Default: None - Activation function in key/value projections
- **init_adaptive_step_bias** (float | None) - Optional - Default: None - Initial bias value for adaptive step linear layer
- **init_momentum_bias** (float | None) - Optional - Default: None - Initial bias value for momentum linear layer
- **init_decay_bias** (float | None) - Optional - Default: None - Initial bias value for decay factor linear layer
- **accept_weight_residual** (bool) - Optional - Default: False - Allow residual connection from previous layer's weights
- **spectral_norm_surprises** (bool) - Optional - Default: False - Apply spectral normalization to weight updates (surprises)
- **gated_transition** (bool) - Optional - Default: False - Gated interpolation between residual and updated weights
- **mem_model_norm_add_residual** (bool) - Optional - Default: True - Wrap memory model with ResidualNorm
- **store_with_lookahead_value** (bool) - Optional - Default: False - Use next timestep's values for gradient computation
- **default_model_kwargs** (dict) - Optional - Default: `depth=2, expansion_factor=4.` - Default kwargs for MemoryMLP if model not provided

### Request Example
```python
# Example usage (not a direct API call, but shows parameter usage)
from torch.nn import Module
from typing import Callable, Tuple, Optional

def default_loss_fn(pred, target):
    # Placeholder for default loss function
    return (pred - target).pow(2).mean()

# Assuming Module and other necessary imports are available
# neural_memory = NeuralMemory(
#     dim=128,
#     chunk_size=(32, 64),
#     batch_size=16,
#     dim_head=64,
#     heads=8,
#     model=None, # or an instance of a memory model
#     store_memory_loss_fn=default_loss_fn,
#     adaptive_step_transform=None,
#     default_step_transform_max_lr=0.5,
#     per_parameter_lr_modulation=True,
#     max_mem_layer_modulation=0.8,
#     per_head_learned_parameters=False,
#     attn_pool_chunks=True,
#     momentum=False,
#     momentum_order=2,
#     learned_momentum_combine=True,
#     learned_combine_include_zeroth=True,
#     num_kv_per_token=2,
#     qkv_receives_diff_views=True,
#     pre_rmsnorm=False,
#     post_rmsnorm=True,
#     qk_rmsnorm=True,
#     max_grad_norm=1.0,
#     use_accelerated_scan=True,
#     activation=None, # or an activation module
#     init_adaptive_step_bias=0.1,
#     init_momentum_bias=0.2,
#     init_decay_bias=0.3,
#     accept_weight_residual=True,
#     spectral_norm_surprises=True,
#     gated_transition=True,
#     mem_model_norm_add_residual=False,
#     store_with_lookahead_value=True,
#     default_model_kwargs={'depth': 3, 'expansion_factor': 5.0}
# )
```

### Response

#### Success Response (200)
This constructor does not return a value directly but initializes an object of the `NeuralMemory` class.

#### Response Example
None
```

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

### Multi-Head Memory Configuration

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

Configuration for NeuralMemory with multiple heads and per-head learned parameters. Ensure dim_head is set correctly to match the total dimension.

```python
mem = NeuralMemory(
    dim=384,
    chunk_size=64,
    heads=8,
    dim_head=48,  # 8 heads * 48 = 384
    per_head_learned_parameters=True  # Each head has own parameters
)
```

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

### MemoryMLP

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Default memory model, a simple feedforward MLP.

```APIDOC
## MemoryMLP

### Description
Default memory model: simple feedforward MLP.

### Parameters
- **dim** (int) - Dimension of the model.
- **depth** (int) - Depth of the MLP.
- **expansion_factor** (int) - Expansion factor for the MLP layers.

### Location
titans_pytorch/memory_models.py:54
```

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

### NestedAttention Constructor

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/attention-variants.md

Initializes the NestedAttention module. Use this to set up the attention mechanism with specified dimensions, heads, and normalization options.

```python
class NestedAttention(Module):
    def __init__(
        self,
        dim: int,
        *,
        dim_head: int = 64,
        heads: int = 8,
        prenorm: bool = True,
        keys_rmsnorm: bool = True
    )
```

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

### Exported from mac_transformer

Source: https://github.com/lucidrains/titans-pytorch/blob/main/_autodocs/module-exports.md

Imports the MemoryAsContextTransformer, a full transformer model augmented with neural memory. Use this for end-to-end language modeling tasks.

```python
from titans_pytorch import MemoryAsContextTransformer
```