### Install SE3 Transformer

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

Install the library via pip.

```bash
pip install se3-transformer-pytorch
```

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

### Install Test Dependencies and Run Tests

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

Install necessary packages like pytest and lie_learn, then execute the test suite using setup.py pytest to verify equivariance properties.

```bash
# Install test dependencies
pip install pytest lie_learn

# Run tests
python setup.py pytest
```

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

### Install sidechainnet

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Install the sidechainnet library using pip.

```bash
pip install sidechainnet
```

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

### Install SE3 Transformer

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Install the package via pip.

```bash
$ pip install se3-transformer-pytorch
```

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

### Initialize SE3 Transformer for Neighbor Masking

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

Setup for using custom neighbor masks to control node-to-node attention.

```python
import torch
from se3_transformer_pytorch.se3_transformer_pytorch import SE3Transformer
```

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

### Initialize and Run SE3 Transformer

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Basic usage example for initializing the model and performing a forward pass with features, coordinates, and a mask.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 512,
    heads = 8,
    depth = 6,
    dim_head = 64,
    num_degrees = 4,
    valid_radius = 10
)

feats = torch.randn(1, 1024, 512)
coors = torch.randn(1, 1024, 3)
mask  = torch.ones(1, 1024).bool()

out = model(feats, coors, mask) # (1, 1024, 512)
```

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

### Train protein backbone denoising model

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

Full training loop example using SideChainNet data to denoise protein backbone coordinates.

```python
import torch
import torch.nn.functional as F
from torch.optim import Adam
from einops import rearrange, repeat
import sidechainnet as scn
from se3_transformer_pytorch.se3_transformer_pytorch import SE3Transformer

torch.set_default_dtype(torch.float64)
BATCH_SIZE = 1
GRADIENT_ACCUMULATE_EVERY = 16

def cycle(loader, len_thres=500):
    while True:
        for data in loader:
            if data.seqs.shape[1] > len_thres:
                continue
            yield data

# Initialize model for coordinate denoising
transformer = SE3Transformer(
    num_tokens = 24,
    dim = 8,
    dim_head = 8,
    heads = 2,
    depth = 2,
    attend_self = True,
    input_degrees = 1,
    output_degrees = 2,
    reduce_dim_out = True,
    differentiable_coors = True,
    num_neighbors = 0,
    attend_sparse_neighbors = True,
    num_adj_degrees = 2,
    adj_dim = 4,
    num_degrees = 2,
).cuda()

# Load protein data
data = scn.load(
    casp_version = 12,
    thinning = 30,
    with_pytorch = 'dataloaders',
    batch_size = BATCH_SIZE,
    dynamic_batching = False
)

dl = cycle(data['train'])
optim = Adam(transformer.parameters(), lr=1e-4)

# Training loop
for _ in range(10000):
    for _ in range(GRADIENT_ACCUMULATE_EVERY):
        batch = next(dl)
        seqs, coords, masks = batch.seqs, batch.crds, batch.msks

        seqs = seqs.cuda().argmax(dim=-1)
        coords = coords.cuda().type(torch.float64)
        masks = masks.cuda().bool()

        l = seqs.shape[1]
        coords = rearrange(coords, 'b (l s) c -> b l s c', s=14)
        coords = coords[:, :, 0:3, :]  # Keep only backbone atoms
        coords = rearrange(coords, 'b l s c -> b (l s) c')

        seq = repeat(seqs, 'b n -> b (n c)', c=3)
        masks = repeat(masks, 'b n -> b (n c)', c=3)

        # Add noise to coordinates
        noised_coords = coords + torch.randn_like(coords).cuda()

        # Create chain adjacency matrix
        i = torch.arange(seq.shape[-1], device=seqs.device)
        adj_mat = (i[:, None] >= (i[None, :] - 1)) & (i[:, None] <= (i[None, :] + 1))

        # Predict denoised coordinates
        out = transformer(seq, noised_coords, mask=masks, adj_mat=adj_mat, return_type=1)
        denoised_coords = noised_coords + out

        loss = F.mse_loss(denoised_coords[masks], coords[masks])
        (loss / GRADIENT_ACCUMULATE_EVERY).backward()

    print('loss:', loss.item())
    optim.step()
    optim.zero_grad()
