---

name: sonnet description: "Use DeepMind Sonnet for TensorFlow 2 modules, layers, training loops, functional transforms, serialization, and distribution workflows." disable-model-invocation: true metadata: disco-role: operating license: Apache 2.0

Sonnet Repo Skill

Use this repo skill when a task involves DeepMind Sonnet (dm-sonnet, imported as sonnet or snt) for TensorFlow 2 neural-network research code: custom modules, built-in layers and nets, manual training loops, RNNs, functional transforms, checkpoints, SavedModel export, mixed precision, or TensorFlow distribution helpers.

Sonnet is intentionally small and unopinionated. It provides composable Module objects, built-in module families, and optimizers, but it does not provide a full training framework, dataset pipeline, experiment runner, or serving stack.

Quick install and import check

python -m pip install tensorflow dm-sonnet

import tensorflow as tf
import sonnet as snt
print("TensorFlow", tf.__version__)
print("Sonnet", snt.__version__)
print(snt.nets.MLP([16, 4])(tf.ones([2, 8])).shape)

Run scripts/check_sonnet_install.py when you need a no-download smoke check for importability, device visibility, lazy variables, a tiny optimizer step, and basic RNN construction.

Route by task

If the task asks about...	Read this sub-skill	Why
Custom snt.Module classes, @snt.once, lazy variables, names, Sequential, Deferred, BatchApply, or snt.build	sub-skills/module-authoring/SKILL.md	Owns Sonnet's core module programming model and composition contracts.
Built-in layers, normalization, initializers, metrics, regularizers, snt.nets.MLP, ResNet, Cifar10ConvNet, or VQ-VAE modules	sub-skills/layers-and-nets/SKILL.md	Owns constructor signatures, shape rules, state behavior, and layer/net smoke checks.
tf.GradientTape training loops, snt.optimizers.*, optimizer.apply, tiny local training checks, or metric accumulation	sub-skills/training-and-optimization/SKILL.md	Owns object-oriented Sonnet optimization and safe training-loop recipes.
RNNCore, LSTM, GRU, DeepRNN, dynamic_unroll, static_unroll, trainable state, or ConvLSTM	sub-skills/sequence-and-rnn/SKILL.md	Owns sequence shape/state conventions and recurrent troubleshooting.
snt.functional.variables, transform, transform_with_state, grad, value_and_grad, jit, device_put/get, or functional optimizers	sub-skills/functional-transforms/SKILL.md	Owns Sonnet's TensorFlow-based Haiku/JAX-like functional API.
TensorFlow checkpoints, SavedModel export/load, Keras/pickle caveats, XLA, mixed precision, snt.distribute.Replicator, TPU, or cross-replica BatchNorm	sub-skills/serialization-and-distribution/SKILL.md	Owns persistence, export, distribution, and backend limitations.

Root references

• references/package-overview.md summarizes Sonnet's public package mental model, dependencies, public imports, and common workflow boundaries.
• references/troubleshooting.md covers cross-cutting install/import, TensorFlow backend, and public-API routing failures before you enter a sub-skill-specific troubleshooting file.
• references/repo-provenance.md records the source snapshot used to build this skill.
• references/repo-routing-metadata.json is structured routing metadata consumed by the managed repo-skill importer.

Operating checklist

1. Import public symbols from sonnet (import sonnet as snt), not private implementation modules.
2. Build modules once with representative inputs before inspecting variables, exporting, or passing parameters to optimizers.
3. Assert small shapes at module boundaries: final feature dimension for Linear, channel dimension for convolution/normalization, and time-major sequence layout for RNN unroll helpers.
4. Use tf.GradientTape plus optimizer.apply(gradients, variables) for object-oriented Sonnet training loops, or use snt.functional only when the task explicitly wants stateless-style init/apply workflows.
5. Treat CUDA, TPU, XLA, and distributed execution as environment-specific. Do not claim accelerator verification unless the current TensorFlow runtime has been probed and the sub-skill's backend guidance is satisfied.

When not to use this skill

• The user is asking for general TensorFlow or Keras usage with no Sonnet API, module, optimizer, or distribution helper involved.
• The user wants to edit or maintain the Sonnet repository itself rather than use Sonnet as a package; route to a repository-maintenance workflow instead.
• The user needs a full experiment framework, trainer, data pipeline, model zoo, or serving platform that Sonnet does not provide.