gpjax 0.13.3__tar.gz → 0.13.5__tar.gz

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/auto-label.yml

@@ -16,7 +16,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Auto-label based on files changed
         if: github.event_name == 'pull_request'

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/build_docs.yml

@@ -27,7 +27,7 @@ jobs:
     steps:
       # Grap the latest commit from the branch
       - name: Checkout the branch
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           persist-credentials: false
 
@@ -61,7 +61,7 @@ jobs:
           uv run mkdocs build
 
       - name: Deploy Page 🚀
-        uses: JamesIves/github-pages-deploy-action@v4.7.4
+        uses: JamesIves/github-pages-deploy-action@v4.8.0
         with:
           branch: gh-pages
           folder: site

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/commit-lint.yml

@@ -13,7 +13,7 @@ jobs:
     
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
 

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/integration.yml

@@ -17,7 +17,7 @@ jobs:
       fail-fast: true
     steps:
       - name: Check out the code
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           fetch-depth: 1
       - name: Set up Python ${{ matrix.python-version }}

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/release.yml

@@ -29,7 +29,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
 
@@ -72,7 +72,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v6
@@ -108,7 +108,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Set up Python
         uses: actions/setup-python@v6
@@ -132,7 +132,7 @@ jobs:
           uv run bandit -r gpjax/ -f json -o bandit-report.json || echo "Bandit scan completed with warnings"
 
       - name: Upload security reports
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         with:
           name: security-reports
           path: |
@@ -146,7 +146,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Set up Python
         uses: actions/setup-python@v6
