AlphaGenome PyTorch

Inference, fine-tuning, and research training with AlphaGenome in PyTorch.

Installation · Quickstart · Documentation · Issues

alphagenome-pt implements AlphaGenome in PyTorch, including its prediction heads, losses, and PyTorch-converted checkpoint loading. It supports the published model and custom configurations for inference and training.

Get Started

Install the package:

python -m pip install alphagenome-pt

Load the published checkpoint and generate predictions:

import torch

from alphagenome_pt import deepmind_model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = deepmind_model(load_state=True, device=device)
model.eval()

sequence = "ACGT" * 2048  # 8,192 bp
predictions, embeddings = model.predict(
    sequence,
    organism_index=0,  # human in the published metadata
    return_embeddings=True,
)

for head, outputs in predictions.items():
    print(head)
    for name, value in outputs.items():
        print(f"  {name}: {tuple(value.shape)}")

atac
  scaled_predictions_1bp: (1, 8192, 256)
  predictions_1bp: (1, 8192, 256)
  scaled_predictions_128bp: (1, 64, 256)
  predictions_128bp: (1, 64, 256)
dnase
  scaled_predictions_1bp: (1, 8192, 384)
  predictions_1bp: (1, 8192, 384)
  scaled_predictions_128bp: (1, 64, 384)
  predictions_128bp: (1, 64, 384)
procap
  scaled_predictions_1bp: (1, 8192, 128)
  predictions_1bp: (1, 8192, 128)
  scaled_predictions_128bp: (1, 64, 128)
  predictions_128bp: (1, 64, 128)
cage
  scaled_predictions_1bp: (1, 8192, 640)
  predictions_1bp: (1, 8192, 640)
  scaled_predictions_128bp: (1, 64, 640)
  predictions_128bp: (1, 64, 640)
rna_seq
  scaled_predictions_1bp: (1, 8192, 768)
  predictions_1bp: (1, 8192, 768)
  scaled_predictions_128bp: (1, 64, 768)
  predictions_128bp: (1, 64, 768)
chip_tf
  scaled_predictions_128bp: (1, 64, 1664)
  predictions_128bp: (1, 64, 1664)
chip_histone
  scaled_predictions_128bp: (1, 64, 1152)
  predictions_128bp: (1, 64, 1152)
contact_maps
  predictions: (1, 4, 4, 28)
splice_sites_classification
  logits: (1, 8192, 5)
  predictions: (1, 8192, 5)
splice_sites_usage
  logits: (1, 8192, 734)
  predictions: (1, 8192, 734)
splice_sites_junction
  predictions: (1, 512, 512, 734)
  splice_site_positions: (1, 4, 512)
  splice_junction_mask: (1, 512, 512, 734)

print(embeddings.embeddings_1bp.shape)
print(embeddings.embeddings_128bp.shape)
print(embeddings.embeddings_pair.shape)

torch.Size([1, 8192, 1536])
torch.Size([1, 64, 3072])
torch.Size([1, 4, 4, 128])

Note

When the artifacts are not already cached, deepmind_model(load_state=True) downloads the published metadata and a roughly 1.8 GB checkpoint. See DeepMind Checkpoints for loading arguments and behavior.

Warning

Sequences must be at least 2,048 bp, divisible by 2,048, and no longer than model.max_seq_len. See Data and Metadata for the complete input contract.

Select Prediction Heads

Run only the heads needed for a prediction:

selected_heads = {"atac", "rna_seq"}

for head_name, head_metadata in model.metadata.metadata["heads"].items():
    head_metadata["enabled"] = head_name in selected_heads

predictions = model.predict(sequence, organism_index=0)

Disabled heads are skipped and omitted from predictions. This is especially useful for the memory-intensive splice_sites_junction head. Changing enabled does not rebuild the model or add heads that were not configured.

Model Construction

Choose a construction path:

Published Model

Load the published metadata, architecture, and checkpoint:

from alphagenome_pt import deepmind_model

published_model = deepmind_model(load_state=True)

Published Architecture with Custom Metadata

Define custom metadata, then flexibly load compatible organism and prediction-head state from the published checkpoint:

from alphagenome_pt import Metadata, deepmind_model

