### Install HRM-pytorch Package

Source: https://github.com/lucidrains/hrm/blob/main/README.md

This snippet shows the command to install the HRM-pytorch library using pip. It is a prerequisite for using the HRM model in Python projects.

```bash
pip install HRM-pytorch
```

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

### HRM Model Usage and Training Example in Python

Source: https://github.com/lucidrains/hrm/blob/main/README.md

Demonstrates how to import and instantiate the HRM model in Python. It includes examples of performing forward passes for training with loss calculation and backpropagation, as well as performing inference after training.

```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(
            dim = 32,
            depth = 2,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        ),
        dict(
            dim = 32,
            depth = 4,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        )
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10
)

seq = torch.randint(0, 256, (3, 1024))
labels = torch.randint(0, 256, (3, 1024))

loss, hiddens, _ = hrm(seq, labels = labels)
loss.backward()

loss, hiddens, _ = hrm(seq, hiddens = hiddens, labels = labels)
loss.backward()

# after much training

pred = hrm(seq, reasoning_steps = 5)
```

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

### HRM Model Initialization

Source: https://context7.com/lucidrains/hrm/llms.txt

Initializes a Hierarchical Reasoning Model with a specified configuration of neural networks at different levels of the hierarchy. This setup allows for parallel processing at varying temporal frequencies to mimic human reasoning.

```APIDOC
## HRM Model Initialization

### Description
Initialize a Hierarchical Reasoning Model with multiple networks arranged in a hierarchy from low to high level processing.

### Method
`__init__` (Python class constructor)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **networks** (list of dict) - Required - Configuration for each hierarchical network.
  - **dim** (int) - Dimension of the model.
  - **depth** (int) - Number of layers in the network.
  - **attn_dim_head** (int) - Dimension of attention heads.
  - **heads** (int) - Number of attention heads.
  - **use_rmsnorm** (bool) - Whether to use RMSNorm.
  - **rotary_pos_emb** (bool) - Whether to use rotary positional embeddings.
  - **pre_norm** (bool) - Whether to use pre-layer normalization.
- **num_tokens** (int) - Required - Vocabulary size.
- **dim** (int) - Required - The main dimension for the model.
- **reasoning_steps** (int) - Required - The number of reasoning steps to perform.
- **causal** (bool) - Optional - Whether to use causal attention (for language modeling).
- **ignore_index** (int) - Optional - Index to ignore for loss calculation.

### Request Example
```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(
            dim = 32,
            depth = 2,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        ),
        dict(
            dim = 32,
            depth = 4,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        )
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10,
    causal = False,
    ignore_index = -1
)
```

### Response
None (This is an initialization function.)
```

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

### Batch Processing and Hidden State Management with HRM

Source: https://context7.com/lucidrains/hrm/llms.txt

Illustrates efficient batch processing with HRM, including managing hidden states across multiple forward passes. This example details the use of `hiddens` parameter, `detach_hiddens`, and `one_step_grad` for stable training.

```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(dim=64, depth=3, attn_dim_head=16, heads=2),
        dict(dim=64, depth=6, attn_dim_head=16, heads=2)
    ],
    num_tokens = 1000,
    dim = 64,
    reasoning_steps = 5
)

# Batch of sequences
batch_size = 16
seq_length = 256
sequences = torch.randint(0, 1000, (batch_size, seq_length))
labels = torch.randint(0, 1000, (batch_size, seq_length))

# Initialize hidden states (optional - defaults to zeros)
hiddens = None

# Process in epochs
for epoch in range(3):
    # Forward pass
    loss, hiddens, predictions = hrm(
        sequences,
        hiddens=hiddens,
        labels=labels,
        detach_hiddens=True,  # Prevent gradient explosion
        one_step_grad=True     # Only backprop through last step
    )

    loss.backward()

    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

    # Hidden states are automatically detached and reused
    # Shape of each hidden: (batch_size, seq_length, dim)
    for i, h in enumerate(hiddens):
        print(f"  Level {i} hidden shape: {h.shape}")

# Final inference without labels
final_predictions, final_hiddens = hrm(sequences, hiddens=hiddens)
accuracy = (final_predictions.argmax(dim=-1) == labels).float().mean()
print(f"Accuracy: {accuracy.item():.4f}")
```

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

