Does Netket support involving vjp function model construction?

[miaoL1]

When constructing probability for metropolis sampling, I used a custom_vjp function. However, when optimizing parameters, error occured:

File D:\Users\mate\anaconda3\envs\netket\lib\site-packages\netket\optimizer\linear_operator.py:100, in LinearOperator.solve(self, solve_fun, y, x0, **kwargs)
     84 def solve(
     85     self, solve_fun: SolverT, y: PyTree, *, x0: Optional[PyTree] = None, **kwargs
     86 ) -> PyTree:
     87     """
     88     Solve the linear system x=⟨S⟩⁻¹⟨y⟩ with the chosen iterative solver.
     89 
   (...)
     98             None if there are no additional information provided.
     99     """
--> 100     return self._solve(jax.tree_util.Partial(solve_fun), y, x0=x0, **kwargs)
File D:\Users\mate\anaconda3\envs\netket\lib\site-packages\netket\optimizer\qgt\qgt_onthefly.py:172, in QGTOnTheFlyT._solve(self, solve_fun, y, x0, **kwargs)
    171 def _solve(self, solve_fun, y: PyTree, *, x0: Optional[PyTree], **kwargs) -> PyTree:
--> 172     return _solve(self, solve_fun, y, x0=x0)

    [... skipping hidden 23 frame]

File D:\Users\mate\anaconda3\envs\netket\lib\site-packages\jax\_src\interpreters\ad.py:751, in raise_custom_vjp_error_on_jvp(*_, **__)
    750 def raise_custom_vjp_error_on_jvp(*_, **__):
--> 751   raise TypeError("can't apply forward-mode autodiff (jvp) to a custom_vjp "
    752                   "function.")

TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.

It seems that jvp function is expected. Does Netket support involving vjp function in autodiff calculation?

[PhilipVinc]

You are using SR with QGTOnTheFly that requires forward mode differentiation.

Use another QGT implementation based on backward mode AD such as QGTJacobian