move away from experimental host callback and use pure_callback on new versions of jax (!1312) · Merge requests · Vicentini Filippo / netket

Vicentini Filippo requested to merge pv/pure into master Aug 31, 2022

Created by: PhilipVinc

jax.experimental.host_callback has a few issues that have been solved in jax.pure_callback. This should work better inside of TDVP when we jit the solvers.

Also, the experimental function will eventually be deprecated.

For now we just add a version switch everywhere. We can remove it if we ever increase the minimum jax version that we require

move away from experimental host callback and use pure_callback on new versions of jax

Merge request reports