Is it possible to have a module without a __call__ or compact?

[dieterichlawson]

I am writing a module which has several methods which use the parameters of the module but where none of them are 'privileged' over the other. So it doesn't really make senses for one of them to be called __call__. Is it possible to write a module without __call__ or compact in this case?

[dieterichlawson]

To answer my own question, it seem that having an empty __call__ solves this:

  @nn.compact
  def __call__(self):
    pass

This is not ideal, however.

In general, it seems that the user experience for modules that use multiple methods, and methods which are not named __call__ isn't ideal (although I am probably missing some stuff). Specifically, I have seen the following used to call other_method on module:

module.apply({'params':params}, arg1, arg2, method=module.other_method) 

This is verbose for my taste. Is there a way to have it be something more like

module.other_method(params, arg1, arg2) 

?

[avital]

Hi Dietrcich, there's nothing special about __call__ when defining modules in Linen. If you want to use nn.compact, you can only specify it on one method. But you can also write modules with no compact methods at all, and define all of your parameters and submodules in setup. You can see this for example on the "Sharing submodules" section in https://docs.google.com/document/d/1hYavTVPaKVVe9Be8pCB7yW7r6dDv3RALVNit8NZca4c/edit" or in our VAE example

As for your question about the call to apply -- yes, you're right that it's a bit odd. The reason that it's helpful to encapsulate calling methods with an explicit apply wrapper is that if you have additional state (e.g. batch norm, autoregressive cache, diagnostic output, intermediate pre-activations) then you need a way to extract those from the output (in that case apply returns two outputs). Another option we considered is making module instance fully mutable, and asking the user to read out the value of those additional "state" collections, but then the problem is that your code is not safe under jit and other JAX transformations. So while this approach is not ideal, we generally find it a reasonable trade-off between other concerns.

Please do share additional perspectives and let us know if you think we've got this wrong -- we are always learning!

[avital]

(Responding concretely to the module.other_method(params, arg1) proposal, that would confuse users because the method is defined without the params argument -- so not sure how that'd work.)

But maybe there's some syntactic sugar we haven't thought of that does work

[dieterichlawson]

Thanks for clarifying about nn.compact. I should have been clearer. What I want is the ability to write a module both without a __call__ method and without nn.compact, the reason being that none of the methods in the module which use the parameters are privileged above the others enough to denote them as the 'default behavior' of the module. Both examples that you gave have a __call__ method, but no nn.compact. When I tried to define a module without either, I got an error.

In regards to the calling syntax, what about something like

module(params).other_method(arg1, arg2)

or

module.apply(params).other_method(arg1, arg2)

?

These take care of your issue with adding arguments to method definitions. I'm not exactly sure what the issues are when dealing with other state like batchnorm, caches, etc... so I'm sure this proposal fails for those cases, but it seems nice to me nonetheless.

[avital]

oh -- what error do you get w/o __call__ or nn.compact? That's a bug -- could you file an issue with a repro? We should fix that.

Let me think a bit more about your module.apply idea. I think it's similar to something that @levskaya floated around a while back.

[avital]

module.apply(params).other_method(arg1, arg2) might be reasonable -- it's just currying our current apply method differently.

The upside is that we no longer have to supply a janky way to decide which method to call.

The downside is that we would have to explain that if you special the mutable argument to apply then the return value of other_method would be output, updated_variables, not just the output.

Curious to hear @jheek and @levskaya's thoughts. @shoyer also liked this.

[dieterichlawson]

Hey just wanted to let you know that I’m on vacation this week so won’t be able to respond until next Monday. Thanks for considering this stuff!

…

On Sun, Oct 11, 2020 at 7:32 AM Avital Oliver ***@***.***> wrote: module.apply(params).other_method(arg1, arg2) might be reasonable -- it's just currying our current apply method differently. The upside is that we no longer have to supply a janky way to decide which method to call. The downside is that we would have to explain that if you special the mutable argument to apply then the return value of other_method would be output, updated_variables, not just the output. Curious to hear @jheek <https://github.com/jheek> and @levskaya <https://github.com/levskaya>'s thoughts. @shoyer <https://github.com/shoyer> also liked this. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#519 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AABD74BEGYWC5RIOT6P5F4TSKGXX7ANCNFSM4SIBYISA> .