### Install Tiny Recursive Model

Source: https://github.com/lucidrains/tiny-recursive-model/blob/main/README.md

Install the package via pip.

```bash
$ pip install tiny-recursive-model
```

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

### Initialize and Configure Trainer

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Set up the Trainer class with a TinyRecursiveModel instance, dataset, and various training hyperparameters. The Trainer handles the recurrent training loop, optimizer, and EMA updates.

```python
# Initialize model
trm = TinyRecursiveModel(
    dim=128,
    num_tokens=256,
    network=MLPMixer1D(dim=128, depth=2, seq_len=256),
    num_refinement_blocks=3,
    num_latent_refinements=6,
)

# Create trainer with full configuration
trainer = Trainer(
    model=trm,
    dataset=dataset,
    learning_rate=1e-4,           # Base learning rate
    muon_learning_rate=1e-3,      # Muon optimizer LR (for transformers)
    weight_decay=1.0,             # L2 regularization
    batch_size=32,                # Batch size
    epochs=10,                    # Training epochs
    halt_prob_thres=0.5,          # Halt threshold during training
    max_recurrent_steps=12,       # Max recurrent steps per sample
    warmup_steps=2000,            # LR warmup steps
    ema_decay_rate=0.999,         # EMA decay for model averaging
    switch_ema_every=10000,       # Switch EMA update frequency
    cpu=False                     # Set True for CPU-only training
)
```

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

### Initialize, Train, and Run Inference with TRM

Source: https://github.com/lucidrains/tiny-recursive-model/blob/main/README.md

Demonstrates model instantiation, training with a custom dataset, and performing inference with refinement steps.

```python
import torch
from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D, Trainer

trm = TinyRecursiveModel(
    dim = 16,
    num_tokens = 256,
    network = MLPMixer1D(
        dim = 16,
        depth = 2,
        seq_len = 256
    ),
)

# mock dataset

from torch.utils.data import Dataset
class MockDataset(Dataset):
    def __len__(self):
        return 16

    def __getitem__(self, idx):
        inp = torch.randint(0, 256, (256,))
        out = torch.randint(0, 256, (256,))
        return inp, out

mock_dataset = MockDataset()

# trainer

trainer = Trainer(
    trm,
    mock_dataset,
    epochs = 1,
    batch_size = 16,
    cpu = True
)

trainer()

# inference

pred_answer, exit_indices = trm.predict(
    torch.randint(0, 256, (1, 256)),
    max_deep_refinement_steps = 12,
    halt_prob_thres = 0.1
)

# save to collection of specialized networks for tool call

torch.save(trm.state_dict(), 'saved-trm.pt')
```

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

### Initialize and Train TinyRecursiveModel

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Demonstrates how to instantiate the model with an MLPMixer1D backbone and execute a recurrent training loop with adaptive halting.

```python
import torch
from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D

# Initialize the model with MLPMixer1D as the shared network
trm = TinyRecursiveModel(
    dim=512,                      # Feature dimension
    num_tokens=256,               # Vocabulary size
    network=MLPMixer1D(
        dim=512,
        depth=4,
        seq_len=1024
    ),
    num_refinement_blocks=3,      # T in paper - deep refinement steps
    num_latent_refinements=6,     # n in paper - latent refinement per output
    halt_loss_weight=1.0,         # Weight for halt prediction loss
    num_register_tokens=0         # Optional register tokens for attention
)

# Training loop example
from torch.optim import AdamW

optim = AdamW(trm.parameters(), lr=1e-4)

# Sample input and target sequences
seq = torch.randint(0, 256, (batch_size := 4, 1024))
labels = torch.randint(0, 256, (batch_size, 1024))

# Get initial output and latent embeddings
outputs, latents = trm.get_initial()

# Recurrent training steps
for step in range(12):
    loss, (main_loss, halt_loss), outputs, latents, pred, halt_prob = trm(
        seq, outputs, latents, labels=labels
    )

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

    print(f"Step {step}: main_loss={main_loss.mean():.4f}, halt_loss={halt_loss.mean():.4f}")

    # Early exit if all samples halted
    halt_mask = halt_prob >= 0.5
    if halt_mask.all():
        break
```

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

### Run Training and Save Model

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Execute the training process by calling the trainer instance and then save the trained model's state dictionary.

```python
# Run training
trainer()

# Save trained model
torch.save(trm.state_dict(), 'trained-trm.pt')
```

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

### Configure MLPMixer1D Network

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Demonstrates the instantiation and usage of the MLPMixer1D network as a standalone component.

