How are parameters passed to call?

[Logon27]

Just to preface I am very new to Flax. But I was kind of curious (from an implementation perspective) how parameters were passed to the classes / functions of the call method?

Given the example in the README...

class CNN(nn.Module):
  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    x = nn.log_softmax(x)
    return x

model = CNN()
batch = jnp.ones((32, 64, 64, 10))  # (N, H, W, C) format
variables = model.init(jax.random.PRNGKey(0), batch)
output = model.apply(variables, batch)

I see that the 'variables' are passed to the apply function. But from an implementation perspective, how exactly are these variables bound to the individual classes defined inside the call method after being passed to the apply function? Not sure if this is a lack of my understanding of python or if this is just non-trivial to do.

[cgarciae]

Hey @Logon27, long story short: when you call init or apply, a Scope object is created that has a references to the variables, the Scope is attached the Module.scope field and, through a somewhat complicated process, submodules also get a scope with their corresponding variables.

[Logon27]

So the scope object is created with the variables. A subclass of the module is created with the scope variables and internal state. Then the computation of the stateful module is performed. And finally the result of the stateful subclass is returned to the parent (pure) Module?

[cgarciae]

And finally the result of the stateful subclass is returned to the parent (pure) Module?

I would not say it like this. Rather, what is returned from init and apply is usually the new state (the variables) and the output of the method called.