gpjax 0.9.2__py3-none-any.whl → 0.9.4__py3-none-any.whl

gpjax/__init__.py

@@ -40,7 +40,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.9.2"
+__version__ = "0.9.4"
 
 __all__ = [
     "base",

gpjax/decision_making/test_functions/non_conjugate_functions.py

@@ -32,7 +32,7 @@ from gpjax.typing import (
 class PoissonTestFunction:
     """
     Test function for GPs utilising the Poisson likelihood. Function taken from
-    https://docs.jaxgaussianprocesses.com/examples/poisson/#dataset.
+    https://docs.jaxgaussianprocesses.com/_examples/poisson/#dataset.
 
     Attributes:
         search_space (ContinuousSearchSpace): Search space for the function.
@@ -77,7 +77,7 @@ class PoissonTestFunction:
     def evaluate(self, x: Float[Array, "N 1"]) -> Int[Array, "N 1"]:
         """
         Evaluate the test function at a set of points. Function taken from
-        https://docs.jaxgaussianprocesses.com/examples/poisson/#dataset.
+        https://docs.jaxgaussianprocesses.com/_examples/poisson/#dataset.
 
         Args:
             x (Float[Array, 'N D']): Points to evaluate the test function at.

gpjax/gps.py

@@ -652,7 +652,8 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         """
         super().__init__(prior=prior, likelihood=likelihood, jitter=jitter)
 
-        latent = latent or jr.normal(key, shape=(self.likelihood.num_datapoints, 1))
+        if latent is None:
+            latent = jr.normal(key, shape=(self.likelihood.num_datapoints, 1))
 
         # TODO: static or intermediate?
         self.latent = latent if isinstance(latent, Parameter) else Real(latent)

gpjax/likelihoods.py

@@ -28,7 +28,6 @@ from gpjax.integrators import (
     GHQuadratureIntegrator,
 )
 from gpjax.parameters import (
-    Parameter,
     PositiveReal,
     Static,
 )
@@ -152,10 +151,9 @@ class Gaussian(AbstractLikelihood):
                 likelihoods. Must be an instance of `AbstractIntegrator`. For the Gaussian likelihood, this defaults to
                 the `AnalyticalGaussianIntegrator`, as the expected log likelihood can be computed analytically.
         """
-        if isinstance(obs_stddev, Parameter):
-            self.obs_stddev = obs_stddev
-        else:
-            self.obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        if not isinstance(obs_stddev, (PositiveReal, Static)):
+            obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        self.obs_stddev = obs_stddev
 
         super().__init__(num_datapoints, integrator)
 

gpjax/scan.py

@@ -22,7 +22,6 @@ from beartype.typing import (
 )
 import jax
 from jax import lax
-from jax.experimental import host_callback as hcb
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jaxtyping import (
@@ -54,7 +53,8 @@ def _callback(cond: ScalarBool, func: Callable, *args: Any) -> None:
 
     def _do_callback(_) -> int:
         """Perform the callback."""
-        return hcb.id_tap(func, *args, result=_dummy_result)
+        jax.debug.callback(func, *args)
+        return _dummy_result
 
     def _not_callback(_) -> int:
         """Do nothing."""
@@ -113,19 +113,19 @@ def vscan(
     _progress_bar = trange(_length)
     _progress_bar.set_description("Compiling...", refresh=True)
 
-    def _set_running(args: Any, transform: Any) -> None:
+    def _set_running(*args: Any) -> None:
         """Set the tqdm progress bar to running."""
         _progress_bar.set_description("Running", refresh=False)
 
-    def _update_tqdm(args: Any, transform: Any) -> None:
+    def _update_tqdm(*args: Any) -> None:
         """Update the tqdm progress bar with the latest objective value."""
         _value, _iter_num = args
-        _progress_bar.update(_iter_num)
+        _progress_bar.update(_iter_num.item())
 
         if log_value and _value is not None:
             _progress_bar.set_postfix({"Value": f"{_value: .2f}"})
 
-    def _close_tqdm(args: Any, transform: Any) -> None:
+    def _close_tqdm(*args: Any) -> None:
         """Close the tqdm progress bar."""
         _progress_bar.close()
 
@@ -145,16 +145,16 @@ def vscan(
         _is_last: bool = iter_num == _length - 1
 
         # Update progress bar, if first of log_rate.
-        _callback(_is_first, _set_running, (y, log_rate))
+        _callback(_is_first, _set_running)
 
         # Update progress bar, if multiple of log_rate.
-        _callback(_is_multiple, _update_tqdm, (y, log_rate))
+        _callback(_is_multiple, _update_tqdm, y, log_rate)
 
         # Update progress bar, if remainder.
-        _callback(_is_remainder, _update_tqdm, (y, _remainder))
+        _callback(_is_remainder, _update_tqdm, y, _remainder)
 
         # Close progress bar, if last iteration.
-        _callback(_is_last, _close_tqdm, (y, None))
+        _callback(_is_last, _close_tqdm)
 
         return carry, y
 

gpjax/variational_families.py

@@ -108,10 +108,17 @@ class AbstractVariationalGaussian(AbstractVariationalFamily[L]):
     def __init__(
         self,
         posterior: AbstractPosterior[P, L],
-        inducing_inputs: Float[Array, "N D"],
+        inducing_inputs: tp.Union[
+            Float[Array, "N D"],
+            Real,
+            Static,
+        ],
         jitter: ScalarFloat = 1e-6,
     ):
-        self.inducing_inputs = Static(inducing_inputs)
+        if not isinstance(inducing_inputs, (Real, Static)):
+            inducing_inputs = Real(inducing_inputs)
+
+        self.inducing_inputs = inducing_inputs
         self.jitter = jitter
 
         super().__init__(posterior)
@@ -142,12 +149,14 @@ class VariationalGaussian(AbstractVariationalGaussian[L]):
     ):
         super().__init__(posterior, inducing_inputs, jitter)
 
-        self.variational_mean = Real(
-            variational_mean or jnp.zeros((self.num_inducing, 1))
-        )
-        self.variational_root_covariance = LowerTriangular(
-            variational_root_covariance or jnp.eye(self.num_inducing)
-        )
+        if variational_mean is None:
+            variational_mean = jnp.zeros((self.num_inducing, 1))
+
+        if variational_root_covariance is None:
+            variational_root_covariance = jnp.eye(self.num_inducing)
+
+        self.variational_mean = Real(variational_mean)
+        self.variational_root_covariance = LowerTriangular(variational_root_covariance)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the prior KL divergence.
@@ -371,12 +380,14 @@ class NaturalVariationalGaussian(AbstractVariationalGaussian[L]):
     ):
         super().__init__(posterior, inducing_inputs, jitter)
 
-        self.natural_vector = Static(
-            natural_vector or jnp.zeros((self.num_inducing, 1))
-        )
-        self.natural_matrix = Static(
-            natural_matrix or -0.5 * jnp.eye(self.num_inducing)
-        )
+        if natural_vector is None:
+            natural_vector = jnp.zeros((self.num_inducing, 1))
+
+        if natural_matrix is None:
+            natural_matrix = -0.5 * jnp.eye(self.num_inducing)
+
+        self.natural_vector = Static(natural_vector)
+        self.natural_matrix = Static(natural_matrix)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Compute the KL-divergence between our current variational approximation
@@ -533,13 +544,14 @@ class ExpectationVariationalGaussian(AbstractVariationalGaussian[L]):
     ):
         super().__init__(posterior, inducing_inputs, jitter)
 
-        # must come after super().__init__
-        self.expectation_vector = Static(
-            expectation_vector or jnp.zeros((self.num_inducing, 1))
-        )
-        self.expectation_matrix = Static(
-            expectation_matrix or jnp.eye(self.num_inducing)
-        )
+        if expectation_vector is None:
+            expectation_vector = jnp.zeros((self.num_inducing, 1))
+
+        if expectation_matrix is None:
+            expectation_matrix = jnp.eye(self.num_inducing)
+
+        self.expectation_vector = Static(expectation_vector)
+        self.expectation_matrix = Static(expectation_matrix)
 
     def prior_kl(self) -> ScalarFloat:
         r"""Evaluate the prior KL-divergence.

{gpjax-0.9.2.dist-info → gpjax-0.9.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: gpjax
-Version: 0.9.2
+Version: 0.9.4
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -19,7 +19,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: <3.13,>=3.10
 Requires-Dist: beartype>0.16.1
 Requires-Dist: cola-ml==0.0.5
-Requires-Dist: flax>=0.8.5
+Requires-Dist: flax<0.10.0
 Requires-Dist: jax<0.4.28
 Requires-Dist: jaxlib<0.4.28
 Requires-Dist: jaxopt==0.8.2
@@ -103,23 +103,23 @@ helped to shape GPJax into the package it is today.
 
 ## Notebook examples
 
-> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/examples/regression/)
-> - [**Classification**](https://docs.jaxgaussianprocesses.com/examples/classification/)
-> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/collapsed_vi/)
-> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/examples/uncollapsed_vi/)
-> - [**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/)
-> - [**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/)
-> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/examples/bayesian_optimisation/)
+> - [**Conjugate Inference**](https://docs.jaxgaussianprocesses.com/_examples/regression/)
+> - [**Classification**](https://docs.jaxgaussianprocesses.com/_examples/classification/)
+> - [**Sparse Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/collapsed_vi/)
+> - [**Stochastic Variational Inference**](https://docs.jaxgaussianprocesses.com/_examples/uncollapsed_vi/)
+> - [**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/)
+> - [**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/)
+> - [**Bayesian Optimisation**](https://docs.jaxgaussianprocesses.com/_examples/bayesian_optimisation/)
 
 ## Guides for customisation
 >
-> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/examples/constructing_new_kernels/#custom-kernel)
-> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/examples/yacht/)
+> - [**Custom kernels**](https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels/#custom-kernel)
+> - [**UCI regression**](https://docs.jaxgaussianprocesses.com/_examples/yacht/)
 
 ## Conversion between `.ipynb` and `.py`
 Above examples are stored in [examples](docs/examples) directory in the double
@@ -180,7 +180,7 @@ optimiser = ox.adam(learning_rate=1e-2)
 # Obtain Type 2 MLEs of the hyperparameters
 opt_posterior, history = gpx.fit(
     model=posterior,
-    objective=gpx.objectives.conjugate_mll,
+    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
     train_data=D,
     optim=optimiser,
     num_iters=500,

{gpjax-0.9.2.dist-info → gpjax-0.9.4.dist-info}/RECORD

@@ -1,18 +1,18 @@
-gpjax/__init__.py,sha256=Bx5JFaveeVk3qJMTzbmrKOFy0U7fNcQ_JnVo5m0ACGA,1697
+gpjax/__init__.py,sha256=f1Sl-8Oz6YuEueKxvzIAL0iH_9b9xGzQv07tddS5wto,1697
 gpjax/citation.py,sha256=R4Pmvjt0ndA0avEDSvIbxDxKapkRRYXWX7RRWBvZCRQ,5306
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
 gpjax/distributions.py,sha256=zxkSEZIlTg0PHvvgj0BQuIFEg-ugx6_NkEwSsbqWUM0,9325
 gpjax/fit.py,sha256=OHv8jUHxa1ndpqMERSDRtYtUDzubk9rMPVIhfCiIH5Q,11551
-gpjax/gps.py,sha256=NO18geRfcjo4mA3PGkuGont_Mj_yRqfvWzJqYmoKwiY,31225
+gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
-gpjax/likelihoods.py,sha256=Uh4kgLTod8ODw178L--G3w4olpm9XvCdcAZ8l7FwkF4,9255
+gpjax/likelihoods.py,sha256=DOyV1L0ompkpeImMTiOOiWLJfqSqvDX_acOumuFqPEc,9234
 gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
 gpjax/mean_functions.py,sha256=et2HzlsYJNViBvTohF2wZYgCWQfDX4KboYeO7egMR1c,6420
 gpjax/objectives.py,sha256=XwkPyL_iovTNKpKGVNt0Lt2_OMTJitSPhuyCtUrJpbc,15383
 gpjax/parameters.py,sha256=Z4Wy3gEzPZG23-dtqC437_ZWnd_sPe9LcLCKn21ZBvA,4886
-gpjax/scan.py,sha256=mtMsg8yLdkVuOYeTHLnATPGfGDnCMAQNdUA-FJlpfLs,5475
+gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=Eik5CCU7qH7_7cacpZ-1lIXm4tElELwSYfVw-n0rI20,27742
+gpjax/variational_families.py,sha256=s1rk7PtNTjQPabmVu-jBsuJBoqsxAAXwKFZJOEswkNQ,28161
 gpjax/decision_making/__init__.py,sha256=SDuPQl80lJ7nhfRsiB_7c22wCMiQO5ehSNohxUGnB7w,2170
 gpjax/decision_making/decision_maker.py,sha256=S4pOXrWcEHy0NDA0gfWzhk7pG0NJfaPpMXvq03yTy0g,13915
 gpjax/decision_making/posterior_handler.py,sha256=UgXf1Gu7GMh2YDSmiSWJIzmWlFW06KTS44HYz3mazZQ,5905
@@ -21,7 +21,7 @@ gpjax/decision_making/utility_maximizer.py,sha256=VT2amwSJbB64IL_MiWNl9ZgjcqO757
 gpjax/decision_making/utils.py,sha256=5j1GO5kcmG2laZR39NjhqgEjRekAWWzrnREv_5Zct_Y,2367
 gpjax/decision_making/test_functions/__init__.py,sha256=GDCY9_kaAnxDWwzo1FkdxnDx-80MErAHchbGybT9xYs,1109
 gpjax/decision_making/test_functions/continuous_functions.py,sha256=oL5ZQkvmbC3u9rEvSYI2DRAN3r7Ynf7wRZQlUWjKjt0,5612
-gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=cfo3xQOzB5ajMjjl0YFfNlJClkAcY7ZbT23UyBYEofQ,2955
+gpjax/decision_making/test_functions/non_conjugate_functions.py,sha256=eJpCnTS9dRieLxpjH4L6OTsP-w9JM3XhjnzCfk2Xqn8,2957
 gpjax/decision_making/utility_functions/__init__.py,sha256=xXI-4JKWAfTJ7XZ1vRDpqtb91MNzSPD0lP6xo0tOc7o,1445
 gpjax/decision_making/utility_functions/base.py,sha256=FOqrsRDmtHiCVl6IHr12-AEYBLStzMT5EBs-F92e1Og,3882
 gpjax/decision_making/utility_functions/expected_improvement.py,sha256=H6hjC-lj1oiHf2BomeQqroORQ7vtcOngiDAWxRwkNbg,4481
@@ -56,7 +56,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.9.2.dist-info/METADATA,sha256=JWT3cDW7onuKnTYUGqa15WxG4L7oEboJKPHYyAggYZ0,9976
-gpjax-0.9.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-gpjax-0.9.2.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
-gpjax-0.9.2.dist-info/RECORD,,
+gpjax-0.9.4.dist-info/METADATA,sha256=Qx_Qv91sE7_Y-c9CGuF40QBFk2FjLW0Fo2SHqFAgQFQ,10010
+gpjax-0.9.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.9.4.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
+gpjax-0.9.4.dist-info/RECORD,,

{gpjax-0.9.2.dist-info → gpjax-0.9.4.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.25.0
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any