Feat: Support the current JAX release (0.10.x), test Python 3.9–3.12, and drop jaxopt #680
Workflow file for this run
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
| name: DMFF's python tests. | |
| on: | |
| push: | |
| pull_request: | |
| # DMFF supports Python 3.9-3.12. The two axes are not independent: upstream JAX | |
| # drops old Pythons over time, so each Python has a different newest usable JAX | |
| # (3.9 -> 0.4.30, 3.10 -> 0.6.2, 3.11/3.12 -> current). Pinning "latest" on 3.9 | |
| # would silently resolve an old JAX and stop testing what the job claims to. | |
| jobs: | |
| build: | |
| runs-on: ubuntu-latest | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| include: | |
| # Floor of both axes: oldest supported Python, newest JAX it can run. | |
| - name: py3.9-jax0.4.30 | |
| python-version: "3.9" | |
| jax-spec: "jax==0.4.30 jaxlib==0.4.30 jaxopt 'numpy<2'" | |
| # Oldest JAX we pin against on a modern Python. jaxopt is archived and | |
| # optional; test_qeq runs everywhere, but only the two jobs that | |
| # install jaxopt exercise its jitted jaxopt-backed solver path. | |
| - name: py3.11-jax0.4.35 | |
| python-version: "3.11" | |
| jax-spec: "jax==0.4.35 jaxlib==0.4.35 jaxopt 'numpy<2'" | |
| # Newest JAX each remaining Python can install, with numpy 2. | |
| - name: py3.10-jax-latest | |
| python-version: "3.10" | |
| jax-spec: "jax jaxlib" | |
| - name: py3.11-jax-latest | |
| python-version: "3.11" | |
| jax-spec: "jax jaxlib" | |
| - name: py3.12-jax-latest | |
| python-version: "3.12" | |
| jax-spec: "jax jaxlib" | |
| name: ${{ matrix.name }} | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| # setuptools_scm needs full history to derive the version. | |
| fetch-depth: 0 | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: actions/setup-python@v5 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Install Dependencies | |
| run: | | |
| source $CONDA/bin/activate | |
| conda create -n dmff -y python=${{ matrix.python-version }} pip numpy openmm pytest rdkit openbabel mdtraj ambertools -c conda-forge | |
| conda activate dmff | |
| # Always drive pip through the interpreter we are about to test with. | |
| # A bare `pip` can resolve to the setup-python installation ahead of | |
| # the conda env on PATH, silently installing into a different | |
| # interpreter than `python` imports from. | |
| which python | |
| python -m pip install --upgrade pip | |
| python -m pip install ${{ matrix.jax-spec }} networkx parmed pymbar optax tqdm ase | |
| - name: Install DMFF | |
| run: | | |
| source $CONDA/bin/activate dmff | |
| python -m pip install . | |
| - name: Report resolved versions | |
| run: | | |
| source $CONDA/bin/activate dmff | |
| which python | |
| python -c "import sys, jax, numpy; print('python', sys.version.split()[0], '| jax', jax.__version__, '| numpy', numpy.__version__)" \ | |
| || { echo '--- jax missing; installed packages: ---'; python -m pip list; exit 1; } | |
| - name: Run Tests | |
| run: | | |
| source $CONDA/bin/activate dmff | |
| python -m pytest -vs tests/test_classical/test_* | |
| python -m pytest -vs tests/test_common/test_* | |
| python -m pytest -vs tests/test_admp/test_* | |
| python -m pytest -vs tests/test_utils.py | |
| python -m pytest -vs tests/test_mbar/test_* | |
| python -m pytest -vs tests/test_sgnn/test_* | |
| python -m pytest -vs tests/test_difftraj/test_* | |
| python -m pytest -vs tests/test_frontend/test_* | |
| python -m pytest -vs tests/test_dimer/test_* | |
| python -m pytest -vs tests/test_energy.py | |
| python -m pytest -vs tests/test_ase/test_* |