### Install gotennet-pytorch

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

Install the gotennet-pytorch library using pip.

```bash
$ pip install gotennet-pytorch
```

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

### Verify SE(3) Equivariance with GotenNet

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

This script defines rotation helper functions and compares model outputs for original and rotated coordinates to confirm invariance and equivariance.

```python
import torch
from torch import sin, cos, stack
from einops import rearrange
from gotennet_pytorch import GotenNet, torch_default_dtype

# Helper functions to create rotation matrices
def rot_z(gamma):
    c, s = cos(gamma), sin(gamma)
    z, o = torch.zeros_like(gamma), torch.ones_like(gamma)
    return rearrange(stack((c, -s, z, s, c, z, z, z, o), dim=-1), '... (r1 r2) -> ... r1 r2', r1=3)

def rot_y(beta):
    c, s = cos(beta), sin(beta)
    z, o = torch.zeros_like(beta), torch.ones_like(beta)
    return rearrange(stack((c, z, s, z, o, z, -s, z, c), dim=-1), '... (r1 r2) -> ... r1 r2', r1=3)

def random_rotation():
    angles = torch.randn(3)
    return rot_z(angles[0]) @ rot_y(angles[1]) @ rot_z(angles[2])

with torch_default_dtype(torch.float64):
    model = GotenNet(
        dim=256,
        max_degree=2,
        depth=4,
        heads=2,
        dim_head=32,
        dim_edge_refinement=256,
        return_coors=True
    )

    R = random_rotation()

    atom_ids = torch.randint(0, 14, (1, 12))
    coors = torch.randn(1, 12, 3)
    adj_mat = torch.randint(0, 2, (1, 12, 12)).bool()

    # Original coordinates
    inv1, coors1 = model(atom_ids, coors=coors, adj_mat=adj_mat)

    # Rotated coordinates
    inv2, coors2 = model(atom_ids, coors=coors @ R, adj_mat=adj_mat)

    # Invariant features should be identical
    invariance_check = torch.allclose(inv1, inv2, atol=1e-5)
    print(f"Invariance preserved: {invariance_check}")  # True

    # Output coordinates should be equivariant (rotate together)
    equivariance_check = torch.allclose(coors1 @ R, coors2, atol=1e-5)
    print(f"Equivariance preserved: {equivariance_check}")  # True
```

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

### GotenNet with Cutoff Radius and Neighbor Selection

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

Configure GotenNet with `cutoff_radius` and `max_neighbors` for efficient processing of large molecular systems. This limits interactions to nearby atoms and caps the number of neighbors considered per atom.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=4,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    cutoff_radius=5.0,          # Only consider atoms within 5 Angstroms
    max_neighbors=32,           # Maximum neighbors per atom
    return_coors=True
)

batch_size = 1
num_atoms = 100  # Larger molecule

atom_ids = torch.randint(0, 14, (batch_size, num_atoms))
coors = torch.randn(batch_size, num_atoms, 3) * 10  # Spread coordinates

adj_mat = torch.randint(0, 2, (batch_size, num_atoms, num_atoms)).bool()

invariant_features, updated_coors = model(
    atom_ids,
    coors=coors,
    adj_mat=adj_mat
)

print(f"Invariant features shape: {invariant_features.shape}")  # [1, 100, 256]
print(f"Updated coordinates shape: {updated_coors.shape}")      # [1, 100, 3]
```

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

### Initialize and Use GotenNet Model

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

Initialize the GotenNet model with specified dimensions and depth, then use it for prediction. Recommended to set default dtype to float64 for training.

```python
import torch
torch.set_default_dtype(torch.float64) # recommended for equivariant network training

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim = 256,
    max_degree = 2,
    depth = 1,
    heads = 2,
    dim_head = 32,
    dim_edge_refinement = 256,
    return_coors = False
)

atom_ids = torch.randint(0, 14, (1, 12)) # negative atom indices will be assumed to be padding - length of molecule is thus `(atom_ids >= 0).sum(dim = -1)`
coors = torch.randn(1, 12, 3)
adj_mat = torch.randint(0, 2, (1, 12, 12)).bool()

