Repository navigation
Optimize memory usage - #2
Merged
Merged
Conversation
Anya497
force-pushed
the
optimize_memory_usage
branch
3 times, most recently
from
April 25, 2026 18:31
61ca8ac to
516885a
Compare
Anya497
force-pushed
the
optimize_memory_usage
branch
2 times, most recently
from
July 2, 2026 14:08
7eaf76d to
0749356
Compare
Anya497
force-pushed
the
optimize_memory_usage
branch
from
August 22, 2026 16:25
a7cbdbc to
b279797
Compare
…ith shape [grid, orbitals, orbitals].
Anya497
force-pushed
the
optimize_memory_usage
branch
from
October 7, 2026 07:26
33de640 to
aa49333
Compare
Closed
CI on PR #2 failed all 24 parametrizations of test_kinetic_density_equivalence with AttributeError: module 'jax.numpy' has no attribute 'testing'. jax.numpy exposes no `testing` submodule on any jax in the supported range (verified on 0.10.2: hasattr(jnp, "testing") is False), so the assertion could never resolve and every case errored before comparing anything. numpy.testing.assert_allclose has the identical rtol/atol/err_msg surface this test uses and accepts JAX arrays through __array__ conversion; numpy already ships with jax, so no new dependency. Rejected: a try/except import shim around jnp.testing (hides the API's real home) and jnp.allclose + bare assert (loses the value-diff diagnostics). Also style-clean this touched file (black formatting, module/function docstrings, third-party-before-first-party import order, trailing whitespace) because the quality gate covers changed files; the repo-wide style debt stays tracked by #7. Verified locally under the CI environment (PYSCF_CONFIG_FILE set): all 24 parametrizations pass, the three CI unit files pass (29 tests), and all eight integration files pass (118 tests) when run per-file like the workflow does. Closes #12
After the unit-test fix, CI advanced to the integration step and failed on tests/integration/solids/test_functional_implementations.py::test_pw92[sol1] with ValueError: negative dimensions are not allowed. PySCF's FFTDF get_eri computes the remaining budget as max_memory - current RSS (pyscf/pbc/df/fft_ao2mo.py:59); the default MAX_MEMORY is 4000 MB (pyscf/__config__.py), but the CI interpreter carrying JAX + TensorFlow + PySCF already sat at ~4297 MB RSS, so the remaining budget came out at -297.44 MB and the ERI block size went negative. Reproduced locally with PYSCF_MAX_MEMORY=100 (same traceback, max_memory = -853.71), which also confirms the env var reaches pyscf.lib.param. Set MAX_MEMORY = 8000 in .github/workflows/pyscf_conf.py rather than in the workflow env: that file is already the job's PYSCF_CONFIG_FILE, it is exec'd after the environment read so it wins over PYSCF_MAX_MEMORY, and it keeps one knob for PySCF behavior in CI and in local CI-reproduction runs. 8000 MB clears the observed RSS with ~3.7 GB of headroom; pyscf sizes its ERI buffer from the remaining budget, so total RSS stays ~8 GB within the runner's RAM. The failure is allocator-variance sensitive (today's green agent_settings run landed just under the 4000 MB line), so the fix removes the near-4000 margin rather than hoping RSS stays low. Closes #12
This was referenced Oct 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR
einsumoperator;