alphagenome-pt 0.3.0

An implementation of AlphaGenome in PyTorch.

AlphaGenome_PyTorch

*Image generated by Nano Banana 2*

Quick Start

From-Scratch

from alphagenome_pt import deepmind_model

model = deepmind_model()

Finetuned

Load the converted DeepMind state into a model. By default this downloads from Hugging Face cache if the state is not already local:

from alphagenome_pt import deepmind_model

model = deepmind_model(load_state=True)

Get Predictions/Embeddings

from alphagenome_pt import synthetic_batch
batch = synthetic_batch(model.metadata, batch_size=1, seq_len=model.max_seq_len)
predictions, embeddings = model(batch)

print(
  embeddings.embeddings_1bp.shape,
  embeddings.embeddings_128bp.shape,
  embeddings.embeddings_pair.shape,
)
print(predictions.keys())

For finetuning with custom downstream heads, build custom metadata:

from alphagenome_pt import (
    HeadName,
    deepmind_model,
    synthetic_batch,
    synthetic_metadata,
)

metadata = synthetic_metadata(
    heads=(HeadName.RNA_SEQ, HeadName.ATAC, HeadName.SPLICE_SITES_CLASSIFICATION),
    num_organisms=2,
    num_tracks=8,
)
# NOTE: can set heads=False to skip loading heads and/or organisms=False to skip loading organisms.
# They are prefix-loaded by default.
model = deepmind_model(
    metadata=metadata.metadata,
    load_state=True,
)
batch = synthetic_batch(metadata, batch_size=1, seq_len=8192)
loss, scalars, predictions = model.loss(batch)
loss.backward()

Metadata

Metadata defines organisms, output heads, track counts, masks, and track means. For examples and smoke tests, generate synthetic metadata:

from alphagenome_pt import HeadName, synthetic_metadata

metadata = synthetic_metadata(
    heads=(HeadName.RNA_SEQ, HeadName.CONTACT_MAPS),
    num_organisms=2,
    num_tracks=8,
)

To use the converted official DeepMind metadata:

from alphagenome_pt import deepmind_metadata

metadata_dict = deepmind_metadata()

To download all converted metadata files and all checkpoint folds ahead of time:

alphagenome-pt-download --local-dir checkpoints

Docs

The docs directory contains information on environment setup, explanations of the data structure and model architecture, and example trainings. To train on your own data, it's strongly recommended to read the model.md and data.md markdown files in the */AlphaGenome_PyTorch/docs/guides directory so that you understand the structure of data batches and metadata.

Environment

See */AlphaGenome_PyTorch/docs/environment for instructions on how to set up a UV environment to run AlphaGenome_PyTorch.

Official JAX Checkpoint Conversion

Most users should download the converted PyTorch checkpoint files from Hugging Face:

alphagenome-pt-download --local-dir checkpoints

The package also includes optional JAX-to-PyTorch conversion tooling for recreating the PyTorch checkpoint from the official public JAX weights:

pip install "alphagenome-pt[jax2pt]"
python -m alphagenome_pt.jax2pt.convert_state

By default this converts every supported official checkpoint. To convert only one checkpoint, pass --fold all_folds, --fold fold_0, --fold fold_1, --fold fold_2, or --fold fold_3; outputs default to alphagenome_{fold_name}.pt.

The converter also requires the public alphagenome_research checkout used by the official loader:

git clone https://github.com/google-deepmind/alphagenome_research.git
pip install -e alphagenome_research

Guides

See */AlphaGenome_PyTorch/docs/guides for explanations on the AlphaGenome model and its data structure (very helpful for understanding examples).

Examples

See */AlphaGenome_PyTorch/docs/examples for examples of:

• Masked Language Modeling (MLM) training (train_mlm.py)
• Training on Downstream Tasks (RNA-Seq, CAGE, ATAC, Splice Sites Classification/Usage/Junction) (train_downstream.py)
• MLM Pretraining --> Training on Downstream Tasks (train_downstream_from_pretrained.py)

Acknowledgements

This repository is a reimplementation of the AlphaGenome model in PyTorch, with an added option for Masked Language Modeling (MLM).

Within the alphagenome_pt directory, some components are direct ports of the released AlphaGenome code Link1 Link2 (licensed under Apache License 2.0), some are reimplementations based on pseudocode from the BioArXiV paper, and others are original additions (e.g., the MLM head). Attribution is made clear at the top of each file in the alphagenome_pt directory. The docs and tests directories are original work (with LLM coding assistance).

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.

Intended Audience

This intended audience of this repo is for those who might want to train the AlphaGenome architecture in PyTorch with maximal flexibility to hyperparameters and new organisms/contexts. You should be able to prepare a batch of tensor data and set up a train/val/test loop, but don't want the hassle of replicating the architecture.

Other Implementations

There is another AlphaGenome PyTorch implementation out here by Phillip Wang (a.k.a. LucidRains) which is quite good. The GitHub page is down as of March 2nd, 2026, but the PyPi package remains. The main advantages of that implementation (as of version 0.2.8) are in evaluation (loading the published weights and running variant scoring). The main advantage of this implementation is research training (an MLM head and track masks that can vary by batch in training). This implementation also has a .loss() function in the model to compute multi-resolution losses for you, and one head per task with dense weight tensor of shape [O, D, T] rather than separate weights tensors of shape [D, T] for each [organism x task], which is mathematically equivalent but more in-line with the original AlphaGenome implementation.

Reaching Out

Want a new feature or find a bug? Feel free to leave an issue on the GitHub repository.