```python
import torch
from tiny_recursive_model import MLPMixer1D

# Create a standalone MLPMixer1D network
mixer = MLPMixer1D(
    dim=512,                    # Feature dimension
    depth=4,                    # Number of mixer blocks
    seq_len=1024,               # Sequence length (fixed)
    expansion_factor=4,         # Expansion for spatial mixing
    expansion_factor_token=0.5, # Expansion for channel mixing
    dropout=0.1                 # Dropout rate
)

# Process token embeddings
tokens = torch.randn(batch_size := 2, 1024, 512)  # [batch, seq_len, dim]
output = mixer(tokens)

print(f"Input shape: {tokens.shape}")   # [2, 1024, 512]
print(f"Output shape: {output.shape}")  # [2, 1024, 512]
```

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

### Initialize TinyRecursiveModel

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Instantiate the TinyRecursiveModel with specified dimensions and a shared network module.

```python
from tiny_recursive_model import TinyRecursiveModel

trm = TinyRecursiveModel(
    dim=512,
    num_tokens=256,
    network=mixer,  # Pass the mixer as the shared network
)
```

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

### GPT-style Training Dataset

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Create a PyTorch Dataset for GPT-style next token prediction tasks. Each item provides an input sequence and a shifted target sequence.

```python
# GPT-style training dataset (next token prediction)
class GPTDataset(Dataset):
    def __init__(self, num_samples=1000, seq_len=256):
        self.num_samples = num_samples
        self.seq_len = seq_len

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        seq = torch.randint(0, 256, (self.seq_len + 1,))
        return seq[:-1], seq[1:]  # Input and shifted target
```

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

### Perform Inference with TinyRecursiveModel.predict()

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Shows how to run inference using the predict method, which supports adaptive halting based on a confidence threshold.

```python
import torch
from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D

# Create and load a trained model
trm = TinyRecursiveModel(
    dim=512,
    num_tokens=256,
    network=MLPMixer1D(dim=512, depth=4, seq_len=1024),
)

# Load pretrained weights
trm.load_state_dict(torch.load('saved-trm.pt'))
trm.eval()

# Prepare input sequence
input_seq = torch.randint(0, 256, (1, 1024))

# Run inference with adaptive halting
predictions, exit_indices = trm.predict(
    input_seq,
    halt_prob_thres=0.5,          # Halt when confidence >= 50%
    max_deep_refinement_steps=12  # Maximum iterations before forced halt
)

print(f"Predictions shape: {predictions.shape}")  # [batch, seq_len]
print(f"Exit indices: {exit_indices}")            # Step at which each sample halted

# Batch inference example
batch_input = torch.randint(0, 256, (8, 1024))
batch_preds, batch_exits = trm.predict(
    batch_input,
    halt_prob_thres=0.3,           # Lower threshold = earlier halting
    max_deep_refinement_steps=20   # Allow more iterations for complex inputs
)

# Analyze computation distribution
for i, exit_step in enumerate(batch_exits.tolist()):
    print(f"Sample {i}: halted at step {exit_step}")
```

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

### Model Persistence: Saving and Loading TRM

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Shows how to save and load trained Tiny Recursive Model (TRM) models using PyTorch's state_dict for weights or the entire model object. Includes loading weights into a new instance and loading a complete model for inference.

```python
import torch
from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D

# Create and train model
trm = TinyRecursiveModel(
    dim=256,
    num_tokens=256,
    network=MLPMixer1D(dim=256, depth=4, seq_len=1024),
    num_refinement_blocks=3,
    num_latent_refinements=6,
)

# ... training code ...

# Save model weights
torch.save(trm.state_dict(), 'trm-checkpoint.pt')

# Save complete model (for inference without knowing architecture)
torch.save(trm, 'trm-complete.pt')

# Load weights into new model instance
new_trm = TinyRecursiveModel(
    dim=256,
    num_tokens=256,
    network=MLPMixer1D(dim=256, depth=4, seq_len=1024),
    num_refinement_blocks=3,
    num_latent_refinements=6,
)
new_trm.load_state_dict(torch.load('trm-checkpoint.pt'))
new_trm.eval()

# Load complete model
loaded_trm = torch.load('trm-complete.pt')
loaded_trm.eval()

# Verify inference works
test_input = torch.randint(0, 256, (1, 1024))
preds, exits = new_trm.predict(test_input, max_deep_refinement_steps=12)
print(f"Loaded model predictions: {preds.shape}")
```

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

### Initialize TinyRecursiveModel with x-transformers Decoder

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Instantiate TinyRecursiveModel using an x-transformers Decoder as the shared network for causal self-attention, suitable for autoregressive tasks.

```python
# Using Decoder (causal self-attention for autoregressive tasks)
decoder_trm = TinyRecursiveModel(
    dim=256,
    num_tokens=256,
    network=Decoder(
        dim=256,
        depth=4,
        heads=8,
        attn_dropout=0.1,
        ff_dropout=0.1
    ),
)
```

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

### Initialize TinyRecursiveModel with x-transformers Encoder

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Instantiate TinyRecursiveModel using an x-transformers Encoder as the shared network for bidirectional self-attention.

