Porting AlphaGenome to PyTorch

Danila Bredikhin, Alejandro Buendia, Martin Kjellberg, Christopher Zou, Xinming Tu, Anshul Kundaje · 2026-03-10

This post introduces alphagenome-pytorch, a faithful port of Google DeepMind's AlphaGenome model to PyTorch. We reproduce the full AlphaGenome architecture in PyTorch and release weights for fold-specific and distilled models. We also verify numerical equivalence of predictions across all output tracks with the original JAX checkpoint from DeepMind. Our package exposes a simple inference API that slots into any PyTorch project without requiring JAX, XLA, or TPU-specific tooling. In this post, we review the initial release for the package and walk through two use cases it enables: 1. Inference within existing PyTorch pipelines, which allows for genome-wide inference across tracks 2. Variant effect prediction and in silico mutagenesis (ISM) Contributions: Initial Development Team: Danila Bredikhin (Lead), Martin Kjellberg, Christopher Zou Finetuning and Validation: Danila Bredikhin, Alejandro Buendia, Xinming Tu Blog Post: Alejandro Buendia, Danila Bredikhin Code: alphagenome-pytorch