### Install fast-weight-product-key-memory

Source: https://github.com/lucidrains/fast-weight-product-key-memories/blob/main/README.md

Install the library using pip. This command is used to add the package to your Python environment.

```bash
pip install fast-weight-product-key-memory
```

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

### Perform low-level memory retrieval and storage

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Demonstrates initializing the fwPKM module and using the retrieve and store methods to perform memory lookups and updates.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim=256,
    heads=2,
    num_memories=64 * 64,
    dim_queries_keys=128,
    dim_values=128,
    topk=4,
    chunk_size=32
)

tokens = torch.randn(1, 32, 256)

# Low-level retrieve call
output, intermediates = mem.retrieve(
    tokens,
    past_memories=None,
    idw_eps=1e-3  # Inverse distance weighting epsilon
)

print(f"Retrieved output shape: {output.shape}")
print(f"Intermediate keys: {list(intermediates.keys())}")
# Keys include: q1, q2, dist1, dist2, indices1, indices2, scores,
#               sub_indices, final_indices, final_scores, gates,
#               target_values, values, memories, k1, k2

# Low-level store call - computes memory updates
memory_values, keys = mem.store(
    intermediates,
    past_memories=None,
    detach_next_memories=False,
    idw_eps=1e-3
)

print(f"Updated memory values shape: {memory_values.shape}")
print(f"Updated keys shape: {keys.shape}")
```

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

### Initialize and use fwPKM module

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Initialize the fwPKM memory module and perform a forward pass with input tokens.

```python
import torch
from fast_weight_product_key_memory import fwPKM

# Initialize the memory module
mem = fwPKM(
	dim=512,                    # Input/output dimension
	heads=4,                    # Number of attention heads
	num_memories=256 * 256,     # Total memory slots (must have integer sqrt)
	dim_queries_keys=512,       # Dimension for queries and keys
	dim_values=512,             # Dimension for stored values
	topk=8,                     # Number of top memories to retrieve
	learning_rate=1.0,          # Fast weight update learning rate
	chunk_size=256,             # Chunk size for memory updates
	addressing_loss_weight=10., # Weight for addressing loss (entropy regularization)
	mse_loss_weight_to_keys=1.  # MSE loss weight for key updates
)

# Create sample input tokens
tokens = torch.randn(2, 256, 512)  # (batch, sequence, dim)

# Forward pass - retrieves from memory
output = mem(tokens)
print(f"Output shape: {output.shape}")  # torch.Size([2, 256, 512])
```

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

### Implement Sequence Memorization Task

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Shows a minimal model using fwPKM to recall the first half of a sequence to predict the second half.

```python
import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import Adam
from fast_weight_product_key_memory import fwPKM

class MemorizingModel(nn.Module):
    def __init__(self, num_tokens, dim, chunk_size=4):
        super().__init__()
        self.embed = nn.Embedding(num_tokens, dim)

        self.memory = fwPKM(
            dim=dim,
            heads=1,
            num_memories=256,
            dim_queries_keys=64,
            dim_values=dim,
            learning_rate=1.,
            topk=4,
            chunk_size=chunk_size,
            addressing_loss_weight=1e-3
        )

        self.head = nn.Linear(dim, num_tokens)

    def forward(self, x):
        h = self.embed(x)
        h = h + self.memory(h)
        return self.head(h)

# Training setup
num_tokens, dim, half_len = 32, 128, 8
model = MemorizingModel(num_tokens, dim, chunk_size=4)
optim = Adam(model.parameters(), lr=1e-3)

# Training loop
for step in range(100):
    # Create memorization task: repeat sequence
    half = torch.randint(0, num_tokens, (1, half_len))
    seq = torch.cat((half, half), dim=-1)
    x, labels = seq[:, :-1], seq[:, 1:]

    logits = model(x)

    # Loss only on recall portion (second half)
    loss = F.cross_entropy(
        logits[:, (half_len - 1):].reshape(-1, num_tokens),
        labels[:, (half_len - 1):].reshape(-1)
    )

    loss.backward()
    optim.step()
    optim.zero_grad()

    if step % 20 == 0:
        print(f"Step {step}, Loss: {loss.item():.4f}")