```python
import torch
from torch.utils.data import Dataset
from x_transformers import Encoder, Decoder
from tiny_recursive_model import TinyRecursiveModel, Trainer

# Using Encoder (bidirectional self-attention)
encoder_trm = TinyRecursiveModel(
    dim=256,
    num_tokens=256,
    network=Encoder(
        dim=256,
        depth=4,
        heads=8,
        attn_dropout=0.1,
        ff_dropout=0.1
    ),
    num_register_tokens=4,  # Register tokens enhance attention
)
```

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

### Train with Decoder and Trainer

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Train a TinyRecursiveModel configured with an x-transformers Decoder using the Trainer class. The MuonAdamAtan2 optimizer is automatically used.

```python
# Train with Decoder (uses MuonAdamAtan2 optimizer automatically)
trainer = Trainer(
    decoder_trm,
    GPTDataset(),
    epochs=5,
    batch_size=16,
    learning_rate=1e-4,
    muon_learning_rate=1e-3,
    cpu=True
)

trainer()
```

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

### Custom Dataset for Training

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Define a custom PyTorch Dataset for training the TinyRecursiveModel. Ensure input and output sequences are within the specified sequence length.

```python
import torch
from torch.utils.data import Dataset
from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D, Trainer

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, data_pairs, seq_len=256):
        self.data_pairs = data_pairs
        self.seq_len = seq_len

    def __len__(self):
        return len(self.data_pairs)

    def __getitem__(self, idx):
        input_seq, output_seq = self.data_pairs[idx]
        return input_seq[:self.seq_len], output_seq[:self.seq_len]

# Create mock training data
mock_data = [
    (torch.randint(0, 256, (256,)), torch.randint(0, 256, (256,)))
    for _ in range(1000)
]
dataset = TextDataset(mock_data)
```

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

### TRM Forward Pass and Inference

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Demonstrates a forward pass of the TRM with labels for training and a subsequent inference pass without labels. Shows how to access loss components and output shapes.

```python
total_loss, (main_loss, halt_loss), outputs, latents, pred, halt_prob = trm(
    input_seq, outputs, latents, labels=labels
)

print(f"Total loss: {total_loss.item():.4f}")
print(f"Main loss (per sample): {main_loss}")
print(f"Halt loss (per sample): {halt_loss}")
print(f"Predictions shape: {pred.shape}")  # [8, 512, 256]
print(f"Halt probabilities: {halt_prob}")  # [8]

# Forward pass without labels (inference mode in training)
outputs, latents, pred, halt_prob = trm(input_seq, outputs, latents)
print(f"Inference pred shape: {pred.shape}")
```

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

### Forward Pass with Labels for Training

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Demonstrates the forward pass of TinyRecursiveModel when labels are provided, which computes predictions, halt probabilities, and losses. Returns detached outputs and latents for the next recurrent step.

```python
import torch
from tiny_recursive_model import TinyRecursiveModel, MLPMixer1D

trm = TinyRecursiveModel(
    dim=128,
    num_tokens=256,
    network=MLPMixer1D(dim=128, depth=2, seq_len=512),
    num_refinement_blocks=3,
    halt_loss_weight=1.0,
)

# Prepare batch data
batch_size = 8
seq_len = 512
input_seq = torch.randint(0, 256, (batch_size, seq_len))
labels = torch.randint(0, 256, (batch_size, seq_len))

# Initialize outputs and latents
outputs, latents = trm.get_initial()

# Forward pass with labels returns:
# - total_loss: Combined main loss + halt loss
# - losses: Tuple of (main_loss, halt_loss) per sample
# - outputs: Detached output embeddings for next step
```

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

### TRM Citation

Source: https://github.com/lucidrains/tiny-recursive-model/blob/main/README.md

BibTeX citation for the Tiny Recursive Model paper.

```bibtex
@misc{jolicoeurmartineau2025morerecursivereasoningtiny,
    title   = {Less is More: Recursive Reasoning with Tiny Networks}, 
    author  = {Alexia Jolicoeur-Martineau},
    year    = {2025},
    eprint  = {2510.04871},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2510.04871}, 
}
```

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

### Decoder Inference

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Perform inference with a TinyRecursiveModel configured with an x-transformers Decoder. Prints the shape of the predicted sequence.

```python
# Inference
pred, exits = decoder_trm.predict(torch.randint(0, 256, (1, 256)))
print(f"Predicted sequence shape: {pred.shape}")
```

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

### Inference with Trained Model

Source: https://context7.com/lucidrains/tiny-recursive-model/llms.txt

Perform inference using the trained TinyRecursiveModel. The EMA model is automatically used for predictions after training.

```python
# Use EMA model for inference (automatically copied after training)
predictions, exits = trm.predict(torch.randint(0, 256, (1, 256)))
```

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