Flax/nnx backend

.github/workflows/build_docs.yml

@@ -47,22 +47,16 @@ jobs:
       - name: Install and configure Poetry
         uses: snok/install-poetry@v1
         with:
-          version: 1.2.2
+          version: 1.5.1
           virtualenvs-create: false
           virtualenvs-in-project: false
           installer-parallel: true
 
-      - name: Install LaTex
-        run: |
-          sudo apt-get update
-          sudo apt-get install texlive-fonts-recommended texlive-fonts-extra texlive-latex-extra dvipng cm-super
-
       - name: Build the documentation with MKDocs
         run: |
-          cp docs/examples/gpjax.mplstyle .
           poetry install --all-extras --with docs
           conda install pandoc
-          poetry run mkdocs build
+          poetry run python docs/scripts/gen_examples.py --execute && poetry run mkdocs build
 
       - name: Deploy Page 🚀
         uses: JamesIves/github-pages-deploy-action@v4.4.1

.github/workflows/integration.yml

@@ -29,7 +29,7 @@ jobs:
       - name: Install Poetry
         uses: snok/install-poetry@v1.3.3
         with:
-          version: 1.4.0
+          version: 1.5.1
 
       # Configure Poetry to use the virtual environment in the project
       - name: Setup Poetry
@@ -39,7 +39,7 @@ jobs:
       # Install the dependencies
       - name: Install Package
         run: |
-          poetry install --all-extras --with docs
+          poetry install --with docs
 
       # Run the unit tests and build the coverage report
       - name: Run Integration Tests

.github/workflows/test_docs.yml

@@ -33,32 +33,17 @@ jobs:
           auto-update-conda: true
           python-version: ${{ matrix.python-version }}
 
-      # Install katex for math support
-      - name: Install NPM
-        uses: actions/setup-node@v3
-        with:
-          node-version: 16
-      - name: Install KaTeX
-        run: |
-          npm install katex
-
-      - name: Install LaTex
-        run: |
-          sudo apt-get update
-          sudo apt-get install texlive-fonts-recommended texlive-fonts-extra texlive-latex-extra dvipng cm-super
-
       # Install Poetry and build the documentation
       - name: Install and configure Poetry
         uses: snok/install-poetry@v1
         with:
-          version: 1.2.2
+          version: 1.5.1
           virtualenvs-create: false
           virtualenvs-in-project: false
           installer-parallel: true
 
       - name: Build the documentation with MKDocs
         run: |
-          cp docs/examples/gpjax.mplstyle .
           poetry install --all-extras --with docs
           conda install pandoc
-          poetry run mkdocs build
+          poetry run python docs/scripts/gen_examples.py --execute && poetry run mkdocs build

.github/workflows/tests.yml

@@ -26,10 +26,13 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       # Install Poetry
-      - name: Install Poetry
-        uses: snok/install-poetry@v1.3.3
+      - name: Install and configure Poetry
+        uses: snok/install-poetry@v1
         with:
-          version: 1.4.0
+          version: 1.5.1
+          virtualenvs-create: false
+          virtualenvs-in-project: false
+          installer-parallel: true
 
       # Configure Poetry to use the virtual environment in the project
       - name: Setup Poetry
@@ -39,7 +42,7 @@ jobs:
       # Install the dependencies
       - name: Install Package
         run: |
-          poetry install --with tests
+          poetry install --with dev
 
       - name: Check docstrings
         run: |

.gitignore

@@ -152,3 +152,4 @@ package-lock.json
 node_modules/
 
 docs/api
+docs/_examples

.pre-commit-config.yaml

@@ -46,14 +46,14 @@ repos:
         language: system
         types: [python]
         exclude: examples/
-  - repo: https://github.com/econchick/interrogate
-    rev: 1.5.0
-    hooks:
-      - id: interrogate
-        args:
-          [
-            "gpjax",
-            "--config",
-            "pyproject.toml",
-          ]
-        pass_filenames: false
+  # - repo: https://github.com/econchick/interrogate
+  #   rev: 1.5.0
+  #   hooks:
+  #     - id: interrogate
+  #       args:
+  #         [
+  #           "gpjax",
+  #           "--config",
+  #           "pyproject.toml",
+  #         ]
+  #       pass_filenames: false

README.md

@@ -72,10 +72,9 @@ helped to shape GPJax into the package it is today.
 ## Notebook examples
 
 > - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification with MCMC**](https://docs.jaxgaussianprocesses.com/examples/classification/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
 > - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
 > - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**BlackJax Integration**](https://docs.jaxgaussianprocesses.com/examples/classification/#mcmc-inference)
 > - [**Laplace Approximation**](https://docs.jaxgaussianprocesses.com/examples/classification/#laplace-approximation)
 > - [**Inference on Non-Euclidean Spaces**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
 > - [**Inference on Graphs**](https://docs.jaxgaussianprocesses.com/examples/graph_kernels/)
@@ -146,13 +145,10 @@ posterior = prior * likelihood
 # Define an optimiser
 optimiser = ox.adam(learning_rate=1e-2)
 
-# Define the marginal log-likelihood
-negative_mll = jit(gpx.objectives.ConjugateMLL(negative=True))
-
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=negative_mll,
+    objective=gpx.objectives.conjugate_mll,
     train_data=D,
     optim=optimiser,
     num_iters=500,