```

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

### Refine Coordinates for Alphafold2

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Example usage for coordinate refinement with differentiable coordinates.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 2,
    input_degrees = 1,
    num_degrees = 2,
    output_degrees = 2,
    reduce_dim_out = True,
    differentiable_coors = True
)

atom_feats = torch.randn(2, 32, 64)
coors = torch.randn(2, 32, 3)
mask  = torch.ones(2, 32).bool()

refined_coors = coors + model(atom_feats, coors, mask, return_type = 1) # (2, 32, 3)
```

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

### Define Adjacency Matrix for Chain Sequences

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Example of creating a chain-like adjacency matrix for input into the model.

```python
i = torch.arange(128)
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

out = model(feats, coors, mask, adj_mat = adj_mat, return_type = 1)
```

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

### Run Tests

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Execute the test suite using setup.py.

```bash
$ python setup.py pytest
```

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

### Initialize and Run Basic SE3 Transformer

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

Create a standard SE3Transformer model and perform a forward pass with node features and 3D coordinates.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Create a basic SE3 Transformer model
model = SE3Transformer(
    dim = 512,          # Feature dimension
    heads = 8,          # Number of attention heads
    depth = 6,          # Number of transformer layers
    dim_head = 64,      # Dimension per attention head
    num_degrees = 4,    # Number of SE(3) degrees (higher = more expressive)
    valid_radius = 10   # Radius for neighbor selection
)

# Input data
feats = torch.randn(1, 1024, 512)  # (batch, num_nodes, dim)
coors = torch.randn(1, 1024, 3)     # (batch, num_nodes, 3D coordinates)
mask = torch.ones(1, 1024).bool()   # (batch, num_nodes) attention mask

# Forward pass - returns type 0 features by default
out = model(feats, coors, mask)  # Output shape: (1, 1024, 512)
```

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

### Initialize SE3Transformer with custom fiber structure

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

Configures the model with specific fiber dimensions and coordinate update clamping for stability.

```python
model = SE3Transformer(
    dim = 32,
    num_neighbors = 8,
    hidden_fiber_dict = {0: 32, 1: 16, 2: 8, 3: 4},  # Custom dimensions
    use_egnn = True,
    depth = 4,
    egnn_hidden_dim = 64,
    egnn_weights_clamp_value = 2,  # Clamp coordinate updates for stability
    reduce_dim_out = True
).cuda()

feats = torch.randn(2, 32, 32).cuda()
coors = torch.randn(2, 32, 3).cuda()
mask = torch.ones(2, 32).bool().cuda()

refinement = model(feats, coors, mask, return_type=1)  # (2, 32, 3)
coors = coors + refinement
```

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

### Initialize SE3 Transformer with Edge Tokens

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Pass num_edge_tokens and edge_dim during initialization to incorporate discrete edge information like bond types.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    num_tokens = 28,
    dim = 64,
    num_edge_tokens = 4,       # number of edge type, say 4 bond types
    edge_dim = 16,             # dimension of edge embedding
    depth = 2,
    input_degrees = 1,
    num_degrees = 3,
    output_degrees = 1,
    reduce_dim_out = True
)

atoms = torch.randint(0, 28, (2, 32))
bonds = torch.randint(0, 4, (2, 32, 32))
coors = torch.randn(2, 32, 3)
mask  = torch.ones(2, 32).bool()

pred = model(atoms, coors, mask, edges = bonds, return_type = 0) # (2, 32, 1)
```

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

### Use Linear Projected Keys

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Enable linear_proj_keys for potential memory savings in specific tasks.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 1,
    num_degrees = 4,
    num_neighbors = 8,
    valid_radius = 10,
    splits = 4,
    linear_proj_keys = True # set this to True
).cuda()

feats = torch.randn(1, 32, 64).cuda()
coors = torch.randn(1, 32, 3).cuda()
mask  = torch.ones(1, 32).bool().cuda()

out = model(feats, coors, mask, return_type = 0)
```

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

### SE3Transformer with Linear Projected Keys (Memory Optimization)

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

Reduce memory usage by enabling linear projections for keys with `linear_proj_keys = True`. This offers significant memory savings.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Memory-efficient model with linear key projection
model = SE3Transformer(
    dim = 64,
    depth = 1,
    num_degrees = 4,
    num_neighbors = 8,
    valid_radius = 10,
    splits = 4,
    linear_proj_keys = True  # 25% memory savings
).cuda()

feats = torch.randn(1, 32, 64).cuda()
coors = torch.randn(1, 32, 3).cuda()
mask = torch.ones(1, 32).bool().cuda()

out = model(feats, coors, mask, return_type=0)
```

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

