nn.compact should be default

[juliuskunze]

Nice to see flax progress, getting processed into linen! (:

I guess the overwhelming majority of user module definitions can and should use nn.compact. Can we remove the attribute and have something like nn.noncompact instead for the few remaining cases where it's necessary? Otherwise, I think this would clutter most use cases with an extra line. In particular, linen would become more friendly for the (important) use case of defining tiny wrappers around existing parametrized functions, like a loss around a classifier etc.

[avital]

Hi @juliuskunze -- nice to "see" you here :) I turned this into a discussion -- if we end up with consensus on a change let's file an issue then.

During our Linen design discussions, at first we did make @nn.compact default (and always on __call__ -- note that the implementation of @nn.compact requires that only one method is annotated -- this is to avoid footguns about which submodules are shared and how things are named -- note in comparison to Haiku where "everything is compact" where parameter names are prefixed by method name and you need to (sometimes) use hk.transparant if you call one method from another. We hoped to avoid this complexity altogether).

So in our intermediate design, we had a distinction between Module (where __call__ is wrapped with nn.compact) and MultiModule (where it isn't). We did an internal alpha user research group and multiple people found this confusing. I suppose we could consider annotating setup with nn.noncompact but I'm pretty sure people would find that similarly confusing)

Another way to think about it, while I agree that the vast majority of modules benefit from using nn.compact that it's a bit of an educational device -- first, people learn that Modules look like they have seen before in PyTorch. And then we say, well, to co-locate concerns and make your code more concise you can use nn.compact.

Your point about wrapping existing functions is a valid one, and I've heard of people who want that kind of behavior to interoperate with existing numerical computation library. For that we could add a small wrapper e.g. module(lambda parent, x: parent.param('foo', ...)). @levskaya has been thinking about this. (and in the past @shoyer and @adarob have also asked)

[peterjliu]

I think the pytorch pattern has shown to be durable and popular, so I'd be very careful about making @nn.compact the default.