netket.jax.apply_chunked#
- netket.jax.apply_chunked(f, in_axes=0, *, chunk_size, axis_0_is_sharded=False)[source]#
Takes an implicitly vmapped function over the axis 0 and uses scan to do the computations in smaller chunks over the 0-th axis of all input arguments.
For this to work, the function f should be vectorized along the in_axes of the arguments. This means that the function f should respect the following condition:
assert f(x) == jnp.concatenate([f(x_i) for x_i in x], axis=0)
which is automatically satisfied if f is obtained by vmapping a function, such as:
f = jax.vmap(f_orig)
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
) – A function that satisfies the condition abovein_axes – The axes that should be scanned along. Only supports 0 or None
chunk_size (
int
|None
) – The maximum size of the chunks to be used. If it is None, chunking is disabledaxis_0_is_sharded (
bool
) – 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: