Sonnet 2: A Module Abstraction for TensorFlow 2 Research Code

What Sonnet Actually Adds to TensorFlow 2

TensorFlow 2 gives you layers, variables, and an eager execution model, but it does not prescribe how you organize a model. Sonnet's answer is a single class, snt.Module, described in the README as the concept the whole programming model is centered around. A module can hold references to parameters, to other modules, and to methods that apply some function to an input. Everything else in the library is either a predefined module (snt.Linear, snt.Conv2D, snt.BatchNorm) or a predefined composition of modules (snt.nets.MLP). The target user is a researcher at an organization like DeepMind who wants composable pieces and no framework telling them how to train. The README says as much directly: Sonnet does not ship with a training framework, and users are encouraged to build their own or adopt one built by others. That is the whole pitch, and it is a narrow one. If you want a trainer, an experiment config system, or a model zoo, Sonnet is not offering any of that.

Lazy Parameter Creation and the @snt.once Decorator

The mechanism that distinguishes Sonnet from plain Keras layers is when parameters come into existence. Most Sonnet modules create their parameters the first time they are called with an input, because in most cases the shape of the parameters is a function of the input shape. The README's MyLinear example makes this explicit: the _initialize method is decorated with @snt.once and creates self.w and self.b from x.shape[1] and output_size. The __call__ method then invokes _initialize before doing the matmul. This means you can construct a module without knowing the input dimension, which is convenient for research code where shapes change between experiments. It also means the first call has a side effect. If you call the module inside a tf.function with a traced shape that later changes, the initialization path and the traced graph interact in ways the README does not discuss. The documentation presents lazy initialization as the normal case and does not walk through what happens when the same module is first called with two different input shapes. That is a gap worth testing yourself before you build on it.

Variables, Trainable Variables, and What You Pass to an Optimizer

Sonnet modules expose two properties. The variables property returns all tf.Variable objects referenced by the module, and the trainable_variables property returns the subset marked trainable. The README is careful about the distinction and gives a concrete reason: tf.Variable is not only used for model parameters, it also holds state in metrics used by snt.BatchNorm. If you hand the full variables list to an optimizer, you are handing it batch norm statistics too, which are updated by a different mechanism. The README's own guidance is that trainable_variables is probably what you want to pass to an optimizer. This is a small design decision with real consequences, and it is the kind of thing that catches people migrating from frameworks where the optimizer is wired up for you. There is no optimizer in Sonnet. You collect the list and you decide what to do with it.

Getting It Running: Install, Import, and a First Module

Installation is two pip commands from the README. First pip install tensorflow tensorflow-probability, then pip install dm-sonnet. The package name on PyPI is dm-sonnet, not sonnet, which is worth noting because the import name is sonnet. Verification is a short script: import tensorflow as tf, import sonnet as snt, then print tf.__version__ and snt.__version__. A minimal model uses snt.Sequential with snt.Linear and tf.nn.relu, for example snt.Sequential([snt.Linear(1024), tf.nn.relu, snt.Linear(10)]), called as mlp(tf.random.normal([batch_size, input_size])). To write your own module you subclass snt.Module, call super().__init__(name=name), decorate an _initialize method with @snt.once, and define __call__. Subclassing gives you a generated __repr__ that prints constructor arguments, which the README shows as MyLinear(output_size=10), and it enters the module name scope on every method call so TensorBoard groups operations under the module name. Those two properties are the practical payoff of using snt.Module instead of a bare Python class.

Serialization: TensorFlow Checkpoints Over Pickle

