Thomaspinder/replace cola (#529)

* Deprecate cola

* Fix flaky test

* Upper bound flax

* Upper bound flax

* Test doctest

* Fix flaky test

* Migrate to UV and fix parameter tags

Author: thomaspinder

.github/pull_request_template.md

@@ -1,6 +1,6 @@
 ## Checklist
 
-- [ ] I've formatted the new code by running `hatch run dev:format` before committing.
+- [ ] I've formatted the new code by running `uv run poe format` before committing.
 - [ ] I've added tests for new code.
 - [ ] I've added docstrings for the new code.
 

.github/workflows/build_docs.yml

@@ -22,7 +22,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10"]
+        python-version: ["3.11"]
 
     steps:
       # Grap the latest commit from the branch
@@ -47,14 +47,18 @@ jobs:
         run: |
           npm install katex
 
-      # Install Hatch
-      - name: Install Hatch
-        uses: pypa/hatch@install
+      # Install uv
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "latest"
 
       - name: Build the documentation with MKDocs
         run: |
           conda install pandoc
-          hatch run docs:build
+          uv sync --extra docs
+          uv run python docs/scripts/gen_examples.py --execute
+          uv run mkdocs build
 
       - name: Deploy Page 🚀
         uses: JamesIves/[email protected]

.github/workflows/integration.yml

@@ -13,7 +13,7 @@ jobs:
       matrix:
         # Select the Python versions to test against
         os: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
       fail-fast: true
     steps:
       - name: Check out the code
@@ -25,10 +25,14 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      # Install Hatch
-      - name: Install Hatch
-        uses: pypa/hatch@install
+      # Install uv
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "latest"
 
       # Run the unit tests and build the coverage report
       - name: Run Integration Tests
-        run: hatch run docs:integration
+        run: |
+          uv sync --extra docs
+          uv run python tests/integration_tests.py

.github/workflows/stale_prs.yml

@@ -17,7 +17,7 @@ jobs:
     strategy:
       matrix:
         os: ["ubuntu-latest"]
-        python-version: ["3.10"]
+        python-version: ["3.11"]
 
     steps:
       # Grap the latest commit from the branch
@@ -33,11 +33,15 @@ jobs:
           auto-update-conda: true
           python-version: ${{ matrix.python-version }}
 
-      # Install Hatch
-      - name: Install Hatch
-        uses: pypa/hatch@install
+      # Install uv
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "latest"
 
       - name: Build the documentation with MKDocs
         run: |
           conda install pandoc
-          hatch run docs:build
+          uv sync --extra docs
+          uv run python docs/scripts/gen_examples.py --execute
+          uv run mkdocs build

.github/workflows/test_docs.yml

@@ -8,13 +8,12 @@ on:
 jobs:
   unit-tests:
     name: Run Tests
-    runs-on: ubuntu-latest
     strategy:
       matrix:
-        # Select the Python versions to test against
         os: ["ubuntu-latest", "macos-latest"]
-        python-version: ["3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
       fail-fast: true
+    runs-on: ${{ matrix.os }}
     steps:
       - name: Check out the code
         uses: actions/[email protected]
@@ -25,18 +24,22 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      # Install Hatch
-      - name: Install Hatch
-        uses: pypa/hatch@install
+      # Install uv
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "latest"
 
       # Install the dependencies
       - name: Check docstrings
         run: |
-          hatch run dev:docstrings
+          uv sync --extra dev
+          uv run xdoctest ./gpjax
 
       # Run the unit tests and build the coverage report
       - name: Run Tests
-        run: hatch run dev:coverage
+        run: uv run pytest . -v --cov=./gpjax --cov-report=xml:./coverage.xml
+
 
       - name: Upload code coverage
         uses: codecov/codecov-action@v3

.github/workflows/tests.yml

@@ -72,16 +72,16 @@ you through every detail!
   Always use a `feature` branch. It's good practice to avoid
   work on the ``main`` branch of any repository.
 
-4.  We use [Hatch](https://hatch.pypa.io/latest/) for packaging and dependency management. Project requirements are in ``pyproject.toml``. To install GPJax into a Hatch virtual environment, run:
+4.  We use [uv](https://docs.astral.sh/uv/) for packaging and dependency management. Project requirements are in ``pyproject.toml``. To install GPJax with uv, run:
 
   ```bash
-  $ hatch env create
+  $ uv sync --extra dev
   ```
 
   At this point we recommend you check your installation passes the supplied unit tests:
 
   ```bash
-  $ hatch run dev:all-tests
+  $ uv run poe all-tests
   ```
 
 5.  Add changed files using `git add` and then `git commit` files to record your
@@ -142,7 +142,7 @@ request, we recommend you check the following:
   accepted. Test coverage can be checked with:
 
     ```bash
-    $ hatch run dev:coverage
+    $ uv run poe coverage
     ```
 
   Navigate to the newly created folder `htmlcov` and open `index.html` to view

docs/contributing.md

@@ -31,7 +31,7 @@ hardware acceleration support as detailed in the
     ```bash
     git clone https://github.com/thomaspinder/GPJax.git
     cd GPJax
-    hatch shell create
+    uv sync --extra dev
     ```
 
 !!! tip
@@ -45,5 +45,5 @@ hardware acceleration support as detailed in the
     and recommend you check your installation passes the supplied unit tests:
 
     ```bash
-    hatch run dev:all-tests
+    uv run poe all-tests
     ```

docs/installation.md

@@ -8,7 +8,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.7
+#       jupytext_version: 1.17.3
 #   kernelspec:
 #     display_name: gpjax
 #     language: python

examples/barycentres.py

@@ -8,9 +8,9 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.7
+#       jupytext_version: 1.11.2
 #   kernelspec:
-#     display_name: gpjax
+#     display_name: .venv
 #     language: python
 #     name: python3
 # ---
@@ -22,7 +22,6 @@
 # with non-Gaussian likelihoods via maximum a posteriori (MAP). We focus on a classification task here.
 
 # %%
-import cola
 from flax import nnx
 import jax
 
@@ -41,7 +40,7 @@
 import optax as ox
 
 from examples.utils import use_mpl_style
-from gpjax.lower_cholesky import lower_cholesky
+from gpjax.linalg import lower_cholesky, PSD, solve
 
 config.update("jax_enable_x64", True)
 
@@ -219,7 +218,7 @@
 # Compute (latent) function value map estimates at training points:
 Kxx = opt_posterior.prior.kernel.gram(x)
 Kxx += identity_matrix(D.n) * jitter
-Kxx = cola.PSD(Kxx)
+Kxx = PSD(Kxx)
 Lx = lower_cholesky(Kxx)
 f_hat = Lx @ opt_posterior.latent.value
 
@@ -267,10 +266,10 @@ def construct_laplace(test_inputs: Float[Array, "N D"]) -> npd.MultivariateNorma
     Kxt = opt_posterior.prior.kernel.cross_covariance(x, test_inputs)
     Kxx = opt_posterior.prior.kernel.gram(x)
     Kxx += identity_matrix(D.n) * jitter
-    Kxx = cola.PSD(Kxx)
+    Kxx = PSD(Kxx)
 
     # Kxx⁻¹ Kxt
-    Kxx_inv_Kxt = cola.solve(Kxx, Kxt)
+    Kxx_inv_Kxt = solve(Kxx, Kxt)
 
     # Ktx Kxx⁻¹[ H⁻¹ ] Kxx⁻¹ Kxt
     laplace_cov_term = jnp.matmul(jnp.matmul(Kxx_inv_Kxt.T, H_inv), Kxx_inv_Kxt)