invariant, coors_out = model(atom_ids, adj_mat = adj_mat, coors = coors)
```

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

### GotenNet with Variable Length Sequences (Negative IDs)

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

Handle batches with different molecule sizes by using negative atom IDs as padding. Negative IDs are treated as padding by the model.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=2,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    return_coors=True
)

batch_size = 2
max_atoms = 16

# Method 3: Using negative atom IDs as padding
atom_ids_padded = torch.randint(0, 14, (batch_size, max_atoms))
atom_ids_padded[0, 12:] = -1  # Pad first molecule
atom_ids_padded[1, 8:] = -1   # Pad second molecule

# Assuming coors is defined as in the previous example
# inv_feats, out_coors = model(atom_ids_padded, coors=coors)

# print(f"Output invariant shape: {inv_feats.shape}")  # [2, 16, 256]
```

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

### GotenNet Main Model Class Initialization and Usage

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

Initialize and use the main GotenNet model class with atom IDs. Ensure float64 dtype is set for training. The model processes atom IDs and coordinates to return invariant features and updated coordinates.

```python
import torch
torch.set_default_dtype(torch.float64)  # Recommended for equivariant network training

from gotennet_pytorch import GotenNet

# Initialize the model with basic configuration
model = GotenNet(
    dim=256,                    # Model dimension
    max_degree=2,               # Maximum degree of spherical harmonics (L_max)
    depth=4,                    # Number of transformer layers
    heads=2,                    # Number of attention heads
    dim_head=32,                # Dimension per attention head
    dim_edge_refinement=256,    # Edge refinement hidden dimension
    return_coors=True           # Return updated coordinates
)

# Prepare input data
batch_size = 2
num_atoms = 12

# Atom IDs: negative values are treated as padding
atom_ids = torch.randint(0, 14, (batch_size, num_atoms))

# 3D coordinates for each atom
coors = torch.randn(batch_size, num_atoms, 3)

# Adjacency matrix (optional, defines molecular bonds)
adj_mat = torch.randint(0, 2, (batch_size, num_atoms, num_atoms)).bool()

# Forward pass
invariant_features, updated_coors = model(
    atom_ids,
    coors=coors,
    adj_mat=adj_mat
)

print(f"Invariant features shape: {invariant_features.shape}")  # [2, 12, 256]
print(f"Updated coordinates shape: {updated_coors.shape}")      # [2, 12, 3]
```

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

### GotenNet with Custom Atom Embeddings

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

Initialize and use GotenNet accepting pre-computed atom embeddings instead of IDs by setting `accept_embed=True`. This allows integration with custom featurization pipelines.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=4,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    accept_embed=True,          # Accept embeddings instead of atom IDs
    return_coors=True
)

batch_size = 1
num_atoms = 12

# Pre-computed atom embeddings
atom_feats = torch.randn(batch_size, num_atoms, 256)

# 3D coordinates
coors = torch.randn(batch_size, num_atoms, 3)

# Adjacency matrix
adj_mat = torch.randint(0, 2, (batch_size, num_atoms, num_atoms)).bool()

# Forward pass with embeddings
invariant_features, updated_coors = model(
    atom_feats,
    coors=coors,
    adj_mat=adj_mat
)

print(f"Invariant features shape: {invariant_features.shape}")  # [1, 12, 256]
print(f"Updated coordinates shape: {updated_coors.shape}")      # [1, 12, 3]
```

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

### GotenNet with Variable Length Sequences (Mask)

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

Handle batches with different molecule sizes using an explicit mask tensor. The mask should be `True` for valid atoms and `False` for padding.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=2,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    return_coors=True
)

batch_size = 2
max_atoms = 16

# Method 1: Using explicit mask
atom_ids = torch.randint(0, 14, (batch_size, max_atoms))
coors = torch.randn(batch_size, max_atoms, 3)
mask = torch.tensor([
    [True] * 12 + [False] * 4,  # First molecule: 12 atoms
    [True] * 8 + [False] * 8    # Second molecule: 8 atoms
])

inv_feats, out_coors = model(atom_ids, coors=coors, mask=mask)
```

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

### GotenNet with Variable Length Sequences (Lengths)

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

