gpjax 0.11.1__tar.gz → 0.12.0__tar.gz

gpjax-0.12.0/.claude/settings.local.json

@@ -0,0 +1,32 @@
+{
+  "permissions": {
+    "allow": [
+      "WebFetch(domain:raw.githubusercontent.com)",
+      "Bash(hatch run dev:test:*)",
+      "Bash(hatch run dev:python:*)",
+      "Bash(hatch run dev:lint:*)",
+      "Bash(hatch env:*)",
+      "Bash(hatch run dev:format:*)",
+      "Bash(hatch run docs:ipython:*)",
+      "Bash(timeout 30s hatch run docs:ipython examples/orthogonal_additive_kernel.py)",
+      "Bash(gtimeout:*)",
+      "mcp__fetch__imageFetch",
+      "Bash(hatch run:*)",
+      "Bash(mkdir:*)",
+      "Bash(find:*)",
+      "Bash(grep:*)",
+      "WebFetch(domain:flax.readthedocs.io)",
+      "WebSearch",
+      "WebFetch(domain:docs.jax.dev)",
+      "mcp__github__search_pull_requests",
+      "mcp__github__get_me",
+      "mcp__github__get_pull_request_status",
+      "mcp__github__list_workflows",
+      "mcp__github__list_workflow_runs",
+      "mcp__github__get_job_logs",
+      "Bash(uv run:*)"
+    ],
+    "deny": [],
+    "ask": []
+  }
+}

{gpjax-0.11.1 → gpjax-0.12.0}/.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.
 

{gpjax-0.11.1 → gpjax-0.12.0}/.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/github-pages-deploy-action@v4.4.1

{gpjax-0.11.1 → gpjax-0.12.0}/.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

{gpjax-0.11.1 → gpjax-0.12.0}/.github/workflows/test_docs.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

{gpjax-0.11.1 → gpjax-0.12.0}/.github/workflows/tests.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/checkout@v3.5.2
@@ -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

gpjax-0.12.0/CLAUDE.md

