### Install RIM-pytorch

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

Install the library using pip.

```bash
pip install RIM-pytorch
```

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

### DepthlessTransformer with Sequential Routing Schedule

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

Utilize a custom routing schedule to control block activation for sequential processing. This example demonstrates activating one block at a time for attention and then feedforward operations.

```python
import torch
from RIM_pytorch import DepthlessTransformer

model = DepthlessTransformer(
    dim=256,
    num_tokens=1000,
    num_blocks=3,
    num_message_exchanges=6  # Will be overridden by schedule length
)

# Sequential routing: one block at a time, attention then feedforward
sequential_schedule = (
    (('attn', (0,)),),    # Round 1: attention on block 0
    (('ff', (0,)),),      # Round 2: feedforward on block 0
    (('attn', (1,)),),    # Round 3: attention on block 1
    (('ff', (1,)),),      # Round 4: feedforward on block 1
    (('attn', (2,)),),    # Round 5: attention on block 2
    (('ff', (2,)),),      # Round 6: feedforward on block 2
)

tokens = torch.randint(0, 1000, (2, 32))
logits = model(tokens, routing_schedule=sequential_schedule)
print(f"Sequential routing output: {logits.shape}")  # (2, 32, 1000)
```

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

### Initialize Feedforward Module in PyTorch

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

Set up a SwiGLU feedforward network with pre-normalization.

```python
import torch
from RIM_pytorch.depth_less_transformer import Feedforward

ff = Feedforward(
    dim=256,
    expansion_factor=4.   # Inner dimension = dim * expansion * 2/3
)

tokens = torch.randn(2, 32, 256)
output = ff(tokens)
print(f"Feedforward output: {output.shape}")  # (2, 32, 256)
```

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

### Configure Attention Module in PyTorch

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

Initialize a standalone attention module supporting causal masking, key RMSNorm, and cross-attention.

```python
import torch
from RIM_pytorch.depth_less_transformer import Attention

attn = Attention(
    dim=256,
    dim_head=64,
    heads=8,
    causal=True,           # Causal masking for autoregressive
    key_rmsnorm=True,      # Apply RMSNorm to keys
    dropout=0.1            # Attention dropout
)

# Self-attention
tokens = torch.randn(2, 32, 256)
output = attn(tokens)
print(f"Self-attention output: {output.shape}")  # (2, 32, 256)

# Cross-attention with context
context = torch.randn(2, 64, 256)
output = attn(tokens, context=context)
print(f"Cross-attention output: {output.shape}")  # (2, 32, 256)
```

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

### Implement Multi-Message Modules in PyTorch

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

Create custom modules that return multiple messages by stacking tensors along a leading dimension.

```python
import torch
from torch import nn
from RIM_pytorch import EnsemblesWithMessagePassing
from RIM_pytorch.depth_less_transformer import Feedforward

class TwoMessageFF(nn.Module):
    """Custom module that produces two messages per forward pass."""
    def __init__(self, dim):
        super().__init__()
        self.ff = Feedforward(dim)

    def forward(self, x):
        out = self.ff(x)
        # Return two messages by stacking along dim 0
        return torch.stack((out, out * 0.5), dim=0)

model = EnsemblesWithMessagePassing(
    modules=TwoMessageFF(dim=256),
    ensemble_size=3,
    dim=256,
    num_message_exchanges=1
)

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

messages = model(
    tokens,
    repeat_input_for_ensemble=True,
    return_all_messages=True
)

# Each block produces 2 messages
print(f"Total messages: {len(messages)}")  # 3: init + 2 from TwoMessageFF
print(f"Message shapes: {[m.shape for m in messages]}")
```

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

### DepthlessTransformer for Language Modeling

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

Initialize and use the DepthlessTransformer for autoregressive language modeling with token indices. It supports returning intermediate messages for recurrent processing.

