gpjax 0.10.0__py3-none-any.whl → 0.10.2__py3-none-any.whl

gpjax/__init__.py

@@ -39,7 +39,7 @@ __license__ = "MIT"
 __description__ = "Didactic Gaussian processes in JAX"
 __url__ = "https://github.com/JaxGaussianProcesses/GPJax"
 __contributors__ = "https://github.com/JaxGaussianProcesses/GPJax/graphs/contributors"
-__version__ = "0.10.0"
+__version__ = "0.10.2"
 
 __all__ = [
     "base",

gpjax/kernels/base.py

@@ -32,6 +32,7 @@ from gpjax.kernels.computations import (
 from gpjax.parameters import (
     Parameter,
     Real,
+    Static,
 )
 from gpjax.typing import (
     Array,
@@ -220,7 +221,9 @@ class Constant(AbstractKernel):
     def __init__(
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
-        constant: tp.Union[ScalarFloat, Parameter[ScalarFloat]] = jnp.array(0.0),
+        constant: tp.Union[
+            ScalarFloat, Parameter[ScalarFloat], Static[ScalarFloat]
+        ] = jnp.array(0.0),
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),
     ):
         if isinstance(constant, Parameter):

gpjax/kernels/nonstationary/polynomial.py

@@ -46,7 +46,7 @@ class Polynomial(AbstractKernel):
         self,
         active_dims: tp.Union[list[int], slice, None] = None,
         degree: int = 2,
-        shift: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 0.0,
+        shift: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 1.0,
         variance: tp.Union[ScalarFloat, nnx.Variable[ScalarArray]] = 1.0,
         n_dims: tp.Union[int, None] = None,
         compute_engine: AbstractKernelComputation = DenseKernelComputation(),

gpjax/mean_functions.py

@@ -28,6 +28,7 @@ from jaxtyping import (
 from gpjax.parameters import (
     Parameter,
     Real,
+    Static
 )
 from gpjax.typing import (
     Array,
@@ -130,9 +131,9 @@ class Constant(AbstractMeanFunction):
     """
 
     def __init__(
-        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter] = 0.0
+        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0
     ):
-        if isinstance(constant, Parameter):
+        if isinstance(constant, Parameter) or isinstance(constant, Static):
             self.constant = constant
         else:
             self.constant = Real(jnp.array(constant))
@@ -158,7 +159,7 @@ class Zero(Constant):
     """
 
     def __init__(self):
-        super().__init__(constant=jnp.array(0.0))
+        super().__init__(constant=Static(jnp.array(0.0)))
 
 
 class CombinationMeanFunction(AbstractMeanFunction):

gpjax/parameters.py

@@ -1,6 +1,7 @@
 import typing as tp
 
 from flax import nnx
+from jax.experimental import checkify
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jax.typing import ArrayLike
@@ -84,8 +85,7 @@ class PositiveReal(Parameter[T]):
 
     def __init__(self, value: T, tag: ParameterTag = "positive", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
-
-        _check_is_positive(self.value)
+        _safe_assert(_check_is_positive, self.value)
 
 
 class Real(Parameter[T]):
@@ -101,7 +101,17 @@ class SigmoidBounded(Parameter[T]):
     def __init__(self, value: T, tag: ParameterTag = "sigmoid", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
 
-        _check_in_bounds(self.value, 0.0, 1.0)
+        # Only perform validation in non-JIT contexts
+        if (
+            not isinstance(value, jnp.ndarray)
+            or not getattr(value, "aval", None) is None
+        ):
+            _safe_assert(
+                _check_in_bounds,
+                self.value,
+                low=jnp.array(0.0),
+                high=jnp.array(1.0),
+            )
 
 
 class Static(nnx.Variable[T]):
@@ -120,8 +130,13 @@ class LowerTriangular(Parameter[T]):
     def __init__(self, value: T, tag: ParameterTag = "lower_triangular", **kwargs):
         super().__init__(value=value, tag=tag, **kwargs)
 
-        _check_is_square(self.value)
-        _check_is_lower_triangular(self.value)
+        # Only perform validation in non-JIT contexts
+        if (
+            not isinstance(value, jnp.ndarray)
+            or not getattr(value, "aval", None) is None
+        ):
+            _safe_assert(_check_is_square, self.value)
+            _safe_assert(_check_is_lower_triangular, self.value)
 
 
 DEFAULT_BIJECTION = {
@@ -132,36 +147,83 @@ DEFAULT_BIJECTION = {
 }
 
 
-def _check_is_arraylike(value: T):
+def _check_is_arraylike(value: T) -> None:
+    """Check if a value is array-like.
+
+    Args:
+        value: The value to check.
+
+    Raises:
+        TypeError: If the value is not array-like.
+    """
     if not isinstance(value, (ArrayLike, list)):
         raise TypeError(
             f"Expected parameter value to be an array-like type. Got {value}."
         )
 
 
-def _check_is_positive(value: T):
-    if jnp.any(value < 0):
-        raise ValueError(
-            f"Expected parameter value to be strictly positive. Got {value}."
-        )
+@checkify.checkify
+def _check_is_positive(value):
+    checkify.check(
+        jnp.all(value > 0), "value needs to be positive, got {value}", value=value
+    )
 
 
-def _check_is_square(value: T):
-    if value.shape[0] != value.shape[1]:
-        raise ValueError(
-            f"Expected parameter value to be a square matrix. Got {value}."
-        )
+@checkify.checkify
+def _check_is_square(value: T) -> None:
+    """Check if a value is a square matrix.
 
+    Args:
+        value: The value to check.
 
-def _check_is_lower_triangular(value: T):
-    if not jnp.all(jnp.tril(value) == value):
-        raise ValueError(
-            f"Expected parameter value to be a lower triangular matrix. Got {value}."
-        )
+    Raises:
+        ValueError: If the value is not a square matrix.
+    """
+    checkify.check(
+        value.shape[0] == value.shape[1],
+        "value needs to be a square matrix, got {value}",
+        value=value,
+    )
 
 
-def _check_in_bounds(value: T, low: float, high: float):
-    if jnp.any((value < low) | (value > high)):
-        raise ValueError(
-            f"Expected parameter value to be bounded between {low} and {high}. Got {value}."
-        )
+@checkify.checkify
+def _check_is_lower_triangular(value: T) -> None:
+    """Check if a value is a lower triangular matrix.
+
+    Args:
+        value: The value to check.
+
+    Raises:
+        ValueError: If the value is not a lower triangular matrix.
+    """
+    checkify.check(
+        jnp.all(jnp.tril(value) == value),
+        "value needs to be a lower triangular matrix, got {value}",
+        value=value,
+    )
+
+
+@checkify.checkify
+def _check_in_bounds(value: T, low: T, high: T) -> None:
+    """Check if a value is bounded between low and high.
+
+    Args:
+        value: The value to check.
+        low: The lower bound.
+        high: The upper bound.
+
+    Raises:
+        ValueError: If any element of value is outside the bounds.
+    """
+    checkify.check(
+        jnp.all((value >= low) & (value <= high)),
+        "value needs to be bounded between {low} and {high}, got {value}",
+        value=value,
+        low=low,
+        high=high,
+    )
+
+
+def _safe_assert(fn: tp.Callable[[tp.Any], None], value: T, **kwargs) -> None:
+    error, _ = fn(value, **kwargs)
+    checkify.check_error(error)

{gpjax-0.10.0.dist-info → gpjax-0.10.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.10.0
+Version: 0.10.2
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues

{gpjax-0.10.0.dist-info → gpjax-0.10.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-gpjax/__init__.py,sha256=LeAdMRx9XYvLf6csLhCIv6IHnDbAFB9rP--TYXECgz0,1654
+gpjax/__init__.py,sha256=F9GVk18tdmvwiDEHZNo_4Wr0TkmPhWIEwl3KzEWQcaQ,1654
 gpjax/citation.py,sha256=f2Hzj5MLyCE7l0hHAzsEQoTORZH5hgV_eis4uoBiWvE,3811
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
 gpjax/distributions.py,sha256=X48FJr3reop9maherdMVt7-XZOm2f26T8AJt_IKM_oE,9339
@@ -7,14 +7,14 @@ gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
 gpjax/likelihoods.py,sha256=DOyV1L0ompkpeImMTiOOiWLJfqSqvDX_acOumuFqPEc,9234
 gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
-gpjax/mean_functions.py,sha256=et2HzlsYJNViBvTohF2wZYgCWQfDX4KboYeO7egMR1c,6420
+gpjax/mean_functions.py,sha256=BpeFkR3Eqa3O_FGp9BtSu9HKNSYZ8M08VtyfPfWbwRg,6479
 gpjax/objectives.py,sha256=XwkPyL_iovTNKpKGVNt0Lt2_OMTJitSPhuyCtUrJpbc,15383
-gpjax/parameters.py,sha256=Z4Wy3gEzPZG23-dtqC437_ZWnd_sPe9LcLCKn21ZBvA,4886
+gpjax/parameters.py,sha256=6VKq6wBzEUtx-GXniC8fEqjTNrTC1YwIOw66QguW6UM,6457
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
 gpjax/variational_families.py,sha256=s1rk7PtNTjQPabmVu-jBsuJBoqsxAAXwKFZJOEswkNQ,28161
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
-gpjax/kernels/base.py,sha256=abkj3zidsBs7YSkYEfjeJ5jTs1YyDCPoBM2ZzqaqrgI,11561
+gpjax/kernels/base.py,sha256=wXsrpm5ofy9S5MNgUkJk4lx2umcIJL6dDNhXY7cmTGk,11616
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
 gpjax/kernels/approximations/rff.py,sha256=4kD1uocjHmxkLgvf4DxB4_Gy7iefdPgnWiZB3jDiExI,4126
 gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
@@ -30,7 +30,7 @@ gpjax/kernels/non_euclidean/utils.py,sha256=z42aw8ga0zuREzHawemR9okttgrAUPmq-aN5
 gpjax/kernels/nonstationary/__init__.py,sha256=YpWQfOy_cqOKc5ezn37vqoK3Z6jznYiJz28BD_8F7AY,930
 gpjax/kernels/nonstationary/arccosine.py,sha256=UCTVJEhTZFQjARGFsYMImLnTDyTyxobIL5f2LiAHkPI,5822
 gpjax/kernels/nonstationary/linear.py,sha256=UKDHFCQzKWDMYo76qcb5-ujjnP2_iL-1tcN017xjK48,2562
-gpjax/kernels/nonstationary/polynomial.py,sha256=yTGobMPbCnKlj4PiQPSXEkWNrj2sjg_x9zFsnFa_j4E,3257
+gpjax/kernels/nonstationary/polynomial.py,sha256=7SDMfEcBCqnRn9xyj4iGcYLNvYJZiveN3uLZ_h12p10,3257
 gpjax/kernels/stationary/__init__.py,sha256=j4BMTaQlIx2kNAT1Dkf4iO2rm-f7_oSVWNrk1bN0tqE,1406
 gpjax/kernels/stationary/base.py,sha256=pQNkMo-E4bIT4tNfb7JvFJZC6fIIXNErsT1iQopFlAA,7063
 gpjax/kernels/stationary/matern12.py,sha256=b2vQCUqhd9NJK84L2RYjpI597uxy_7xgwsjS35Gc958,1807
@@ -42,7 +42,7 @@ gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UO
 gpjax/kernels/stationary/rbf.py,sha256=G13gg5phO7ite7D9QgoCy7gB2_y0FM6GZhgFW4RL6Xw,1734
 gpjax/kernels/stationary/utils.py,sha256=Xa9EEnxgFqEi08ZSFAZYYHhJ85_3Ac-ZUyUk18B63M4,2225
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-gpjax-0.10.0.dist-info/METADATA,sha256=wZyZSD1p2t_K5m25TrvGJr6lTlfaUxoB12F-0f1d9Co,9970
-gpjax-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.10.0.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.10.0.dist-info/RECORD,,
+gpjax-0.10.2.dist-info/METADATA,sha256=mqIBMOMKKiI9qkM_uFHSuPEXY17Jd6bOL5EM2hPiaok,9970
+gpjax-0.10.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.10.2.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.10.2.dist-info/RECORD,,