Creating named parameter collections inside @nn.compact

[minqi]

When using @nn.compact with __call__, is it possible to have the initialized parameters be nested under a provided name key inside the self.variables dictionary?

I noticed that when initializing a module that is bound to self, the resulting parameters are nested under self.variables['params'][<variable name of module>], but this is not the case for a module that is first defined inline inside __call__.

[cgarciae]

Hey @minqi, you can set the name for any Module as it would appear in the variables structure using the name constructor argument e.g:

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

class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(10, name='bar')(x)
    
foo = Foo()
variables = foo.init(jax.random.PRNGKey(0), jnp.ones((1, 10)))
print(jax.tree_map(jnp.shape, variables))

FrozenDict({
    params: {
        bar: {
            bias: (10,),
            kernel: (10, 10),
        },
    },
})

[minqi]

Thanks, @cgarciae. My issue was slightly different, as the submodule being initialized in __call__ was a nn.Sequential. In this case, it seems that flax ignores any name assigned to the Sequential submodule. However the Sequential's children module names are included as usual inside the params object. For example, consider this variation of your snippet:

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

class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Sequential([
        	nn.Dense(10, name='dense1'),
        	nn.Dense(10, name='dense2'),
        ], name='bar')(x)
    
foo = Foo()
variables = foo.init(jax.random.PRNGKey(0), jnp.ones((1, 10)))
print(jax.tree_map(jnp.shape, variables))

This would output

FrozenDict({
    params: {
        dense1: {
            bias: (10,),
            kernel: (10, 10),
        },
        dense2: {
            bias: (10,),
            kernel: (10, 10),
        },
    },
})

Notice that the name bar assigned to the Sequential instance is missing.

[cgarciae]

Ok, took me a while to figure out. Because you are using compact, the Dense modules get attached Foo first and then they are shared with Sequential, thus their path appears directly under Foo, Sequential gets ignored path-wise.