Improving JIT and tree_map compatibility: separating state and structure

[theo-brown]

Hi! This is a suggestion crossed with a question.

In Flax, the 'correct' way of doing things is to separate the state and the structure. For example,

model = nn.Dense(features=5)
params = model.init(key, x)
model.apply(params, x)

This enables better JIT compatibility, as functions can be JITted once based on the structure of the model (e.g. a GP with constant mean and a Matern52 kernel), and then reused even if the parameters change (e.g. a modified lengthscale from marginal likelihood fitting to some new data). It also might fix the problem I encountered in #516.

So the question/suggestion is: why is this not done in GPJax? It would be fairly straightforward to modify some of the APIs (especially the fit functions) so that the nnx.GraphDef and nnx.State are handled as separate arguments. In fact, I've ended up doing this in my project that uses GPjax, so I'd be happy to propagate some of the changes over.

While this may not make a difference to the actual runtime of performing GP inference - because presumably the matrix inversions etc are all JITted anyway? - it would make it a lot easier to integrate GPJax GPs with more complex routines where GPs must be passed around between JITted functions (e.g. in some more complex BayesOpt workflows).

[thomaspinder]

Hi @theo-brown . I 100% agree with your proposal here. To be frank, the reason we didn't do this is that NNX was still evolving as we migrated to it, and likely this detail passed us by. If you'd be willing to lead the API change, then I'd most happily support and integrate when ready.

[theo-brown]

Awesome! Yes, I thought it would be something like this, as NNX is still fairly new. I'll start putting together a PR over the next few weeks.