- exception netket.errors.JaxOperatorSetupDuringTracingError#
Illegal attempt to use a Jax-operator Numba-operators constructed inside of a Jax function transformation with non-constant data.
This happens when building a
DiscreteJaxOperatorinside of a function that is being transformed by jax with transformations such as
jax.grad(), and the operator is not compatible with Jax.
DiscreteJaxOperatorcan be used inside of jax function transformations, but NetKet is currently limited in that you cannot build them inside of Jax transformations.
To avoid this error you should build your operators outside of the jax context.
(i) Building a Jax operator outside of a jax context#
Build the operator outside of a jax context.
import netket as nk import jax import jax.numpy as jnp N = 2 ham = nk.operator.PauliStringsJax(['XI', 'IX'], jnp.array([0.3, 0.4])) samples = ham.hilbert.all_states() @jax.jit def compute_values(ham, s): return ham.get_conn_padded(s) compute_values(ham, samples)
This limitation is not systematic, and it could be lifted in the future by some interested coder. If at that moment Jax will support dynamic shape, this feature could be implemented at no additional runtime cost. If Jax won’t support yet dynamic shapes then it should be implemented as a secondary path (instead of this error) that is only taken if the operator is constructed inside of a jax context. This path will lead to a slightly less optimized, higher computational cost operators.
If you are really interested in contributing to NetKet and find yourself in need of building operators in a jax context (for example because you’re doing optimal control), get in touch with us by opening an issue.
Most operators lazily initialise the fields used to compute the connected elements only when needed. To check whether an operator was initialized you can probe the boolean flag
operator._initializedis True, you can safely call
get_conn_padded()and similar methods. If it is False, then the setup procedure will be handled by an internal method usually called
operator._setup(). If you see this error, it means that this method internally uses dynamically determined shapes, and it is what should be converted to be jax-friendly.