netket.jax.vmap_chunked#
- netket.jax.vmap_chunked(f, in_axes=0, *, chunk_size, axis_0_is_sharded=False)[source]#
Behaves like jax.vmap but uses scan to chunk the computations in smaller chunks.
This function is essentially equivalent to:
nk.jax.apply_chunked(jax.vmap(f, in_axes), in_axes, chunk_size)
Some limitations to in_axes apply.
Note
If netket_experimental_sharding is enabled, this function assumes that chunked in_axes are sharded by default. This can be overridden by specifying axis_0_is_sharded=False.
- Parameters:
f (
Callable
) – The function to be vectorised.in_axes – The axes that should be scanned along. Only supports 0 or None
chunk_size (
Optional
[int
]) – The maximum size of the chunks to be used. If it is None, chunking is disabledaxis_0_is_sharded – specifies if axis 0 of the arrays scanned is sharded among multiple devices, The function is then computed in chunks of size chunk_size on every device. Defaults True if config.netket_experimental_sharding, oterhwise defaults to False.
- Return type:
- Returns:
A vectorised and chunked function