@@ -0,0 +1,100 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Development Commands
+
+### Testing
+- **Run all tests**: `hatch run dev:test` or `pytest . -v -n 4 --beartype-packages='gpjax'`
+- **Run tests with coverage**: `hatch run dev:coverage`
+- **Run a single test file**: `pytest tests/test_file.py -v`
+- **Run docstring tests**: `hatch run dev:docstrings`
+- **Full test suite**: `hatch run dev:all-tests` (includes format check, docstrings, and tests)
+
+### Code Formatting and Linting
+- **Format all code**: `hatch run dev:format` (runs black, isort, and ruff format)
+- **Check formatting**: `hatch run dev:check` (black-check, imports-check, lint-check)
+- **Format with black**: `hatch run dev:black-format`
+- **Format imports**: `hatch run dev:imports-format`
+- **Lint and auto-fix**: `hatch run dev:lint-check`
+- **Remove unused imports**: `hatch run dev:remove-unused`
+
+### Documentation
+- **Build docs**: `hatch run docs:build`
+- **Serve docs locally**: `hatch run docs:serve`
+- **Run integration tests**: `hatch run docs:integration`
+
+### Build and Installation
+- **Install for development**: `hatch env create` then `hatch shell`
+- **Install stable version**: `pip install gpjax`
+
+## High-Level Architecture
+
+GPJax is a Gaussian Process library built on JAX and Flax (nnx), designed with a modular architecture that closely mirrors mathematical abstractions:
+
+### Core Components
+
+1. **Gaussian Processes (`gpjax.gps`)**
+   - `AbstractPrior`: Base class for GP priors combining kernels and mean functions
+   - `Prior`: Standard GP prior implementation  
+   - `AbstractPosterior`: Base class for posterior inference
+   - `ConjugatePosterior`: Exact inference for Gaussian likelihoods
+   - `NonConjugatePosterior`: Approximate inference for non-Gaussian likelihoods
+   - Uses Flax nnx.Module for parameter management
+
+2. **Kernels (`gpjax.kernels`)**
+   - `AbstractKernel`: Base kernel class with composition support (+, *)
+   - Stationary kernels: RBF, Matern12/32/52, RationalQuadratic, Periodic, White
+   - Non-stationary kernels: Linear, Polynomial, ArcCosine
+   - Non-Euclidean kernels: Graph kernels for structured data
+   - Kernel computations: Dense, Diagonal, Eigen decomposition strategies
+   - Approximations: Random Fourier Features (RFF)
+
+3. **Likelihoods (`gpjax.likelihoods`)**
+   - `AbstractLikelihood`: Base class defining the observation model
+   - `Gaussian`: Standard Gaussian likelihood with observation noise
+   - Non-Gaussian: Bernoulli, Poisson for classification and count data
+   - Links prediction and observation spaces
+
+4. **Variational Inference (`gpjax.variational_families`)**
+   - `AbstractVariationalFamily`: Base for variational approximations
+   - `VariationalGaussian`: Mean-field Gaussian variational distribution
+   - `WhitenedVariationalGaussian`: Whitened parameterization
+   - `NaturalVariationalGaussian`: Natural gradient parameterization
+   - Supports both collapsed and uncollapsed VI
+
+5. **Objectives (`gpjax.objectives`)**
+   - `AbstractObjective`: Base class for optimization objectives
+   - `ConjugateMLL`: Marginal likelihood for exact inference
+   - `NonConjugateMLL`: Laplace approximation for non-Gaussian likelihoods
+   - `ELBO`: Evidence lower bound for variational inference
+   - `CollapsedELBO`: Analytically integrated ELBO
+
+6. **Optimization (`gpjax.fit`)**
+   - `fit()`: General optimizer using Optax optimizers
+   - `fit_lbfgs()`: L-BFGS optimization via Optax
+   - `fit_scipy()`: Interface to scipy optimizers via JAXopt
+   - Supports custom loss functions and stopping criteria
+
+### Key Design Patterns
+
+- **Functional Design**: Functions are first-class, mirroring mathematical notation
+- **Composability**: Kernels support arithmetic operations for easy combination
+- **JAX Integration**: Full support for JIT compilation, automatic differentiation, and vectorization
+- **Type Safety**: Extensive use of jaxtyping and beartype for runtime type checking
+- **CoLA Integration**: Uses CoLA (Compositional Linear Algebra) for efficient linear algebra operations
+- **Parameter Management**: Uses Flax nnx for trainable parameters with PyTree support
+
+### Data Structures
+
+- `Dataset`: Simple dataclass for (X, y) pairs with optional y
+- `GaussianDistribution`: Represents multivariate Gaussians with efficient sampling
+- Parameters use `nnx.Param` with transformation support (e.g., `SoftplusTransformation`)
+
+### Important Implementation Details
+
+- All kernels must implement `__call__(x, y)` for single point evaluation
+- Kernel matrices are computed via `compute_engine.gram()` or `compute_engine.cross_covariance()`
+- Jitter (small diagonal noise) is added for numerical stability, default 1e-6
+- Uses `cola.PSD` annotations for positive semi-definite matrices
+- Cholesky decompositions use custom `lower_cholesky` for better gradients

{gpjax-0.11.1 → gpjax-0.12.0}/Makefile

@@ -21,11 +21,16 @@ black:  ## Format code in-place using black.
 isort:  ## Format imports in-place using isort.
 		isort ${PKGROOT}/ tests/
 
-format: ## Code styling - black, isort
-		black ${PKGROOT}/ tests/ -l 88 .
+format: ## Code styling - black, isort, ruff
+		uv run black ${PKGROOT}/ tests/
 		@printf "\033[1;34mBlack passes!\033[0m\n\n"