```python
import torch
from RIM_pytorch import DepthlessTransformer

# Create a DepthlessTransformer for language modeling
model = DepthlessTransformer(
    dim=256,                      # Model dimension
    num_tokens=1000,              # Vocabulary size (enables embedding layer)
    num_blocks=3,                 # Number of parallel ensemble blocks
    num_message_exchanges=3,      # Rounds of message passing between blocks
    dim_head=64,                  # Attention head dimension
    heads=8,                      # Number of attention heads
    causal=True,                  # Causal attention for autoregressive modeling
    ff_expansion_factor=4.,       # Feedforward expansion factor
    use_pope=False                # Use PoPE positional encoding
)

# Forward pass with token indices
tokens = torch.randint(0, 1000, (2, 32))  # (batch, seq_len)
logits = model(tokens)
print(f"Output logits shape: {logits.shape}")  # (2, 32, 1000)

# Get messages for recurrent processing
logits, messages = model(tokens, return_messages=True)
print(f"Number of messages: {len(messages)}")  # 7: init + 6 (3 rounds * 2 modules)

# Recurrent message passing - feed messages back for continued processing
logits_next, messages_next = model(tokens, return_messages=True, messages=messages)
print(f"Messages after recurrent pass: {len(messages_next)}")  # 14
```

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

### EnsemblesWithMessagePassing with Single Module

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

Create an ensemble of a single module type (e.g., Attention) that replicates across ensemble members. Input can be repeated for each block, and message history can be returned.

```python
import torch
from RIM_pytorch import EnsemblesWithMessagePassing
from RIM_pytorch.depth_less_transformer import Attention

# Create ensemble from a single attention module
attn = Attention(dim=256)

model = EnsemblesWithMessagePassing(
    dim=256,
    modules=attn,                 # Single module replicated across ensemble
    ensemble_size=3,              # Number of parallel blocks
    num_message_exchanges=2       # Rounds of message exchange
)

# Input tokens (batch, seq_len, dim)
tokens = torch.randn(2, 32, 256)

# Forward with ensemble replication
output = model(
    tokens,
    repeat_input_for_ensemble=True,   # Replicate input for each block
    return_all_messages=True          # Return message history
)

print(f"Number of messages: {len(output)}")  # 3
```

```python
# Using routing schedule with single module
routing_schedule = (
    ((0, 1),),   # Round 1: blocks 0 and 1
    (2,),        # Round 2: block 2 only
)

messages = model(
    tokens,
    repeat_input_for_ensemble=True,
    return_all_messages=True,
    routing_schedule=routing_schedule
)
print(f"Message shapes: {[m.shape for m in messages]}")
```

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

### EnsemblesWithMessagePassing with Multiple Modules

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

Construct ensembles with multiple named modules (e.g., Attention and Feedforward) for complex architectures. The input shape for multiple modules should include the ensemble dimension first.

```python
import torch
from RIM_pytorch import EnsemblesWithMessagePassing
from RIM_pytorch.depth_less_transformer import Attention, Feedforward

dim = 256
ensemble_size = 4

# Create ensemble with multiple module types
model = EnsemblesWithMessagePassing(
    modules=dict(
        attn=Attention(dim, heads=8, dim_head=32),
        ff=Feedforward(dim, expansion_factor=4.)
    ),
    ensemble_size=ensemble_size,
    dim=dim,
    voting_attn_kwargs=dict(dim_head=32, heads=8),
    num_message_exchanges=3
)

tokens = torch.randn(ensemble_size, 2, 32, dim)  # (blocks, batch, seq, dim)

# Forward pass processes both attn and ff modules in each round
output = model(tokens)
print(f"Output shape: {output.shape}")  # (4, 2, 32, 256)
```

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

### DepthlessTransformer with Rotational Routing Schedule

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

Implement a rotational routing schedule to alternate block operations, allowing blocks to perform attention in one round and feedforward in the next. This enables flexible and dynamic module activation patterns.

