Inconsistent shapes between value and initializer for linen.Module

[PgLoLo]

I have a module, one of the parameters of which is "shape-agnostic", i.e. it may slightly vary during inference (actually, it firstly doubles its size during Riemannian optimization step and then shrink back to its initial size). And during this shape increase it fails to run Module.apply method with "Inconsistent shapes between value and initializer" error.

Consider the following simple example: we have a module and "shape-agnostic" parameter:

from flax import linen


class A(linen.Module):
    size: int
        
    def setup(self):
        self.array = self.param('array', lambda _: jnp.zeros(self.size))
        
    def __call__(self):
        return self.array.mean()
    
    
model = A(10)
params = model.init(jax.random.PRNGKey(0))

print(model.apply(params))

As expected, this works perfectly well. Please note, that __call__ method can work with array of any shape without troubles.

Next, we want to do some model-surgery and call __call__ method once again:

from flax.core import freeze, unfreeze

params = unfreeze(params)
params['params']['array'] = jnp.concatenate([params['params']['array'], params['params']['array']])
parms = freeze(params)

print(model.apply(params))

That would raise ValueError: Inconsistent shapes between value and initializer for parameter "array" in "/": (20,), (10,) exception.

As far, as I understand, to rebuild the model from params, flax firstly reevaluates A.setup() function, where the array is instantiated with shape [10], and then tries to replace its values with corresponding values from params argument. How can I overcome this behaviour? Is there a way to specify parameters in setup method without shape information?

[marcvanzee]

Is there a way to specify parameters in setup method without shape information?

I'm not sure if I understand this question. If your model doubles in size, why don't you just construct it with the new size? So your last line looks as follows:

print(A(20).apply(params))

Or you derive the desired size from the input params:

size = params['params']['array'].shape[0]
print(A(size).apply(params))

Also, as a general rule: I would suggest to name the return value of init() to variables since params["params"] looks a bit odd, but variables["params"] make much more sense, at least to me 😄

[PgLoLo]

why don't you just construct it with the new size?

Basically, because of incapsulation reasons, I just don't want to know anything about the underlying structure. Consider a case, where this module A is used inside some different module. Say, it is a parameterization of prior distribution in GAN or something like that. During GAN inference I don't want to know anything about this underlying parameterization, I just want to call apply of it.

Here is the cause of the problem -- there is an explicit check that the shapes of incoming parameters and "initialized" parameters are the same. And I don't really understand why we need such a check here: firstly, as far, as I understand, everything would work even without this check. Secondly, I don't see reasons why models should not change (theoretically speaking) their shapes during their lifetime.

[jheek]

For params we use shape inference to check that the initialiser and it's value have the same shape. This avoids a lot of issues with hyper paramaters and params being out of sync for example after restoring a checkpoint. We might at some point at a keyword arg to disable this check but for now an easy workaround is:

class A(linen.Module):
    size: int
        
    def setup(self):
        self.array = self.variable('params', 'array', jnp.zeros, self.size)
        
    def __call__(self):
        return self.array.mean()

btw I would consider putting such a variable in a separate collection than "params" anyway. Quite often you need to enforce shape invariance outside of the model as well for example when the params are being optimized.

[PgLoLo]

implementing such arguments as variables really solves the problem, thank you very much for such a nice solution!