### Install self-reasoning-tokens-pytorch

Source: https://github.com/lucidrains/self-reasoning-tokens-pytorch/blob/main/README.md

Install the library using pip. No specific setup is required beyond this.

```bash
pip install self-reasoning-tokens-pytorch
```

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

### Complete Transformer Training Example

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Provides a full example of setting up, training, and performing inference with the Transformer model, including optimizer setup and a basic training loop. Ensure `stop_grad_next_tokens_to_reason` is set appropriately for gradient control.

```python
import torch
import torch.optim as optim
from self_reasoning_tokens_pytorch import Transformer

# Model configuration
model = Transformer(
    dim=256,
    depth=6,
    num_tokens=1000,
    max_seq_len=512,
    max_reason_seq_len=8,
    dim_head=32,
    heads=8,
    stop_grad_next_tokens_to_reason=True
)

# Optimizer
optimizer = optim.AdamW(model.parameters(), lr=1e-4)

# Training loop
for step in range(100):
    # Generate random training data
    batch = torch.randint(0, 1000, (4, 64))

    # Forward pass with reasoning tokens
    loss = model(
        batch,
        num_reason_tokens=4,
        num_steps_future_can_use_reason=8,
        return_loss=True
    )

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

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

# Inference mode
model.eval()
with torch.no_grad():
    test_input = torch.randint(0, 1000, (1, 32))
    logits = model(
        test_input,
        num_reason_tokens=4,
        remove_reason_tokens_at_end=True
    )
    predictions = logits.argmax(dim=-1)
    print(f"Predictions shape: {predictions.shape}")
```

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

### Initialize and use Transformer with reasoning tokens

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Initialize the Transformer model with reasoning token support and use it for training or inference. Ensure correct parameters for reasoning token handling and loss computation.

```python
import torch
from self_reasoning_tokens_pytorch import Transformer

# Initialize transformer with reasoning token support
model = Transformer(
    dim=512,                              # Model dimension
    depth=4,                              # Number of transformer layers
    num_tokens=256,                       # Vocabulary size
    max_seq_len=2048,                     # Maximum sequence length
    max_reason_seq_len=4,                 # Maximum reasoning tokens per timestep
    dim_head=64,                          # Dimension per attention head
    heads=8,                              # Number of attention heads
    ignore_index=-1,                      # Index to ignore in loss computation
    stop_grad_next_tokens_to_reason=True  # Apply stop gradients to reasoning tokens
)

# Create input sequence
x = torch.randint(0, 256, (1, 4))

# Training: Forward pass with loss computation
loss = model(
    x,
    num_reason_tokens=4,                  # Number of reasoning tokens per time step
    num_steps_future_can_use_reason=16,   # Tokens 16 steps ahead can attend to reason tokens
    return_loss=True                      # Return cross-entropy loss
)

# Backward pass for training
loss.backward()

# Inference: Get logits without loss
logits = model(x, num_reason_tokens=4)
print(f"Logits shape: {logits.shape}")  # Shape: (batch, seq * (1 + num_reason_tokens), num_tokens)

# Inference with reasoning tokens removed from output
logits_clean = model(
    x,
    num_reason_tokens=4,
    remove_reason_tokens_at_end=True
)
print(f"Clean logits shape: {logits_clean.shape}")  # Shape: (batch, seq, num_tokens)
```

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

### Initialize and use Transformer with reasoning tokens

Source: https://github.com/lucidrains/self-reasoning-tokens-pytorch/blob/main/README.md

Initialize a Transformer model with reasoning token capabilities and calculate loss. Ensure to import torch and Transformer. The model can be used for both loss calculation and logit generation.

```python
import torch
from self_reasoning_tokens_pytorch import Transformer

model = Transformer(
    dim = 512,
    depth = 4,
    num_tokens = 256,
    stop_grad_next_tokens_to_reason = True
)

x = torch.randint(0, 256, (1, 4))

loss = model(
    x,
    num_reason_tokens = 4,                # number of reasoning tokens per time step
    num_steps_future_can_use_reason = 16, # say you wish for reason tokens to be only attended to by tokens 16 time steps into the future
    return_loss = True
)

loss.backward()

logits = model(x, num_reason_tokens = 4)
```

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

### Stop-Graddable Attention with Gradient Scaling

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Illustrates using float masks for gradient scaling instead of binary stop gradients. Float values between 0.0 and 1.0 scale the gradient accordingly.

```python
# Using float masks for gradient scaling instead of binary stop
# Float values scale the gradient (0.0 = no gradient, 1.0 = full gradient)
gradient_scale_mask = torch.rand(num_heads, seq_len, seq_len)
output_scaled = stop_graddable_attn(
    q, k, v,
    causal=True,
    q_stop_grad_mask=gradient_scale_mask,
    k_stop_grad_mask=gradient_scale_mask,
    v_stop_grad_mask=gradient_scale_mask
)
```

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

### Initialize and use CausalAttention with masks

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Initialize the CausalAttention module and use it for standard attention, masked attention, or with stop gradient masks for fine-grained gradient control.