-		isort ${PKGROOT}/ tests/
+		uv run jupytext --pipe black examples/*.py
+		@printf "\033[1;34mJupytext passes!\033[0m\n\n"
+		uv run isort ${PKGROOT}/ tests/
+		uv run isort examples/*.py --treat-comment-as-code '# %%' --float-to-top
 		@printf "\033[1;34misort passes!\033[0m\n\n"
+		uv run ruff format ${PKGROOT}/ tests/ examples/
+		@printf "\033[1;34mRuff passes!\033[0m\n\n"
 
 ##@ Testing
 test:  ## Test code using pytest.

{gpjax-0.11.1 → gpjax-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.11.1
+Version: 0.12.0
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -14,11 +14,11 @@ Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
-Requires-Python: <3.13,>=3.10
+Requires-Python: <=3.13,>=3.10
 Requires-Dist: beartype>0.16.1
-Requires-Dist: cola-ml>=0.0.7
 Requires-Dist: flax>=0.10.0
 Requires-Dist: jax>=0.5.0
 Requires-Dist: jaxlib>=0.5.0
@@ -27,6 +27,47 @@ Requires-Dist: numpy>=2.0.0
 Requires-Dist: numpyro
 Requires-Dist: optax>0.2.1
 Requires-Dist: tqdm>4.66.2
+Provides-Extra: dev
+Requires-Dist: absolufy-imports>=0.3.1; extra == 'dev'
+Requires-Dist: autoflake; extra == 'dev'
+Requires-Dist: black; extra == 'dev'
+Requires-Dist: codespell>=2.2.4; extra == 'dev'
+Requires-Dist: coverage>=7.2.2; extra == 'dev'
+Requires-Dist: interrogate>=1.5.0; extra == 'dev'
+Requires-Dist: isort; extra == 'dev'
+Requires-Dist: jupytext; extra == 'dev'
+Requires-Dist: mktestdocs>=0.2.1; extra == 'dev'
+Requires-Dist: networkx; extra == 'dev'
+Requires-Dist: pre-commit>=3.2.2; extra == 'dev'
+Requires-Dist: pytest-beartype; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest-pretty>=1.1.1; extra == 'dev'
+Requires-Dist: pytest-xdist>=3.2.1; extra == 'dev'
+Requires-Dist: pytest>=7.2.2; extra == 'dev'
+Requires-Dist: ruff>=0.6; extra == 'dev'
+Requires-Dist: xdoctest>=1.1.1; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: blackjax>=0.9.6; extra == 'docs'
+Requires-Dist: ipykernel>=6.22.0; extra == 'docs'
+Requires-Dist: ipython>=8.11.0; extra == 'docs'
+Requires-Dist: ipywidgets>=8.0.5; extra == 'docs'
+Requires-Dist: jupytext>=1.14.5; extra == 'docs'
+Requires-Dist: markdown-katex>=202406.1035; extra == 'docs'
+Requires-Dist: matplotlib>=3.7.1; extra == 'docs'
+Requires-Dist: mkdocs-gen-files>=0.5.0; extra == 'docs'
+Requires-Dist: mkdocs-git-authors-plugin>=0.7.0; extra == 'docs'
+Requires-Dist: mkdocs-jupyter>=0.24.3; extra == 'docs'
+Requires-Dist: mkdocs-literate-nav>=0.6.0; extra == 'docs'
+Requires-Dist: mkdocs-material>=9.5.12; extra == 'docs'
+Requires-Dist: mkdocs>=1.5.3; extra == 'docs'
+Requires-Dist: mkdocstrings[python]<0.28.0; extra == 'docs'
+Requires-Dist: nbconvert>=7.16.2; extra == 'docs'
+Requires-Dist: networkx>=3.0; extra == 'docs'
+Requires-Dist: pandas>=1.5.3; extra == 'docs'
+Requires-Dist: pymdown-extensions>=10.7.1; extra == 'docs'
+Requires-Dist: scikit-learn>=1.5.1; extra == 'docs'
+Requires-Dist: seaborn>=0.12.2; extra == 'docs'
+Requires-Dist: watermark>=2.3.1; extra == 'docs'
 Description-Content-Type: text/markdown
 
 <!-- <h1 align='center'>GPJax</h1>
@@ -41,12 +82,12 @@ Description-Content-Type: text/markdown
 [![PyPI version](https://badge.fury.io/py/GPJax.svg)](https://badge.fury.io/py/GPJax)
 [![DOI](https://joss.theoj.org/papers/10.21105/joss.04455/status.svg)](https://doi.org/10.21105/joss.04455)
 [![Downloads](https://pepy.tech/badge/gpjax)](https://pepy.tech/project/gpjax)
-[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
+[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
 
 [**Quickstart**](#simple-example)
 | [**Install guide**](#installation)
 | [**Documentation**](https://docs.jaxgaussianprocesses.com/)
-| [**Slack Community**](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
+| [**Slack Community**](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
 
 GPJax aims to provide a low-level interface to Gaussian process (GP) models in
 [Jax](https://github.com/google/jax), structured to give researchers maximum
@@ -81,7 +122,7 @@ behaviours through [this form](https://jaxgaussianprocesses.com/contact/) or rea
 one of the project's [_gardeners_](https://docs.jaxgaussianprocesses.com/GOVERNANCE/#roles).
 
 Feel free to join our [Slack
-Channel](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw),
+Channel](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA),
 where we can discuss the development of GPJax and broader support for Gaussian
 process modelling.
 
@@ -177,14 +218,13 @@ configuration in development mode.
 ```bash
 git clone https://github.com/JaxGaussianProcesses/GPJax.git
 cd GPJax
-hatch env create
-hatch shell
+uv sync --extra dev
 ```
 
 > We recommend you check your installation passes the supplied unit tests:
 >
 > ```python
-> hatch run dev:test
+> uv run pytest --beartype-packages='gpjax'
 > ```
 
 # Citing GPJax

{gpjax-0.11.1 → gpjax-0.12.0}/README.md

@@ -10,12 +10,12 @@
 [![PyPI version](https://badge.fury.io/py/GPJax.svg)](https://badge.fury.io/py/GPJax)
 [![DOI](https://joss.theoj.org/papers/10.21105/joss.04455/status.svg)](https://doi.org/10.21105/joss.04455)
 [![Downloads](https://pepy.tech/badge/gpjax)](https://pepy.tech/project/gpjax)
-[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
+[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
 
 [**Quickstart**](#simple-example)
 | [**Install guide**](#installation)
 | [**Documentation**](https://docs.jaxgaussianprocesses.com/)
-| [**Slack Community**](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
+| [**Slack Community**](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA)
 
 GPJax aims to provide a low-level interface to Gaussian process (GP) models in
 [Jax](https://github.com/google/jax), structured to give researchers maximum
@@ -50,7 +50,7 @@ behaviours through [this form](https://jaxgaussianprocesses.com/contact/) or rea
 one of the project's [_gardeners_](https://docs.jaxgaussianprocesses.com/GOVERNANCE/#roles).
 
 Feel free to join our [Slack
-Channel](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw),
+Channel](https://join.slack.com/t/gpjax/shared_invite/zt-3cesiykcx-nzajjRdnV3ohw7~~eMlCYA),
 where we can discuss the development of GPJax and broader support for Gaussian
 process modelling.
 
@@ -146,14 +146,13 @@ configuration in development mode.
 ```bash
 git clone https://github.com/JaxGaussianProcesses/GPJax.git
 cd GPJax
-hatch env create
-hatch shell
+uv sync --extra dev
 ```
 
 > We recommend you check your installation passes the supplied unit tests:
 >
 > ```python
-> hatch run dev:test
+> uv run pytest --beartype-packages='gpjax'
 > ```
 
 # Citing GPJax

{gpjax-0.11.1 → gpjax-0.12.0}/docs/contributing.md

@@ -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

{gpjax-0.11.1 → gpjax-0.12.0}/docs/installation.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
     ```

{gpjax-0.11.1 → gpjax-0.12.0}/examples/barycentres.py

@@ -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

{gpjax-0.11.1 → gpjax-0.12.0}/examples/classification.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 numpyro.distributions as npd
 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 @@ jitter = 1e-6
 # 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)

{gpjax-0.11.1 → gpjax-0.12.0}/examples/collapsed_vi.py

@@ -7,9 +7,9 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.16.7
+#       jupytext_version: 1.11.2
 #   kernelspec:
-#     display_name: gpjax_beartype
+#     display_name: .venv
 #     language: python
 #     name: python3
 # ---

{gpjax-0.11.1 → gpjax-0.12.0}/examples/constructing_new_kernels.py

@@ -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
@@ -95,7 +95,7 @@ meanf = gpx.mean_functions.Zero()
 for k, ax, c in zip(kernels, axes.ravel(), cols, strict=False):
     prior = gpx.gps.Prior(mean_function=meanf, kernel=k)
     rv = prior(x)
-    y = rv.sample(seed=key, sample_shape=(10,))
+    y = rv.sample(key=jr.key(22), sample_shape=(10,))
     ax.plot(x, y.T, alpha=0.7, color=c)
     ax.set_title(k.name)
 
@@ -326,7 +326,7 @@ opt_posterior, history = gpx.fit_scipy(
 # %%
 posterior_rv = opt_posterior.likelihood(opt_posterior.predict(angles, train_data=D))
 mu = posterior_rv.mean
-one_sigma = posterior_rv.stddev()
+one_sigma = jnp.sqrt(posterior_rv.variance)
 
 # %%
 fig = plt.figure(figsize=(7, 3.5))

{gpjax-0.11.1 → gpjax-0.12.0}/examples/deep_kernels.py

@@ -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

{gpjax-0.11.1 → gpjax-0.12.0}/examples/graph_kernels.py

@@ -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

{gpjax-0.11.1 → gpjax-0.12.0}/examples/intro_to_gps.py

@@ -7,7 +7,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

{gpjax-0.11.1 → gpjax-0.12.0}/examples/intro_to_kernels.py

@@ -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

{gpjax-0.11.1 → gpjax-0.12.0}/examples/likelihoods_guide.py

@@ -7,7 +7,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