custom_metadata = Metadata({
    "organisms": ["human"],
    "heads": {
        "rna_seq": {
            "num_tracks": [2],
            "means": [[2.1, 0.8]],  # Per-track divisor in target scaling (1.0 has no effect)
        },
    },
})

adapted_model = deepmind_model(
    metadata=custom_metadata,
    load_state=True,
    organisms=True,
    heads=True,
)

Custom Architecture

Define custom metadata and initialize the model from scratch:

from alphagenome_pt import AlphaGenome, AlphaGenomeConfig, Metadata

custom_metadata = Metadata({
    "organisms": ["dragon"],
    "heads": {
        "rna_seq": {
            "num_tracks": [2],
            "means": [[2.1, 0.8]],  # Per-track divisor in target scaling (1.0 has no effect)
        },
    },
})

custom_model = AlphaGenome(
    AlphaGenomeConfig(
        max_seq_len=8_192,
        num_channels=96,
        transformer_layers=3,
        metadata=custom_metadata,
    )
)

Save and Load a Model

Save and load the model configuration, metadata, parameters, and persistent buffers in a directory:

from alphagenome_pt import AlphaGenome

model.save("checkpoints/model")
loaded_model = AlphaGenome.load("checkpoints/model", device=device)

See Model Construction, DeepMind Checkpoints, and Configuration for complete options.

Model Inputs and Forward Calls

Prediction and embedding calls accept several input forms. Built-in loss calls require a DataBatch with targets.

DNA Inputs

Pass raw strings or one-hot tensors directly, or provide either representation through DataBatch:

from alphagenome_pt import DataBatch, DNAOneHotEncoder

sequences = ["ACGT" * 2048, "TGCA" * 2048]
one_hot = DNAOneHotEncoder().encode(sequences)  # [2, 8192, 4]
batch = DataBatch(dna_sequence=sequences, organism_index=[0, 1])

# Raw DNA and organism indices passed separately
string_predictions = model.predict(sequences, organism_index=[0, 1])

# One-hot DNA and organism indices passed separately
one_hot_predictions = model.predict(one_hot, organism_index=[0, 1])

# DNA and organism indices read from DataBatch
batch_predictions = model.predict(batch)

A single raw sequence or [S, 4] tensor receives a leading batch dimension. One-hot channels are ordered A, C, G, T.

Organism Indices

Omit the index, share one index across the batch, or provide one per sequence:

import torch

sequences = ["ACGT" * 2048, "TGCA" * 2048]

# Omitted indices default every sequence to organism 0
default_predictions = model.predict(sequences)

# A scalar index is shared across the batch
shared_predictions = model.predict(sequences, organism_index=1)

# A length-B tensor selects one organism per sequence
per_sequence_predictions = model.predict(
    sequences,
    organism_index=torch.tensor([0, 1]),
)

The index may instead be stored in DataBatch.organism_index. Index meanings follow the order of metadata["organisms"].

Forward Calls

Choose the call based on its return value and gradient behavior:

Warning

During DDP training, run parameter-using forwards through ddp_model(...), including mode="embed" and mode="loss". Calling the underlying module bypasses DDP's hooks and can cause reducer errors when find_unused_parameters=True.

Use direct model calls to define custom objectives from differentiable predictions or embeddings.

For a target-bearing DataBatch, backpropagate the built-in loss total with:

from alphagenome_pt import LossOutput

# Contains a target for every enabled head
output: LossOutput = model(batch, mode="loss")
loss, tree = output.total, output.tree
loss.backward()

Note

The built-in loss method remains under active development and testing.

See Training, Losses and the Metric Tree, and Predictions and Embeddings for complete behavior.

Command-Line Interface

Inspect the install:

alphagenome-pt --version
alphagenome-pt --help

Download the converted all-folds checkpoint and metadata to a local directory:

alphagenome-pt download --local-dir checkpoints --fold all_folds

Load the downloaded artifacts into a model:

model = deepmind_model(
    load_state=True,
    local_dir="checkpoints",
    device=device,
)

See the Command-Line Interface for all download options.

Learn More

• Understand: Data and Metadata · Architecture
• Construct: Model Construction · Checkpoints
• Run and Train: Running the Model

Development and Support

git clone https://github.com/RylieWeaver/AlphaGenome_PyTorch.git
cd AlphaGenome_PyTorch
python -m pip install -e ".[dev]"
python -m pytest tests