### Multi-Network Hierarchy Configuration with HRM

Source: https://context7.com/lucidrains/hrm/llms.txt

Demonstrates configuring HRM with a multi-level hierarchy of networks, allowing for varying depths and computation frequencies. This section explains setting up the network configurations and relative periods for hierarchical processing.

```python
import torch
from HRM import HRM
from x_transformers import Encoder

# Three-level hierarchy: shallow -> medium -> deep networks
hrm = HRM(
    networks = [
        dict(dim=32, depth=2, attn_dim_head=8, heads=1),   # Low: runs every step
        dict(dim=32, depth=4, attn_dim_head=8, heads=1),   # Mid: runs every 2 steps
        Encoder(dim=32, depth=8, attn_dim_head=8, heads=1) # High: runs every 4 steps
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 8,
    relative_period = (1, 2, 2),  # Explicit period specification
)

seq = torch.randint(0, 256, (2, 256))
labels = torch.randint(0, 256, (2, 256))

# Train with hierarchical processing
loss, hiddens, predictions = hrm(seq, labels=labels)
print(f"Number of hierarchy levels: {len(hiddens)}")  # 3 levels
print(f"Loss: {loss.item()}")

loss.backward()

# Verify evaluate_networks_at periods
print(f"Network evaluation periods: {hrm.evaluate_networks_at}")  # [1, 2, 4]
print(f"Total steps per reasoning step: {hrm.lowest_steps_per_reasoning_step}")  # 4
```

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

### Adaptive Computation Time (ACT) with HRM_ACT

Source: https://context7.com/lucidrains/hrm/llms.txt

Shows how to initialize and train the HRM_ACT variant, which learns to exit early based on prediction confidence. It covers calculating both the main prediction loss and the ACT loss, and performing inference with adaptive computation.

```python
import torch
from HRM import HRM_ACT

# Initialize HRM with adaptive computation time
hrm_act = HRM_ACT(
    networks = [
        dict(dim=32, depth=2, attn_dim_head=8, heads=1, use_rmsnorm=True),
        dict(dim=32, depth=4, attn_dim_head=8, heads=1, use_rmsnorm=True)
    ],
    num_tokens = 256,
    dim = 32,
    max_reasoning_steps = 10,
    min_reasoning_steps_epsilon_prob = 0.5,
    act_loss_weight = 1.0,
    discount_factor = 1.0,
    causal = False
)

# Training with ACT - learns Q(halt|continue) values
seq = torch.randint(0, 256, (4, 512))
labels = torch.randint(0, 256, (4, 512))

# Returns total loss combining main prediction loss and ACT loss
total_loss, hiddens, (logits, loss_breakdown) = hrm_act(
    seq,
    labels=labels,
    compute_loss_across_reasoning_steps=False
)

print(f"Total loss: {total_loss.item()}")
print(f"Main loss: {loss_breakdown.main.item()}")
print(f"ACT loss: {loss_breakdown.act.item()}")

total_loss.backward()

# Inference with adaptive computation - exits early when confident
seq_test = torch.randint(0, 256, (16, 512))

predictions, exit_hiddens, reasoning_steps_used = hrm_act(
    seq_test,
    max_reasoning_steps=10,
    adaptive_compute=True
)

print(f"Predictions shape: {predictions.shape}")  # (16, 512, 256)
print(f"Reasoning steps per sample: {reasoning_steps_used}")  # (16,) - varies per sample
print(f"Average reasoning steps: {reasoning_steps_used[reasoning_steps_used > 0].float().mean().item():.2f}")
```

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

### Autoregressive Text Generation with HRM

Source: https://context7.com/lucidrains/hrm/llms.txt

