gpjax 0.9.2__py3-none-any.whl → 0.9.3__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.3"
 
 __all__ = [
     "base",

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)

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

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

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

@@ -1,18 +1,18 @@
-gpjax/__init__.py,sha256=Bx5JFaveeVk3qJMTzbmrKOFy0U7fNcQ_JnVo5m0ACGA,1697
+gpjax/__init__.py,sha256=UNfvnpRhJEZfz9qYjzWYNSSh3xg1FqNFst_A6xl_nfE,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/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=JO78dywHNH9__hjJkrP2ASb1L3C9aEBOW7fd0run-e4,27918
 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
@@ -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.3.dist-info/METADATA,sha256=A-tRR4wxz_YizCra4tZdOfRD1aUGJ_5sKgzL7Ax81B0,9976
+gpjax-0.9.3.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+gpjax-0.9.3.dist-info/licenses/LICENSE,sha256=tAkwu8-AdEyGxGoSvJ2gVmQdcicWw3j1ZZueVV74M-E,11357
+gpjax-0.9.3.dist-info/RECORD,,