Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add jac_chunk_size keyword argument to ObjectiveFunction to reduce memory usage of forward mode Jacobian calculation #1052

Merged
merged 74 commits into from
Sep 26, 2024
Merged

Add jac_chunk_size keyword argument to ObjectiveFunction to reduce memory usage of forward mode Jacobian calculation #1052

merged 74 commits into from
Sep 26, 2024

Commits on Jun 12, 2024

  1. add netket to use vmap_chunked, change the vectorize calls in objecti… 

    …ve jvp_scaled to use custom vectorize
    @dpanici
    dpanici committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    e1ffec8 View commit details
    Browse the repository at this point in the history

  2. add chunk_size arg

    @dpanici
    dpanici committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    5aef59c View commit details
    Browse the repository at this point in the history

  3. add chunk_size to objectivefunction as well

    @dpanici
    dpanici committed Jun 12, 2024
    Configuration menu
    Copy the full SHA
    67844a7 View commit details
    Browse the repository at this point in the history

Commits on Jun 13, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Jun 13, 2024
    Configuration menu
    Copy the full SHA
    60e142c View commit details
    Browse the repository at this point in the history

Commits on Jun 14, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Jun 14, 2024
    Configuration menu
    Copy the full SHA
    c27bf9b View commit details
    Browse the repository at this point in the history

  2. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Jun 14, 2024
    Configuration menu
    Copy the full SHA
    f3b6683 View commit details
    Browse the repository at this point in the history

  3. set default chunk size to None to disable chunking by default:

    @dpanici
    dpanici committed Jun 14, 2024
    Configuration menu
    Copy the full SHA
    517a519 View commit details
    Browse the repository at this point in the history

  4. fix default chunk size in backend

    @dpanici
    dpanici committed Jun 14, 2024
    Configuration menu
    Copy the full SHA
    6b786d0 View commit details
    Browse the repository at this point in the history

Commits on Jun 24, 2024

  1. add a comment

    @dpanici
    dpanici committed Jun 24, 2024
    Configuration menu
    Copy the full SHA
    135adce View commit details
    Browse the repository at this point in the history

Commits on Jul 2, 2024

  1. add chunk_size arg to solve continuation

    @dpanici
    dpanici committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    0366043 View commit details
    Browse the repository at this point in the history

  2. fix error

    @dpanici
    dpanici committed Jul 2, 2024
    Configuration menu
    Copy the full SHA
    131ae78 View commit details
    Browse the repository at this point in the history

Commits on Jul 3, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Jul 3, 2024
    Configuration menu
    Copy the full SHA
    8cfcf8d View commit details
    Browse the repository at this point in the history

Commits on Jul 15, 2024

  1. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Jul 15, 2024
    Configuration menu
    Copy the full SHA
    8f0515c View commit details
    Browse the repository at this point in the history

Commits on Jul 23, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Jul 23, 2024
    Configuration menu
    Copy the full SHA
    c3ba7ac View commit details
    Browse the repository at this point in the history

Commits on Aug 16, 2024

  1. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Aug 16, 2024
    Configuration menu
    Copy the full SHA
    ce5742f View commit details
    Browse the repository at this point in the history

Commits on Aug 22, 2024

  1. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    1f32705 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    02ba37e View commit details
    Browse the repository at this point in the history

  3. remove netket dependence, vendor only needed parts from their vmap_ch… 

    …unked implementation
    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    feb006e View commit details
    Browse the repository at this point in the history

  4. add chunk_size arg to every Obective

    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    eda29ee View commit details
    Browse the repository at this point in the history

  5. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Aug 22, 2024
    Configuration menu
    Copy the full SHA
    95e25cc View commit details
    Browse the repository at this point in the history

  6. put chunk size in a few tests

    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    7aca7ad View commit details
    Browse the repository at this point in the history

  7. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    fcdf492 View commit details
    Browse the repository at this point in the history

  8. add chunk size arg to blocked test

    @dpanici
    dpanici committed Aug 22, 2024
    Configuration menu
    Copy the full SHA
    2239746 View commit details
    Browse the repository at this point in the history

