netket.jax.vmap_chunked#

netket.jax.vmap_chunked(f, in_axes=0, *, chunk_size)[source]#

Behaves like jax.vmap but uses scan to chunk the computations in smaller chunks.

Parameters