Sonnet supports multiple serialization formats and the README is blunt about which one to avoid. Pickle is described as the simplest format and all built-in modules are tested to round-trip through it in the same Python process, but the README says it discourages pickle in general because it is not well supported by many parts of TensorFlow and can be brittle. The recommended path is TensorFlow checkpointing, which saves parameter values periodically during training so a crashed or stopped job can resume. The README's example builds a save_prefix from a checkpoint_root and a checkpoint_name, then constructs a Checkpoint object around a module that extends snt.Module. The material here is truncated mid-example, so the exact Checkpoint constructor arguments are not visible in what was supplied. If you are evaluating Sonnet for a long-running training job, read the full checkpointing section in the documentation at sonnet.readthedocs.io rather than relying on the README excerpt, because the truncated portion is exactly the part that determines whether your restore path works.

Where Sonnet Is the Wrong Tool

The README states plainly that Sonnet does not ship with a training framework. That is a design choice, not an oversight, but it has a cost. If you are a small team without an existing training harness, adopting Sonnet means writing the loop, the checkpoint cadence, the logging, and the distributed strategy wiring yourself, or pulling in another library to do it. The README points to an snt.distribute example for distributed CIFAR-10 training, so distribution support exists, but it is demonstrated rather than abstracted into a trainer. A second limitation is release cadence. The v2 branch's most recent releases are v2.0.1 and v2.0.2, both dated 2024-01-02, and v2.0.0 dates to 2020-03-27. The last push to the repository is 2026-07-07, so the project is not abandoned, but the tagged releases are sparse. If your project needs a library with frequent versioned releases and a documented deprecation policy, Sonnet's release history is thin evidence that you will get one. Neither of these is a defect in the code. They are reasons the library fits some teams and not others.

How Sonnet Differs from Keras and PyTorch

The nearest comparison is tf.keras, which ships with TensorFlow and includes layers, a fit/evaluate training API, model saving, and callbacks. Keras decides more for you: it manages parameter creation at build time, it owns the training loop, and it has a serialization format tied to the framework. Sonnet decides almost nothing for you. Parameters appear on first call rather than at a build step, there is no fit method, and serialization is delegated to TensorFlow's checkpoint machinery. The trade is control against convenience. If your research involves nonstandard parameter shapes, custom initialization tied to input dimensions, or a training procedure that does not map onto fit, Sonnet's model is a better fit than Keras. If your work is standard supervised training on fixed-shape inputs, Keras will get you there with less code. The other comparison worth naming is PyTorch, which has its own module abstraction and its own eager execution model. Choosing between them is mostly a choice of ecosystem and team familiarity, not of module design, since both let you define a class with parameters and a forward pass. Sonnet's specific contribution is that it keeps that abstraction inside TensorFlow 2 while staying unopinionated about everything around it.

Maintenance Cost and Licence

Sonnet is licensed under Apache-2.0, which permits commercial and private use and requires that you retain the licence and attribution notices when you redistribute. This is a permissive licence and does not impose copyleft obligations on your own code. It is not legal advice; read the LICENSE file in the repository if the terms matter to your organization. On maintenance: the dependency surface is TensorFlow 2 plus tensorflow-probability, both of which move quickly. The README's install line pins nothing, so pip will resolve the latest versions of each, and a TensorFlow major version bump is the most likely thing to break an existing Sonnet codebase. The @snt.once decorator and the lazy initialization pattern are the parts most exposed to changes in how TensorFlow traces and retraces functions. If you adopt Sonnet, pin your TensorFlow and dm-sonnet versions together in your requirements file rather than tracking latest, and re-run your own module tests after any TensorFlow upgrade. The library's own test suite is not a substitute for testing your subclasses, because your subclasses are where the lazy initialization logic lives.

Editorial conclusion

Adopt Sonnet if your team already writes TensorFlow 2 code, wants explicit control over when parameters are created, and is comfortable supplying its own training loop. Do not adopt it if you need a bundled trainer, a large model zoo, or a framework that is actively releasing new versions. Before committing, verify that the dm-sonnet release you install matches your TensorFlow version, that your checkpointing path uses the TensorFlow checkpoint format rather than pickle, and that your own snt.Module subclasses behave correctly under the @snt.once decorator.

Official sources

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