### SE3Transformer with Reversible Networks

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Enable reversible networks by setting `reversible = True`. This technique allows for greater network depth by reducing memory usage, preserving equivariance.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    num_tokens = 20,
    dim = 32,
    dim_head = 32,
    heads = 4,
    depth = 12,             # 12 layers
    input_degrees = 1,
    num_degrees = 3,
    output_degrees = 1,
    reduce_dim_out = True,
    reversible = True       # set reversible to True
).cuda()

atoms = torch.randint(0, 4, (2, 32)).cuda()
coors = torch.randn(2, 32, 3).cuda()
mask  = torch.ones(2, 32).bool().cuda()

pred = model(atoms, coors, mask = mask, return_type = 0)

loss = pred.sum()
loss.backward()
```

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

### Run Protein Backbone Denoising Task

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Execute the protein backbone denoising task using the provided script.

```bash
python denoise.py
```

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

### Manage basis vector cache

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

Configures cache behavior for basis vectors using environment variables.

```bash
# Clear cache on startup
CLEAR_CACHE=1 python train.py

# Use custom cache directory
CACHE_PATH=./my_cache python train.py

# Default cache location
rm -rf ~/.cache.equivariant_attention
```

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

### SE3Transformer with Neighbor Mask Support

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

Initialize the SE3Transformer model with neighbor mask support. The `num_neighbors` parameter should be set to the maximum number of neighbors in your mask. A custom `neighbor_mask` can be provided to control attention between nodes.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 16,
    dim_head = 16,
    attend_self = True,
    num_degrees = 4,
    output_degrees = 2,
    num_edge_tokens = 4,
    num_neighbors = 8,  # Set to max neighbors in your mask
    edge_dim = 2,
    depth = 3
)

feats = torch.randn(1, 32, 16)
coors = torch.randn(1, 32, 3)
mask = torch.ones(1, 32).bool()
bonds = torch.randint(0, 4, (1, 32, 32))

# Custom neighbor mask (True = can attend, False = masked out)
neighbor_mask = torch.ones(1, 32, 32).bool()
# Example: mask out certain interactions
neighbor_mask[:, 0, 5:10] = False

out = model(
    feats, coors, mask,
    edges=bonds,
    neighbor_mask=neighbor_mask,
    return_type=1
)
```

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

### SE3Transformer with Tied Key/Value (Memory Optimization)

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

Reduce memory usage by 50% on key/value projections by tying them together with `tie_key_values = True` (K=V).

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model with tied key/values
model = SE3Transformer(
    dim = 64,
    depth = 8,
    num_degrees = 4,
    num_neighbors = 8,
    valid_radius = 10,
    splits = 4,
    tie_key_values = True  # K = V for memory savings
).cuda()

feats = torch.randn(1, 32, 64).cuda()
coors = torch.randn(1, 32, 3).cuda()
mask = torch.ones(1, 32).bool().cuda()

out = model(feats, coors, mask, return_type=0)
```

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

### Manage Cache via Environment Variables

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Use environment variables to control cache behavior during script execution.

```bash
$ CLEAR_CACHE=1 python train.py
```

```bash
CACHE_PATH=./path/to/my/cache python train.py
```

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

### SE3Transformer with Tied Key/Value Heads

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Set `tie_key_values = True` to ensure key and value heads are identical, effectively halving memory usage for these components.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 8,
    num_degrees = 4,
    num_neighbors = 8,
    valid_radius = 10,
    splits = 4,
    tie_key_values = True # set this to True
).cuda()

feats = torch.randn(1, 32, 64).cuda()
coors = torch.randn(1, 32, 3).cuda()
mask  = torch.ones(1, 32).bool().cuda()

out = model(feats, coors, mask, return_type = 0)
```

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

### SE3Transformer with Shared Key/Value Heads (Memory Optimization)

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