See Development Setup for contributor instructions. To report a bug or request a feature, open an issue.

Citation

If AlphaGenome PyTorch supports your work, cite the published AlphaGenome paper and link this repository:

Avsec, Ž., Latysheva, N., Cheng, J. et al. Advancing regulatory variant effect prediction with AlphaGenome. Nature 649, 1206–1218 (2026). https://doi.org/10.1038/s41586-025-10014-0

Use the Nature article when citing AlphaGenome. The original bioRxiv entry is included below for historical reference.

@article{avsec_advancing_2026,
	title = {Advancing regulatory variant effect prediction with {AlphaGenome}},
	volume = {649},
	issn = {1476-4687},
	url = {https://www.nature.com/articles/s41586-025-10014-0},
	doi = {10.1038/s41586-025-10014-0},
	number = {8099},
	journal = {Nature},
	publisher = {Nature Publishing Group},
	author = {Avsec, Žiga and Latysheva, Natasha and Cheng, Jun and Novati, Guido and Taylor, Kyle R. and Ward, Tom and Bycroft, Clare and Nicolaisen, Lauren and Arvaniti, Eirini and Pan, Joshua and Thomas, Raina and Dutordoir, Vincent and Perino, Matteo and De, Soham and Karollus, Alexander and Gayoso, Adam and Sargeant, Toby and Mottram, Anne and Wong, Lai Hong and Drotár, Pavol and Kosiorek, Adam and Senior, Andrew and Tanburn, Richard and Applebaum, Taylor and Basu, Souradeep and Hassabis, Demis and Kohli, Pushmeet},
	month = jan,
	year = {2026},
	pages = {1206--1218},
}

@article{avsec_alphagenome_2025,
	title = {{AlphaGenome}: advancing regulatory variant effect prediction with a unified {DNA} sequence model},
	url = {https://www.biorxiv.org/content/early/2025/06/27/2025.06.25.661532},
	doi = {10.1101/2025.06.25.661532},
	journal = {bioRxiv},
	publisher = {Cold Spring Harbor Laboratory},
	author = {Avsec, Žiga and Latysheva, Natasha and Cheng, Jun and Novati, Guido and Taylor, Kyle R. and Ward, Tom and Bycroft, Clare and Nicolaisen, Lauren and Arvaniti, Eirini and Pan, Joshua and Thomas, Raina and Dutordoir, Vincent and Perino, Matteo and De, Soham and Karollus, Alexander and Gayoso, Adam and Sargeant, Toby and Mottram, Anne and Wong, Lai Hong and Drotár, Pavol and Kosiorek, Adam and Senior, Andrew and Tanburn, Richard and Applebaum, Taylor and Basu, Souradeep and Hassabis, Demis and Kohli, Pushmeet},
	year = {2025},
}

Acknowledgements

This implementation bases its model components off of AlphaGenome package and AlphaGenome research code repositories. Some components are direct ports of the code, others are paper-based reimplementations, and some are original additions (e.g. the Masked Language Modeling Head). Source files identify applicable provenance.

We acknowledge the independent genomicsxai/alphagenome-pytorch implementation from the Kundaje Lab at Stanford and the authors named in its copyright notice: Danila Bredikhin, Martin Kjellberg, Christopher Zou, Alejandro Buendia, Xinming Tu, and Anshul Kundaje.

We also acknowledge earlier independent AlphaGenome PyTorch work by Phil Wang, Miquel Anglada-Girotto, and Xinming Tu.

This package was developed with assistance from LLM coding agents, including OpenAI Codex. The repository banner was generated with Nano Banana 2.

Developing this project used resources of the Oak Ridge Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC05-00OR22725.

License and Model Terms

Important

AlphaGenome PyTorch is an independent research implementation, not an official Google DeepMind package. It has not been designed or validated for direct clinical use.

Portions of this project are derived from Google DeepMind's AlphaGenome research code, which is licensed under the Apache License 2.0:

Copyright 2026 Google LLC

Google DeepMind's published AlphaGenome model parameters, outputs produced from them, and related derivatives remain subject to the AlphaGenome Model Terms, including restrictions on commercial use.

The AlphaGenome PyTorch repository and its alphagenome-pt Python package are available under the Apache License 2.0:

Copyright 2026 Rylie Weaver, Gomathi Lakshmanan, and John Lagergren