@@ -166,7 +166,7 @@ jobs:
           uv run twine check dist/*
 
       - name: Upload build artifacts
-        uses: actions/upload-artifact@v5
+        uses: actions/upload-artifact@v6
         with:
           name: dist-packages
           path: dist/
@@ -181,7 +181,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0
 
@@ -261,10 +261,10 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
 
       - name: Download build artifacts
-        uses: actions/download-artifact@v6
+        uses: actions/download-artifact@v7
         with:
           name: dist-packages
           path: dist/
@@ -294,7 +294,7 @@ jobs:
 
     steps:
       - name: Download build artifacts
-        uses: actions/download-artifact@v6
+        uses: actions/download-artifact@v7
         with:
           name: dist-packages
           path: dist/

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/ruff.yml

@@ -8,5 +8,5 @@ jobs:
   ruff:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: chartboost/ruff-action@v1

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/security-analysis.yml

@@ -20,7 +20,7 @@ jobs:
     
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v5
+      uses: actions/checkout@v6
 
     - name: Set up Python
       uses: actions/setup-python@v6
@@ -47,7 +47,7 @@ jobs:
         uv run bandit -r gpjax/ -f json -o bandit-report.json || true
 
     - name: Upload dependency scan results
-      uses: actions/upload-artifact@v5
+      uses: actions/upload-artifact@v6
       if: always()
       with:
         name: security-scan-results
@@ -62,7 +62,7 @@ jobs:
     
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v5
+      uses: actions/checkout@v6
       with:
         fetch-depth: 0  # Fetch full history for comprehensive scanning
 

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/test_docs.yml

@@ -18,7 +18,7 @@ jobs:
     steps:
       # Grap the latest commit from the branch
       - name: Checkout the branch
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           persist-credentials: false
 

{gpjax-0.13.3 → gpjax-0.13.5}/.github/workflows/tests.yml

@@ -16,7 +16,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     steps:
       - name: Check out the code
-        uses: actions/checkout@v5
+        uses: actions/checkout@v6
         with:
           fetch-depth: 1
       - name: Set up Python ${{ matrix.python-version }}

{gpjax-0.13.3 → gpjax-0.13.5}/.gitignore

@@ -153,3 +153,8 @@ node_modules/
 
 docs/api
 docs/_examples
+local_libs/
+local_papers/
+GEMINI.md
+AGENTS.md
+plans/

{gpjax-0.13.3 → gpjax-0.13.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.13.3
+Version: 0.13.5
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/thomaspinder/GPJax/issues
@@ -25,6 +25,7 @@ Requires-Dist: jaxtyping>0.2.10
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: numpyro
 Requires-Dist: optax>0.2.1
+Requires-Dist: tensorstore!=0.1.76; sys_platform == 'darwin'
 Requires-Dist: tqdm>4.66.2
 Provides-Extra: dev
 Requires-Dist: absolufy-imports>=0.3.1; extra == 'dev'
@@ -59,7 +60,7 @@ 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.31.0; extra == 'docs'
+Requires-Dist: mkdocstrings[python]<1.1.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'
@@ -141,7 +142,7 @@ GPJax into the package it is today.
 > - [**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/)
-> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/_examples/spatial/)
+> - [**Heteroscedastic Inference**](https://docs.jaxgaussianprocesses.com/_examples/heteroscedastic_inference/)
 > - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/_examples/barycentres/)
 > - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/_examples/deep_kernels/)
 > - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/_examples/poisson/)

{gpjax-0.13.3 → gpjax-0.13.5}/README.md

@@ -70,7 +70,7 @@ GPJax into the package it is today.
 > - [**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/)
-> - [**Pathwise Sampling**](https://docs.jaxgaussianprocesses.com/_examples/spatial/)
+> - [**Heteroscedastic Inference**](https://docs.jaxgaussianprocesses.com/_examples/heteroscedastic_inference/)
 > - [**Learning Gaussian Process Barycentres**](https://docs.jaxgaussianprocesses.com/_examples/barycentres/)
 > - [**Deep Kernel Regression**](https://docs.jaxgaussianprocesses.com/_examples/deep_kernels/)
 > - [**Poisson Regression**](https://docs.jaxgaussianprocesses.com/_examples/poisson/)

{gpjax-0.13.3 → gpjax-0.13.5}/docs/sharp_bits.md

@@ -178,12 +178,19 @@ points. We demonstrate its use in
 
 ## JIT compilation
 
-There are a subset of operations in GPJax that are not JIT compatible by default. This
-is because we have assertions in place to check the properties of the parameters. For
-example, we check that the lengthscale parameter that a user provides is positive. This
-makes for a better user experience as we can provide more informative error messages;
-however, JIT compiling functions wherein these assertions are made will break the code.
-As an example, consider the following code:
+GPJax validates parameters at construction time using two kinds of checks:
+
+1. **Type checks** — plain Python `isinstance` checks that verify values are array-like.
+2. **Value checks** — JAX-compatible assertions (via `checkify`) that verify constraints
+   like positivity or bounds.
+
+During JIT tracing, concrete values are replaced by abstract tracers. The type checks
+use `isinstance`, which is a pure Python operation that cannot be intercepted by JAX's
+`checkify` transformation. This means that constructing GPJax objects (kernels, mean
+functions, likelihoods, etc.) **inside** a JIT boundary will fail.
+
+As an example, consider the following code that constructs a kernel inside a
+JIT-compiled function:
 
 ```python
 import jax
@@ -192,43 +199,40 @@ import gpjax as gpx
 
 x = jnp.linspace(0, 1, 10)[:, None]
 
-def compute_gram(lengthscale):
+def compute_gram_bad(lengthscale):
     k = gpx.kernels.RBF(active_dims=[0], lengthscale=lengthscale, variance=jnp.array(1.0))
     return k.gram(x)
 
-compute_gram(1.0)
+compute_gram_bad(1.0)  # works fine outside JIT
 ```
 
-so far so good. However, if we try to JIT compile this function, we will get an error:
+If we try to JIT compile this function, we get a `TypeError` because the kernel
+constructor receives a JAX tracer instead of a concrete array:
 
 ```python
-jit_compute_gram = jax.jit(compute_gram)
+jit_compute_gram_bad = jax.jit(compute_gram_bad)
 try:
-    jit_compute_gram(1.0)
+    jit_compute_gram_bad(1.0)
 except Exception as e:
     print(e)
 ```
 
-This error is due to the fact that the `RBF` kernel contains an assertion that checks
-that the lengthscale is positive. It does not matter that the assertion is satisfied;
-the very presence of the assertion will break JIT compilation.
+### The fix: construct objects outside JIT
 
-To resolve this, we can use the `checkify` decorator to remove the assertion. This will
-allow the function to be JIT compiled.
+The solution is to construct GPJax objects **outside** the JIT boundary and only JIT the
+computation itself. This follows the standard JAX pattern of keeping object construction
+separate from traced computation:
 
 ```python
-from jax.experimental import checkify
+k = gpx.kernels.RBF(active_dims=[0], lengthscale=1.0, variance=jnp.array(1.0))
 
-jit_compute_gram = jax.jit(checkify.checkify(compute_gram))
-error, value = jit_compute_gram(1.0)
-```
-By virtue of the `checkify.checkify`, a tuple is returned where the first element is the
-output of the assertion, and the second element is the value of the function. 
+@jax.jit
+def compute_gram(x):
+    return k.gram(x)
 
-This design is not perfect, and in an ideal world we would not enforce the user to wrap
-their code in `checkify.checkify`. We are actively looking into cleaner ways to provide
-guardrails in a less intrusive manner. However, for now, should you try to JIT compile
-a component of GPJax wherein there is an assertion, you will need to wrap the function
-in `checkify.checkify` as shown above.
+result = compute_gram(x)
+```
 
-For more on `checkify`, please see the [JAX Checkify Doc](https://docs.jax.dev/en/latest/debugging/checkify_guide.html).
+More generally, any GPJax object should be constructed outside of `jax.jit`, `jax.vmap`,
+or `jax.grad` boundaries. Once constructed, their methods can be freely used inside
+these JAX transformations.