Handle batches with different molecule sizes using a length tensor. The `lens` tensor specifies the number of valid atoms in each molecule.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=2,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    return_coors=True
)

batch_size = 2
max_atoms = 16

# Method 2: Using lengths
lens = torch.tensor([12, 8])  # Number of atoms per molecule
# Assuming atom_ids and coors are defined as in the previous example
# inv_feats, out_coors = model(atom_ids, coors=coors, lens=lens)
```

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

### GotenNet with Hyper-Connections

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

Configure the number of residual streams for hyper-connections to improve gradient flow in deep networks. `add_value_residual=True` enables value residual learning, and `layernorm_input=True` adds layernorm before the feedforward layer.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=8,                      # Deeper network
    heads=4,
    dim_head=32,
    dim_edge_refinement=256,
    num_residual_streams=4,       # Number of hyper-connection streams
    add_value_residual=True,      # Enable value residual learning
    return_coors=True,
    ff_kwargs=dict(
        layernorm_input=True      # Add layernorm before feedforward
    )
)

atom_ids = torch.randint(0, 14, (1, 12))
coors = torch.randn(1, 12, 3)
adj_mat = torch.randint(0, 2, (1, 12, 12)).bool()

invariant_features, updated_coors = model(
    atom_ids,
    coors=coors,
    adj_mat=adj_mat
)

print(f"Invariant features shape: {invariant_features.shape}")  # [1, 12, 256]
```

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

### GotenNet with Output Projection

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

Project the output invariant features to a different dimension using `proj_invariant_dim`. This is useful for task-specific heads. The model returns invariant features and updated coordinates.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=4,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    proj_invariant_dim=64,        # Project to 64-dimensional output
    return_coors=True
)

atom_ids = torch.randint(0, 14, (1, 12))
coors = torch.randn(1, 12, 3)

invariant_features, updated_coors = model(atom_ids, coors=coors)

print(f"Projected invariant shape: {invariant_features.shape}")  # [1, 12, 64]
```

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

### GotenNet with Invariant Full Attention

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

Enable full self-attention on invariant features alongside the equivariant attention mechanism. Use `invariant_attn_use_flash=True` for efficiency. This configuration returns only invariant features and high-degree tensors.

```python
import torch
torch.set_default_dtype(torch.float64)

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim=256,
    max_degree=2,
    depth=4,
    heads=2,
    dim_head=32,
    dim_edge_refinement=256,
    invariant_full_attn=True,       # Enable full attention on invariants
    invariant_attn_use_flash=True,  # Use Flash Attention for efficiency
    return_coors=False              # Only return invariant features
)

atom_ids = torch.randint(0, 14, (1, 12))
coors = torch.randn(1, 12, 3)
adj_mat = torch.randint(0, 2, (1, 12, 12)).bool()

# Returns invariant features and high-degree tensors (not coordinates)
invariant_features, high_degree_tensors = model(
    atom_ids,
    coors=coors,
    adj_mat=adj_mat
)

print(f"Invariant features shape: {invariant_features.shape}")  # [1, 12, 256]
print(f"Number of high-degree tensors: {len(high_degree_tensors)}")  # 2
```

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

### torch_default_dtype Context Manager

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

A context manager for temporarily setting PyTorch's default dtype. Essential for equivariant network training which benefits from float64 precision. The default dtype is restored upon exiting the context.

```python
import torch
from gotennet_pytorch import torch_default_dtype, GotenNet

# Default dtype is float32
print(f"Default dtype: {torch.get_default_dtype()}")  # torch.float32

# Temporarily switch to float64
with torch_default_dtype(torch.float64):
    print(f"Inside context: {torch.get_default_dtype()}")  # torch.float64

    model = GotenNet(
        dim=128,
        max_degree=2,
        depth=2,
        heads=2,
        dim_head=32,
        dim_edge_refinement=128,
        return_coors=True
    )

    atom_ids = torch.randint(0, 14, (1, 8))
    coors = torch.randn(1, 8, 3)  # Created as float64

    inv, out_coors = model(atom_ids, coors=coors)
    print(f"Output dtype: {inv.dtype}")  # torch.float64

# Back to default
print(f"After context: {torch.get_default_dtype()}")  # torch.float32
```

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