# Evaluation
model.eval()
with torch.no_grad():
    test_half = torch.randint(0, num_tokens, (1, half_len))
    test_seq = torch.cat((test_half, test_half), dim=-1)
    test_x, test_labels = test_seq[:, :-1], test_seq[:, 1:]

    preds = model(test_x).argmax(dim=-1)
    acc = (preds[:, (half_len - 1):] == test_labels[:, (half_len - 1):]).float().mean()
    print(f"Recall accuracy: {acc.item():.1%}")
```

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

### Run Enwik8 Training Script

Source: https://github.com/lucidrains/fast-weight-product-key-memories/blob/main/README.md

Execute the training script for a character-level language model using `fwPKM` on the enwik8 dataset.

```shell
uv run train_enwik8.py
```

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

### Initialize and Use fwPKM in Python

Source: https://github.com/lucidrains/fast-weight-product-key-memories/blob/main/README.md

Initialize the fwPKM memory module and process a chunk of tokens. The module can also chain memories by passing `past_memories`.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim = 512,
    num_memories = 256 * 256,
    dim_queries_keys = 512,
    dim_values = 512,
    topk = 8,
    learning_rate = 1.,
    chunk_size = 256
)

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

# forward a chunk of tokens for retrieved and the fast weight episodic memories

retrieved, next_memories = mem(tokens, return_next_memories = True)

# chain memories

retrieved, next_memories = mem(tokens, return_next_memories = True, past_memories = next_memories)
retrieved, next_memories = mem(tokens, return_next_memories = True, past_memories = next_memories)
retrieved, next_memories = mem(tokens, return_next_memories = True, past_memories = next_memories)
```

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

### Integrate fwPKM into Transformer Architecture

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Demonstrates how to conditionally insert fwPKM layers into a standard transformer module and manage memory state across forward passes.

```python
import torch
from torch import nn
from fast_weight_product_key_memory import fwPKM

class TransformerWithMemory(nn.Module):
    def __init__(
        self,
        dim=512,
        depth=6,
        heads=8,
        memory_layers=(2, 4),  # Add memory at layers 2 and 4
    ):
        super().__init__()
        self.layers = nn.ModuleList()

        for layer_idx in range(1, depth + 1):
            # Add fwPKM at specified layers
            memory = None
            if layer_idx in memory_layers:
                memory = fwPKM(
                    dim=dim,
                    heads=4,
                    num_memories=512 * 512,
                    dim_queries_keys=512,
                    dim_values=512,
                    learning_rate=1.,
                    topk=8,
                    chunk_size=256,
                    addressing_loss_weight=5.
                )

            self.layers.append(nn.ModuleDict({
                'memory': memory,
                'attention': nn.MultiheadAttention(dim, heads, batch_first=True),
                'ff': nn.Sequential(
                    nn.Linear(dim, dim * 4),
                    nn.GELU(),
                    nn.Linear(dim * 4, dim)
                ),
                'norm1': nn.LayerNorm(dim),
                'norm2': nn.LayerNorm(dim)
            }))

    def forward(self, x, cache=None):
        new_cache = []
        cache = cache or [None] * len(self.layers)

        for layer, layer_cache in zip(self.layers, cache):
            # Apply episodic memory if present
            next_mem_cache = layer_cache
            if layer['memory'] is not None:
                mem_out, next_mem_cache = layer['memory'](
                    x,
                    past_memories=layer_cache,
                    return_next_memories=True
                )
                x = x + mem_out

            # Standard transformer operations
            attn_out, _ = layer['attention'](
                layer['norm1'](x),
                layer['norm1'](x),
                layer['norm1'](x)
            )
            x = x + attn_out
            x = x + layer['ff'](layer['norm2'](x))

            new_cache.append(next_mem_cache)

        return x, new_cache

# Usage
model = TransformerWithMemory()
tokens = torch.randn(2, 256, 512)

output, cache = model(tokens)
print(f"Output shape: {output.shape}")  # torch.Size([2, 256, 512])

# Continue with next chunk using cached memories
tokens_2 = torch.randn(2, 256, 512)
output_2, cache = model(tokens_2, cache=cache)
```

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

### Configure token-by-token processing

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Set chunk_size to 1 to enable autoregressive inference, allowing the model to process sequences one token at a time while maintaining state.

