Parameters not available as instance properties after setup()

[ingmarschuster]

Hi,

I'm using the following minimal example for two methods that use module parameters:

import jax.numpy as np
from jax.numpy import log, exp
import jax.random as rand
import flax.linen as ln

class MultipleForw(ln.Module):
    def setup(self):
        print("Setting up")
        self.s1 = self.param("s1", ln.initializers.ones, (1,))
        self.s2 = log(1 + exp(self.param("s2", ln.initializers.ones, (1,)) - 1))
        print(self.s1, self.s2)

    def __call__(self, X, ):
        return X * log(1 + exp(self.s1 - 1))

    def call2(self, X):
        return X * self.s2


def test():
    mf = MultipleForw()
    X = np.arange(5)
    mf.init(rand.PRNGKey(0), X)
    mf(X)
    mf.call2(X)

test()

this prints out

Setting up
[1.] [0.6931472]

as expected, then fails at the call mf(X) with
AttributeError: 'MultipleForw' object has no attribute 's1'
Any ideas? What am I missing? Btw, is the form of self.s2 = log(1 + exp(self.param("s2", ln.initializers.ones, (1,)) - 1)) for getting a parameter that can be optimized in an unconstrained way but remains positive supported or should I rather go the way self.s1 is used?

[ingmarschuster]

Self-answering: I forgot to use apply rather than calling the method directly. Calling it directly meant the parameters where not set. corrected function:

def test():
    mf = MultipleForw()
    X = np.arange(5)
    par = mf.init(rand.PRNGKey(0), X)
    print(mf.apply(par, X))
    print(mf.apply(par, X, method = mf.call2))

The self.s2 version works as expected and is equivalent to the self.s1 version.

[avital]

Yeah we should make it raise an error in case you call a method directly on an unbound module (that is -- not via apply). Do you mind filing an issue?

[ingmarschuster]

done in issue #1014