Package: dm-haiku
Haiku:SonnetforJAXOverview|Why Haiku?|Quickstart|Installation|Examples|User manual|Documentation|Citing Haiku[!IMPORTANT]
📣As of July 2023Google DeepMindrecommends that new projects adoptFlaxinstead of Haiku.Flaxis a neural network library originally
developed byGoogle Brainand now byGoogle DeepMind.📣At the time of writingFlaxhas superset of the features available in Haiku,
alargerandmore activedevelopment team and
more adoption with users outside of Alphabet.Flaxhasmore extensive documentation,examplesand anactive communitycreating end
to end examples.Haiku will remain best-effort supported, however the project will entermaintenance mode, meaning
that development efforts will be focussed on bug fixes and compatibility with
new releases of JAX.New releases will be made to keep Haiku working with newer versions of Python
andJAX, however we will not be adding (or accepting PRs for) new features.We have significant usage of Haiku internally atGoogle DeepMindand
currently plan to support Haiku in this mode indefinitely.What is Haiku?Haiku is a toolFor building neural networksThink: "SonnetforJAX"Haiku is a simple neural network library forJAXdeveloped by some of the
authors ofSonnet, a neural network library forTensorFlow.Documentation on Haiku can be found athttps://dm-haiku.readthedocs.io/.Disambiguation:if you are looking for Haiku the operating system then
please seehttps://haiku-os.org/.OverviewJAXis a numerical computing library that combines NumPy, automatic
differentiation, and first-class GPU/TPU support.Haiku is a simple neural network library for JAX that enables users to use
familiarobject-oriented programming modelswhile allowing full access to
JAX's pure function transformations.Haiku provides two core tools: a module abstraction,hk.Module, and a simple
function transformation,hk.transform.hk.Modules are Python objects that hold references to their own parameters,
other modules, and methods that apply functions on user inputs.hk.transformturns functions that use these object-oriented, functionally
"impure" modules into pure functions that can be used withjax.jit,jax.grad,jax.pmap, etc.Why Haiku?There are a number of neural network libraries for JAX. Why should you choose
Haiku?Haiku has been tested by researchers at DeepMind at scale.DeepMind has reproduced a number of experiments in Haiku and JAX with relative
ease. These include large-scale results in image and language processing,
generative models, and reinforcement learning.Haiku is a library, not a framework.Haiku is designed to make specific things simpler: managing model parameters
and other model state.Haiku can be expected to compose with other libraries and work well with the
rest of JAX.Haiku otherwise is designed to get out of your way - it does not define custom
optimizers, checkpointing formats, or replication APIs.Haiku does not reinvent the wheel.Haiku builds on the programming model and APIs of Sonnet, a neural network
library with near universal adoption at DeepMind. It preserves Sonnet'sModule-based programming model for state management while retaining access
to JAX's function transformations.Haiku APIs and abstractions are as close as reasonable to Sonnet. Many users
have found Sonnet to be a productive programming model in TensorFlow; Haiku
enables the same experience in JAX.Transitioning to Haiku is easy.By design, transitioning from TensorFlow and Sonnet to JAX and Haiku is easy.Outside of new features (e.g.hk.transform), Haiku aims to match the API of
Sonnet 2. Modules, methods, argument names, defaults, and initialization
schemes should match.Haiku makes other aspects of JAX simpler.Haiku offers a trivial model for working with random numbers. Within a
transformed function,hk.next_rng_key()returns a unique rng key.These unique keys are deterministically derived from an initial random key
passed into the top-level transformed function, and are thus safe to use with
JAX program transformations.QuickstartLet's take a look at an example neural network, loss function, and training
loop. (For more examples, see ourexamples directory.
TheMNIST exampleis a good place to start.)importhaikuashkimportjax.numpyasjnpdefsoftmax_cross_entropy(logits,labels):one_hot=jax.nn.one_hot(labels,logits.shape[-1])return-jnp.sum(jax.nn.log_softmax(logits)*one_hot,axis=-1)defloss_fn(images,labels):mlp=hk.Sequential([hk.Linear(300),jax.nn.relu,hk.Linear(100),jax.nn.relu,hk.Linear(10),])logits=mlp(images)returnjnp.mean(softmax_cross_entropy(logits,labels))loss_fn_t=hk.transform(loss_fn)loss_fn_t=hk.without_apply_rng(loss_fn_t)rng=jax.random.PRNGKey(42)dummy_images,dummy_labels=next(input_dataset)params=loss_fn_t.init(rng,dummy_images,dummy_labels)defupdate_rule(param,update):returnparam-0.01*updateforimages,labelsininput_dataset:grads=jax.grad(loss_fn_t.apply)(params,images,labels)params=jax.tree_util.tree_map(update_rule,params,grads)The core of Haiku ishk.transform. Thetransformfunction allows you to
write neural network functions that rely on parameters (here the weights of theLinearlayers) without requiring you to explicitly write the boilerplate
for initialising those parameters.transformdoes this by transforming the
function into a pair of functions that arepure(as required by JAX)initandapply.initTheinitfunction, with signatureparams = init(rng, ...)(where...are
the arguments to the untransformed function), allows you tocollectthe
initial value of any parameters in the network. Haiku does this by running your
function, keeping track of any parameters requested throughhk.get_parameter(called by e.g.hk.Linear) and returning them to you.Theparamsobject returned is a nested data structure of all the
parameters in your network, designed for you to inspect and manipulate.
Concretely, it is a mapping of module name to module parameters, where a module
parameter is a mapping of parameter name to parameter value. For example:{'linear': {'b': ndarray(..., shape=(300,), dtype=float32),
            'w': ndarray(..., shape=(28, 300), dtype=float32)},
 'linear_1': {'b': ndarray(..., shape=(100,), dtype=float32),
              'w': ndarray(..., shape=(1000, 100), dtype=float32)},
 'linear_2': {'b': ndarray(..., shape=(10,), dtype=float32),
              'w': ndarray(..., shape=(100, 10), dtype=float32)}}applyTheapplyfunction, with signatureresult = apply(params, rng, ...), allows
