stochastic reconfiguration with jvp and vjp
Created by: inailuig
Adds the option to do stochastic reconfiguration where the gradients are never stored, but (re-)calculated using jvp and vjp for every matrix-vector product in the iterative solver.
There are several advantages of this approach:
- works for large numbers of samples and parameters where its impossible to store the gradients
- non-homogeneous parameters can be treated more easily
- can be faster for certain kinds of machines