Option to switch of jitting for parts of netket

[jwnys]

Hi,

I'm wondering whether it would be feasible to switch off jax jitting in certain sections of netket with some environment variables.
For example, it's my understanding that currently any model that we develop needs to be jittable.
This also makes debugging rather challenging...
It would be also useful though, to allow people to jit only subroutines in their model.
Is there currently an option to do this, without just switching off jax jitting overall?
What would be necessary to implement such a thing, and is this feasible at all?

Jannes

[PhilipVinc]

I'm wondering whether it would be feasible to switch off jax jitting in certain sections of netket with some environment variables.

For debugging purposes that's easily achieved (taken from jax-ml/jax#252 ):

import jax
jax.config.update('jax_disable_jit', True)
import netket

However for proper usage (jitting only part of the model) that is going to be quite hard: the whole of netket is written assuming a functional model.
We call in several places functional transformations on the model itself like jax.vmap (local energy) and jax.lax.while and jax.lax.fori_loop (sampling).
If your model is not jittable, those won't work either.

Why do you want to jit only parts of the model?

[jwnys]

Jitting part of the model is mainly relevant when you try something new, and e.g. use external packages, or custom python implementations etc. However, I came across this today: jax-ml/jax#766 which might be a good enough hack to accomplish something like that.

[femtobit]

Indeed, host_callback should work for the purpose of calling non-jitable code in a model. Note that there are actually a few places where this workaround is used in netket at the moment, see https://github.com/netket/netket/search?q=hcb.call.

[PhilipVinc]

ah yeah. I had forgotten about this.
It's fairly slow however