you toinjectparameter values into your function. Wheneverhk.get_parameteris called, the value returned will come from theparamsyou
provide as input toapply:loss=loss_fn_t.apply(params,rng,images,labels)Note that since the actual computation performed by our loss function doesn't
rely on random numbers, passing in a random number generator is unnecessary, so
we could also pass inNonefor therngargument. (Note that if your
computationdoesuse random numbers, passing inNoneforrngwill cause
an error to be raised.) In our example above, we ask Haiku to do this for us
automatically with:loss_fn_t=hk.without_apply_rng(loss_fn_t)Sinceapplyis a pure function we can pass it tojax.grad(or any of JAX's
other transforms):grads=jax.grad(loss_fn_t.apply)(params,images,labels)TrainingThe training loop in this example is very simple. One detail to note is the use
ofjax.tree_util.tree_mapto apply thesgdfunction across all matching
entries inparamsandgrads. The result has the same structure as the
previousparamsand can again be used withapply.InstallationHaiku is written in pure Python, but depends on C++ code via JAX.Because JAX installation is different depending on your CUDA version, Haiku does
not list JAX as a dependency inrequirements.txt.First, followthese instructionsto install JAX with the relevant accelerator support.Then, install Haiku using pip:$pipinstallgit+https://github.com/deepmind/dm-haikuAlternatively, you can install via PyPI:$pipinstall-Udm-haikuOur examples rely on additional libraries (e.g.bsuite). You can install the full set of additional requirements using pip:$pipinstall-rexamples/requirements.txtUser manualWriting your own modulesIn Haiku, all modules are a subclass ofhk.Module. You can implement any
method you like (nothing is special-cased), but typically modules implement__init__and__call__.Let's work through implementing a linear layer:classMyLinear(hk.Module):def__init__(self,output_size,name=None):super().__init__(name=name)self.output_size=output_sizedef__call__(self,x):j,k=x.shape[-1],self.output_sizew_init=hk.initializers.TruncatedNormal(1./np.sqrt(j))w=hk.get_parameter("w",shape=[j,k],dtype=x.dtype,init=w_init)b=hk.get_parameter("b",shape=[k],dtype=x.dtype,init=jnp.zeros)returnjnp.dot(x,w)+bAll modules have a name. When nonameargument is passed to the module, its
name is inferred from the name of the Python class (for exampleMyLinearbecomesmy_linear). Modules can have named parameters that are accessed
usinghk.get_parameter(param_name, ...). We use this API (rather than just
using object properties) so that we can convert your code into a pure function
usinghk.transform.When using modules you need to define functions and transform them into a pair
of pure functions usinghk.transform. See ourquickstartfor
more details about the functions returned fromtransform:defforward_fn(x):model=MyLinear(10)returnmodel(x)# Turn `forward_fn` into an object with `init` and `apply` methods. By default,# the `apply` will require an rng (which can be None), to be used with# `hk.next_rng_key`.forward=hk.transform(forward_fn)x=jnp.ones([1,1])# When we run `forward.init`, Haiku will run `forward_fn(x)` and collect initial# parameter values. Haiku requires you pass a RNG key to `init`, since parameters# are typically initialized randomly:key=hk.PRNGSequence(42)params=forward.init(next(key),x)# When we run `forward.apply`, Haiku will run `forward_fn(x)` and inject parameter# values from the `params` that are passed as the first argument.  Note that# models transformed using `hk.transform(f)` must be called with an additional# `rng` argument: `forward.apply(params, rng, x)`. Use# `hk.without_apply_rng(hk.transform(f))` if this is undesirable.y=forward.apply(params,None,x)Working with stochastic modelsSome models may require random sampling as part of the computation.
For example, in variational autoencoders with the reparametrization trick,
a random sample from the standard normal distribution is needed. For dropout we
need a random mask to drop units from the input. The main hurdle in making this
work with JAX is in management of PRNG keys.In Haiku we provide a simple API for maintaining a PRNG key sequence associated
with modules:hk.next_rng_key()(ornext_rng_keys()for multiple keys):classMyDropout(hk.Module):def__init__(self,rate=0.5,name=None):super().__init__(name=name)self.rate=ratedef__call__(self,x):key=hk.next_rng_key()p=jax.random.bernoulli(key,1.0-self.rate,shape=x.shape)returnx*p/(1.0-self.rate)forward=hk.transform(lambdax:MyDropout()(x))key1,key2=jax.random.split(jax.random.PRNGKey(42),2)params=forward.init(key1,x)prediction=forward.apply(params,key2,x)For a more complete look at working with stochastic models, please see ourVAE example.Note:hk.next_rng_key()is not functionally pure which means you should
avoid using it alongside JAX transformations which are insidehk.transform.
For more information and possible workarounds, please consult the docs onHaiku transformsand availablewrappers for JAX transforms inside Haiku networks.Working with non-trainable stateSome models may want to maintain some internal, mutable state. For example, in
batch normalization a moving average of values encountered during training is
maintained.In Haiku we provide a simple API for maintaining mutable state that is
associated with modules:hk.set_stateandhk.get_state. When using these
functions you need to transform your function usinghk.transform_with_statesince the signature of the returned pair of functions is different:defforward(x,is_training):net=hk.nets.ResNet50(1000)returnnet(x,is_training)forward=hk.transform_with_state(forward)# The `init` function now returns parameters **and** state. State contains# anything that was created using `hk.set_state`. The structure is the same as# params (e.g. it is a per-module mapping of named values).params,state=forward.init(rng,x,is_training=True)# The apply function now takes both params **and** state. Additionally it will# return updated values for state. In the resnet example this will be the# updated values for moving averages used in the batch norm layers.logits,state=forward.apply(params,state,rng,x,is_training=True)If you forget to usehk.transform_with_statedon't worry, we will print a
clear error pointing you tohk.transform_with_staterather than silently
dropping your state.Distributed training withjax.pmapThe pure functions returned fromhk.transform(orhk.transform_with_state)
are fully compatible withjax.pmap. For more details on SPMD programming withjax.pmap,look here.One common use ofjax.pmapwith Haiku is for data-parallel training on many
accelerators, potentially across multiple hosts. With Haiku, that might look
like this:defloss_fn(inputs,labels):logits=hk.nets.MLP([8,4,2])(x)returnjnp.mean(softmax_cross_entropy(logits,labels))loss_fn_t=hk.transform(loss_fn)loss_fn_t=hk.without_apply_rng(loss_fn_t)# Initialize the model on a single device.rng=jax.random.PRNGKey(428)sample_image,sample_label=next(input_dataset)params=loss_fn_t.init(rng,sample_image,sample_label)# Replicate params onto all devices.num_devices=jax.local_device_count()params=jax.tree_util.tree_map(lambdax:np.stack([x]*num_devices),params)defmake_superbatch():"""Constructs a superbatch, i.e. one batch of data per device."""# Get N batches, then split into list-of-images and list-of-labels.superbatch=[next(input_dataset)for_inrange(num_devices)]superbatch_images,superbatch_labels=zip(*superbatch)# Stack the superbatches to be one array with a leading dimension, rather than# a python list. This is what `jax.pmap` expects as input.superbatch_images=np.stack(superbatch_images)superbatch_labels=np.stack(superbatch_labels)returnsuperbatch_images,superbatch_labelsdefupdate(params,inputs,labels,axis_name='i'):"""Updates params based on performance on inputs and labels."""grads=jax.grad(loss_fn_t.apply)(params,inputs,labels)# Take the mean of the gradients across all data-parallel replicas.grads=jax.lax.pmean(grads,axis_name)# Update parameters using SGD or Adam or ...new_params=my_update_rule(params,grads)returnnew_params# Run several training updates.for_inrange(10):superbatch_images,superbatch_labels=make_superbatch()params=jax.pmap(update,axis_name='i')(params,superbatch_images,superbatch_labels)For a more complete look at distributed Haiku training, take a look at ourResNet-50 on ImageNet example.Citing HaikuTo cite this repository:@software{haiku2020github,
  author = {Tom Hennigan and Trevor Cai and Tamara Norman and Lena Martens and Igor Babuschkin},
  title = {{H}aiku: {S}onnet for {JAX}},
  url = {http://github.com/deepmind/dm-haiku},
  version = {0.0.10},
  year = {2020},
}In this bibtex entry, the version number is intended to be fromhaiku/__init__.py,
and the year corresponds to the project's open-source release.