Repository navigation
Agent settings. Tests fixes - #8
Merged
Merged
Conversation
Anya497
commented
Oct 6, 2026
Anya497
commented
Owner
- Add basic agent settings
- Support newer jax
- Fix linter mistakes
The workflow-management skill, AGENTS.md, and the six remaining skills were copied from another project. They pointed at facts and files that do not exist here, so following them literally was impossible: every one of them referenced `docs/developer.rst`, `docs/reference/` autosummary lists, `.github/workflows/ python_tests.yaml`, `python_linting.yaml`, `docs.yaml`, `poetry`, component directories (`AIAgent/`, `tools/compstrat/`), the ONNX/path-selector worked example, and the `todowrite` tool. All of that was replaced with what this repository actually has. Facts corrected to match the repository: - CI workflow is `.github/workflows/install_and_test.yml` (not `.yaml`) and runs `pytest -v <file>` per file from the root, with `PYSCF_CONFIG_FILE` set to `.github/workflows/pyscf_conf.py`. That env var is now called out because the functional-prediction tests silently differ from CI without it. - Style gate is `.pre-commit-config.yaml` (`black` + `pylint -rn -sn --rcfile=.pylintrc`). There is no linting CI workflow. - Test layout is `tests/unit/`, `tests/integration/molecules/`, `tests/integration/solids/`. - Installation is `pip install -e .` / `pip install -e ".[examples]"`. - The docs build gate is gone: this repository has no Sphinx site, so the quality gate is tests + style only. Documentation surface replaces the Sphinx docs: `README.md`, `examples/` (`basic_notebooks`, `intermediate_notebooks`, `advanced_scripts`, `article_experiments`), and numpydoc-style docstrings. The mapping is now defined once, in the `planning` skill, and referenced from `subtask-loop` and `code-review` so it is not restated. Dangling skill references resolved by pointing at the owning artifact instead of loading a skill that does not exist here: - `code-style` -> `.pre-commit-config.yaml` + `.pylintrc` - `documentation` -> the documentation surface section in `planning` - `reusing` -> an explicit grep/glob reuse search over `grad_dft/`, `tests/`, `examples/`, seeded with `grad_dft/utils/` - `git-workflow` -> the commit format is now the single source of truth in `subtask-loop` step 6 (`(<issue>-S<n>): <summary>`); branching and integration stay in `workflow-management` - `todowrite` -> per-subtask cycle checklists inside `tasks/detailed_plan.md`, since that file is already the single source of truth and the plan is mirrored to the issue comment AGENTS.md listed 12 skills but only 7 exist; its table now lists exactly the 7 present skills, and a short "Repository Workflow Configuration" section records the integration branch so a fresh clone does not depend on local git config. Mode simplified to stacked: the integration branch is the personal branch `agent_settings`, so the direct-mode/PR-per-task branches were removed and every reference now resolves `INTEGRATION` as `git config --get graddft.integrationBranch || echo agent_settings`. That snippet is stated once in `workflow-management` and reused verbatim in `code-review` and `subtask-loop`. The local repo config was set to match. Worked examples in the skills are now GradDFT-specific (`utils/types.py` aliases, `utils/chunk.py`, a DM21 prediction test mirroring `test_predict_B88.py`) instead of ONNX export and graph-algorithm blockers. Not committed: untracked `.vscode/` and `grad_dft.egg-info/` build/editor artifacts, which are out of scope and not covered by `.gitignore`. Closes #4
`.vscode/` and `grad_dft.egg-info/` showed up as untracked noise in every `git status`, because `.gitignore` covered neither. The existing entries (`*__pycache__/`, `data/training/dissociation/`, `runs/`, `*/__pycache__/`) are all narrow, so these follow the same style: `.vscode/` for the editor directory and the generalized `*.egg-info/` rather than the single `grad_dft.egg-info/`, since `pip install -e .` regenerates that directory under whatever the package is named and would otherwise reappear as untracked noise. Also drop a trailing blank line from `.gitignore`; the file ended without a newline, which makes every future diff on it noisier than it needs to be. Closes #4
`from jax.config import config` no longer resolves: JAX removed the
`jax.config` submodule, so every module that used it raised
`ModuleNotFoundError: No module named 'jax.config'` at import. That took out
`test_predict_B88.py`, which CI *does* run, so the suite could not pass at all on
a fresh install from `requirements.txt` (verified on jax/jaxlib 0.10.2).
`jax.config` still exists as an attribute of the `jax` package, which is why
`grad_dft/utils/types.py` (`jax.config.x64_enabled`) keeps working and needs no
change. That makes the supported spelling `from jax import config` -- already used
by 12 of the 17 files in this repo that call `config.update("jax_enable_x64", ...)`,
including `tests/integration/molecules/test_Harris.py` and
`tests/unit/test_eigenproblem.py`. This change adopts that existing convention
rather than inventing one, so the call sites are untouched.
Fixed in the three `test_predict_*.py` modules and, deliberately, in
`examples/advanced_scripts/{train_scf_loop,train_complex_model}.py` as well:
those are the user-facing documentation surface and carry the identical defect,
so following the README-adjacent example scripts would fail the same way.
`test_predict_B3LYP.py` stays out of CI (Grad-DFT's B3LYP differs from PySCF's,
as that file documents) and `test_predict_DM21.py` stays out because it is broken
independently of JAX; both are fixed only so they are not left importing a module
that no longer exists.
Measured, full-tree `pytest --collect-only`: 7 collection errors / 60 tests
collected before, 4 / 69 after. The 4 that remain are pre-existing and untouched
by this change -- 24 removed `jnp.clip(a_min=...)` call sites and PySCF's
`get_init_guess(max_cycle=...)`. They are recorded as S2's scope with their root
causes, not silently left behind.
`black` was applied to the five touched files (it also cleared 13 pre-existing
findings there: trailing whitespace, a missing final newline, collapsible
imports). Pylint on them went 124 -> 111 messages with no genuinely new message.
The repo-wide state -- 51 of 59 tracked Python files failing `black --check`, with
no lint job in CI -- is recorded in the plan's Design Notes and deliberately not
fixed here, since reformatting ~50 files would bury this compatibility fix.
S1 fixed the import error that stopped the suite collecting at all. This
subtask fixes the API removals it uncovered, so the suite can actually be run
end to end: `pytest --collect-only tests/` goes from 4 collection errors / 69
tests collected to 0 errors / 127 collected.
Three independent causes, none of them JAX's fault:
1. `jnp.clip(..., a_min=..., a_max=...)` was removed. On jax 0.10.2 the keyword
is simply gone -- `jnp.clip(x, a_min=0.5)` raises `TypeError: clip() got an
unexpected keyword argument 'a_min'`. 28 call sites now pass `min=` / `max=`:
24 in `grad_dft/{functional,popular_functionals,train}.py`, 4 in the two
`test_training.py` modules. Pylint, measured against this commit's parent:
20 errors down to 9, all 11 removed being exactly the `E1123
unexpected-keyword-arg 'a_min'` ones, and no genuinely new message (total
197 -> 186; the six apparent additions are pre-existing `W0621` warnings
renumbered by the one line added at each `max_cycle` site).
2. PySCF's `get_init_guess()` no longer takes `max_cycle`, so passing it through
`mf.kernel(max_cycle=...)` leaks it into `get_init_guess(**kwargs)` and
raises `TypeError`. It is an attribute on the SCF class, so the 5 call sites
set it and then call `kernel()`.
3. `test_training.py`, `test_non_xc_energy.py` and
`test_functional_implementations.py` each exist under both `molecules/` and
`solids/`. Pytest derives a module name from the basename, so the second
import collides with the first and fails with "import file mismatch". This is
why the `solids/` errors looked order-dependent: each file collected alone and
only broke once both copies were collected. The `tests/**/__init__.py` files
make each directory a package so the derived names differ. CI never saw this,
because it runs one file per process and never collects both copies together.
The same two renames are applied to the six example notebooks that carry them,
so the documented workflows stop raising on their first cell.
Gate: the 10 per-file invocations from `.github/workflows/install_and_test.yml`,
run with `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py`, are 9/10 green:
2, 3, 35, 12+2failed, 6, 4, 26, 14, 10, 8 tests.
The 2 failures are `test_dm21[mol0]`/`[mol1]`, and they are NOT caused by this
task -- they fail identically on the parent commit, because two upstream breaks
are stacked: tensorflow-hub 0.16 deleted the TF1 `hub.Module` API that the
vendored DeepMind code calls, and PySCF 2.13's `eval_xc_eff` now dispatches
through `eval_xc1`, which `NeuralNumInt` does not override, so its `eval_xc` is
never called at all. Neither is a JAX change. The gate is reported BLOCKED on
this rather than reported green; both root causes, a working prototype for the
first, and the derivative-layout analysis for the second are written up in
#6.
`black` is deliberately not applied here: reformatting the 7 touched modules
costs ~737 changed lines against this 44-line fix. The repo-wide state -- 46 of
64 tracked Python files failing `black --check`, no lint job in CI -- is filed
as #7.
`requirements.txt` declared `jax>=0.4.14`, a floor the code has never been able to satisfy. Two independent reasons, and the higher of the two is what this commit sets: 1. `grad_dft/utils/chunk.py` imports `jax.extend.linear_util`. Inspecting the published sdists, `jax/extend/linear_util.py` is absent in 0.4.14 and 0.4.15 and present from 0.4.16. (`jax/extend/__init__.py` alone already exists in 0.4.14, which is why this is easy to miss.) An install that resolved to the oldest allowed JAX would fail at `import grad_dft`. 2. #5-S2 removed the last `jnp.clip(a_min=...)` call site in favour of `min=`/`max=`. Reading `jnp.clip`'s published signature per release: `min`/`max` do not exist at or below 0.4.26, exist alongside the deprecated `a_min`/`a_max` from 0.4.27 to 0.4.30, and are the only spelling from 0.4.31 on. So the code cannot run below 0.4.27 either. 0.4.27 is also the lowest value that keeps CI's existing matrix installable. The matrix still runs Python 3.9, and 0.4.30 is the last JAX release supporting 3.9, so any floor above 0.4.30 would make that job unresolvable. 0.4.27-0.4.30 all accept both spellings, so the 3.9 job is safe even when pip picks the newest JAX it can install there. `jax` and `jaxlib` move together because they are released in lockstep. No other pin is touched. `README.md` states no JAX version, so no documentation changes.
…plan Tick the remaining cycle steps and fold in what was measured while running the gate: the pylint before/after (20 errors -> 9), the corrected 0.4.27 floor with the per-release jnp.clip signature table, the full 10-file gate table, and the DM21 root-cause analysis with its decision to defer to #6 and the black debt to #7.
The vendored DeepMind code built its graph with the TF1 Hub Module API
(`hub.Module(spec=...)`, `hub.add_signature`, `hub.create_module_spec`), all of
which tensorflow-hub 0.16 removed, so `test_dm21` died at import-time graph
construction with `AttributeError: hub.Module`. Three changes, no behaviour
change intended:
1. `_build_graph` now loads the checkpoint with
`hub.load(self._model_path).signatures["default"]` and calls it as
`self._functional(**tensor_dict)`. The keyword names keep the
`tensor_dict$` prefix (only the TensorSpec `name=` fields are sanitised to
`_`), so the existing `tensor_dict = {f"tensor_dict${k}": ...}` construction
and the `grid_contribution` output key are untouched, as is the whole
`tf.gradients` derivative machinery downstream.
2. The `outputs` dict and the trailing `hub.add_signature(...)` call are gone
with the API. Their only consumer was the spec-based export, which now
raises, so the derivative tensors stay plain graph attributes consumed by
`eval_xc` via `self._session.run`.
3. `export_functional_and_derivatives` gets an explicit outcome instead of a
TF1 reimplementation: a documented `NotImplementedError` (TF-Hub >= 0.16
removed the Module/export API it was built on; a TF2 `tf.saved_model.save`
port would be a new feature). Its only callers are the vendored
`export_saved_model.py` / `neural_numint_test.py`, neither in the suite.
`requirements.txt` raises the floor to `tensorflow-hub>=0.16.0`: 0.14/0.15
import `tensorflow_hub.estimator`, which pulls `tensorflow.compat.v1.estimator`
— gone in TF 2.21 — so older hubs cannot even be imported in this environment
(verified from the published wheels).
Gate for this subtask, `PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py
pytest -v tests/integration/molecules/test_functional_implementations.py`:
12 passed, 2 failed — exactly the pre-#6 baseline, with `test_dm21` now failing
at the second documented break (`ValueError: cannot reshape array of size
240720 into shape (2,1,24072)`, the PySCF 2.13 `eval_xc_eff` routing that S2
fixes) instead of `AttributeError: hub.Module`. Pylint vs the subtask parent:
no new messages (the 4 baseline `E1101 tensorflow_hub no-member` errors are
gone with the removed API). Black is deliberately not applied: the baseline
already fails `black --check`, and a file-wide reformat of this vendored file
is out of scope (task #5 precedent).
PySCF 2.x's nr_rks/nr_uks grid loops call `ni.eval_xc_eff(...)` instead of `ni.eval_xc(...)`. NeuralNumInt only overrode eval_xc, so the call fell through to the base class, which evaluated with libxc and died on the meta-GGA-shaped density (`ValueError: cannot reshape array of size 240720 into shape (2,1,24072)`) — and NeuralNumInt.eval_xc was dead code. The new override routes the whole evaluation back to eval_xc (the DM21 graph run plus the local Hartree-Fock `_vmat_hf` side effects) and applies the chain rule from eval_xc's sigma-derivative convention to the density-parameter convention eval_xc_eff promises: - spin=0 -> (5, N): [dE/d rho, dE/d nabla rho, dE/d tau] - spin=1 -> (2, 5, N): the same rows per spin channel, with the cross term dE/d sigma_ab distributing between the two gradient rows. The layout is not guessed: it reproduces `pyscf.dft.numint.NumInt.eval_xc_eff` exactly (max |diff| = 0.0 on mgga_x_tpss for both spins). The alternative of feeding eval_xc's outputs through `xc_deriv.transform_xc` was prototyped first and produced garbage (its out row order is the xcfun convention), so the explicit chain rule is the design of record; the code comments point at PySCF's MGGA contract sites (nr_rks: wv[0]/wv[4] halving, _scale_ao_sparse on ao[:4]/wv[:4], _tau_dot_sparse on wv[4]; nr_uks the same per spin row). Rho arrives from eval_rho without the laplacian row (MGGA_DENSITY_LAPL is off) while the graph's placeholders are (6, N); a zero laplacian row is inserted before evaluation (the graph discards it). The preamble mirrors the base method (omega fallback, asarray, spin inference from rho.shape[0]), deriv == 0 still evaluates for the _vmat_hf side effect and returns [exc, None, None, None], and deriv > 1 raises NotImplementedError (DM21 has no second derivatives; the base class would fail inside transform_xc). The signature mirrors the base exactly, which is why the new method carries an inline pylint suppression for the two too-many-arguments checks. With this commit `test_dm21`'s PySCF leg — `mf.kernel()` on both parametrized molecules — completes for the first time on these dependencies; the test as a whole still cannot pass until S3, because its second leg (`generate_DM21_weights`) hits an independent break (TF1 SavedModel loading on TF 2.21). Gate for this subtask, with PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py: test_functional_implementations.py gives 12 passed / 2 failed (both failures now at test_dm21's second leg), and the collateral files test_predict_B88.py, test_non_xc_energy.py, test_Harris.py give 45 passed. Pylint vs the S1 parent: message multiset identical (the eval_xc_eff too-many-arguments pair is suppressed inline). Black is not applied to the vendored file, same precedent as S1 (baseline already fails `black --check`).
With the eval_xc_eff shim (S2) in place, test_dm21 moved to its second leg:
`DM21.generate_DM21_weights()` loads the same vendored checkpoint with
`tf.saved_model.load(folder)`. That is a TensorFlow 1 SavedModel, and on
TF >= 2.13 (the declared floor) its V1 restore machinery in
load_v1_in_v2.restore_variables fails inside the lifted restore function
(`TypeError: Binding inputs to tf.function failed due to 'too many
positional arguments'` — the pruned function has an empty signature but is
called with the variables path), so the test could never complete.
Instead of importing the SavedModel, read the variables straight out of its
checkpoint files with `tf.train.load_checkpoint(folder/variables/variables)`
+ `tf.train.list_variables`, appending the ":0" op suffix so the existing
name patterns in `vars_to_params` ("/w:", "/b:", "gamma:", "beta:") still
match. `vars_to_params` now iterates `(name, value)` pairs and
`tf_tensor_to_jax` is a plain `jnp.asarray` (the reader returns numpy arrays
directly, no .numpy() call needed). The method's interface (folder default,
return value) is unchanged, and the whole inspection stays local to the
method.
This completes test_dm21 for the first time on these dependencies: both
parametrized molecules pass end-to-end with the unchanged `atol=1`,
validating the S1+S2 PySCF leg against grad_dft's independently implemented
JAX port of DM21. Gate for this subtask, with
PYSCF_CONFIG_FILE=.github/workflows/pyscf_conf.py:
tests/integration/molecules/test_functional_implementations.py -> 14 passed
(0 failures). Pylint vs the S2 parent: symbol multiset identical (no new
messages; the line-shifted redefined-outer-name/too-many-statements are the
pre-existing ones). Black introduces no new debt: the baseline already fails
`black --check` (task #7's repo-wide reformat).
setuptools >= 81 removed pkg_resources, and a fresh pip install resolves setuptools (an unpinned requirement of pyscf) to 84.0.0, so CI crashed at collection with ModuleNotFoundError: No module named 'pkg_resources'. tensorflow-hub 0.16.1 (the newest release) still imports pkg_resources.parse_version in _ensure_tf_install() at import time, reached via the vendored neural_numint. Provide a stdlib-only stand-in for that one symbol on the import path; no dependency pins or new deps. Closes #10
The agent workflow's plan files (tasks/detailed_plan.md) are personal workflow state whose single source of truth is the plan comment mirrored onto the task issue (<!-- detailed-plan -->). Tracking them in the public repository adds per-subtask churn to every PR diff (e.g. PR #8) and to maintainers' history. tasks/global_plan.md already worked this way (untracked); detailed_plan.md now follows suit. The file stays on disk and keeps being mirrored to issues exactly as before.
The pull_request trigger fired on every push (synchronize/opened/reopened), so non-code changes (.gitignore, the removed tasks/ plan churn) wasted a full ~15-minute suite run. Add a paths whitelist matching exactly the workflow's inputs: grad_dft/**, tests/**, requirements.txt, setup.py, .github/** (the workflow itself plus the pyscf_conf.py it consumes), and the lint configs that #7 will put into CI. Whitelist semantics: docs/notebook/ plan-file-only PRs skip CI entirely. .github/** is in the filter, so this very change re-triggers CI on push, self-validating the filter. Closes #11
This was referenced Oct 7, 2026
Closed
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.