Migrating from Haiku to Flax#
This guide will walk through the process of migrating Haiku models to Flax, and highlight the differences between the two libraries.
Basic Example#
To create custom Modules you subclass from a Module base class in
both Haiku and Flax. However, Haiku classes use a regular __init__ method
whereas Flax classes are dataclasses, meaning you define some class
attributes that are used to automatically generate a constructor. Also,
all Flax Modules accept a name argument without needing to define it,
whereas in Haiku name must be explicitly defined in the constructor
signature and passed to the superclass constructor.
import haiku as hk
class Block(hk.Module):
def __init__(self, features: int, name=None):
super().__init__(name=name)
self.features = features
def __call__(self, x, training: bool):
x = hk.Linear(self.features)(x)
x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)
x = jax.nn.relu(x)
return x
class Model(hk.Module):
def __init__(self, dmid: int, dout: int, name=None):
super().__init__(name=name)
self.dmid = dmid
self.dout = dout
def __call__(self, x, training: bool):
x = Block(self.dmid)(x, training)
x = hk.Linear(self.dout)(x)
return x
import flax.linen as nn
class Block(nn.Module):
features: int
@nn.compact
def __call__(self, x, training: bool):
x = nn.Dense(self.features)(x)
x = nn.Dropout(0.5, deterministic=not training)(x)
x = jax.nn.relu(x)
return x
class Model(nn.Module):
dmid: int
dout: int
@nn.compact
def __call__(self, x, training: bool):
x = Block(self.dmid)(x, training)
x = nn.Dense(self.dout)(x)
return x
The __call__ method looks very similar in both libraries, however, in Flax
you have to use the @nn.compact decorator in order to be able to define
submodules inline. In Haiku, this is the default behavior.
Now, a place where Haiku and Flax differ substantially is in how you construct
the model. In Haiku, you use hk.transform over a function
that calls your Module, transform will return an object with init
and apply methods. In Flax, you simply instantiate your Module.
def forward(x, training: bool):
return Model(256, 10)(x, training)
model = hk.transform(forward)
...
model = Model(256, 10)
To get the model parameters in both libraries you use the init method
with a random.key plus some inputs to run the model. The main difference here is
that Flax returns a mapping from collection names to nested array dictionaries,
params is just one of these possible collections. In Haiku, you get the params
structure directly.
sample_x = jax.numpy.ones((1, 784))
params = model.init(
random.key(0),
sample_x, training=False # <== inputs
)
...
sample_x = jax.numpy.ones((1, 784))
variables = model.init(
random.key(0),
sample_x, training=False # <== inputs
)
params = variables["params"]
One very important thing to note is that in Flax the parameters structure is
hierarchical, with one level per nested module and a final level for the
parameter name.
In Haiku the parameters structure is a python dictionary with a two level hierarchy:
the fully qualified module name mapping to the parameter name. The module name
consists of a / separated string path of all the nested Modules.
...
{
'model/block/linear': {
'b': (256,),
'w': (784, 256),
},
'model/linear': {
'b': (10,),
'w': (256, 10),
}
}
...
FrozenDict({
Block_0: {
Dense_0: {
bias: (256,),
kernel: (784, 256),
},
},
Dense_0: {
bias: (10,),
kernel: (256, 10),
},
})
During training in both frameworks you pass the parameters structure to the
apply method to run the forward pass. Since we are using dropout, in
both cases we must provide a key to apply in order to generate
the random dropout masks.
def train_step(key, params, inputs, labels):
def loss_fn(params):
logits = model.apply(
params,
key,
inputs, training=True # <== inputs
)
return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
grads = jax.grad(loss_fn)(params)
params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)
return params
def train_step(key, params, inputs, labels):
def loss_fn(params):
logits = model.apply(
{'params': params},
inputs, training=True, # <== inputs
rngs={'dropout': key}
)
return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
grads = jax.grad(loss_fn)(params)
params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)
return params
The most notable differences is that in Flax you have to
pass the parameters inside a dictionary with a params key, and the
key inside a dictionary with a dropout key. This is because in Flax
you can have many types of model state and random state. In Haiku, you
just pass the parameters and the key directly.
Handling State#
Now let’s see how mutable state is handled in both libraries. We will take the same model as before, but now we will replace Dropout with BatchNorm.
class Block(hk.Module):
def __init__(self, features: int, name=None):
super().__init__(name=name)
self.features = features
def __call__(self, x, training: bool):
x = hk.Linear(self.features)(x)
x = hk.BatchNorm(
create_scale=True, create_offset=True, decay_rate=0.99
)(x, is_training=training)
x = jax.nn.relu(x)
return x
class Block(nn.Module):
features: int
@nn.compact
def __call__(self, x, training: bool):
x = nn.Dense(self.features)(x)
x = nn.BatchNorm(
momentum=0.99
)(x, use_running_average=not training)
x = jax.nn.relu(x)
return x
The code is very similar in this case as both libraries provide a BatchNorm
layer. The most notable difference is that Haiku uses is_training to
control whether or not to update the running statistics, whereas Flax uses
use_running_average for the same purpose.
To instantiate a stateful model in Haiku you use hk.transform_with_state,
which changes the signature for init and apply to accept and return
state. As before, in Flax you construct the Module directly.
def forward(x, training: bool):
return Model(256, 10)(x, training)
model = hk.transform_with_state(forward)
...
model = Model(256, 10)
To initialize both the parameters and state you just call the init method
as before. However, in Haiku you now get state as a second return value, and
in Flax you get a new batch_stats collection in the variables dictionary.
Note that since hk.BatchNorm only initializes batch statistics when
is_training=True, we must set training=True when initializing parameters
of a Haiku model with an hk.BatchNorm layer. In Flax, we can set
training=False as usual.
sample_x = jax.numpy.ones((1, 784))
params, state = model.init(
random.key(0),
sample_x, training=True # <== inputs
)
...
sample_x = jax.numpy.ones((1, 784))
variables = model.init(
random.key(0),
sample_x, training=False # <== inputs
)
params, batch_stats = variables["params"], variables["batch_stats"]
In general, in Flax you might find other state collections in the variables
dictionary such as cache for auto-regressive transformers models,
intermediates for intermediate values added using Module.sow, or other
collection names defined by custom layers. Haiku only makes a distinction
between params (variables which do not change while running apply) and
state (variables which can change while running apply).
Now, training looks very similar in both frameworks as you use the same
apply method to run the forward pass. In Haiku, now pass the state
as the second argument to apply, and get the new state as the second
return value. In Flax, you instead add batch_stats as a new key to the
input dictionary, and get the updates variables dictionary as the second
return value.
def train_step(params, state, inputs, labels):
def loss_fn(params):
logits, new_state = model.apply(
params, state,
None, # <== rng
inputs, training=True # <== inputs
)
loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
return loss, new_state
grads, new_state = jax.grad(loss_fn, has_aux=True)(params)
params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)
return params, new_state
def train_step(params, batch_stats, inputs, labels):
def loss_fn(params):
logits, updates = model.apply(
{'params': params, 'batch_stats': batch_stats},
inputs, training=True, # <== inputs
mutable='batch_stats',
)
loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
return loss, updates["batch_stats"]
grads, batch_stats = jax.grad(loss_fn, has_aux=True)(params)
params = jax.tree_util.tree_map(lambda p, g: p - 0.1 * g, params, grads)
return params, batch_stats
One major difference is that in Flax a state collection can be mutable or immutable.
During init all collections are mutable by default, however, during apply
you have to explicitly specify which collections are mutable. In this example,
we specify that batch_stats is mutable. Here a single string is passed but a list
can also be given if there are more mutable collections. If this is not done an
error will be raised at runtime when trying to mutate batch_stats.
Also, when mutable is anything other than False, the updates
dictionary is returned as the second return value of apply, else only the
model output is returned.
Haiku makes the mutable/immutable distinction through having params
(immutable) and state (mutable) and using either hk.transform or
hk.transform_with_state
Using Multiple Methods#
In this section we will take a look at how to use multiple methods in Haiku and Flax.
As an example, we will implement an auto-encoder model with three methods:
encode, decode, and __call__.
In Haiku, we can just define the submodules that encode and decode need
directly in __init__, in this case each will just use a Linear layer.
In Flax, we will define an encoder and a decoder Module ahead of time
in setup, and use them in the encode and decode respectively.
class AutoEncoder(hk.Module):
def __init__(self, embed_dim: int, output_dim: int, name=None):
super().__init__(name=name)
self.encoder = hk.Linear(embed_dim, name="encoder")
self.decoder = hk.Linear(output_dim, name="decoder")
def encode(self, x):
return self.encoder(x)
def decode(self, x):
return self.decoder(x)
def __call__(self, x):
x = self.encode(x)
x = self.decode(x)
return x
class AutoEncoder(nn.Module):
embed_dim: int
output_dim: int
def setup(self):
self.encoder = nn.Dense(self.embed_dim)
self.decoder = nn.Dense(self.output_dim)
def encode(self, x):
return self.encoder(x)
def decode(self, x):
return self.decoder(x)
def __call__(self, x):
x = self.encode(x)
x = self.decode(x)
return x
Note that in Flax setup doesn’t run after __init__, instead it runs
when init or apply are called.
Now, we want to be able to call any method from our AutoEncoder model. In Haiku we
can define multiple apply methods for a module through hk.multi_transform. The
function passed to multi_transform defines how to initialize the module and which
different apply methods to generate.
def forward():
module = AutoEncoder(256, 784)
init = lambda x: module(x)
return init, (module.encode, module.decode)
model = hk.multi_transform(forward)
...
model = AutoEncoder(256, 784)
To initialize the parameters of our model, init can be used to trigger the
__call__ method, which uses both the encode and decode
method. This will create all the necessary parameters for the model.
params = model.init(
random.key(0),
x=jax.numpy.ones((1, 784)),
)
...
variables = model.init(
random.key(0),
x=jax.numpy.ones((1, 784)),
)
params = variables["params"]
This generates the following parameter structure.
{
'auto_encoder/~/decoder': {
'b': (784,),
'w': (256, 784)
},
'auto_encoder/~/encoder': {
'b': (256,),
'w': (784, 256)
}
}
FrozenDict({
decoder: {
bias: (784,),
kernel: (256, 784),
},
encoder: {
bias: (256,),
kernel: (784, 256),
},
})
Finally, let’s explore how we can employ the apply function to invoke the encode method:
encode, decode = model.apply
z = encode(
params,
None, # <== rng
x=jax.numpy.ones((1, 784)),
)
...
z = model.apply(
{"params": params},
x=jax.numpy.ones((1, 784)),
method="encode",
)
Because the Haiku apply function is generated through
hk.multi_transform, it’s a tuple of two functions which we can unpack into
an encode and decode function which correspond to the methods on the
AutoEncoder module. In Flax we call the encode method through passing
the method name as a string.
Another noteworthy distinction here is that in Haiku, rng needs to be
explicitly passed, even though the module does not use any stochastic
operations during apply. In Flax this is not necessary (check out
Randomness and PRNGs in Flax).
The Haiku rng is set to None here, but you could also use
hk.without_apply_rng on the apply function to remove the rng argument.
Lifted Transforms#
Both Flax and Haiku provide a set of transforms, which we will refer to as lifted transforms,
that wrap JAX transformations in such a way that they can be used with Modules and sometimes
provide additional functionality. In this section we will take a look at how to use the
lifted version of scan in both Flax and Haiku to implement a simple RNN layer.
To begin, we will first define a RNNCell module that will contain the logic for a single
step of the RNN. We will also define a initial_state method that will be used to initialize
the state (a.k.a. carry) of the RNN. Like with jax.lax.scan, the RNNCell.__call__
method will be a function that takes the carry and input, and returns the new
carry and output. In this case, the carry and the output are the same.
class RNNCell(hk.Module):
def __init__(self, hidden_size: int, name=None):
super().__init__(name=name)
self.hidden_size = hidden_size
def __call__(self, carry, x):
x = jnp.concatenate([carry, x], axis=-1)
x = hk.Linear(self.hidden_size)(x)
x = jax.nn.relu(x)
return x, x
def initial_state(self, batch_size: int):
return jnp.zeros((batch_size, self.hidden_size))
class RNNCell(nn.Module):
hidden_size: int
@nn.compact
def __call__(self, carry, x):
x = jnp.concatenate([carry, x], axis=-1)
x = nn.Dense(self.hidden_size)(x)
x = jax.nn.relu(x)
return x, x
def initial_state(self, batch_size: int):
return jnp.zeros((batch_size, self.hidden_size))
Next, we will define a RNN Module that will contain the logic for the entire RNN.
In Haiku, we will first initialze the RNNCell, then use it to construct the carry,
and finally use hk.scan to run the RNNCell over the input sequence. In Flax its
done a bit differently, we will use nn.scan to define a new temporary type that wraps
RNNCell. During this process we will also specify instruct nn.scan to broadcast
the params collection (all steps share the same parameters) and to not split the
params rng stream (so all steps intialize with the same parameters), and finally
we will specify that we want scan to run over the second axis of the input and stack
the outputs along the second axis as well. We will then use this temporary type immediately
to create an instance of the lifted RNNCell and use it to create the carry and
the run the __call__ method which will scan over the sequence.
class RNN(hk.Module):
def __init__(self, hidden_size: int, name=None):
super().__init__(name=name)
self.hidden_size = hidden_size
def __call__(self, x):
cell = RNNCell(self.hidden_size)
carry = cell.initial_state(x.shape[0])
carry, y = hk.scan(cell, carry, jnp.swapaxes(x, 1, 0))
y = jnp.swapaxes(y, 0, 1)
return y
class RNN(nn.Module):
hidden_size: int
@nn.compact
def __call__(self, x):
rnn = nn.scan(RNNCell, variable_broadcast='params', split_rngs={'params': False},
in_axes=1, out_axes=1)(self.hidden_size)
carry = rnn.initial_state(x.shape[0])
carry, y = rnn(carry, x)
return y
In general, the main difference between lifted transforms between Flax and Haiku is that
in Haiku the lifted transforms don’t operate over the state, that is, Haiku will handle the
params and state in such a way that it keeps the same shape inside and outside of the
transform. In Flax, the lifted transforms can operate over both variable collections and rng
streams, the user must define how different collections are treated by each transform
according to the transform’s semantics.
Finally, let’s quickly view how the RNN Module would be used in both Haiku and Flax.
def forward(x):
return RNN(64)(x)
model = hk.without_apply_rng(hk.transform(forward))
params = model.init(
random.key(0),
x=jax.numpy.ones((3, 12, 32)),
)
y = model.apply(
params,
x=jax.numpy.ones((3, 12, 32)),
)
...
model = RNN(64)
variables = model.init(
random.key(0),
x=jax.numpy.ones((3, 12, 32)),
)
params = variables['params']
y = model.apply(
{'params': params},
x=jax.numpy.ones((3, 12, 32)),
)
The only notable change with respect to the examples in the previous sections is that
this time around we used hk.without_apply_rng in Haiku so we didn’t have to
pass the rng argument as None to the apply method.
Scan over layers#
One very important application of scan is apply a sequence of layers iteratively
over an input, passing the output of each layer as the input to the next layer. This
is very useful to reduce compilation time for big models. As an example we will create
a simple Block Module, and then use it inside an MLP Module that will apply
the Block Module num_layers times.
In Haiku, we define the Block Module as usual, and then inside MLP we will
use hk.experimental.layer_stack over a stack_block function to create a stack
of Block Modules. In Flax, the definition of Block is a little different,
__call__ will accept and return a second dummy input/output that in both cases will
be None. In MLP, we will use nn.scan as in the previous example, but
by setting split_rngs={'params': True} and variable_axes={'params': 0}
we are telling nn.scan create different parameters for each step and slice the
params collection along the first axis, effectively implementing a stack of
Block Modules as in Haiku.
class Block(hk.Module):
def __init__(self, features: int, name=None):
super().__init__(name=name)
self.features = features
def __call__(self, x, training: bool):
x = hk.Linear(self.features)(x)
x = hk.dropout(hk.next_rng_key(), 0.5 if training else 0, x)
x = jax.nn.relu(x)
return x
class MLP(hk.Module):
def __init__(self, features: int, num_layers: int, name=None):
super().__init__(name=name)
self.features = features
self.num_layers = num_layers
def __call__(self, x, training: bool):
@hk.experimental.layer_stack(self.num_layers)
def stack_block(x):
return Block(self.features)(x, training)
stack = hk.experimental.layer_stack(self.num_layers)
return stack_block(x)
class Block(nn.Module):
features: int
training: bool
@nn.compact
def __call__(self, x, _):
x = nn.Dense(self.features)(x)
x = nn.Dropout(0.5)(x, deterministic=not self.training)
x = jax.nn.relu(x)
return x, None
class MLP(nn.Module):
features: int
num_layers: int
@nn.compact
def __call__(self, x, training: bool):
ScanBlock = nn.scan(
Block, variable_axes={'params': 0}, split_rngs={'params': True},
length=self.num_layers)
y, _ = ScanBlock(self.features, training)(x, None)
return y
Notice how in Flax we pass None as the second argument to ScanBlock and ignore
its second output. These represent the inputs/outputs per-step but they are None
because in this case we don’t have any.
Initializing each model is the same as in previous examples. In this case,
we will be specifying that we want to use 5 layers each with 64 features.
def forward(x, training: bool):
return MLP(64, num_layers=5)(x, training)
model = hk.transform(forward)
sample_x = jax.numpy.ones((1, 64))
params = model.init(
random.key(0),
sample_x, training=False # <== inputs
)
...
...
model = MLP(64, num_layers=5)
sample_x = jax.numpy.ones((1, 64))
variables = model.init(
random.key(0),
sample_x, training=False # <== inputs
)
params = variables['params']
When using scan over layers the one thing you should notice is that all layers
are fused into a single layer whose parameters have an extra “layer” dimension on
the first axis. In this case, the shape of all parameters will start with (5, ...)
as we are using 5 layers.
...
{
'mlp/__layer_stack_no_per_layer/block/linear': {
'b': (5, 64),
'w': (5, 64, 64)
}
}
...
FrozenDict({
ScanBlock_0: {
Dense_0: {
bias: (5, 64),
kernel: (5, 64, 64),
},
},
})
Top-level Haiku functions vs top-level Flax modules#
In Haiku, it is possible to write the entire model as a single function by using the raw hk.{get,set}_{parameter,state} to define/access model parameters and states. It very common to write the top-level “Module” as a function instead:
The Flax team recommends a more Module-centric approach that uses __call__ to define the forward function. The corresponding accessor will be nn.module.param and nn.module.variable (go to Handling State for an explanaion on collections).
def forward(x):
counter = hk.get_state('counter', shape=[], dtype=jnp.int32, init=jnp.ones)
multiplier = hk.get_parameter('multiplier', shape=[1,], dtype=x.dtype, init=jnp.ones)
output = x + multiplier * counter
hk.set_state("counter", counter + 1)
return output
model = hk.transform_with_state(forward)
params, state = model.init(random.key(0), jax.numpy.ones((1, 64)))
class FooModule(nn.Module):
@nn.compact
def __call__(self, x):
counter = self.variable('counter', 'count', lambda: jnp.ones((), jnp.int32))
multiplier = self.param('multiplier', nn.initializers.ones_init(), [1,], x.dtype)
output = x + multiplier * counter.value
if not self.is_initializing(): # otherwise model.init() also increases it
counter.value += 1
return output
model = FooModule()
variables = model.init(random.key(0), jax.numpy.ones((1, 64)))
params, counter = variables['params'], variables['counter']