Demonstrates how to use the HRM model for autoregressive text generation. It details how to compute the loss, backpropagate, and then generate new text autoregressively using the model's hidden states.

```python
loss, hiddens, predictions = hrm(seq, return_autoreg_loss=True)
loss.backward()

print(f"Autoregressive loss: {loss.item()}")

# Generate text autoregressively
with torch.no_grad():
    context = torch.randint(0, 50257, (1, 10))  # Starting context
    generated = context.clone()

    for _ in range(50):  # Generate 50 tokens
        logits, hiddens = hrm(generated, hiddens=hiddens if _ > 0 else None)
        next_token = logits[:, -1:].argmax(dim=-1)
        generated = torch.cat([generated, next_token], dim=1)

    print(f"Generated sequence length: {generated.shape[1]}")
```

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

### Initialize Hierarchical Reasoning Model (PyTorch)

Source: https://context7.com/lucidrains/hrm/llms.txt

Initializes a Hierarchical Reasoning Model (HRM) in PyTorch. It configures multiple transformer networks in a hierarchy, specifying dimensions, depth, attention heads, and other parameters. The model can be used for tasks requiring multi-step reasoning.

```python
import torch
from HRM import HRM

# Define hierarchical network configuration
hrm = HRM(
    networks = [
        dict(
            dim = 32,
            depth = 2,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        ),
        dict(
            dim = 32,
            depth = 4,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        )
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10,
    causal = False,
    ignore_index = -1
)

# Input: token sequences (batch_size, sequence_length)
seq = torch.randint(0, 256, (3, 1024))
labels = torch.randint(0, 256, (3, 1024))

# Forward pass with labels for training
loss, hiddens, predictions = hrm(seq, labels=labels)
loss.backward()

print(f"Loss: {loss.item()}")
print(f"Hidden states: {len(hiddens)} layers")
print(f"Predictions shape: {predictions.shape}")
```

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

### Inference with Variable Reasoning Steps

Source: https://context7.com/lucidrains/hrm/llms.txt

Performs inference using the HRM model, allowing for flexible control over the number of reasoning steps to adapt to varying input complexities and desired output speeds.

```APIDOC
## Inference with Variable Reasoning Steps

### Description
Generate predictions with flexible reasoning step counts for different complexity levels.

### Method
Forward pass of the HRM model for inference.

### Endpoint
`hrm(seq, reasoning_steps=None)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **seq** (torch.Tensor) - Required - Input token sequences (batch_size, sequence_length).
- **reasoning_steps** (int) - Optional - Number of reasoning steps to perform. Defaults to the model's initialized value.

### Request Example
```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(dim=32, depth=2, attn_dim_head=8, heads=1),
        dict(dim=32, depth=4, attn_dim_head=8, heads=1)
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10  # Default reasoning steps
)

# Load trained model weights
# hrm.load_state_dict(torch.load('hrm_checkpoint.pt'))

# Inference with default reasoning steps
seq = torch.randint(0, 256, (1, 1024))
predictions, hiddens = hrm(seq)

print(f"Predictions shape: {predictions.shape}")  # (1, 1024, 256)
print(f"Hidden states: {len(hiddens)}")

# Inference with custom reasoning steps for faster/slower thinking
fast_predictions, fast_hiddens = hrm(seq, reasoning_steps=3)
slow_predictions, slow_hiddens = hrm(seq, reasoning_steps=20)

# Argmax to get predicted tokens
predicted_tokens = predictions.argmax(dim=-1)
print(f"Predicted tokens: {predicted_tokens.shape}")  # (1, 1024)
```

### Response
#### Success Response (200)
- **predictions** (torch.Tensor) - The model's predictions (batch_size, sequence_length, num_tokens).
- **hiddens** (list of torch.Tensor) - The hidden states after the forward pass.

