google-deepmind/sonnet: a TensorFlow 2 module library for research code

What Sonnet 2 is for, and who it is actually for

Sonnet is a library built on top of TensorFlow 2 that provides what its README calls "simple, composable abstractions for machine learning research." The abstraction is a single class, snt.Module. A module can hold references to parameters, to other modules, and to methods that apply some function to the input. Sonnet ships predefined modules such as snt.Linear, snt.Conv2D and snt.BatchNorm, plus a few predefined networks such as snt.nets.MLP, and the README explicitly encourages users to build their own.

The intended audience is narrow and stated plainly. Sonnet was designed and built by researchers at DeepMind, and the README says the team finds it a successful abstraction for its organization. That is a research-lab framing, not a product framing. If your work involves un/supervised learning or reinforcement learning experiments where you want to swap layer implementations quickly, the module model fits. If you want a framework that decides how training happens, Sonnet is the wrong shape, and the README says so: it does not ship with a training framework, and users are encouraged to build their own or adopt those built by others.

The unopinionated stance cuts both ways. Modules are designed to be self contained and entirely decoupled from one another, so you can compose them without inheriting a lifecycle you did not ask for. You also inherit the work of wiring optimizers, metrics, checkpointing schedules and distributed strategies yourself. For a research group with existing TensorFlow infrastructure that is a feature. For a small team trying to get a first model into production, it is unpaid work.

How snt.Module works: lazy parameters and name scoping

The mechanism worth understanding before adopting Sonnet is lazy parameter creation. Most modules create their parameters the first time they are called with some input, because in most cases the shape of the parameters is a function of the input shape. The README's MyLinear example makes this concrete: the constructor stores output_size, and an _initialize method decorated with @snt.once builds the weight and bias variables on the first call.

That decorator is the whole trick. You call the module, the initializer runs once, and subsequent calls reuse the variables. It also means a module has no parameters until you have fed it a tensor, which affects how you inspect or restore a model. Sonnet exposes two properties for that. The variables property returns all tf.Variables referenced by the module, and the README warns that this includes non-parameter state, such as the state held by metrics in snt.BatchNorm. The trainable_variables property returns the subset TensorFlow marks as trainable, and the README says this is probably what you want to pass to an optimizer.

Calling a module also enters its name scope, so variables get a prefix like my_linear/w:0 and operations appear grouped in TensorBoard. Subclassing snt.Module gives you a default __repr__ that prints constructor arguments, which matters more than it sounds when you are debugging a deep stack of composed modules. None of this is novel machinery; it is a consistent convention applied across a library, and consistency is the actual product.

Installing Sonnet and building a first MLP

Installation is two pip commands from the README. The first pulls TensorFlow 2 and TensorFlow Probability, the second installs the library under its PyPI name dm-sonnet.

pip install tensorflow tensorflow-probability
pip install dm-sonnet

The README then suggests verifying the install by importing both packages and printing their versions. If the imports succeed you should see the TensorFlow version and the Sonnet version printed to stdout.

import tensorflow as tf
import sonnet as snt

print("TensorFlow version {}".format(tf.__version__))
print("Sonnet version {}".format(snt.__version__))

A first real use is an MLP built from snt.Sequential, which calls a sequence of modules and passes the output of each as the input to the next. Here snt.Linear and tf.nn.relu are mixed freely, which is the composability claim in practice.

mlp = snt.Sequential([
    snt.Linear(1024),
    tf.nn.relu,
    snt.Linear(10),
])
logits = mlp(tf.random.normal([batch_size, input_size]))

Because parameters are created on first call, mlp.variables is empty until that call happens. After the call, mlp.trainable_variables holds the weights and biases you would hand to an optimizer. The README also links Colab notebooks for predicting MNIST with an MLP, training a little GAN on MNIST and distributed training with snt.distribute; those are the fastest way to see a complete loop rather than a fragment.

Serialization is where the design shows its edges

Sonnet supports several serialization formats and is candid about the weakest one. Pickle is described as the simplest format supported, and built in modules are tested to save and load via pickle in the same Python process. The README then discourages it: pickle is not well supported by many parts of TensorFlow and in the team's experience can be quite brittle. That is an unusually direct warning in a project README, and it should be read as a constraint rather than a footnote.

The recommended path is TensorFlow checkpointing, which Sonnet is designed to work with. The README's example builds a save_prefix from a checkpoint root and name, notes that the module can be anything extending snt.Module, and shows a Checkpoint object being constructed. The README text available here stops mid-example, so the exact restore call is not documented in what I can see; check the TensorFlow checkpointing guide linked from that section, which the README cites as a reference.

