netket.jax.vmap_chunked

netket.jax.vmap_chunked(f, in_axes=0, *, chunk_size)

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

Parameters

• f (Callable) –

• chunk_size (Optional[int]) –