Commits on Aug 26, 2024

  1. add info about chunk_size to docs

    @dpanici
    dpanici committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    74e9d0c View commit details
    Browse the repository at this point in the history

  2. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    d9b34d1 View commit details
    Browse the repository at this point in the history

  3. change attribution statements

    @dpanici
    dpanici committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    9ad2688 View commit details
    Browse the repository at this point in the history

  4. fix typo in attribution statements

    @dpanici
    dpanici committed Aug 26, 2024
    Configuration menu
    Copy the full SHA
    41171c8 View commit details
    Browse the repository at this point in the history

Commits on Aug 27, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Aug 27, 2024
    Configuration menu
    Copy the full SHA
    4dc0c98 View commit details
    Browse the repository at this point in the history

Commits on Aug 28, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Aug 28, 2024
    Configuration menu
    Copy the full SHA
    4b09c02 View commit details
    Browse the repository at this point in the history

Commits on Sep 3, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 3, 2024
    Configuration menu
    Copy the full SHA
    3e38d25 View commit details
    Browse the repository at this point in the history

Commits on Sep 6, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici committed Sep 6, 2024
    Configuration menu
    Copy the full SHA
    0d59bd9 View commit details
    Browse the repository at this point in the history

Commits on Sep 8, 2024

  1. add warnings

    @dpanici
    dpanici committed Sep 8, 2024
    Configuration menu
    Copy the full SHA
    5d15de3 View commit details
    Browse the repository at this point in the history

  2. set default to be dim_x/4 for chunk_size

    @dpanici
    dpanici committed Sep 8, 2024
    Configuration menu
    Copy the full SHA
    9aca31f View commit details
    Browse the repository at this point in the history

Commits on Sep 9, 2024

  1. fix warning causing test fail

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    6f2bbdb View commit details
    Browse the repository at this point in the history

  2. change chunk_size to jac_chunk_size, change docstring

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    f3b5f25 View commit details
    Browse the repository at this point in the history

  3. remove unused functions

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    e4e27a2 View commit details
    Browse the repository at this point in the history

  4. simplify chunking fxn

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    51e3e40 View commit details
    Browse the repository at this point in the history

  5. place batcehd_vectorize utils in own file

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    da359d5 View commit details
    Browse the repository at this point in the history

  6. add file

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    fbebfb8 View commit details
    Browse the repository at this point in the history

  7. remove jac_looped, deprecate 'looped' deriv_mode

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    fc8c98c View commit details
    Browse the repository at this point in the history

  8. re-implement scan_append, need to be more careful in replacing it

    @dpanici
    dpanici committed Sep 9, 2024
    Configuration menu
    Copy the full SHA
    98c93d5 View commit details
    Browse the repository at this point in the history

Commits on Sep 10, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 10, 2024
    Configuration menu
    Copy the full SHA
    0787b84 View commit details
    Browse the repository at this point in the history

  2. simplify function slightly

    @dpanici
    dpanici committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    bc1d498 View commit details
    Browse the repository at this point in the history

  3. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Sep 10, 2024
    Configuration menu
    Copy the full SHA
    1fe12df View commit details
    Browse the repository at this point in the history

  4. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 10, 2024
    Configuration menu
    Copy the full SHA
    5b245f4 View commit details
    Browse the repository at this point in the history

Commits on Sep 11, 2024

  1. remove print statements

    @dpanici
    dpanici committed Sep 11, 2024
    Configuration menu
    Copy the full SHA
    f088a06 View commit details
    Browse the repository at this point in the history

  2. simplify scan further

    @dpanici
    dpanici committed Sep 11, 2024
    Configuration menu
    Copy the full SHA
    7c804a1 View commit details
    Browse the repository at this point in the history

  3. add doc for reducing memory usage

    @dpanici
    dpanici committed Sep 11, 2024
    Configuration menu
    Copy the full SHA
    480a4a7 View commit details
    Browse the repository at this point in the history

  4. change constraint wrapper to use batched

    @dpanici
    dpanici committed Sep 11, 2024
    Configuration menu
    Copy the full SHA
    15a7f48 View commit details
    Browse the repository at this point in the history

Commits on Sep 12, 2024

  1. address some comments

    @dpanici
    dpanici committed Sep 12, 2024
    Configuration menu
    Copy the full SHA
    ea74b24 View commit details
    Browse the repository at this point in the history

  2. update chunk size to auto as default

    @dpanici
    dpanici committed Sep 12, 2024
    Configuration menu
    Copy the full SHA
    83c39b4 View commit details
    Browse the repository at this point in the history

  3. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 12, 2024
    Configuration menu
    Copy the full SHA
    83e96c0 View commit details
    Browse the repository at this point in the history

