gpjax 0.13.4__tar.gz → 0.13.5__tar.gz

{gpjax-0.13.4 → 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.4 → 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.4 → 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.4 → 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.4 → 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.4 → 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.4 → 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.4 → 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.4 → 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.4 → gpjax-0.13.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.13.4
+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'

{gpjax-0.13.4 → 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.

{gpjax-0.13.4 → gpjax-0.13.5}/examples/heteroscedastic_inference.py

@@ -16,7 +16,7 @@
 # ---
 
 # %% [markdown]
-# # Heteroscedastic inference for regression and classification
+# # Heteroscedastic Inference
 #
 # This notebook shows how to fit a heteroscedastic Gaussian processes (GPs) that
 # allows one to perform regression where there exists non-constant, or
@@ -293,13 +293,13 @@ posterior_adv = mean_prior * likelihood_adv
 # The signal requires a richer inducing set to capture its oscillations, whereas the
 # noise process can be summarised with fewer points because the burst is localised.
 z_signal = jnp.linspace(-2.0, 2.0, 30)[:, None]
-z_noise = jnp.linspace(-2.0, 2.0, 15)[:, None]
+z_noise = jnp.linspace(-2.0, 2.0, 20)[:, None]
 
 # Use VariationalGaussianInit to pass specific configurations
 q_init_f = VariationalGaussianInit(inducing_inputs=z_signal)
 q_init_g = VariationalGaussianInit(inducing_inputs=z_noise)
 
-q_adv = HeteroscedasticVariationalFamily(
+q_sparse = HeteroscedasticVariationalFamily(
     posterior=posterior_adv,
     signal_init=q_init_f,
     noise_init=q_init_g,
@@ -315,19 +315,19 @@ q_adv = HeteroscedasticVariationalFamily(
 # Optimize
 objective_adv = lambda model, data: -gpx.objectives.heteroscedastic_elbo(model, data)
 optimiser_adv = ox.adam(1e-2)
-q_adv_trained, _ = gpx.fit(
-    model=q_adv,
+q_sparse_trained, _ = gpx.fit(
+    model=q_sparse,
     objective=objective_adv,
     train_data=data_adv,
     optim=optimiser_adv,
-    num_iters=8000,
+    num_iters=10000,
     verbose=False,
 )
 
 # %%
 # Plotting
-xtest = jnp.linspace(-2.2, 2.2, 200)[:, None]
-pred = q_adv_trained.predict(xtest)
+xtest = jnp.linspace(-2.2, 2.2, 300)[:, None]
+pred = q_sparse_trained.predict(xtest)
 
 # Unpack the named tuple
 mf = pred.mean_f
@@ -339,36 +339,41 @@ vg = pred.variance_g
 # The likelihood expects the *latent* noise distribution to compute the predictive
 # but here we can just use the transformed expected variance for plotting.
 # For accurate predictive intervals, we should use likelihood.predict.
-signal_dist, noise_dist = q_adv_trained.predict_latents(xtest)
+signal_dist, noise_dist = q_sparse_trained.predict_latents(xtest)
 predictive_dist = likelihood_adv.predict(signal_dist, noise_dist)
 predictive_mean = predictive_dist.mean
 predictive_std = jnp.sqrt(jnp.diag(predictive_dist.covariance_matrix))
 
-fig, ax = plt.subplots()
-ax.plot(x, y, "o", color="black", alpha=0.3, label="Data")
-ax.plot(xtest, mf, color="C0", label="Signal Mean")
+fig, ax = plt.subplots(figsize=(6, 2.5))
+ax.plot(x, y, "x", color="black", alpha=0.5, label="Data")
+
+# Plot total uncertainty (signal + noise)
+ax.plot(xtest, predictive_mean, "--", color=cols[1], linewidth=2)
 ax.fill_between(
     xtest.squeeze(),
-    mf.squeeze() - 2 * jnp.sqrt(vf.squeeze()),
-    mf.squeeze() + 2 * jnp.sqrt(vf.squeeze()),
-    color="C0",
-    alpha=0.2,
-    label="Signal Uncertainty",
+    predictive_mean - predictive_std,
+    predictive_mean + predictive_std,
+    color=cols[1],
+    alpha=0.3,
+    label="One std. dev.",
 )
-
-# Plot total uncertainty (signal + noise)
-ax.plot(xtest, predictive_mean, "--", color="C1", alpha=0.5)
+ax.plot(xtest.squeeze(), predictive_mean - predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
+ax.plot(xtest.squeeze(), predictive_mean + predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
 ax.fill_between(
     xtest.squeeze(),
     predictive_mean - 2 * predictive_std,
     predictive_mean + 2 * predictive_std,
-    color="C1",
+    color=cols[1],
     alpha=0.1,
-    label="Predictive Uncertainty (95%)",
+    label="Two std. dev.",
 )
+ax.plot(xtest.squeeze(), predictive_mean - 2 * predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
+ax.plot(xtest.squeeze(), predictive_mean + 2 * predictive_std, "--", color=cols[1], alpha=0.5, linewidth=0.75)
 
-ax.set_title("Heteroscedastic Regression with Custom Inducing Points")
-ax.legend(loc="upper left", fontsize="small")
+ax.set_title("Sparse Heteroscedastic Regression")
+ax.legend(loc="best", fontsize="small")
+ax.set_xlabel("$x$")
+ax.set_ylabel("$y$")
 
 # %% [markdown]
 # ## Takeaways

{gpjax-0.13.4 → gpjax-0.13.5}/gpjax/__init__.py

@@ -40,7 +40,7 @@ __license__ = "MIT"
 __description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/thomaspinder/GPJax"
 __contributors__ = "https://github.com/thomaspinder/GPJax/graphs/contributors"
-__version__ = "0.13.4"
+__version__ = "0.13.5"
 
 __all__ = [
     "gps",

{gpjax-0.13.4 → gpjax-0.13.5}/gpjax/parameters.py

@@ -1,6 +1,7 @@
 import typing as tp
 
 from flax import nnx
+import jax
 from jax.experimental import checkify
 import jax.numpy as jnp
 import jax.tree_util as jtu
@@ -162,7 +163,7 @@ def _check_is_arraylike(value: T) -> None:
     Raises:
         TypeError: If the value is not array-like.
     """
-    if not isinstance(value, (ArrayLike, list)):
+    if not isinstance(value, (jax.Array, ArrayLike, list)):
         raise TypeError(
             f"Expected parameter value to be an array-like type. Got {value}."
         )

{gpjax-0.13.4 → gpjax-0.13.5}/mkdocs.yml

@@ -118,10 +118,9 @@ plugins:
       handlers:
         python:
           paths: ["gpjax"]
-          rendering:
+          options:
             show_symbol_type_toc: true
             show_signature_annotations: true
-          options:
             members_order: source
             inherited_members: true
             show_source: false

{gpjax-0.13.4 → gpjax-0.13.5}/pyproject.toml

@@ -30,13 +30,14 @@ dependencies = [
   "beartype>0.16.1",
   "flax>=0.12.0",
   "numpy>=2.0.0",
+  "tensorstore!=0.1.76; sys_platform == 'darwin'",
 ]
 
 [project.optional-dependencies]
 docs = [
   "mkdocs>=1.5.3",
   "mkdocs-material>=9.5.12",
-  "mkdocstrings[python]<0.31.0",
+  "mkdocstrings[python]<1.1.0",
   "mkdocs-jupyter>=0.24.3",
   "mkdocs-gen-files>=0.5.0",
   "mkdocs-literate-nav>=0.6.0",
@@ -78,6 +79,7 @@ dev = [
 ]
 
 [tool.uv]
+exclude-newer = "7 days"
 managed = true
 dev-dependencies = [
   "ruff>=0.6",
@@ -260,12 +262,11 @@ convention = "numpy"
 
 [tool.ruff.lint.per-file-ignores]
 "gpjax/__init__.py" = ['I', 'F401', 'E402', 'D104']
-"gpjax/progress_bar.py" = ["TCH004"]
 "gpjax/scan.py" = ["PLR0913"]
 "gpjax/citation.py" = ["F811"]
-"tests/test_base/test_module.py" = ["PLR0915"]
 "tests/test_objectives.py" = ["PLR0913"]
-"docs/examples/barycentres.py" = ["PLR0913"]
+"examples/barycentres.py" = ["PLR0913"]
+"tests/*.py" = ["PLW0108"]
 
 [tool.isort]
 profile = "black"

{gpjax-0.13.4 → gpjax-0.13.5}/tests/test_parameters.py

@@ -1,3 +1,4 @@
+import jax
 from flax import nnx
 from jax import jit
 from jax.experimental import checkify
@@ -109,3 +110,22 @@ def test_check_in_bounds():
         _safe_assert(
             _check_in_bounds, jnp.array(1.5), low=jnp.array(0.0), high=jnp.array(1.0)
         )
+
+
+@pytest.mark.parametrize(
+    "param_cls, value",
+    [
+        (PositiveReal, jnp.array(1.0)),
+        (PositiveReal, jnp.array([1.0, 2.0])),
+        (Real, jnp.array(1.0)),
+        (NonNegativeReal, jnp.array(1.0)),
+    ],
+)
+def test_parameter_construction_under_grad(param_cls, value):
+    """Regression test for #592: parameter construction must accept JAX tracers."""
+
+    def f(x):
+        return param_cls(x).value.sum()
+
+    grad = jax.grad(f)(value)
+    assert grad.shape == value.shape