How to get access of the parameters within a model?

[wztdream]

Hi,
Bellow is an simple example, I need to do some calculation based on the value of parameters, how to get access of the parameters within a flax model?

import flax.linen as nn


class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=64, kernel_size=(3, 3), strides=(1, 1), padding="SAME")(x)
        x = nn.Conv(
            features=1,
            kernel_size=(1, 1),
            strides=(1, 1),
            padding="SAME",
            name="need_kernel_here",
        )(x)
        # here wee need to do calculation based on the kernel of "need_kernel_here"
        w, b = self.get_paramers(
            param_name="need_kernel_here"
        )  # how to do this in flax?
        my_calculation = (
            w.mean() + b.mean()
        )  # just a simple example, in my real project the calculation is more complex but the calculation only need read the value of w and b, not change them
        return x * my_calculation

[andsteing]

Inside your __call__() function the module is in a "bound" state and the variables of submodules can simply be accessed via sub_module.variables:

class Foo(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Conv(features=64, kernel_size=(3, 3), strides=(1, 1), padding="SAME")(x)
        conv = nn.Conv(
            features=1,
            kernel_size=(1, 1),
            strides=(1, 1),
            padding="SAME",
            name="need_kernel_here",
        )
        x = conv(x)
        w = conv.variables['params']['kernel']
        b = conv.variables['params']['bias']
        my_calculation = (
            w.mean() + b.mean()
        )  # just a simple example, in my real project the calculation is more complex but the calculation only need read the value of w and b, not change them
        return x * my_calculation


foo = Foo()
x = jnp.zeros([1, 2, 2, 3])
vs = foo.init(jax.random.PRNGKey(0), x)
foo.apply(vs, x)

[wztdream]

Nice! so simple, thank you!