The practical consequence is that lazy initialization and checkpointing interact. A module that has never been called has no variables to restore into, so the order of construction, first call and restore matters. Sonnet handles this cleanly when you follow the pattern, but the pattern is not optional. If your pipeline constructs modules and immediately tries to restore, you will hit an empty-variable situation rather than a helpful error. That is a design consequence of lazy parameters, not a bug, and it is the first thing to verify in your own code.

Where Sonnet is the wrong tool

The clearest limitation is stated by the project itself: Sonnet does not ship a training framework. There is no fit loop, no metrics aggregation layer, no learning-rate schedule manager, no data pipeline. If your team's bottleneck is getting a model trained end to end with minimal code, Sonnet adds a layer rather than removing one, and you will write the same loop you would have written against raw TensorFlow 2, only with snt.Module objects inside it.

The second limitation is the TensorFlow 2 dependency. Sonnet 2 is built on top of TensorFlow 2 and the install instructions pull tensorflow and tensorflow-probability. If your stack is on PyTorch, or on a TensorFlow 1 codebase, Sonnet 2 is not a migration path. The v2 branch name and the v2.0.0 release in 2020 mark the break from the original Sonnet; there is no indication in the README that the v1 API is supported here.

The third is maturity signalling from the packaging itself. setup.py classifies the project as Development Status :: 3 - Alpha despite the library being used inside DeepMind. That classifier is a statement about API stability expectations, and it is worth weighing against the fact that the most recent release listed is v2.0.2 from January 2024, with v2.0.0 dating to March 2020. A long gap between major versions plus an alpha classifier means you should pin versions and read release notes rather than assume the surface is frozen.

Sonnet against Keras, which ships with TensorFlow

The obvious alternative is tf.keras, which is bundled with TensorFlow and occupies the same conceptual slot: layers, models, and a fit loop. The difference in approach is what each library owns. Keras owns the training step. It gives you compile and fit, callbacks, metric objects and a serialization format designed for deployment, and it makes the common case short. Sonnet owns nothing beyond the module abstraction and expects you to write the training step.

That difference shows up in composition. In Sonnet, modules are designed to be self contained and entirely decoupled from one another, and the README frames this as being extremely unopinionated about how you use them. You can put a raw tf.nn.relu inside an snt.Sequential alongside snt.Linear, as the README example does, without wrapping it in a layer class. In Keras the equivalent requires a Lambda layer or a functional API call, because the framework wants every step to be a layer it can inspect and serialize.

The trade is inspection against freedom. Keras can serialize a whole model graph because it constrains what a model is. Sonnet's modules are easier to rearrange inside a custom research loop, and correspondingly harder to hand to a serving stack that expects a Keras artifact. If you are choosing today, pick Keras when you want the framework to own training, and Sonnet when you already have training code and want a cleaner layer vocabulary inside it. There is no version of this where Sonnet replaces Keras wholesale, and the README does not claim otherwise.

Maintenance, licensing and upgrade cost

The repository is not archived, and the last push was on 2026-07-07, which is recent enough that the codebase is being touched. That said, the release list tells a different story about the API: v2.0.0 landed in March 2020, and the two releases after it, v2.0.1 and v2.0.2, both landed on 2024-01-02. Two patch releases on the same day suggests packaging or compatibility fixes rather than new surface area. Anyone planning a dependency on Sonnet should treat the module API as stable by neglect rather than by policy, and pin the version in requirements.

Upgrade cost is mostly TensorFlow's, not Sonnet's. The install instructions require TensorFlow 2 and tensorflow-probability, so a TensorFlow major bump is the event that will force you to revisit Sonnet. The runtime dependencies listed in requirements.txt are small and ordinary: absl-py, numpy, dm-tree, wrapt and tabulate. None of those are likely to pin you to an old TensorFlow, which is good news for the upgrade path.

Licensing is Apache-2.0, stated in the repository LICENSE file and in setup.py, with the PyPI package name dm-sonnet. Apache-2.0 is a permissive licence with a patent grant and requires attribution and notice retention. It does not impose copyleft obligations on your own code. This is a description of the licence text, not legal advice; if you are redistributing Sonnet inside a product, have your own counsel read the NOTICE requirements. The package metadata also declares requires_python of at least 3.6, which is only relevant if you are maintaining something very old.

Editorial conclusion

Adopt Sonnet if you already write TensorFlow 2 training loops and want small, composable layer objects with predictable variable handling; the pip install is two commands and the module model is easy to reason about. Skip it if you want a framework that also owns training, data pipelines and checkpoint policy, or if you are not on TensorFlow 2 at all. Before committing, verify that the modules you need exist in your installed version, check that snt.Module subclasses pickle or checkpoint the way your pipeline expects, and read the release notes for v2.0.2 to see what changed since v2.0.0. The project's own README points at Colab notebooks for MNIST and CIFAR-10, so run one of those against your TensorFlow build first.

Frequently asked questions

Official sources

1. google-deepmind/sonnet on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases