@compact is great and I'm looking forward to more development/moving it out of experimental. 🙏

[kailukowiak]

Thanks for developing @compact it makes writing neural nets so much easier. This is partially because I'm coming from Python frameworks which are more similar to that but also because it introduces so much flexibility and readability.

To me this is a huge improvement from defining models only with Chain. While Chain is great and gives a lot of improved usability over PyTorch and Flax's nn.Sequential it is still limiting to use on more complex models.

As it currently stands, it's totally fine but what do you think about trying to move the layer definition into the do block like Flax's @nn.compact?

Thanks again 🤗

[avik-pal]

Can you elaborate on what you mean by the last part (maybe an example_?

[kailukowiak]

For sure. In Flax you can either explicitly define the layers in setup:

class MLP(nn.Module):
  def setup(self):
    # Submodule names are derived by the attributes you assign to. In this
    # case, "dense1" and "dense2". This follows the logic in PyTorch.
    self.dense1 = nn.Dense(32)
    self.dense2 = nn.Dense(32)

  def __call__(self, x):
    x = self.dense1(x)
    x = nn.relu(x)
    x = self.dense2(x)
    return x

Or define them more implicitly in the actual call:

class MLP(nn.Module):

  @nn.compact
  def __call__(self, x):
    x = nn.Dense(32, name="dense1")(x)
    x = nn.relu(x)
    x = nn.Dense(32, name="dense2")(x)
    return x

Currently Lux works like:

my_model = @compact(dense1 = Dense(32, 32),
                    dense2 = Dense(32, 32),
                    act = relu) do x
    x = act(dense1(x))
    x = dense2(x)
    return x

but could potentially work like:

my_model = @compact do x
    x = Dense(32, 32)(x)
    x = relu(x)
    x = Dense(32, 32)(x)
    return x

I'm not the best at meta programming so I'm not sure about the feasibility of this but having the option to keep the layer definitions and control flow of the inputs in the same area can be nice.

[avik-pal]

I see what you mean. What you are proposing used to be what @compact was initially going to do. But supporting that is quite messy because you need to parse the user code and modify the function. Doing that doesn't really justify the code complexity vs usability tradeoff.

[avik-pal]

#584 moves compact out of experimental!