Multiple forward pass methods

[nico-bohlinger]

Hello,

I want to implement a simple Linen module that consists of two methods for the forward pass. A main __call__(self, x) method that calls a torso(self, x) method and further uses its outputs. During training I want to be able to call __call__(self, x) as well as torso(self, x). Until now I did everything with the @nn.compact syntax but for what I want to do I seem to need to use the setup method, at least that's what the documentation tells me.

Now I'm still confused how to call both methods correctly. Before training I normally call mymodule.apply = jax.jit(mymodule.apply) to jit the apply and later I would call mymodule.apply(module_params, x). How would the call look like for my torso method? I can pass the torso method to the apply mymodule.apply(module_params, x, method=mymodule.torso) but this does give me a <class 'method'> is not a valid JAX type error and I'm also worried that the jitting of the apply only applies to the usage with the __call__(self, x) method.

I googled a bit but didn't find a great example that shows this use case :(

[cgarciae]

Hey @nico-bohlinger, I don't think you should jit your Modules apply function directly, a better option is to create a helper function that is easier to jit e.g:

@jax.jit
def forward_torso(params, x):
  return mymodule.apply({'params': params}, x, method=mymodule.torso)

You can create a similar one for the normal call. You can even create a more general one that accepts the method as a static argument e.g:

from functools import partial

@partial(jax.jit, static_argnums=(2,))
def forward(params, x, method):
  return mymodule.apply({'params': params}, x, method=method)