From ed7779d1264b39e85e3440471cde30b0179870c0 Mon Sep 17 00:00:00 2001
From: michalk8 <46717574+michalk8@users.noreply.github.com>
Date: Tue, 21 Feb 2023 20:58:01 +0100
Subject: [PATCH] Add CI tests for upstream JAX (#313)

* Add `jax-latest` to `tox`

* Update test.yml

* Fix `tox` environment names

* Add `continue-on-error` for `jax-latest`

* Fix typo in `tests.yml`
---
 .github/workflows/tests.yml | 16 ++++++++++++----
 CONTRIBUTING.md             |  6 +++---
 pyproject.toml              | 10 ++++++----
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 885f5e719..4bc8fc4fc 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -14,14 +14,22 @@ concurrency:
 
 jobs:
   fast-tests:
-    name: Fast tests Python 3.9 on ubuntu-latest
+    name: Fast tests Python ${{ matrix.python-version }} ${{ matrix.jax-version }}
     runs-on: ubuntu-latest
+    # allow tests using latest Jax to fail
+    continue-on-error: ${{ matrix.jax-version == 'jax-latest' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.9']
+        jax-version: [jax-default, jax-latest]
+
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python 3.9
       uses: actions/setup-python@v4
       with:
-        python-version: '3.9'
+        python-version: ${{ matrix.python-version }}
 
     - name: Install apt dependencies
       if: ${{ runner.os == 'Linux' }}
@@ -40,11 +48,11 @@ jobs:
 
     - name: Setup environment
       run: |
-        tox -e py39 --notest -v
+        tox -e py${{ matrix.python-version }}-${{ matrix.jax-version }} --notest -v
 
     - name: Run tests
       run: |
-        tox -e py39 --skip-pkg-install -- -m fast --memray -n auto -vv
+        tox -e py${{ matrix.python-version }}-${{ matrix.jax-version }} --skip-pkg-install -- -m fast --memray -n auto -vv
 
   gpu-tests:
     name: Fast GPU tests Python 3.8 on ubuntu-20.04
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 5cedfa56d..df1cd79d0 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -25,9 +25,9 @@ In order to run tests, we utilize [tox](https://tox.wiki/):
 ```shell
 tox run  # run linter and all tests on all available Python versions
 tox run -- -n auto -m fast  # run linter and fast tests in parallel
-tox -e py38  # run all tests on Python3.8
-tox -e py39 -- -k "test_euclidean_point_cloud"  # run tests matching the expression on Python3.9
-tox -e py310 -- --memray  # test also memory on Python3.10
+tox -e py3.8  # run all tests on Python3.8
+tox -e py3.9 -- -k "test_euclidean_point_cloud"  # run tests matching the expression on Python3.9
+tox -e py3.10 -- --memray  # test also memory on Python3.10
 ```
 Alternatively, tests can be also run using the [pytest](https://docs.pytest.org/):
 ```shell
diff --git a/pyproject.toml b/pyproject.toml
index 729a2ffbf..27c390ef4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -169,15 +169,17 @@ ignore_path = ['docs/**/_autosummary']
 legacy_tox_ini = """
     [tox]
     min_version = 4.0
-    env_list = lint-code,py{38,39,310,311}
+    env_list = lint-code,py{3.8,3.9,3.10,3.11}
     skip_missing_interpreters = true
 
     [testenv]
     extras = test
     pass_env = CUDA_*,PYTEST_*,CI
-    deps = py{38,39,310}: scikit-sparse>=0.4.6
+    deps = py{3.8,3.9,3.10}: scikit-sparse>=0.4.6
+    commands_pre =
+        gpu: python -I -m pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+        jax-latest: python -I -m pip install 'git+https://github.com/google/jax@main'
     commands =
-        gpu: pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
         python -m pytest {tty:--color=yes} {posargs: \
             --cov={env_site_packages_dir}{/}ott --cov-config={tox_root}{/}pyproject.toml \
             --no-cov-on-fail --cov-report=xml --cov-report=term-missing:skip-covered}
@@ -215,7 +217,7 @@ legacy_tox_ini = """
     commands =
         sphinx-build -b html {tox_root}/docs {tox_root}/docs/_build/html {posargs}
     commands_post =
-        python -c 'import pathlib; print(f"Documentation is under:", pathlib.Path(f"{tox_root}") / "docs" / "_build" / "html" / "index.html")'
+        python -c 'import pathlib; print(pathlib.Path(f"{tox_root}") / "docs" / "_build" / "html" / "index.html")'
 
     [testenv:clean-docs]
     description = Remove the documentation.