When exactly is setup() called for linen modules?

[davisyoshida]

I have a module which uses jax.jit in its setup(), under the assumption that setup will only be called when I call model.init. However, by logging the calls to setup() I found that it was called not just when I initialize the model, but also later when I use its apply(). I found this very unexpected, is there some documentation on what all will trigger a call to setup()?

(The reason I was jitting the function inside setup is that it depends on submodules, and jax.jit() doesn't like modules as static args since they can't be hashed, so I figured I'd make it in a closure where the submodules were accessible)