Creating module attributes with @nn.compact

[matthias-wright]

Hi all,

I have noticed that it is no longer possible to create module attributes in the call method when using the @nn.compact decorator.
See attached code for an example.

This used to work in version 0.2.2 but no longer works in version 0.3.0.

I am wondering why this is no longer possible. I think there are many cases, where this can be helpful.

Thanks!

import flax.nn as nn
from jax import random
import jax.numpy as jnp

class Net(nn.Module):

    @nn.compact
    def __call__(self, x):
        self.n = 1
        return x

class Net2(nn.Module):

    def setup(self):
        self.n = 1

    def __call__(self, x):
        return x

key = random.PRNGKey(0)
x = random.normal(key, shape=(1,1))

# This does not work in version 0.3.0
model = Net()
params = model.init(key, x)
y = model.apply(params, x)

# This still works
model = Net2()
params = model.init(key, x)
y = model.apply(params, x)

Error message:
TypeError: Module instance is frozen outside of setup method.

Edit:
My use case: I ported some models from PyTorch and Tensorflow to Flax and stored the weights and parameters in a nested dictionary that can be loaded from a file (see Net for a simplified example).
The reason why I don't pass this dictionary to the model during initialization is that I want to "hide" the code for loading this dictionary inside the model. Just because I think it is neat.

I did not want to use the setup method for this, because I absolutely love using the compact method decorator. It makes the model look so tidy.

I also know that I can just declare the dictionary as a normal variable inside the call method and then pass it to the resblock. However, for the larger resnets, when there are many blocks, it just looks cleaner when you don't have to do that.

I guess a good compromise is to just use the setup method for initializing the dictionary and still using the compact decorator to define the submodules in the call method (see Net2 for a simplified example).

# does not work
class Net(nn.Module):

    @nn.compact
    def __call__(self, x):
        self.param_dict = {} # dict with params
        x = self.resblock(x)
        return x
    
    def resblock(self, x):
        # resblock code here
        w = lambda *_ : jnp.array(self.param_dict['dense1']['weight'])
        b = lambda *_ : jnp.array(self.param_dict['dense1']['bias'])
        x = nn.Conv(features=1, kernel_init=w, bias_init=b, name='dense1')(x)
        return x

# works
class Net2(nn.Module):

    def setup(self):
        self.param_dict = {} # dict with params

    @nn.compact
    def __call__(self, x):
         x = self.resblock(x)
        return x

    def resblock(self, x):
        # resblock code here
        w = lambda *_ : jnp.array(self.param_dict['dense1']['weight'])
        b = lambda *_ : jnp.array(self.param_dict['dense1']['bias'])
        x = nn.Conv(features=1, kernel_init=w, bias_init=b, name='dense1')(x)
        return x

[jheek]

TL;DR it was never safe to do this but now we actually raise an error.

The practical argument is that like all Jax/Flax APIs we want to avoid internal state so we typically freeze things after construction.

The theoretical argument is that in the case of Module we actually need the instance to be Frozen after setup otherwise we cannot guarantee correct behaviour in all cases. Basically, we need to clone Module instances when using transformations (vmap, jit) or Module.apply. Instead of trying to do a magical deep clone we reconstruct the Module from the construction arguments. This is clean and simple as long as there is no internal state like attribute assignment in methods.

[avital]

@matthias-wright can you share more about the particular use-case you are looking at? Per @jheek's comment it's fundamentally unsafe to mutate anything in the context of JAX transformations. That's why we expose mutable variables via Module.variable(). Would this work for your use-case?

[marcvanzee]

It does seem to me we could improve the error message here, I think it is a bit too short. I've filed an issue for this: #944

[matthias-wright]

Hey everyone, thanks for the quick responses!
@jheek That makes sense, thank you for the explanation!
@avital I have edited my original post with my use case. I know that it is very specific and there are many workarounds. It is just a matter of taste and I found a way of doing things that I liked and was surprised that it did not work in the newer versions.
Keep up the great work! :)