How to init BatchNorm in the setup function

[davidshen84]

Hi,

The example only shows how to initialize the BatchNorm module in the compact form. What if I need to initialize it in the setup form?

In the compact form, the BatchNorm constructor takes a use_running_average variable from the __call__ method. But in the setup form, the BatchNorm is initialized in the setup function, which does not have the use_running_average argument at runtime.

[chiamp]

You can pass in use_running_average at call time:

import flax.linen as nn
import jax, jax.numpy as jnp

class Model(nn.Module):
  def setup(self):
    self.dense = nn.Dense(3)
    self.batchnorm = nn.BatchNorm()
  def __call__(self, x, train):
    return self.dense(self.batchnorm(x, use_running_average=not train))

x = jnp.ones((1, 2))
model = Model()
variables = model.init(jax.random.key(0), x, train=False)
out, batch_stats = model.apply(variables, x, train=True, mutable=['batch_stats'])

You could also pass in use_running_average at setup time, but it will lock that value in:

class Model(nn.Module):
  def setup(self):
    self.dense = nn.Dense(3)
    self.batchnorm = nn.BatchNorm(use_running_average=True) # will always use running average
  ...

If you want control over when to use running average or not, it's best to pass in use_running_average at call time.