What is the difference between Module.param() and Module.variable()?

[marcvanzee]

[marcvanzee]

Module.param() is a special case of Module.variable that assumes 3 things:

1. the variable is immutable, so it returns the parameter directly, rather than a mutable reference to it
2. the param init() fn accepts a PRNGKey as the first arg (in addition to any additionally provided ones)

3. the variable collection name is params

In code:

p = self.param('param_name', init_fn, shape, dtype)
# is a convenient shorthand for this:
p = self.variable('params', 'param_name', lambda s, d: init_fn(self.make_rng('params'), s, d), shape, dtype).value

So self.param is simply a a convenience, because 95% or more of the time user want to define NN parameters, but the self.variable call lets you do whatever you want.

self.variable has a much more general pattern than self.param - we don't presume that a general variable init fn takes an rng. To do that:

from flax.linen.initializers import lecun_normal, variance_scaling, zeros, normal
class Foo(nn.Module):
  features: int
  @nn.compact
  def __call__(self):
    key = self.make_rng('spectral_norm_stats')
    u0_variable = self.variable('spectral_norm_stats', 'u0', normal(), key, (1, self.features))
    return u0_variable.value

foovars = Foo(3).init({'params': random.PRNGKey(0), 'spectral_norm_stats': random.PRNGKey(1)})
Foo(3).apply(foovars, rngs={'spectral_norm_stats': random.PRNGKey(1)})  
# --> DeviceArray([[0.00711277, 0.0107195 , 0.019903  ]], dtype=float32)

We create a new PRNG sequence inline with self.make_rng (the 'params' rng stream always exist, but we create a new one for use with spectral_norm_stats - note the PRNG stream doesn't -have- to share the same name as the variable collection, it's just a convention.). For every PRNGKey stream used, we have to pass in an initializing seed key for each stream as shown here.

If you only need it once in the initializer, you can do this:

class Foo(nn.Module):
  features: int
  @nn.compact
  def __call__(self):
    u0_variable = self.variable('spectral_norm_stats', 'u0', 
                                lambda s: normal()(self.make_rng('params'), s), (1, self.features))
    return u0_variable.value

foovars = Foo(3).init({'params': random.PRNGKey(0)})
Foo(3).apply(foovars)