Equinox compilation retnet

Follow the full discussion on Reddit.
I'm working on replicating the retnet model in Equinox + Jax. For the model there are 2 representations: parallel and recurrent (ignoring chunkwise). They use the recurrent at inference and parallel at training. In equinox, should i build 2 seperate models for each representation sharing parameters or build one with if statements. The reason I'm asking is if I use one model with if statements for inference and training, will the model re-compile whenever switching between representations or not? If so then wouldn't building 2 models each compiled originally with their representation make it faster.