### Install assoc-scan

Source: https://github.com/lucidrains/assoc-scan/blob/main/README.md

Install the assoc-scan library using pip.

```bash
pip install assoc-scan
```

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

### Install Assoc Scan

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Commands to install the core library and the optional accelerated GPU backend.

```bash
pip install assoc-scan
```

```bash
pip install accelerated-scan
```

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

### Basic AssocScan Usage in PyTorch

Source: https://github.com/lucidrains/assoc-scan/blob/main/README.md

Demonstrates the basic usage of the AssocScan class with PyTorch tensors. Ensure torch is installed.

```python
import torch
from assoc_scan import AssocScan

scan = AssocScan()

gates = torch.randn(1, 1024, 512).sigmoid()
inputs = torch.randn(1, 1024, 512)

out = scan(gates, inputs) # (1, 1024, 512)
```

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

### AssocScan with Accelerated Backend

Source: https://github.com/lucidrains/assoc-scan/blob/main/README.md

Enables the accelerated version of AssocScan by passing `use_accelerated = True`. Requires the 'accelerated-scan' library to be installed.

```python
scan = AssocScan(use_accelerated = True)

out = scan(gates, inputs)
```

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

### Enable Accelerated GPU Backend

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Initialize the scanner with the accelerated backend for CUDA-enabled devices.

```python
import torch
from assoc_scan import AssocScan

# Initialize with accelerated backend (requires accelerated-scan package)
scan = AssocScan(use_accelerated=True)

# Move tensors to GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
gates = torch.randn(4, 2048, 256, device=device).sigmoid()
inputs = torch.randn(4, 2048, 256, device=device)
```

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

### Initialize and Run AssocScan

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Basic usage of the AssocScan module to perform cumulative gated operations on batched tensors.

```python
import torch
from assoc_scan import AssocScan

# Initialize the scanner
scan = AssocScan(
    use_accelerated=False,  # Set True for GPU-optimized version
    reverse=False           # Set True to scan in reverse direction
)

# Create sample inputs: gates should be in range [0, 1]
batch_size, seq_len, dim = 2, 1024, 512
gates = torch.randn(batch_size, seq_len, dim).sigmoid()  # (2, 1024, 512)
inputs = torch.randn(batch_size, seq_len, dim)           # (2, 1024, 512)

# Perform associative scan
output = scan(gates, inputs)
print(f"Output shape: {output.shape}")  # Output shape: torch.Size([2, 1024, 512])
```

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

### Basic AssocScan Usage

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Demonstrates the basic usage of the scan function with gates and inputs. This operation can be accelerated on GPU using Triton.

```python
output = scan(gates, inputs)
print(f"Accelerated output shape: {output.shape}")  # torch.Size([4, 2048, 256])
```

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

### Perform State Continuation

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Use the forward method with a previous state to process sequences in chunks.

```python
import torch
from assoc_scan import AssocScan

scan = AssocScan()

# Process sequence with previous state continuation
batch_size, seq_len, dim = 1, 10, 4
gates = torch.randn(batch_size, seq_len, dim).sigmoid()
inputs = torch.randn(batch_size, seq_len, dim)

# Previous state from a prior chunk
prev_state = torch.randn(batch_size, dim)  # (1, 4)

# Continue scan from previous state
output = scan(
    gates,
    inputs,
    prev=prev_state,        # Previous hidden state
    remove_prev=True,       # Remove prepended prev from output (default when prev is given)
    return_next=False       # Whether to return next state for chaining
)
print(f"Output shape: {output.shape}")  # Output shape: torch.Size([1, 10, 4])
```

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

### Execute Reverse Scanning

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Configure the scanner to process sequences from end to beginning.

```python
import torch
from assoc_scan import AssocScan

# Initialize reverse scanner
reverse_scan = AssocScan(reverse=True)

# Create inputs
gates = torch.randn(1, 10, 4).sigmoid()
inputs = torch.randn(1, 10, 4)

# Scan in reverse direction
output = reverse_scan(gates, inputs)
print(f"Reverse scan output shape: {output.shape}")  # torch.Size([1, 10, 4])
```

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

### Handle Unbatched Inputs

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

The module supports 2D tensors without batch dimensions and 1D broadcastable gates.

```python
import torch
from assoc_scan import AssocScan

scan = AssocScan()

# Unbatched inputs: just (seq_len, dim)
seq_len, dim = 10, 4
gates = torch.randn(seq_len, dim).sigmoid()  # (10, 4)
inputs = torch.randn(seq_len, dim)           # (10, 4)

# Works without batch dimension
output = scan(gates, inputs)
print(f"Unbatched output shape: {output.shape}")  # torch.Size([10, 4])

# Also supports 1D gates that broadcast across dimensions
gates_1d = torch.randn(seq_len).sigmoid()  # (10,)
inputs_1d = torch.randn(seq_len)           # (10,)
output_1d = scan(gates_1d, inputs_1d)
print(f"1D output shape: {output_1d.shape}")  # torch.Size([10])
```

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

### Linear Attention Layer with AssocScan

Source: https://context7.com/lucidrains/assoc-scan/llms.txt

Integrates AssocScan into a PyTorch nn.Module for building linear attention mechanisms. Ensure `use_accelerated=True` for production deployments on CUDA.

```python
import torch
import torch.nn as nn
from assoc_scan import AssocScan

class LinearAttentionLayer(nn.Module):
    def __init__(self, dim, use_accelerated=False):
        super().__init__()
        self.dim = dim
        self.scan = AssocScan(use_accelerated=use_accelerated)
        self.to_gates = nn.Linear(dim, dim)
        self.to_values = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)

    def forward(self, x):
        # x: (batch, seq_len, dim)
        gates = self.to_gates(x).sigmoid()  # Decay factors
        values = self.to_values(x)          # Input values

        # Associative scan computes cumulative gated sum
        hidden = self.scan(gates, values)

        return self.out_proj(hidden)

# Usage
batch_size, seq_len, dim = 2, 128, 64
model = LinearAttentionLayer(dim)
x = torch.randn(batch_size, seq_len, dim)
output = model(x)
print(f"Layer output shape: {output.shape}")  # torch.Size([2, 128, 64])
```

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