Simplify submodule access from top-level Module

[cgarciae]

Its common for users to want to get access to a submodule defined in setup from the outside, however they soon find out that .setup doesn't run upon construction and then they have to figure out how to get access to this. E.g:

class MyModule(nn.Module):
  def setup(self):
    self.submodule = MySubModule()
  ...
  
module = MyModule()
submodule = module.submodule  #<== ERROR

Our current solution as proposed in the Transfer Learning guide is to use nn.apply with a helper function that extracts the needed objects (since the Module is bounded at this point):

def extract_submodule(module):
    submodule = module.submodule.clone() # clone to avoid leakage
    submodule_variables = module.submodule.variables
    return submodule, submodule_variables

submodule, submodule_variables = nn.apply(extract_submodule, module)(variables)

Since this pattern is both common and has some rough edges (users need to remember to use .clone which is really more of an internal API) it would be good to create a helper function to automate this use case.

[marcvanzee]

Closing this since it seems to be a duplicate of #2478