FLIP: Support __init__ in Modules about flax HOT 7 CLOSED

In the proposed syntax, why can’t the user write z = self.encoder(x) in apply? What happens if they do? And why do the encoder and decoder modules have to be shared? What happens if they’re not?

from flax.

why can’t the user write z = self.encoder(x) in apply?

That will work too. In this case it also doesn't matter because the function is a one-liner. If it is more complicated one might want to use the module method from within apply for code reuse

And why do the encoder and decoder modules have to be shared? What happens if they’re not?

In the example we use the encoder and decoder at most once. So in this case partial would also work.
The idea is to promote the usage of shared in the constructor because using a module multiple times would bite you otherwise.

from flax.

This would be useful for normalizing flows as well. In the meantime, there seems to be a namespace issue with the _shared_module approach. By modifying this line, to use cls.__name__, my implementation is working again.

def compose_transforms(transforms):

    class TransformSequence(Transform):
        
        def _shared_modules(self):
            return [t.shared() for t in transforms]

        @flax.nn.module_method
        def transform(self, x):

            transforms = self._shared_modules()

            for t in transforms:
                x = t.transform(x)
            return x

        @flax.nn.module_method
        def inverse_and_log_det_jac(self, y):

            transforms = self._shared_modules()

            log_det_jac = 0.0
            for t in reversed(transforms):
                y, term = t.inverse_and_log_det_jac(y)
                log_det_jac += term

            return y, log_det_jac

    return TransformSequence

Is there a better workaround for now?

from flax.

My main worry about supporting __init__ is that it could lead to incorrect assumptions about how Flax works, due to apparent similarity with normal class syntax. If I saw a call like the proposed AutoEncoder without knowing anything about Flax, I would expect the only valid argument to AutoEncoder(*args, **kwargs) to be those that appear explicitly on __init__, but that isn't how Flax works.

Some ideas:

• Use a different name from __init__, e.g., setup (we already use init for variables).
• Consider (conditionally?) switching to an explicit setup/call split like Keras or Haiku: AutoEncoder(**params)(x). This would probably be more pervasive than you want.

The magic separation of arguments between __init__ and apply also worries me a little bit. I don't know if there is a good way to do this, but I do think passing all **kwargs to __init__ (or setup) is a better alternative than using introspection.

from flax.

This would be useful for normalizing flows as well. In the meantime, there seems to be a namespace issue with the _shared_module approach. By modifying this line, to use cls.__name__, my implementation is working again.

def compose_transforms(transforms):

    class TransformSequence(Transform):
        
        def _shared_modules(self):
            return [t.shared() for t in transforms]

        @flax.nn.module_method
        def transform(self, x):

            transforms = self._shared_modules()

            for t in transforms:
                x = t.transform(x)
            return x

        @flax.nn.module_method
        def inverse_and_log_det_jac(self, y):

            transforms = self._shared_modules()

            log_det_jac = 0.0
            for t in reversed(transforms):
                y, term = t.inverse_and_log_det_jac(y)
                log_det_jac += term

            return y, log_det_jac

    return TransformSequence

Is there a better workaround for now?

@mattwescott I think your issue was introduced in a recent change to the default name policy. This is probably a bug that should be fixed. I'll look into it ASAP

from flax.

• Use a different name from __init__, e.g., setup (we already use init for variables).

I agree with this proposal

• Consider (conditionally?) switching to an explicit setup/call split like Keras or Haiku: AutoEncoder(**params)(x). This would probably be more pervasive than you want.

That's a very big change to make and takes away a key advantage of calling modules as functions.

The magic separation of arguments between __init__ and apply also worries me a little bit. I don't know if there is a good way to do this, but I do think passing all **kwargs to __init__ (or setup) is a better alternative than using introspection.

I do really dislike the introspection as well. Passing all kwargs to setup is probably fine the main downside is that they are also passed to all the module methods which then probably need something like an **unused_kwargs.

from flax.

This is no longer relevant since Linen has landed, so I'm closing this for now.

from flax.