$ account google-deepmind
Sonnet is a library for building neural networks in TensorFlow 2, developed by DeepMind. It provides a modular abstraction (snt.Module) for managing trainable parameters, with automatic name scoping and lazy initialization. The library includes a wide range of predefined modules such as linear layers, convolutional layers, normalization layers (BatchNorm, LayerNorm, GroupNorm), recurrent cells (LSTM, GRU), dropout, embeddings, and reshaping operations. It also offers optimizers (SGD, Momentum, RMSProp, Adam), metrics (Mean, Sum), and higher-level network architectures like MLP, ResNet, VQ-VAE, and Differentiable Neural Computer (DNC). Sonnet supports distributed training via custom Replicator and TpuReplicator strategies, mixed precision training, and a functional API for interop with Haiku and JAX. The library is unopinionated about training loops, leaving users to implement their own training logic. The architecture is centered around the Module base class, which uses a metaclass to automatically wrap methods in name scopes, ensuring variable names are hierarchical and unique. Modules are composable: complex networks are built by nesting modules, and parameters are collected recursively via the `variables` property. Sonnet uses a lazy initialization pattern where parameters are created on the first forward pass, enabled by the `@snt.once` decorator. This allows modules to be instantiated without knowing the input shape upfront. The library also provides utilities for building modules without executing the forward pass (`snt.build`), applying modules across batch dimensions (`snt.BatchApply`), and deferring module construction (`snt.Deferred`). Key entry points include the top-level `sonnet/__init__.py` which re-exports the public API. Users typically import `sonnet as snt` and use `snt.Module`, `snt.Linear`, `snt.Conv2D`, etc. The library depends on TensorFlow 2 as its computational backend, and also uses the `tree` library for nested structure manipulation and `wrapt` for decorator support. Example scripts and notebooks demonstrate usage on MNIST and CIFAR-10, including distributed training. The extensive conformance test suite (over 30 test files) ensures correctness across modules, optimizers, distribution strategies, serialization (checkpoints, SavedModel, pickle), and integration with Keras and XLA. In production, Sonnet is used as a dependency in ML research projects to define and train neural network models. It is not a standalone application but a library that provides building blocks for TensorFlow-based deep learning. Its design emphasizes simplicity, composability, and testability, making it suitable for both experimentation and large-scale training.
AlphaFold3 is Google DeepMind's official command-line application for running AlphaFold 3 biomolecular structure prediction inference. The dominant character is a production-grade scientific application, though it is heavily intertwined with research data pipelines: it ingests a JSON job specification describing protein, RNA, DNA, and ligand chains, runs Jackhmmer/Nhmmer searches over large genetic databases, builds multiple sequence alignments, searches structural templates with HMMER, featurizes the input into JAX/Haiku tensors, samples structures with a diffusion model, and writes predicted structures as mmCIF files plus confidence JSON (pLDDT, PAE, pTM, ipTM). The codebase is not a training repository; it provides inference only and requires separately downloaded model parameters under restrictive legal terms. Architecturally, the system is a multi-stage pipeline driven by the `run_alphafold.py` entry point. Core modules include the data pipeline (`src/alphafold3/data/pipeline.py`), MSA/template tools (`src/alphafold3/data/tools/`), a structure representation layer (`src/alphafold3/structure/`), JAX geometric primitives (`src/alphafold3/jax/geometry/`), the neural network model (`src/alphafold3/model/`), and performance-critical C++ pybind11 extensions for mmCIF parsing, MSA feature extraction, and string/array operations. It depends on external binaries (Jackhmmer, Nhmmer, HMMER suite), RDKit for ligand conformer generation, JAX/Haiku for model execution, and ~630 GB of downloadable sequence and structure databases. The repository is actively maintained, with recent commits and CI running CPU-only data tests. Technical drift is low for a research release: dependencies are pinned via `pyproject.toml` and `uv.lock`, and active maintenance keeps it aligned with its ecosystem, though the C++ build, GPU-specific requirements, and external database URLs remain fragility points.
This repository reproduces the AlphaGeometry system (DD+AR and neural language model) from the Nature 2024 paper on automated theorem proving for olympiad-level geometry. It combines a symbolic deduction engine (DD) with algebraic reasoning (AR) and a transformer-based language model that suggests auxiliary constructions. The codebase includes problem files in a custom DSL, a large set of deduction rules, and a 150M-parameter transformer model. While not a production application, it provides a well-structured research codebase with clean module separation, comprehensive unit tests, and several reusable components such as a generic beam search implementation, algebraic reasoning tables, and proof traceback utilities. The repository is somewhat stale (last updated in 2023) and has many open issues; dependency pinning is strict but external dependencies like Meliad and a model checkpoint introduce reproducibility risks.
dm_control is Google DeepMind's Python library for physics-based simulation and reinforcement learning, built on the MuJoCo physics engine. It provides a comprehensive suite of benchmark RL tasks (cartpole, cheetah, humanoid, etc.), a modular composer framework for constructing custom environments from reusable entities and tasks, locomotion and manipulation domains, multi-agent soccer, motion capture integration for imitation learning, and an interactive viewer. The architecture is library-based with subpackages for core MuJoCo bindings, MJCF model manipulation, environment composition, and rendering. Key entry points include suite.load() for standard tasks, composer.Environment for custom environments, and viewer.launch() for interactive visualization. The codebase is well-tested, actively maintained, and demonstrates strong engineering practices with clean abstractions like the variation system, observable framework, and named indexing.
AlphaGenome is a Python client library and SDK for interacting with Google DeepMind's AlphaGenome cloud service, a deep learning model that predicts molecular properties (gene expression, splicing, chromatin accessibility, histone marks, contact maps) from DNA sequence. The library provides gRPC-based clients for sequence, interval, and variant prediction, along with precomputed variant scores via the Atlas service. It includes rich data containers (TrackData, JunctionData), modular visualization components (tracks, transcripts, sashimi plots), and interpretation tools like in-silico mutagenesis. The repository also contains a set of Colab notebooks demonstrating end-to-end workflows for variant scoring, haplotype analysis, splicing quantification, and batch processing. The architecture is a thin client layer over a remote API, with protobuf definitions for service contracts and tensor serialization. Core modules handle genomic data types (Interval, Variant, Transcript), gene annotation parsing, ontology mapping, and efficient tensor packing. The variant scoring framework is extensible via dataclass-based scorers (center mask, gene mask LFC, splicing, polyadenylation) and includes utilities to tidy AnnData results into pandas DataFrames. Visualization is component-based, allowing composition of tracks, contact maps, and annotations into publication-quality figures. While primarily an application library, the repo has a strong data research character due to its focus on genomic analysis and the included notebooks. The code is well-structured, with extensive unit tests and CI/CD pipelines. However, it is tightly coupled to the AlphaGenome API and protobuf definitions, limiting standalone reuse. The most reusable components are the data containers (TrackData, JunctionData), tensor packing utilities, and the variant scorer dataclass pattern, which could be extracted for other genomic scoring tasks.
This repository from Google DeepMind implements a family of point tracking models (TAP-Net, TAPIR, BootsTAP, TAPNext, TAPNext++, TRAJAN) and the associated TAP-Vid and TAPVid-3D benchmarks. It is primarily a research codebase for tracking arbitrary points in video, with both JAX/Haiku and PyTorch implementations. The code includes training pipelines (Jaxline), evaluation scripts, data generation for synthetic (Kubric) and real datasets (DAVIS, Kinetics, ADT, DriveTrack, Panoptic Studio), and interactive Colab demos. While it has application-like components such as live webcam demos and a VOT challenge tracker, the dominant character is research: the code is organized around benchmarking and model experimentation rather than production deployment. The repository is actively maintained with recent commits and provides extensive documentation for reproducing results.