```python
import torch
from RIM_pytorch import DepthlessTransformer

model = DepthlessTransformer(
    dim=256,
    num_tokens=1000,
    num_blocks=3,
    num_message_exchanges=6  # Will be overridden by schedule length
)

# Rotational routing: alternate which blocks run attention vs feedforward
rotational_schedule = (
    (('attn', (0, 1)), ('ff', (2,))),   # Round 1: attn on 0,1; ff on 2
    (('attn', (2,)), ('ff', (0, 1))),   # Round 2: attn on 2; ff on 0,1
)

logits = model(tokens, routing_schedule=rotational_schedule)
print(f"Rotational routing output: {logits.shape}")  # (2, 32, 1000)
```

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

### Implicit Routing Schedule

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

Use implicit routing schedules with module names to control message passing flow. 'attn' and 'ff' refer to attention and feedforward blocks, respectively. None indicates all indices.

```python
implicit_schedule = (
    ('attn', 'ff'),                     # All blocks for attn and ff
    (('attn', None), ('ff', None)),     # Explicit None = all indices
)

logits, messages = model(tokens, return_messages=True, routing_schedule=implicit_schedule)
print(f"Messages from implicit schedule: {len(messages)}")  # 5
```

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

### Nested Ensembles

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

Create hierarchical architectures by nesting `EnsemblesWithMessagePassing` modules. The input tensor dimensions should correspond to the nested ensemble structure.

```python
import torch
from RIM_pytorch import EnsemblesWithMessagePassing
from RIM_pytorch.depth_less_transformer import Feedforward

dim = 64

# Inner ensemble: 4 parallel feedforward blocks
ff = Feedforward(dim)
inner_model = EnsemblesWithMessagePassing(
    dim=dim,
    modules=ff,
    ensemble_size=4,
    num_message_exchanges=2
)

# Outer ensemble: 4 parallel copies of the inner ensemble
outer_model = EnsemblesWithMessagePassing(
    dim=dim,
    modules=inner_model,
    ensemble_size=4,
    num_message_exchanges=2
)

# Input with nested block dimensions
x = torch.randn(4, 4, 3, dim)  # (outer_blocks, inner_blocks, batch, dim)
tokens = outer_model(x)
print(f"Nested ensemble output: {tokens.shape}")  # (4, 4, 3, 64)
```

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

### Ensemble Component

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

The low-level Ensemble component parallelizes any PyTorch module using `vmap`. Input tensors must have a leading dimension for the ensemble blocks. Specific blocks can be selected using the `indices` argument.

```python
import torch
from torch import nn
from RIM_pytorch import Ensemble

# Create ensemble from any module
linear = nn.Linear(64, 64)
ensemble = Ensemble(
    net=linear,
    ensemble_size=4
)

# Input must have leading block dimension
tokens = torch.randn(4, 8, 64)  # (blocks, batch, dim)

# Forward through all blocks
output = ensemble(tokens)
print(f"All blocks output: {output.shape}")  # (4, 8, 64)
```

```python
# Forward through specific blocks only
output_subset = ensemble(tokens, indices=(0, 2))
print(f"Blocks 0,2 output: {output_subset.shape}")  # (2, 8, 64)
```

```python
# Single block
output_single = ensemble(tokens, indices=0)
print(f"Single block output: {output_single.shape}")  # (1, 8, 64)
```

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

### DepthlessTransformer with Continuous Inputs

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

Initialize and use the DepthlessTransformer for processing continuous embeddings when `num_tokens` is not specified. This bypasses the embedding layer and is useful for pre-embedded features. Bidirectional attention is enabled by default when `causal` is False.

```python
import torch
from RIM_pytorch import DepthlessTransformer

# Create model for continuous inputs (no embedding layer)
model = DepthlessTransformer(
    dim=32,
    num_blocks=6,
    num_message_exchanges=6,
    causal=False,                 # Bidirectional attention
    use_pope=True                 # Enable PoPE positional encoding
)

# Forward pass with continuous embeddings
x = torch.randn(1, 7, 32)  # (batch, seq_len, dim)
pooled_messages = model(x)
print(f"Output shape: {pooled_messages.shape}")  # (1, 7, 32)
```

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