#### Response Example
```json
{
  "predictions": tensor_predictions,
  "hiddens": [tensor_hidden_state_1, tensor_hidden_state_2, ...]
}
```
```

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

### HRM Inference with Variable Reasoning Steps (PyTorch)

Source: https://context7.com/lucidrains/hrm/llms.txt

Performs inference using the HRM model, allowing for dynamic adjustment of the number of reasoning steps. This enables controlling the trade-off between computational cost and reasoning depth. It also shows how to extract predictions and hidden states.

```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(dim=32, depth=2, attn_dim_head=8, heads=1),
        dict(dim=32, depth=4, attn_dim_head=8, heads=1)
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10  # Default reasoning steps
)

# Load trained model weights
# hrm.load_state_dict(torch.load('hrm_checkpoint.pt'))

# Inference with default reasoning steps
seq = torch.randint(0, 256, (1, 1024))
predictions, hiddens = hrm(seq)

print(f"Predictions shape: {predictions.shape}")  # (1, 1024, 256)
print(f"Hidden states: {len(hiddens)}")

# Inference with custom reasoning steps for faster/slower thinking
fast_predictions, fast_hiddens = hrm(seq, reasoning_steps=3)
slow_predictions, slow_hiddens = hrm(seq, reasoning_steps=20)

# Argmax to get predicted tokens
predicted_tokens = predictions.argmax(dim=-1)
print(f"Predicted tokens: {predicted_tokens.shape}")  # (1, 1024)
```

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

### Train HRM with Hidden State Continuity (PyTorch)

Source: https://context7.com/lucidrains/hrm/llms.txt

Demonstrates training the HRM model while maintaining and optionally flowing gradients through hidden states across multiple training steps. This is crucial for recurrent reasoning patterns. It shows how to control gradient flow and whether to detach hidden states.

```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(dim=32, depth=2, attn_dim_head=8, heads=1, use_rmsnorm=True),
        dict(dim=32, depth=4, attn_dim_head=8, heads=1, use_rmsnorm=True)
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10
)

# Create sample data
seq = torch.randint(0, 256, (3, 1024))
labels = torch.randint(0, 256, (3, 1024))

# First training step
loss, hiddens, _ = hrm(seq, labels=labels)
loss.backward()

# Continue training with previous hidden states (detached by default)
loss, hiddens, _ = hrm(seq, hiddens=hiddens, labels=labels)
loss.backward()

# Control hidden state behavior
loss, hiddens, _ = hrm(
    seq,
    hiddens=hiddens,
    labels=labels,
    detach_hiddens=False,  # Gradients flow through hiddens
    one_step_grad=False     # Gradients through all reasoning steps
)
loss.backward()
```

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

### Causal Language Modeling with HRM (PyTorch)

Source: https://context7.com/lucidrains/hrm/llms.txt

Utilizes the HRM model for causal language modeling, enabling autoregressive text generation. It involves initializing the HRM with `causal=True` and preparing input sequences and labels with appropriate shifting for predicting the next token.

```python
import torch
from HRM import HRM

# Initialize with causal=True for autoregressive modeling
hrm = HRM(
    networks = [
        dict(dim=64, depth=4, attn_dim_head=16, heads=4, use_rmsnorm=True),
        dict(dim=64, depth=8, attn_dim_head=16, heads=4, use_rmsnorm=True)
    ],
    num_tokens = 50257,  # GPT-2 vocabulary size
    dim = 64,
    reasoning_steps = 5,
    causal = True
)

# Training with shifted labels (autoregressive)
seq = torch.randint(0, 50257, (8, 512))

# Option 1: Manual label shifting
labels = seq[:, 1:]
input_seq = seq[:, :-1]
loss, hiddens, predictions = hrm(input_seq, labels=labels)
```

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

### Training with Hidden State Continuity

Source: https://context7.com/lucidrains/hrm/llms.txt

Facilitates training of the HRM model across multiple steps, allowing for the preservation and utilization of hidden states between steps to capture recurrent reasoning patterns.

```APIDOC
## Training with Hidden State Continuity

### Description
Train the model across multiple steps while maintaining hidden states for recurrent reasoning patterns.

### Method
Forward pass of the HRM model during training.