Achieve significant memory reduction by sharing key/value projections across all query heads using `one_headed_key_values = True`.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model with shared key/value head
model = SE3Transformer(
    dim = 64,
    depth = 8,
    num_degrees = 4,
    num_neighbors = 8,
    valid_radius = 10,
    splits = 4,
    one_headed_key_values = True  # One K/V head shared across all Q heads
).cuda()

feats = torch.randn(1, 32, 64).cuda()
coors = torch.randn(1, 32, 3).cuda()
mask = torch.ones(1, 32).bool().cuda()

out = model(feats, coors, mask, return_type=0)
```

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

### SE3Transformer with Shared Key/Value Head

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Use `one_headed_key_values = True` to share one key/value head across all query heads. This can trade memory for depth or a higher number of degrees, potentially impacting performance.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 8,
    num_degrees = 4,
    num_neighbors = 8,
    valid_radius = 10,
    splits = 4,
    one_headed_key_values = True  # one head of key / values shared across all heads of the queries
).cuda()

feats = torch.randn(1, 32, 64).cuda()
coors = torch.randn(1, 32, 3).cuda()
mask  = torch.ones(1, 32).bool().cuda()

out = model(feats, coors, mask, return_type = 0)
```

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

### Configure Sparse Neighbors with Adjacency Matrix

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Set attend_sparse_neighbors to True and provide an adjacency matrix to restrict attention to specific connectivity.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 32,
    heads = 8,
    depth = 1,
    dim_head = 64,
    num_degrees = 2,
    valid_radius = 10,
    attend_sparse_neighbors = True,  # this must be set to true, in which case it will assert that you pass in the adjacency matrix
    num_neighbors = 0,               # if you set this to 0, it will only consider the connected neighbors as defined by the adjacency matrix. but if you set a value greater than 0, it will continue to fetch the closest points up to this many, excluding the ones already specified by the adjacency matrix
    max_sparse_neighbors = 8         # you can cap the number of neighbors, sampled from within your sparse set of neighbors as defined by the adjacency matrix, if specified
)

feats = torch.randn(1, 128, 32)
coors = torch.randn(1, 128, 3)
mask  = torch.ones(1, 128).bool()

# placeholder adjacency matrix
# naively assuming the sequence is one long chain (128, 128)

i = torch.arange(128)
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

out = model(feats, coors, mask, adj_mat = adj_mat) # (1, 128, 512)
```

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

### Enable reversible residual connections

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

Uses reversible layers to reduce memory footprint during backpropagation, enabling deeper network architectures.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Deep model with reversible layers
model = SE3Transformer(
    num_tokens = 20,
    dim = 32,
    dim_head = 32,
    heads = 4,
    depth = 12,             # Deeper network feasible with reversible
    input_degrees = 1,
    num_degrees = 3,
    output_degrees = 1,
    reduce_dim_out = True,
    reversible = True       # Enable reversible residual connections
).cuda()

atoms = torch.randint(0, 4, (2, 32)).cuda()
coors = torch.randn(2, 32, 3).cuda()
mask = torch.ones(2, 32).bool().cuda()

pred = model(atoms, coors, mask=mask, return_type=0)

# Backpropagation works with reduced memory footprint
loss = pred.sum()
loss.backward()
```

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

### SE3Transformer with Global Features

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

Configure the SE3Transformer to accept global features. The `global_feats_dim` parameter enables this functionality. Global features can be used for conditioning on graph-level properties.

```python
import torch
from torch import nn
from se3_transformer_pytorch import SE3Transformer

# Model accepting global features
model = SE3Transformer(
    dim = 64,
    depth = 1,
    num_degrees = 2,
    num_neighbors = 4,
    valid_radius = 10,
    global_feats_dim = 32  # Enable global features
)

feats = torch.randn(1, 32, 64)
coors = torch.randn(1, 32, 3)
mask = torch.ones(1, 32).bool()

# Create global features (e.g., by pooling and projecting)
global_feats = nn.Linear(64, 32)(feats.mean(dim=1, keepdim=True))  # (1, 1, 32)

# All nodes attend to global features
out = model(feats, coors, mask, return_type=0, global_feats=global_feats)
```

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

### Clear Cache Directory

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Manually remove the cache directory if necessary.

```bash
$ rm -rf ~/.cache.equivariant_attention
```

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

### Perform Coordinate Refinement

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

