Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hannes-holey wants to merge 3 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d588dac
Replace deprecated jax.experimental.enable_x64 contextmanager
hannes-holey Jan 27, 2026
9cb6245
Avoid None attributes in JIT-wrapped functions
hannes-holey Jan 27, 2026
e4a7f69
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jan 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/tinygp/means.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(())
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

Expand Down
27 changes: 27 additions & 0 deletions src/tinygp/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from contextlib import contextmanager
from typing import Any

import jax
Expand Down Expand Up @@ -30,3 +31,29 @@ def assert_pytrees_allclose(calculated: Any, expected: Any, *args: Any, **kwargs
jax.tree_util.tree_map(
lambda a, b: assert_allclose(a, b, *args, **kwargs), calculated, expected
)


def _as_context_manager(obj):
Copy link
Owner

@dfm dfm Jan 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is more complicated than it needs to be. Let's just switch to using jax.enable_x64 everywhere where we were using the experimental version!

Copy link
Author

@hannes-holey hannes-holey Jan 29, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree. If you are okay with dropping Python 3.10 I will make that change.

Copy link
Owner

@dfm dfm Jan 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's do it - thanks!

# If it's already a context manager
if hasattr(obj, "__enter__") and hasattr(obj, "__exit__"):
return obj

# If it's a generator, wrap it
if hasattr(obj, "__iter__") and hasattr(obj, "send"):
return contextmanager(lambda: obj)()

raise TypeError("Object is neither a context manager nor a generator")


@contextmanager
def jax_enable_x64():
if hasattr(jax, "enable_x64"):
cm = jax.enable_x64(True)
else:
# deprecated in jax>=0.9
from jax.experimental import enable_x64 as _enable_x64

cm = _enable_x64()

with _as_context_manager(cm):
yield
4 changes: 2 additions & 2 deletions tests/test_gp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from numpy import random as np_random

from tinygp import GaussianProcess, kernels
from tinygp.test_utils import assert_allclose
from tinygp.test_utils import assert_allclose, jax_enable_x64


@pytest.fixture
Expand All @@ -24,7 +24,7 @@ def data(random):
def test_sample(data):
X, _ = data

with jax.experimental.enable_x64(True):
with jax_enable_x64():
gp = GaussianProcess(
kernels.Matern32(1.5), X, diag=0.01, mean=lambda x: jnp.sum(x)
)
Expand Down
4 changes: 2 additions & 2 deletions tests/test_kernels/test_kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

from tinygp import kernels, noise
from tinygp.solvers import DirectSolver
from tinygp.test_utils import assert_allclose
from tinygp.test_utils import assert_allclose, jax_enable_x64


@pytest.fixture
Expand Down Expand Up @@ -71,7 +71,7 @@ def test_ops(data):

def test_conditioned(data):
x1, x2 = data
with jax.experimental.enable_x64(): # type: ignore
with jax_enable_x64(): # type: ignore
k1 = 1.5 * kernels.Matern32(2.5)
k2 = 0.9 * kernels.ExpSineSquared(scale=1.5, gamma=0.3)
K = k1(x1, x1) + 0.1 * jnp.eye(x1.shape[0])
Expand Down