[P] Equinox compilation retnet

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 …

build compilation inference jax machinelearning parameters reason representation training