### Endpoint
`hrm(seq, hiddens=None, labels=None, detach_hiddens=True, one_step_grad=True)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **seq** (torch.Tensor) - Required - Input token sequences (batch_size, sequence_length).
- **hiddens** (list of torch.Tensor) - Optional - Previous hidden states to continue reasoning from.
- **labels** (torch.Tensor) - Optional - Ground truth labels for loss calculation.
- **detach_hiddens** (bool) - Optional - Whether to detach hidden states from the computation graph (default: True).
- **one_step_grad** (bool) - Optional - Whether to compute gradients for only one step (default: True).

### Request Example
```python
import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(dim=32, depth=2, attn_dim_head=8, heads=1, use_rmsnorm=True),
        dict(dim=32, depth=4, attn_dim_head=8, heads=1, use_rmsnorm=True)
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10
)

seq = torch.randint(0, 256, (3, 1024))
labels = torch.randint(0, 256, (3, 1024))

# First training step
loss, hiddens, _ = hrm(seq, labels=labels)
loss.backward()

# Continue training with previous hidden states (detached by default)
loss, hiddens, _ = hrm(seq, hiddens=hiddens, labels=labels)
loss.backward()

# Control hidden state behavior
loss, hiddens, _ = hrm(
    seq,
    hiddens=hiddens,
    labels=labels,
    detach_hiddens=False,  # Gradients flow through hiddens
    one_step_grad=False     # Gradients through all reasoning steps
)
loss.backward()
```

### Response
#### Success Response (200)
- **loss** (torch.Tensor) - The computed loss for the training step.
- **hiddens** (list of torch.Tensor) - The hidden states after the forward pass.
- **predictions** (torch.Tensor) - The model's predictions.

#### Response Example
```json
{
  "loss": 1.234,
  "hiddens": [tensor_hidden_state_1, tensor_hidden_state_2, ...],
  "predictions": tensor_predictions
}
```
```

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

### Causal Language Modeling

Source: https://context7.com/lucidrains/hrm/llms.txt

Applies the HRM model to causal language modeling tasks by utilizing causal attention masks, enabling autoregressive text generation and prediction.

```APIDOC
## Causal Language Modeling

### Description
Use HRM for autoregressive language modeling with causal attention masks.

### Method
Initialization and forward pass of the HRM model with `causal=True`.

### Endpoint
`HRM(..., causal=True)` and `hrm(input_seq, labels=None)`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **HRM Initialization**: See 'HRM Model Initialization' section. Ensure `causal=True` is set.
- **Forward Pass**: 
  - **input_seq** (torch.Tensor) - Required - Input token sequences (batch_size, sequence_length-1) for autoregressive prediction.
  - **labels** (torch.Tensor) - Optional - Ground truth shifted labels (batch_size, sequence_length-1) for loss calculation.

### Request Example
```python
import torch
from HRM import HRM

# Initialize with causal=True for autoregressive modeling
hrm = HRM(
    networks = [
        dict(dim=64, depth=4, attn_dim_head=16, heads=4, use_rmsnorm=True),
        dict(dim=64, depth=8, attn_dim_head=16, heads=4, use_rmsnorm=True)
    ],
    num_tokens = 50257,  # GPT-2 vocabulary size
    dim = 64,
    reasoning_steps = 5,
    causal = True
)

# Training with shifted labels (autoregressive)
seq = torch.randint(0, 50257, (8, 512))

# Option 1: Manual label shifting
labels = seq[:, 1:]
input_seq = seq[:, :-1]
loss, hiddens, predictions = hrm(input_seq, labels=labels)
```

### Response
#### Success Response (200)
- **loss** (torch.Tensor) - The computed loss for the training step (if labels are provided).
- **hiddens** (list of torch.Tensor) - The hidden states after the forward pass.
- **predictions** (torch.Tensor) - The model's predictions (logits for the next token).

#### Response Example
```json
{
  "loss": 3.57,
  "hiddens": [tensor_hidden_state_1, ...],
  "predictions": tensor_predictions
}
```
```

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