Configure the model to output type-1 features for iterative 3D coordinate refinement.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model configured for coordinate refinement
model = SE3Transformer(
    dim = 64,
    depth = 2,
    input_degrees = 1,          # Input contains type 0 features only
    num_degrees = 2,            # Use type 0 and type 1 internally
    output_degrees = 2,         # Output type 0 and type 1
    reduce_dim_out = True,      # Reduce output dimension to 1
    differentiable_coors = True # Enable gradients through basis computation
)

# Input: atom features and coordinates
atom_feats = torch.randn(2, 32, 64)  # (batch, num_atoms, dim)
coors = torch.randn(2, 32, 3)         # (batch, num_atoms, 3)
mask = torch.ones(2, 32).bool()

# Get coordinate refinement (type 1 output is a 3D vector)
refinement = model(atom_feats, coors, mask, return_type=1)  # (2, 32, 3)

# Apply refinement to coordinates
refined_coors = coors + refinement
```

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

### Configure Custom Fiber Dimensions

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

Define specific feature dimensions for each degree using fiber dictionaries for fine-grained control.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model with custom fiber dimensions
model = SE3Transformer(
    num_tokens = 28,
    dim = 64,
    num_edge_tokens = 4,
    edge_dim = 16,
    depth = 2,
    input_degrees = 1,
    num_degrees = 3,
    output_degrees = 1,
    hidden_fiber_dict = {0: 16, 1: 8, 2: 4},  # Custom dims per degree
    out_fiber_dict = {0: 16, 1: 1},           # Output fiber structure
    reduce_dim_out = False
)

atoms = torch.randint(0, 28, (2, 32))
bonds = torch.randint(0, 4, (2, 32, 32))
coors = torch.randn(2, 32, 3)
mask = torch.ones(2, 32).bool()

# Output is a dictionary of feature types
pred = model(atoms, coors, mask, edges=bonds)

print(pred['0'].shape)  # (2, 32, 16) - type 0 features
print(pred['1'].shape)  # (2, 32, 1, 3) - type 1 features (3D vectors)
```

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

### SE3Transformer as EGNN for Higher Order Types

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Utilize `use_egnn = True` for an experimental EGNN implementation that supports higher-order types and greater dimensionality. This configuration requires `num_edge_tokens` and `edge_dim` if bonds are provided.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 32,
    num_neighbors = 8,
    num_edge_tokens = 4,
    edge_dim = 4,
    num_degrees = 4,       # number of higher order types - will use basis on a TCN to project to these dimensions
    use_egnn = True,       # set this to true to use EGNN instead of equivariant attention layers
    egnn_hidden_dim = 64,  # egnn hidden dimension
    depth = 4,             # depth of EGNN
    reduce_dim_out = True  # will project the dimension of the higher types to 1
).cuda()

feats = torch.randn(2, 32, 32).cuda()
coors = torch.randn(2, 32, 3).cuda()
bonds = torch.randint(0, 4, (2, 32, 32)).cuda()
mask  = torch.ones(2, 32).bool().cuda()

refinement = model(feats, coors, mask, edges = bonds, return_type = 1) # (2, 32, 3)

coors = coors + refinement  # update coors with refinement
```

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

### Apply Fourier encoding to continuous values

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

Generates positional encodings for continuous input values using a specified number of frequency bands.

```python
import torch
from se3_transformer_pytorch.utils import fourier_encode

# Continuous values to encode
values = torch.randn(2, 32, 32, 2)  # (batch, n, n, features)

# Fourier encode with 8 frequency bands
encoded = fourier_encode(
    values,
    num_encodings=8,     # Number of frequency bands
    include_self=True    # Include original values
)
# Output shape: (2, 32, 32, 34)
# = 2 features * (2 * 8 frequencies + 1 self) = 34

print(encoded.shape)  # torch.Size([2, 32, 32, 34])
```

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

### Enable Autoregressive Mode

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Set the causal flag to True to use the model in an autoregressive manner.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 512,
    heads = 8,
    depth = 6,
    dim_head = 64,
    num_degrees = 4,
    valid_radius = 10,
    causal = True          # set this to True
)

feats = torch.randn(1, 1024, 512)
coors = torch.randn(1, 1024, 3)
mask  = torch.ones(1, 1024).bool()

out = model(feats, coors, mask) # (1, 1024, 512)
```

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

### Integrate Global Nodes

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Pass global feature vectors that are accessible by all nodes regardless of local connectivity.

