Chunking apply_fun for MCState.to_array (#1470)
`MCState.to_array()` cannot be used to instantiate the wave function
over the full Hilbert space for systems with fairly large size (but that
should be still accessible) due to a large memory usage. The
modification in `MCState.to_array()` consists in the possibility of
computing the wave function over chunks of the Hilbert space of size
`MCState.chunk_size` that are then combined to get the complete state
vector.
---------
Co-authored-by:
Clemens Giuliani <clemens@inailuig.it>
Showing
- netket/jax/_vmap_chunked.py 37 additions, 40 deletionsnetket/jax/_vmap_chunked.py
- netket/nn/utils.py 38 additions, 9 deletionsnetket/nn/utils.py
- netket/vqs/exact/state.py 2 additions, 0 deletionsnetket/vqs/exact/state.py
- netket/vqs/mc/mc_state/state.py 6 additions, 1 deletionnetket/vqs/mc/mc_state/state.py
- test/nn/test_utils.py 5 additions, 1 deletiontest/nn/test_utils.py
Loading
Please register or sign in to comment