Capture nondiff args before calling nkjax.vjp
Created by: inailuig
Alternative fix for https://github.com/netket/netket/issues/1435 Also see https://github.com/netket/netket/pull/1436
@PhilipVinc
Created by: inailuig
Alternative fix for https://github.com/netket/netket/issues/1435 Also see https://github.com/netket/netket/pull/1436
@PhilipVinc