```python
import torch
from torch import nn
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 1,
    num_degrees = 2,
    num_neighbors = 4,
    valid_radius = 10,
    global_feats_dim = 32 # this must be set to the dimension of the global features, in this example, 32
)

feats = torch.randn(1, 32, 64)
coors = torch.randn(1, 32, 3)
mask  = torch.ones(1, 32).bool()

# naively derive global features
# by pooling features and projecting
global_feats = nn.Linear(64, 32)(feats.mean(dim = 1, keepdim = True)) # (1, 1, 32)

out = model(feats, coors, mask, return_type = 0, global_feats = global_feats)
```

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

### Configure Fiber Dimensions

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Use hidden_fiber_dict and out_fiber_dict to control dimensionality for specific degrees.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    num_tokens = 28,
    dim = 64,
    num_edge_tokens = 4,
    edge_dim = 16,
    depth = 2,
    input_degrees = 1,
    num_degrees = 3,
    output_degrees = 1,
    hidden_fiber_dict = {0: 16, 1: 8, 2: 4},
    out_fiber_dict = {0: 16, 1: 1},
    reduce_dim_out = False
)

atoms = torch.randint(0, 28, (2, 32))
bonds = torch.randint(0, 4, (2, 32, 32))
coors = torch.randn(2, 32, 3)
mask  = torch.ones(2, 32).bool()

pred = model(atoms, coors, mask, edges = bonds)

pred['0'] # (2, 32, 16)
pred['1'] # (2, 32, 1, 3)
```

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

### SE3Transformer in EGNN Mode

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

Utilize E(n) Equivariant Graph Neural Networks instead of SE(3) attention layers by setting `use_egnn = True`. This provides an alternative equivariant architecture.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# EGNN-based model
model = SE3Transformer(
    dim = 32,
    num_neighbors = 8,
    num_edge_tokens = 4,
    edge_dim = 4,
    num_degrees = 4,
    use_egnn = True,           # Use EGNN instead of SE3 attention
    egnn_hidden_dim = 64,
    depth = 4,
    reduce_dim_out = True
).cuda()

feats = torch.randn(2, 32, 32).cuda()
coors = torch.randn(2, 32, 3).cuda()
bonds = torch.randint(0, 4, (2, 32, 32)).cuda()
mask = torch.ones(2, 32).bool().cuda()

refinement = model(feats, coors, mask, edges=bonds, return_type=1)  # (2, 32, 3)
coors = coors + refinement  # Update coordinates
```

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

### Process Multi-type Features

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

Pass multiple feature types as a dictionary to handle both scalar and vector inputs.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model accepting multiple input types
model = SE3Transformer(
    dim = 64,
    depth = 2,
    input_degrees = 2,      # Input contains type 0 and type 1
    num_degrees = 2,
    output_degrees = 2,
    reduce_dim_out = True
)

# Multi-type input features
atom_feats = torch.randn(2, 32, 64, 1)   # Type 0: (batch, n, dim, 1)
coors_feats = torch.randn(2, 32, 64, 3)  # Type 1: (batch, n, dim, 3)

# Pass as dictionary with degree keys
features = {'0': atom_feats, '1': coors_feats}
coors = torch.randn(2, 32, 3)
mask = torch.ones(2, 32).bool()

# Output is equivariant to both input types
refined_coors = coors + model(features, coors, mask, return_type=1)  # (2, 32, 3)
```

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

### Derive Nth-Degree Neighbors Automatically

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Use num_adj_degrees to automatically derive neighbor degrees and adj_dim to embed them as edge information.

```python
import torch
from se3_transformer_pytorch.se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 1,
    attend_self = True,
    num_degrees = 2,
    output_degrees = 2,
    num_neighbors = 0,
    attend_sparse_neighbors = True,
    num_adj_degrees = 2,    # automatically derive 2nd degree neighbors
    adj_dim = 4             # embed 1st and 2nd degree neighbors (as well as null neighbors) with edge embeddings of this dimension
)

feats = torch.randn(1, 32, 64)
coors = torch.randn(1, 32, 3)
mask  = torch.ones(1, 32).bool()