```python
import torch
from self_reasoning_tokens_pytorch import CausalAttention

# Initialize causal attention layer
attn = CausalAttention(
    dim=512,       # Input/output dimension
    dim_head=64,   # Dimension per attention head
    heads=8        # Number of attention heads
)

# Input tensor: (batch, sequence, dim)
x = torch.randn(2, 128, 512)

# Standard causal attention forward pass
output = attn(x)
print(f"Output shape: {output.shape}")  # Shape: (2, 128, 512)

# With custom attention mask
seq_len = 128
custom_attn_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril()
output_masked = attn(x, attn_mask=custom_attn_mask)

# With stop gradient mask (single mask applies to keys and values)
stop_grad_mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)
stop_grad_mask[64:, :64] = True  # Stop gradients from second half to first half
output_stop_grad = attn(x, stop_grad_attn_mask=stop_grad_mask)

# With separate stop gradient masks for queries, keys, and values
q_stop_mask = torch.randint(0, 2, (seq_len, seq_len)).bool()
k_stop_mask = torch.randint(0, 2, (seq_len, seq_len)).bool()
v_stop_mask = torch.randint(0, 2, (seq_len, seq_len)).bool()
output_separate = attn(x, stop_grad_attn_mask=(q_stop_mask, k_stop_mask, v_stop_mask))
print(f"Output with separate masks: {output_separate.shape}")
```

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

### Basic Stop-Graddable Attention Usage

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Demonstrates basic usage of stop_graddable_attn with causal masking and binary stop gradient masks. Ensure stop_grad_mask is broadcastable to the attention scores shape.

```python
import torch
from self_reasoning_tokens_pytorch import stop_graddable_attn

# Shape: (batch, heads, sequence, dim_head)
batch_size, num_heads, seq_len, dim_head = 2, 8, 1024, 64

q = torch.randn(batch_size, num_heads, seq_len, dim_head, requires_grad=True)
k = torch.randn(batch_size, num_heads, seq_len, dim_head, requires_grad=True)
v = torch.randn(batch_size, num_heads, seq_len, dim_head, requires_grad=True)

# Create stop gradient masks (True = stop gradient, False = allow gradient)
# Shape: (heads, seq, seq) or broadcastable
stop_grad_mask = torch.randint(0, 2, (num_heads, seq_len, seq_len)).bool()

# Basic usage with causal masking
output = stop_graddable_attn(
    q, k, v,
    causal=True,
    q_stop_grad_mask=stop_grad_mask,
    k_stop_grad_mask=stop_grad_mask,
    v_stop_grad_mask=stop_grad_mask
)
print(f"Output shape: {output.shape}")  # (2, 8, 1024, 64)

# Backward pass - gradients are masked according to stop_grad_masks
loss = output.sum()
loss.backward()
```

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

### Raw Stop-Graddable Attention Function Call

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Demonstrates the direct usage of the raw autograd function `stop_graddable_attn_` with positional arguments. This provides lower-level control over the attention mechanism.

```python
import torch
from self_reasoning_tokens_pytorch import stop_graddable_attn_

# Tensor setup
batch, heads, seq, dim = 2, 8, 256, 64
q = torch.randn(batch, heads, seq, dim, requires_grad=True)
k = torch.randn(batch, heads, seq, dim, requires_grad=True)
v = torch.randn(batch, heads, seq, dim, requires_grad=True)

# Stop gradient masks
stop_mask = torch.zeros(heads, seq, seq, dtype=torch.bool)
stop_mask[:, 128:, :128] = True  # Stop gradients from later tokens to earlier

# Raw function call with positional arguments:
# (q, k, v, mask, attn_mask, causal, q_stop_grad_mask, k_stop_grad_mask, v_stop_grad_mask)
output = stop_graddable_attn_(
    q, k, v,
    None,        # mask (optional key padding mask)
    None,        # attn_mask (optional attention mask)
    True,        # causal
    stop_mask,   # q_stop_grad_mask
    stop_mask,   # k_stop_grad_mask
    stop_mask    # v_stop_grad_mask
)

print(f"Raw output shape: {output.shape}")  # (2, 8, 256, 64)
```

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

### Stop-Graddable Attention with Attention Mask

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Shows how to use stop_graddable_attn with an additional attention mask for controlling which tokens can attend to others. The attn_mask handles causality when causal is set to False.

```python
# With additional attention mask
attn_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril()
output_with_attn_mask = stop_graddable_attn(
    q, k, v,
    attn_mask=attn_mask,
    causal=False,  # attn_mask handles causality
    q_stop_grad_mask=stop_grad_mask,
    k_stop_grad_mask=stop_grad_mask,
    v_stop_grad_mask=stop_grad_mask
)
```

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

### Use stop_graddable_attn with custom stop gradient masks

Source: https://github.com/lucidrains/self-reasoning-tokens-pytorch/blob/main/README.md

Utilize the stop_graddable_attn function for custom attention with specific stop gradient masks for queries, keys, and values. Ensure causal attention is set appropriately. The output shape is (batch_size, num_heads, sequence_length, head_dim).

```python
import torch
from self_reasoning_tokens_pytorch import stop_graddable_attn

q = torch.randn(2, 8, 1024, 64)
k = torch.randn(2, 8, 1024, 64)
v = torch.randn(2, 8, 1024, 64)

stop_grad_mask = torch.randint(0, 2, (8, 1024, 1024)).bool()

out = stop_graddable_attn(
    q, k, v, causal = True,
    q_stop_grad_mask = stop_grad_mask,
    k_stop_grad_mask = stop_grad_mask,
    v_stop_grad_mask = stop_grad_mask
)

out.shape # (2, 8, 1024, 64)
```

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

### Low-level stop_graddable_attn function

Source: https://context7.com/lucidrains/self-reasoning-tokens-pytorch/llms.txt

Utilize the low-level stop_graddable_attn function for custom attention computation with independent stop gradient masks for queries, keys, and values, enabling fine-grained gradient control.

```python
import torch
from self_reasoning_tokens_pytorch import stop_graddable_attn

# Create query, key, value tensors
```

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