Support chunking on nkjax.expect

[PhilipVinc]

Add support for chunking to nk.jax.expect.

[PhilipVinc]

@inailuig what do you think?

[inailuig]

@inailuig what do you think?

seems ok

[jwnys]

Just a small question about this implementation:

Why is f(x) computed first in the forward separately? I'm guessing it's to compute ΔL_σ for which you need the mean in the backward. Given that f(x) is typically the expensive part to compute (since it typically contains connected elements, think about the Elocs): wouldn't it be better (maybe not feasible?) to compute the vjp once for the second term (thereby computing \nabla f(x)). This should give directly also f(x) for free. Using this on all samples, you can compute ΔL_σ, and with it, compute \nabla \log p(x) ΔL_σ, which is just the derivative of the log_pdf. This would just avoid computing f(x) twice, which I think is the case now: once in forward, once in backward.