# placeholder adjacency matrix
```

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

### Apply Neighbor Masking

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Control node consideration by passing a neighbor mask where False values are ignored.

```python
import torch
from se3_transformer_pytorch.se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 16,
    dim_head = 16,
    attend_self = True,
    num_degrees = 4,
    output_degrees = 2,
    num_edge_tokens = 4,
    num_neighbors = 8,      # make sure you set this value as the maximum number of neighbors set by your neighbor_mask, or it will throw a warning
    edge_dim = 2,
    depth = 3
)

feats = torch.randn(1, 32, 16)
coors = torch.randn(1, 32, 3)
mask  = torch.ones(1, 32).bool()
bonds = torch.randint(0, 4, (1, 32, 32))

neighbor_mask = torch.ones(1, 32, 32).bool() # set the nodes you wish to be masked out as False

out = model(
    feats,
    coors,
    mask,
    edges = bonds,
    neighbor_mask = neighbor_mask,
    return_type = 1
)
```

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

### Embed Discrete Tokens

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

Use the built-in embedding layer to process discrete token inputs like atom types.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model with token embedding
model = SE3Transformer(
    num_tokens = 28,        # Number of unique atom types
    dim = 64,
    depth = 2,
    input_degrees = 1,
    num_degrees = 2,
    output_degrees = 2,
    reduce_dim_out = True
)

# Input: atom type indices instead of features
atoms = torch.randint(0, 28, (2, 32))  # (batch, num_atoms) - integer token indices
coors = torch.randn(2, 32, 3)
mask = torch.ones(2, 32).bool()

# Forward pass with tokens
refinement = model(atoms, coors, mask, return_type=1)  # (2, 32, 3)
refined_coors = coors + refinement
```

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

### Incorporate Bond Types as Edge Information

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

Use discrete edge embeddings to represent chemical bond types in the attention mechanism.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model with edge embeddings
model = SE3Transformer(
    num_tokens = 28,
    dim = 64,
    num_edge_tokens = 4,    # Number of edge types (e.g., 4 bond types)
    edge_dim = 16,          # Dimension of edge embeddings
    depth = 2,
    input_degrees = 1,
    num_degrees = 3,
    output_degrees = 1,
    reduce_dim_out = True
)

# Inputs
atoms = torch.randint(0, 28, (2, 32))
bonds = torch.randint(0, 4, (2, 32, 32))  # Pairwise bond types
coors = torch.randn(2, 32, 3)
mask = torch.ones(2, 32).bool()

# Forward pass with edge information
pred = model(atoms, coors, mask, edges=bonds, return_type=0)  # (2, 32, 1)
```

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

### SE3Transformer in Autoregressive Mode

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

Enable causal attention for autoregressive generation by setting `causal = True`. This ensures each position only attends to previous positions.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Causal/autoregressive model
model = SE3Transformer(
    dim = 512,
    heads = 8,
    depth = 6,
    dim_head = 64,
    num_degrees = 4,
    valid_radius = 10,
    causal = True  # Enable causal masking
)

feats = torch.randn(1, 1024, 512)
coors = torch.randn(1, 1024, 3)
mask = torch.ones(1, 1024).bool()

# Each position only attends to previous positions
out = model(feats, coors, mask)
```

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

### Derive Multi-hop Neighbors Automatically

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

Automatically compute higher-degree neighbors from first-degree connectivity using learnable edge embeddings.

```python
import torch
from se3_transformer_pytorch.se3_transformer_pytorch import SE3Transformer

# Model with automatic neighbor derivation
model = SE3Transformer(
    dim = 64,
    depth = 1,
    attend_self = True,
    num_degrees = 2,
    output_degrees = 2,
    num_neighbors = 0,
    attend_sparse_neighbors = True,
    num_adj_degrees = 2,  # Derive up to 2nd degree neighbors
    adj_dim = 4           # Embed neighbor degree as edge features
)

feats = torch.randn(1, 32, 64)
coors = torch.randn(1, 32, 3)
mask = torch.ones(1, 32).bool()

# First-degree adjacency (sequential chain)
i = torch.arange(128)
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

# Model automatically derives 2nd degree neighbors
out = model(feats, coors, mask, adj_mat=adj_mat, return_type=1)
```

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

### Use Continuous Edge Features with Fourier Encoding

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

Apply Fourier encoding to continuous pairwise features to create expressive edge representations.

