how to define and train a model with mult forward?

[Nightbringers]

here is simple case, one model have mult forward:
class model(nn.Module):

dtype: jnp.dtype = jnp.float32

def setup(self):
    self.encoder1 = nn.Conv(self.dim, (3, 3),use_bias = False,dtype=self.dtype)
    self.encoder2 = nn.Conv(self.dim, (3, 3),use_bias = False,dtype=self.dtype)
    self.encoder3 = nn.Conv(self.dim, (3, 3),use_bias = False,dtype=self.dtype)

  
def e1(self, obs):

    s = self.encoder1(obs)
    return s 

def e2(self,obs):

    s = self.encoder2(obs)
    return s 

def e3(self, obs):

    s = self.encoder3(obs)
    return s 

here is the problem: it don't have call it will raise a error. And when train the model, TrainState only have one apply_fn

[lucaslingle]

Disclaimer: I'm not a Flax developer but I've used it for about a year now.

For apply_fn:

Per the docs, you don't have to use TrainState's apply_fn, it is just for convenience. You can just set it to None.

In general, when you want to apply one of the encoders as your forward method, you can call apply on your instantiated Flax module, and pass in the appropriate forward method required, like so:

Model().apply({"params": params}, inputs, method=Model.e1)

For the __call__ error:

During initialization, I assume you want to obtain the parameters for all encoders wrapped by the Model class, so would be convenient to define the __call__ method to run all of them. When you run init with any given method (e.g., method=Model.e1), it will only return the variables/parameters required for that method. So running all encoders with __call__ seems like convenient way to initialize all encoders' parameters and solve your error for __call__ as well.

Misc.:
If you plan on using the encoders in different circumstances during training, you might want to just call the encoders directly, as you might not need a wrapper class. However, in this case, you would have to manage multiple TrainState objects, which might be inconvenient. This depends on your codebase and requirements.