Issue definig subModules on super()

[alonfnt]

Hi, I am trying to boundle two networks into the same large network, and it was my understanding that something like this would work:

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

class Net1(nn.Module):
    n: int

    @nn.compact
    def __call__(self, x):
        for _ in range(2):
            x = nn.relu(nn.Dense(64)(x))
        x = nn.Dense(self.n)(x)
        return x

class Net2(nn.Module):
    n: int

    def setup(self):
        self.inner = Net1(self.n)
        self.outer = Net1(self.n*2)

    def __call__(self, x):
        x = self.outer(x)
        return x


key = jax.random.key(42)
x1_dummy = jnp.ones(shape=(20))
inner_params = Net2(n=10).init(key, x1_dummy, method='inner')
...
apply_inner_fn = lambda x: Net2(n=10).apply(x, method='inner')

Where I want to have access to inner so that I can merge params, and for a specific gradient calculation, I could do something like

def loss_f(params, x, y):
    z = train_state.apply_inner_fn(params, x)
    y_hat = train_state.apply_fn(params, z)
    return jnp.mean((y-y_hat)**2)
...
# Update params together with the same grads

But running the first snippets, I get

AttributeError: "Net2" object has no attribute "inner". If "inner" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'.

I am calling it inside init, am I not? I am a bit confused about how it is best to proceed.

[chiamp]

Net2(n=10).init(key, x1_dummy, method='inner') is looking for a method called inner, which isn't defined. You could try something like this:

class Net2(nn.Module):
    n: int

    def setup(self):
        self.inner = Net1(self.n)
        self.outer = Net1(self.n*2)

    def inner_call(self, x):
        return self.inner(x)
    
    def outer_call(self, x):
        return self.outer(x)

...

inner_params = Net2(n=10).init(key, x1_dummy, method='inner_call')