```python
import torch
from se3_transformer_pytorch import SE3Transformer
from se3_transformer_pytorch.utils import fourier_encode

# Model with continuous edge features
model = SE3Transformer(
    dim = 64,
    depth = 1,
    attend_self = True,
    num_degrees = 2,
    output_degrees = 2,
    edge_dim = 34  # Must match fourier-encoded edge dimension
)

feats = torch.randn(1, 32, 64)
coors = torch.randn(1, 32, 3)
mask = torch.ones(1, 32).bool()

# Create continuous pairwise features (e.g., 2 features per edge)
pairwise_continuous = torch.randint(0, 4, (1, 32, 32, 2))

# Fourier encode: 2 features * (2 * 8 encodings + 1 self) = 34 dimensions
edges = fourier_encode(
    pairwise_continuous,
    num_encodings=8,
    include_self=True
)  # (1, 32, 32, 34)

out = model(feats, coors, mask, edges=edges, return_type=1)
```

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

### Featurize Positional Encoding

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Passing explicit positional features in space using a dictionary of types.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 64,
    depth = 2,
    input_degrees = 2,
    num_degrees = 2,
    output_degrees = 2,
    reduce_dim_out = True  # reduce out the final dimension
)

atom_feats  = torch.randn(2, 32, 64, 1) # b x n x d x type0
coors_feats = torch.randn(2, 32, 64, 3) # b x n x d x type1

# atom features are type 0, predicted coordinates are type 1
features = {'0': atom_feats, '1': coors_feats}
coors = torch.randn(2, 32, 3)
mask  = torch.ones(2, 32).bool()

refined_coors = coors + model(features, coors, mask, return_type = 1) # (2, 32, 3) - equivariant to input type 1 features and coordinates
```

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

### Define Sparse Connectivity with Adjacency Matrix

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

Restrict attention to specific neighbors using an adjacency matrix for known molecular structures.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

# Model with sparse neighbor attention
model = SE3Transformer(
    dim = 32,
    heads = 8,
    depth = 1,
    dim_head = 64,
    num_degrees = 2,
    valid_radius = 10,
    attend_sparse_neighbors = True,  # Enable sparse attention
    num_neighbors = 0,               # Only use adjacency-defined neighbors
    max_sparse_neighbors = 8         # Cap maximum neighbors
)

feats = torch.randn(1, 128, 32)
coors = torch.randn(1, 128, 3)
mask = torch.ones(1, 128).bool()

# Create adjacency matrix (e.g., sequential chain connectivity)
i = torch.arange(128)
adj_mat = (i[:, None] <= (i[None, :] + 1)) & (i[:, None] >= (i[None, :] - 1))

# Forward pass with adjacency matrix
out = model(feats, coors, mask, adj_mat=adj_mat)  # (1, 128, 32)
```

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

### Embed Atom Tokens

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Usage with tokenized atom inputs where the model handles embedding of type 0 features.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    num_tokens = 28,       # 28 unique atoms
    dim = 64,
    depth = 2,
    input_degrees = 1,
    num_degrees = 2,
    output_degrees = 2,
    reduce_dim_out = True
)

atoms = torch.randint(0, 28, (2, 32))
coors = torch.randn(2, 32, 3)
mask  = torch.ones(2, 32).bool()

refined_coors = coors + model(atoms, coors, mask, return_type = 1) # (2, 32, 3)
```

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

### EGNN with Hidden Fiber Dimensions

Source: https://github.com/lucidrains/se3-transformer-pytorch/blob/main/README.md

Configure EGNN with specific dimensions for higher types using `hidden_fiber_dict`. This allows fine-grained control over the dimensionality of each degree.

```python
import torch
from se3_transformer_pytorch import SE3Transformer

model = SE3Transformer(
    dim = 32,
    num_neighbors = 8,
    hidden_fiber_dict = {0: 32, 1: 16, 2: 8, 3: 4},
    use_egnn = True,
    depth = 4,
    egnn_hidden_dim = 64,
    egnn_weights_clamp_value = 2, 
    reduce_dim_out = True
).cuda()

feats = torch.randn(2, 32, 32).cuda()
coors = torch.randn(2, 32, 3).cuda()
mask  = torch.ones(2, 32).bool().cuda()

refinement = model(feats, coors, mask, return_type = 1) # (2, 32, 3)

coors = coors + refinement  # update coors with refinement
```