Replace deprecated jax context manager and JIT-traced non-array arguments (#260)

* Replace deprecated jax.experimental.enable_x64 contextmanager

* Avoid None attributes in JIT-wrapped functions

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Inline comment

* Drop Python 3.10

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* CI test on Python 3.11

* Python 3.11 for readthedocs

* Fix linting issues

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: hannes-holey

.github/workflows/tests.yml

@@ -13,7 +13,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.11"]
         x64: ["0"]
         include:
           - python-version: "3.13"

.readthedocs.yaml

@@ -5,7 +5,7 @@ build:
   apt_packages:
     - fonts-liberation
   tools:
-    python: "3.10"
+    python: "3.11"
 
 python:
   install:

pyproject.toml

@@ -3,7 +3,7 @@ name = "tinygp"
 description = "The tiniest of Gaussian Process libraries"
 authors = [{ name = "Dan Foreman-Mackey", email = "foreman.mackey@gmail.com" }]
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 license = { text = "MIT" }
 classifiers = [
     "Development Status :: 4 - Beta",
@@ -45,11 +45,11 @@ source = "vcs"
 version-file = "src/tinygp/tinygp_version.py"
 
 [tool.black]
-target-version = ["py39"]
+target-version = ["py312"]
 line-length = 88
 
 [tool.ruff]
-target-version = "py39"
+target-version = "py312"
 line-length = 88
 
 [tool.ruff.lint]
@@ -60,6 +60,7 @@ ignore = [
     "PLR0913", # Allow many arguments to functions
     "PLR0915", # Allow many statements
     "PLR2004", # Allow magic numbers in comparisons
+    "B905",    # Allow zip() without explicit `strict=` parameter
 ]
 exclude = []
 

src/tinygp/gp.py

@@ -2,12 +2,11 @@
 
 __all__ = ["GaussianProcess"]
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     NamedTuple,
 )
 

src/tinygp/kernels/base.py

@@ -12,8 +12,8 @@
 ]
 
 from abc import abstractmethod
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, Union
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any
 
 import equinox as eqx
 import jax
@@ -24,7 +24,7 @@
 if TYPE_CHECKING:
     from tinygp.solvers.solver import Solver
 
-Axis = Union[int, Sequence[int]]
+Axis = int | Sequence[int]
 
 
 class Kernel(eqx.Module):

src/tinygp/means.py

@@ -13,7 +13,7 @@
 __all__ = ["Mean", "Conditioned"]
 
 from abc import abstractmethod
-from typing import Callable
+from collections.abc import Callable
 
 import equinox as eqx
 import jax
@@ -39,18 +39,18 @@ class Mean(MeanBase):
             signature.
     """
 
-    value: JAXArray | None = None
+    value: JAXArray
     func: Callable[[JAXArray], JAXArray] | None = eqx.field(default=None, static=True)
 
     def __init__(self, value: JAXArray | Callable[[JAXArray], JAXArray]):
         if callable(value):
             self.func = value
+            self.value = jax.numpy.zeros(())  # avoids undefined traced values
         else:
             self.value = value
 
     def __call__(self, X: JAXArray) -> JAXArray:
-        if self.value is None:
-            assert self.func is not None
+        if self.func is not None:
             return self.func(X)
         return self.value
 

src/tinygp/solvers/quasisep/core.py

@@ -18,8 +18,9 @@
 
 import dataclasses
 from abc import abstractmethod
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable
+from typing import Any
 
 import equinox as eqx
 import jax

src/tinygp/solvers/quasisep/general.py

@@ -16,8 +16,9 @@
 
 __all__ = ["GeneralQSM"]
 
+from collections.abc import Callable
 from functools import wraps
-from typing import Any, Callable
+from typing import Any
 
 import equinox as eqx
 import jax

src/tinygp/transforms.py

@@ -8,9 +8,9 @@
 
 __all__ = ["Transform", "Linear", "Cholesky", "Subspace"]
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from functools import partial
-from typing import Any, Callable
+from typing import Any
 
 import equinox as eqx
 import jax.numpy as jnp

tests/test_gp.py

@@ -24,7 +24,7 @@ def data(random):
 def test_sample(data):
     X, _ = data
 
-    with jax.experimental.enable_x64(True):
+    with jax.enable_x64(True):
         gp = GaussianProcess(
             kernels.Matern32(1.5), X, diag=0.01, mean=lambda x: jnp.sum(x)
         )