Commits on Sep 14, 2024

  1. fix test

    @dpanici
    dpanici committed Sep 14, 2024
    Configuration menu
    Copy the full SHA
    2df20b5 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Sep 14, 2024
    Configuration menu
    Copy the full SHA
    da92365 View commit details
    Browse the repository at this point in the history

  3. remove errant prints

    @dpanici
    dpanici committed Sep 14, 2024
    Configuration menu
    Copy the full SHA
    e433704 View commit details
    Browse the repository at this point in the history

Commits on Sep 18, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 18, 2024
    Configuration menu
    Copy the full SHA
    5ce88c6 View commit details
    Browse the repository at this point in the history

Commits on Sep 19, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 19, 2024
    Configuration menu
    Copy the full SHA
    e9efd27 View commit details
    Browse the repository at this point in the history

Commits on Sep 20, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 20, 2024
    Configuration menu
    Copy the full SHA
    60e1c74 View commit details
    Browse the repository at this point in the history

Commits on Sep 21, 2024

  1. change auto to base chunk size off of device memiry and jacobian size

    @dpanici
    dpanici committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    4137d36 View commit details
    Browse the repository at this point in the history

  2. update test

    @dpanici
    dpanici committed Sep 21, 2024
    Configuration menu
    Copy the full SHA
    bd34534 View commit details
    Browse the repository at this point in the history

Commits on Sep 22, 2024

  1. update changelog

    @dpanici
    dpanici committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    8b2645d View commit details
    Browse the repository at this point in the history

  2. Merge branch 'dp/jacobian-batched-vmap' of github.com:PlasmaControl/D… 

    …ESC into dp/jacobian-batched-vmap
    @dpanici
    dpanici committed Sep 22, 2024
    Configuration menu
    Copy the full SHA
    1f98694 View commit details
    Browse the repository at this point in the history

Commits on Sep 24, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 24, 2024
    Configuration menu
    Copy the full SHA
    1c1fab4 View commit details
    Browse the repository at this point in the history

  2. remove auto option for _Objective chunk size

    @dpanici
    dpanici committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    3e8e364 View commit details
    Browse the repository at this point in the history

  3. rename utils file

    @dpanici
    dpanici committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    b0b78fe View commit details
    Browse the repository at this point in the history

  4. resolve comments

    @dpanici
    dpanici committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    be57bbc View commit details
    Browse the repository at this point in the history

  5. resolve further comment

    @dpanici
    dpanici committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    cd865d0 View commit details
    Browse the repository at this point in the history

  6. change warnings to errors

    @dpanici
    dpanici committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    e0e513f View commit details
    Browse the repository at this point in the history

  7. add disclaimer to ObjectiveFunction for auto chunk size on HPC CPU

    @dpanici
    dpanici committed Sep 24, 2024
    Configuration menu
    Copy the full SHA
    a6e29d6 View commit details
    Browse the repository at this point in the history

  8. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 24, 2024
    Configuration menu
    Copy the full SHA
    6988660 View commit details
    Browse the repository at this point in the history

Commits on Sep 25, 2024

  1. Merge branch 'master' into dp/jacobian-batched-vmap

    @YigitElma
    YigitElma authored Sep 25, 2024
    Configuration menu
    Copy the full SHA
    f4fbbcf View commit details
    Browse the repository at this point in the history

  2. update changelog

    @dpanici
    dpanici committed Sep 25, 2024
    Configuration menu
    Copy the full SHA
    9100e7c View commit details
    Browse the repository at this point in the history

Commits on Sep 26, 2024

  1. correct changelog

    @dpanici
    dpanici committed Sep 26, 2024
    Configuration menu
    Copy the full SHA
    15d95f1 View commit details
    Browse the repository at this point in the history

  2. Merge branch 'master' into dp/jacobian-batched-vmap

    @dpanici
    dpanici authored Sep 26, 2024
    Configuration menu
    Copy the full SHA
    3e99510 View commit details
    Browse the repository at this point in the history