$ show google-deepmind/sonnet
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.
The LLM identified 4 files worth highlighting for reuse:
Provides the `@once` decorator that ensures a method runs only once per instance, with subsequent calls returning the cached result. This is used extensively in Sonnet for lazy parameter creation (e.g., weight matrices are created on the first forward pass). The decorator is self-contained, thread-safe, and can be extracted for use in any Python project that needs one-time initialization of instance attributes.
Implements `BatchApply`, a module that flattens leading dimensions of the input tensor, applies a wrapped submodule, and restores the original shape. This is useful for applying a module designed for single examples across a batch or time dimension. The pattern is generic and can be reused to apply any callable across arbitrary batch dimensions.
Provides the `build()` function that forces a module to create its variables by calling it with a dummy input, without requiring the full forward pass. This is useful for inspecting module parameters before training or for initializing variables on TPU. The function is simple, well-documented, and can be adapted for any framework that uses lazy parameter creation.
Defines a programmatic descriptor system for all Sonnet modules, specifying input shapes, dtypes, and expected output properties. This is used to generate comprehensive tests (e.g., for checkpointing, XLA, Keras integration) without writing individual test cases. The pattern is highly reusable for any module library that needs exhaustive testing with minimal boilerplate.