```python
import torch
from fast_weight_product_key_memory import fwPKM

# Configure for token-by-token processing
mem = fwPKM(
    dim=32,
    heads=4,
    num_memories=4 * 4,
    dim_queries_keys=8,
    dim_values=8,
    chunk_size=1,  # Process one token at a time
    topk=2
)

# Process sequence token by token (autoregressive inference)
seq = torch.randn(1, 5, 32)
state = None
outputs = []

for i in range(5):
    token = seq[:, i:i+1]  # Single token: (batch, 1, dim)
    out, state = mem(
        token,
        return_next_memories=True,
        past_memories=state
    )
    outputs.append(out)
    print(f"Token {i}: output shape {out.shape}, tokens processed: {state.token_count}")

# Concatenate outputs
full_output = torch.cat(outputs, dim=1)
print(f"Full output shape: {full_output.shape}")  # torch.Size([1, 5, 32])

# Compare with parallel processing
parallel_out, _ = mem(seq, return_next_memories=True)
print(f"Outputs match: {torch.allclose(full_output, parallel_out, atol=1e-6)}")  # True
```

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

### Implement truncated BPTT

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Use detach_next_memories_every to periodically detach gradients during long sequence training to prevent memory overflow.

```python
import torch
from fast_weight_product_key_memory import fwPKM

mem = fwPKM(
    dim=512,
    heads=4,
    num_memories=256 * 256,
    dim_queries_keys=512,
    dim_values=512,
    topk=8,
    chunk_size=256
)

# Long sequence training with truncated BPTT
long_sequence = torch.randn(1, 2048, 512)

# Detach memories every 2 chunks to limit gradient computation
output, memories = mem(
    long_sequence,
    return_next_memories=True,
    detach_next_memories_every=2  # Detach after every 2 chunks
)

print(f"Output shape: {output.shape}")  # torch.Size([1, 2048, 512])
print(f"Memories detached periodically for stable training")
```

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

### Chain memory states across sequences

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Use the return_next_memories parameter to persist and accumulate memory state across multiple forward passes.

```python
import torch
from fast_weight_product_key_memory import fwPKM, Memories

mem = fwPKM(
	dim=512,
	num_memories=256 * 256,
	dim_queries_keys=512,
	dim_values=512,
	topk=8,
	learning_rate=1.,
	chunk_size=256
)

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

# First chunk - initialize memories
retrieved_1, memories = mem(tokens, return_next_memories=True)
print(f"Retrieved shape: {retrieved_1.shape}")
print(f"Memories type: {type(memories)}")  # <class 'Memories'>

# Second chunk - use previous memories
tokens_2 = torch.randn(2, 256, 512)
retrieved_2, memories = mem(tokens_2, return_next_memories=True, past_memories=memories)

# Third chunk - continue chaining
tokens_3 = torch.randn(2, 256, 512)
retrieved_3, memories = mem(tokens_3, return_next_memories=True, past_memories=memories)

# Fourth chunk
tokens_4 = torch.randn(2, 256, 512)
retrieved_4, memories = mem(tokens_4, return_next_memories=True, past_memories=memories)

print(f"Final memory values shape: {memories.memory_values.shape}")
print(f"Final keys shape: {memories.keys.shape}")
```

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

### Inspect memory state structure

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Access and print the shapes and counts of internal memory components for debugging or monitoring.

```python
print(f"memory_values: {memories.memory_values.shape}")  # Fast weight updates to memories
print(f"keys: {memories.keys.shape}")                    # Fast weight updates to keys
print(f"last_token: {memories.last_token.shape if memories.last_token is not None else None}")
print(f"cached_tokens: {memories.cached_tokens.shape if memories.cached_tokens is not None else None}")
print(f"token_count: {memories.token_count}")            # Total tokens processed
print(f"num_cached: {memories.num_cached}")              # Tokens cached for next chunk

# Access initial empty memories
init_memories = mem.init_memories
print(f"Initial memory values: {init_memories.memory_values.shape}")
```

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

### Access Memories named tuple

Source: https://context7.com/lucidrains/fast-weight-product-key-memories/llms.txt

Retrieve the accumulated memory state container from the forward pass.

```python
import torch
from fast_weight_product_key_memory import fwPKM, Memories

mem = fwPKM(
	dim=512,
	heads=4,
	num_memories=64 * 64,
	dim_queries_keys=256,
	dim_values=256,
	topk=8,
	chunk_size=128
)

tokens = torch.randn(1, 128, 512)

# Get memories from forward pass
output, memories = mem(tokens, return_next_memories=True)
```

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