Recommended way to call apply on submodule given supermodule's params

[carlosgmartin]

I have a module as well as params for it, and need to call apply on one of its submodules. Is the third way in the following code the recommended way to do this, i.e. by accessing the submodule's name and putting it inside a dictionary with the key params?

Is the recommended way to do this somewhere in the documentation?

from jax import numpy as jnp, random
from flax import linen as nn


class Autoencoder(nn.Module):
    encoder: nn.Module
    decoder: nn.Module

    def __call__(self, x):
        h = self.encoder(x)
        y = self.decoder(h)
        return y


def main():
    model = Autoencoder(nn.Dense(10), nn.Dense(3))
    key = random.PRNGKey(0)
    x = jnp.ones(2)
    params = model.init(key, x)
    y = model.apply(params, x)
    print(y)

    print()

    try:
        h = model.encoder.apply(params, x)
    except Exception as e:
        print(e)

    print()

    try:
        h = model.encoder.apply(params['params']['encoder'], x)
    except Exception as e:
        print(e)

    print()

    h = model.encoder.apply({'params': params['params']['encoder']}, x)
    print(h)


if __name__ == "__main__":
    main()

[-0.2754686   0.35974982  0.5339251 ]

Could not find parameter named "kernel" in scope "/". (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeParamNotFoundError)

The first argument passed to an apply function should be a dictionary of collections. Each collection should be a dictionary with string keys. (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ApplyScopeInvalidVariablesTypeError)

[-0.33321813 -0.8941687  -0.40589067  0.2368109  -0.81948453  1.2879758
 -0.17161958  1.3635392   1.8489769   0.7455205 ]

[chiamp]

You can pass in a method or string in the apply function's method argument that specifically calls the submodules you want. For example:

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

class Autoencoder(nn.Module):
  encoder: nn.Module
  decoder: nn.Module

  def __call__(self, x):
    return self.decode(self.encode(x))
  
  def encode(self, x):
    return self.encoder(x)
  
  def decode(self, x):
    return self.decoder(x)

model = Autoencoder(nn.Dense(10), nn.Dense(3))
key = jax.random.key(0)
x = jnp.ones(2)
params = model.init(key, x)

# encode and then decode using __call__
out = model.apply(params, x)
assert out.shape == (3,)

# encode only
encode_out = model.apply(params, x, method='encode')
assert encode_out.shape == (10,)

# decode only
decode_out = model.apply(params, encode_out, method=model.decode)
assert decode_out.shape == (3,)

If nothing is passed to the method argument, the method __call__ will be used by apply by default. You can find more details in our API reference documentation.

[carlosgmartin]

Thanks, but what if those encode and decode methods in the parent module aren't available, and only the child modules themselves (encoder and decoder) are available/accessible?

[chiamp]

You can define a parent Module to call your AutoEncoder submodules:

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

class Autoencoder(nn.Module):
  encoder: nn.Module
  decoder: nn.Module

class Model(nn.Module):
  def setup(self):
    self.autoencoder1 = Autoencoder(nn.Dense(10), nn.Dense(3))
    self.autoencoder2 = Autoencoder(nn.Dense(9), nn.Dense(2))

  def __call__(self, x):
    x = self.autoencoder1.encoder(x)
    print(f'encoder1 output shape: {x.shape}')
    x = self.autoencoder1.decoder(x)
    print(f'decoder1 output shape: {x.shape}')
    x = self.autoencoder2.encoder(x)
    print(f'encoder2 output shape: {x.shape}')
    x = self.autoencoder2.decoder(x)
    print(f'decoder2 output shape: {x.shape}')
    return x

model = Model()
key = jax.random.key(0)
x = jnp.ones(2)
params = model.init(key, x)
print(jax.tree_map(jnp.shape, flax.core.freeze(params)))

out = model.apply(params, x)

encoder1 output shape: (10,)
decoder1 output shape: (3,)
encoder2 output shape: (9,)
decoder2 output shape: (2,)
FrozenDict({
    params: {
        autoencoder1: {
            decoder: {
                bias: (3,),
                kernel: (10, 3),
            },
            encoder: {
                bias: (10,),
                kernel: (2, 10),
            },
        },
        autoencoder2: {
            decoder: {
                bias: (2,),
                kernel: (9, 2),
            },
            encoder: {
                bias: (9,),
                kernel: (3, 9),
            },
        },
    },
})
encoder1 output shape: (10,)
decoder1 output shape: (3,)
encoder2 output shape: (9,)
decoder2 output shape: (2,)

Submodules / dataclass fields can only be accessed when your module is bound, meaning the variables are attached to your module. A module is only bound when you are calling .init and .apply. See Flax Module Lifecycle for more details.