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

gpjax/likelihoods.py

@@ -19,7 +19,7 @@ from jax import vmap
 import jax.numpy as jnp
 import jax.scipy as jsp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
 
 from gpjax.distributions import GaussianDistribution
 from gpjax.integrators import (
@@ -28,7 +28,7 @@ from gpjax.integrators import (
     GHQuadratureIntegrator,
 )
 from gpjax.parameters import (
-    PositiveReal,
+    NonNegativeReal,
     Static,
 )
 from gpjax.typing import (
@@ -36,9 +36,6 @@ from gpjax.typing import (
     ScalarFloat,
 )
 
-tfb = tfp.bijectors
-tfd = tfp.distributions
-
 
 class AbstractLikelihood(nnx.Module):
     r"""Abstract base class for likelihoods.
@@ -62,7 +59,7 @@ class AbstractLikelihood(nnx.Module):
         self.num_datapoints = num_datapoints
         self.integrator = integrator
 
-    def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> tfd.Distribution:
+    def __call__(self, *args: tp.Any, **kwargs: tp.Any) -> npd.Distribution:
         r"""Evaluate the likelihood function at a given predictive distribution.
 
         Args:
@@ -76,7 +73,7 @@ class AbstractLikelihood(nnx.Module):
         return self.predict(*args, **kwargs)
 
     @abc.abstractmethod
-    def predict(self, *args: tp.Any, **kwargs: tp.Any) -> tfd.Distribution:
+    def predict(self, *args: tp.Any, **kwargs: tp.Any) -> npd.Distribution:
         r"""Evaluate the likelihood function at a given predictive distribution.
 
         Args:
@@ -85,19 +82,19 @@ class AbstractLikelihood(nnx.Module):
                 `predict` method.
 
         Returns:
-            tfd.Distribution: The predictive distribution.
+            npd.Distribution: The predictive distribution.
         """
         raise NotImplementedError
 
     @abc.abstractmethod
-    def link_function(self, f: Float[Array, "..."]) -> tfd.Distribution:
+    def link_function(self, f: Float[Array, "..."]) -> npd.Distribution:
         r"""Return the link function of the likelihood function.
 
         Args:
             f (Float[Array, "..."]): the latent Gaussian process values.
 
         Returns:
-            tfd.Distribution: The distribution of observations, y, given values of the
+            npd.Distribution: The distribution of observations, y, given values of the
                 Gaussian process, f.
         """
         raise NotImplementedError
@@ -137,7 +134,7 @@ class Gaussian(AbstractLikelihood):
         self,
         num_datapoints: int,
         obs_stddev: tp.Union[
-            ScalarFloat, Float[Array, "#N"], PositiveReal, Static
+            ScalarFloat, Float[Array, "#N"], NonNegativeReal, Static
         ] = 1.0,
         integrator: AbstractIntegrator = AnalyticalGaussianIntegrator(),
     ):
@@ -151,26 +148,26 @@ 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 not isinstance(obs_stddev, (PositiveReal, Static)):
-            obs_stddev = PositiveReal(jnp.asarray(obs_stddev))
+        if not isinstance(obs_stddev, (NonNegativeReal, Static)):
+            obs_stddev = NonNegativeReal(jnp.asarray(obs_stddev))
         self.obs_stddev = obs_stddev
 
         super().__init__(num_datapoints, integrator)
 
-    def link_function(self, f: Float[Array, "..."]) -> tfd.Normal:
+    def link_function(self, f: Float[Array, "..."]) -> npd.Normal:
         r"""The link function of the Gaussian likelihood.
 
         Args:
             f (Float[Array, "..."]): Function values.
 
         Returns:
-            tfd.Normal: The likelihood function.
+            npd.Normal: The likelihood function.
         """
-        return tfd.Normal(loc=f, scale=self.obs_stddev.value.astype(f.dtype))
+        return npd.Normal(loc=f, scale=self.obs_stddev.value.astype(f.dtype))
 
     def predict(
-        self, dist: tp.Union[tfd.MultivariateNormalTriL, GaussianDistribution]
-    ) -> tfd.MultivariateNormalFullCovariance:
+        self, dist: tp.Union[npd.MultivariateNormal, GaussianDistribution]
+    ) -> npd.MultivariateNormal:
         r"""Evaluate the Gaussian likelihood.
 
         Evaluate the Gaussian likelihood function at a given predictive
@@ -179,75 +176,79 @@ class Gaussian(AbstractLikelihood):
         distribution's covariance matrix.
 
         Args:
-            dist (tfd.Distribution): The Gaussian process posterior,
+            dist (npd.Distribution): The Gaussian process posterior,
                 evaluated at a finite set of test points.
 
         Returns:
-            tfd.Distribution: The predictive distribution.
+            npd.Distribution: The predictive distribution.
         """
         n_data = dist.event_shape[0]
-        cov = dist.covariance()
+        cov = dist.covariance_matrix
         noisy_cov = cov.at[jnp.diag_indices(n_data)].add(self.obs_stddev.value**2)
 
-        return tfd.MultivariateNormalFullCovariance(dist.mean(), noisy_cov)
+        return npd.MultivariateNormal(dist.mean, noisy_cov)
 
 
 class Bernoulli(AbstractLikelihood):
-    def link_function(self, f: Float[Array, "..."]) -> tfd.Distribution:
+    def link_function(self, f: Float[Array, "..."]) -> npd.BernoulliProbs:
         r"""The probit link function of the Bernoulli likelihood.
 
         Args:
             f (Float[Array, "..."]): Function values.
 
         Returns:
-            tfd.Distribution: The likelihood function.
+            npd.Bernoulli: The likelihood function.
         """
-        return tfd.Bernoulli(probs=inv_probit(f))
+        return npd.Bernoulli(probs=inv_probit(f))
 
-    def predict(self, dist: tfd.Distribution) -> tfd.Distribution:
+    def predict(
+        self, dist: tp.Union[npd.MultivariateNormal, GaussianDistribution]
+    ) -> npd.BernoulliProbs:
         r"""Evaluate the pointwise predictive distribution.
 
         Evaluate the pointwise predictive distribution, given a Gaussian
         process posterior and likelihood parameters.
 
         Args:
-            dist (tfd.Distribution): The Gaussian process posterior, evaluated
-                at a finite set of test points.
+            dist ([npd.MultivariateNormal, GaussianDistribution].): The Gaussian
+                process posterior, evaluated at a finite set of test points.
 
         Returns:
-            tfd.Distribution: The pointwise predictive distribution.
+            npd.Bernoulli: The pointwise predictive distribution.
         """
-        variance = jnp.diag(dist.covariance())
-        mean = dist.mean().ravel()
+        variance = jnp.diag(dist.covariance_matrix)
+        mean = dist.mean.ravel()
         return self.link_function(mean / jnp.sqrt(1.0 + variance))
 
 
 class Poisson(AbstractLikelihood):
-    def link_function(self, f: Float[Array, "..."]) -> tfd.Distribution:
+    def link_function(self, f: Float[Array, "..."]) -> npd.Poisson:
         r"""The link function of the Poisson likelihood.
 
         Args:
             f (Float[Array, "..."]): Function values.
 
         Returns:
-            tfd.Distribution: The likelihood function.
+            npd.Poisson: The likelihood function.
         """
-        return tfd.Poisson(rate=jnp.exp(f))
+        return npd.Poisson(rate=jnp.exp(f))
 
-    def predict(self, dist: tfd.Distribution) -> tfd.Distribution:
+    def predict(
+        self, dist: tp.Union[npd.MultivariateNormal, GaussianDistribution]
+    ) -> npd.Poisson:
         r"""Evaluate the pointwise predictive distribution.
 
         Evaluate the pointwise predictive distribution, given a Gaussian
         process posterior and likelihood parameters.
 
         Args:
-            dist (tfd.Distribution): The Gaussian process posterior, evaluated
-                at a finite set of test points.
+            dist (tp.Union[npd.MultivariateNormal, GaussianDistribution]): The Gaussian
+                process posterior, evaluated at a finite set of test points.
 
         Returns:
-            tfd.Distribution: The pointwise predictive distribution.
+            npd.Poisson: The pointwise predictive distribution.
         """
-        return self.link_function(dist.mean())
+        return self.link_function(dist.mean)
 
 
 def inv_probit(x: Float[Array, " *N"]) -> Float[Array, " *N"]:

gpjax/mean_functions.py

@@ -28,7 +28,7 @@ from jaxtyping import (
 from gpjax.parameters import (
     Parameter,
     Real,
-    Static
+    Static,
 )
 from gpjax.typing import (
     Array,
@@ -131,7 +131,8 @@ class Constant(AbstractMeanFunction):
     """
 
     def __init__(
-        self, constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0
+        self,
+        constant: tp.Union[ScalarFloat, Float[Array, " O"], Parameter, Static] = 0.0,
     ):
         if isinstance(constant, Parameter) or isinstance(constant, Static):
             self.constant = constant
@@ -206,5 +207,5 @@ SumMeanFunction = ft.partial(
     CombinationMeanFunction, operator=ft.partial(jnp.sum, axis=0)
 )
 ProductMeanFunction = ft.partial(
-    CombinationMeanFunction, operator=ft.partial(jnp.sum, axis=0)
+    CombinationMeanFunction, operator=ft.partial(jnp.prod, axis=0)
 )

gpjax/numpyro_extras.py

@@ -0,0 +1,106 @@
+import math
+
+import jax
+import jax.numpy as jnp
+from numpyro.distributions.transforms import Transform
+
+# -----------------------------------------------------------------------------
+# Implementation: FillTriangularTransform
+# -----------------------------------------------------------------------------
+
+
+class FillTriangularTransform(Transform):
+    """
+    Transform that maps a vector of length n(n+1)/2 to an n x n lower triangular matrix.
+    The ordering is assumed to be:
+       (0,0), (1,0), (1,1), (2,0), (2,1), (2,2), ..., (n-1, n-1)
+    """
+
+    # Note: The base class provides `inv` through _InverseTransform wrapping _inverse.
+
+    def __call__(self, x):
+        """
+        Forward transformation.
+
+        Parameters
+        ----------
+        x : array_like, shape (..., L)
+            Input vector with L = n(n+1)/2 for some integer n.
+
+        Returns
+        -------
+        y : array_like, shape (..., n, n)
+            Lower-triangular matrix (with zeros in the upper triangle) filled in
+            row-major order (i.e. [ (0,0), (1,0), (1,1), ... ]).
+        """
+        L = x.shape[-1]
+        # Use static (Python) math.sqrt to compute n. This avoids tracer issues.
+        n = int((-1 + math.sqrt(1 + 8 * L)) // 2)
+        if n * (n + 1) // 2 != L:
+            raise ValueError("Last dimension must equal n(n+1)/2 for some integer n.")
+
+        def fill_single(vec):
+            out = jnp.zeros((n, n), dtype=vec.dtype)
+            row, col = jnp.tril_indices(n)
+            return out.at[row, col].set(vec)
+
+        if x.ndim == 1:
+            return fill_single(x)
+        else:
+            batch_shape = x.shape[:-1]
+            flat_x = x.reshape((-1, L))
+            out = jax.vmap(fill_single)(flat_x)
+            return out.reshape(batch_shape + (n, n))
+
+    def _inverse(self, y):
+        """
+        Inverse transformation.
+
+        Parameters
+        ----------
+        y : array_like, shape (..., n, n)
+            Lower triangular matrix.
+
+        Returns
+        -------
+        x : array_like, shape (..., n(n+1)/2)
+            The vector containing the elements from the lower-triangular portion of y.
+        """
+        if y.ndim < 2:
+            raise ValueError("Input to inverse must be at least two-dimensional.")
+        n = y.shape[-1]
+        if y.shape[-2] != n:
+            raise ValueError(
+                "Input matrix must be square; got shape %s" % str(y.shape[-2:])
+            )
+
+        row, col = jnp.tril_indices(n)
+
+        def inv_single(mat):
+            return mat[row, col]
+
+        if y.ndim == 2:
+            return inv_single(y)
+        else:
+            batch_shape = y.shape[:-2]
+            flat_y = y.reshape((-1, n, n))
+            out = jax.vmap(inv_single)(flat_y)
+            return out.reshape(batch_shape + (n * (n + 1) // 2,))
+
+    def log_abs_det_jacobian(self, x, y, intermediates=None):
+        # Since the transform simply reorders the vector into a matrix, the Jacobian determinant is 1.
+        return jnp.zeros(x.shape[:-1])
+
+    @property
+    def sign(self):
+        # The reordering transformation has a positive derivative everywhere.
+        return 1.0
+
+    # Implement tree_flatten and tree_unflatten because base Transform expects them.
+    def tree_flatten(self):
+        # This transform is stateless.
+        return (), {}
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        return cls()

gpjax/objectives.py

@@ -13,7 +13,7 @@ from jax import vmap
 import jax.numpy as jnp
 import jax.scipy as jsp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
 import typing_extensions as tpe
 
 from gpjax.dataset import Dataset
@@ -29,8 +29,6 @@ from gpjax.typing import (
 )
 from gpjax.variational_families import AbstractVariationalFamily
 
-tfd = tfp.distributions
-
 VF = TypeVar("VF", bound=AbstractVariationalFamily)
 
 
@@ -175,7 +173,7 @@ def conjugate_loocv(posterior: ConjugatePosterior, data: Dataset) -> ScalarFloat
     loocv_means = mx + (y - mx) - Sigma_inv_y / Sigma_inv_diag
     loocv_stds = jnp.sqrt(1.0 / Sigma_inv_diag)
 
-    loocv_posterior = tfd.Normal(loc=loocv_means, scale=loocv_stds)
+    loocv_posterior = npd.Normal(loc=loocv_means, scale=loocv_stds)
     return jnp.sum(loocv_posterior.log_prob(y))
 
 
@@ -232,7 +230,7 @@ def log_posterior_density(
     likelihood = posterior.likelihood.link_function(fx)
 
     # Whitened latent function values prior, p(wx | θ) = N(0, I)
-    latent_prior = tfd.Normal(loc=0.0, scale=1.0)
+    latent_prior = npd.Normal(loc=0.0, scale=1.0)
     return likelihood.log_prob(y).sum() + latent_prior.log_prob(wx).sum()
 
 
@@ -305,7 +303,7 @@ def variational_expectation(
     # inputs, x
     def q_moments(x):
         qx = q(x)
-        return qx.mean().squeeze(), qx.covariance().squeeze()
+        return qx.mean.squeeze(), qx.covariance().squeeze()
 
     mean, variance = vmap(q_moments)(x[:, None])
 

gpjax/parameters.py

@@ -5,7 +5,9 @@ from jax.experimental import checkify
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jax.typing import ArrayLike
-import tensorflow_probability.substrates.jax.bijectors as tfb
+import numpyro.distributions.transforms as npt
+
+from gpjax.numpyro_extras import FillTriangularTransform
 
 T = tp.TypeVar("T", bound=tp.Union[ArrayLike, list[float]])
 ParameterTag = str
@@ -13,7 +15,7 @@ ParameterTag = str
 
 def transform(
     params: nnx.State,
-    params_bijection: tp.Dict[str, tfb.Bijector],
+    params_bijection: tp.Dict[str, npt.Transform],
     inverse: bool = False,
 ) -> nnx.State:
     r"""Transforms parameters using a bijector.
@@ -22,7 +24,7 @@ def transform(
     ```pycon
         >>> from gpjax.parameters import PositiveReal, transform
         >>> import jax.numpy as jnp
-        >>> import tensorflow_probability.substrates.jax.bijectors as tfb
+        >>> import numpyro.distributions.transforms as npt
         >>> from flax import nnx
         >>> params = nnx.State(
         >>>     {
@@ -30,7 +32,7 @@ def transform(
         >>>         "b": PositiveReal(jnp.array([2.0])),
         >>>     }
         >>> )
-        >>> params_bijection = {'positive': tfb.Softplus()}
+        >>> params_bijection = {'positive': npt.SoftplusTransform()}
         >>> transformed_params = transform(params, params_bijection)
         >>> print(transformed_params["a"].value)
          [1.3132617]
@@ -47,11 +49,11 @@ def transform(
     """
 
     def _inner(param):
-        bijector = params_bijection.get(param._tag, tfb.Identity())
+        bijector = params_bijection.get(param._tag, npt.IdentityTransform())
         if inverse:
-            transformed_value = bijector.inverse(param.value)
+            transformed_value = bijector.inv(param.value)
         else:
-            transformed_value = bijector.forward(param.value)
+            transformed_value = bijector(param.value)
 
         param = param.replace(transformed_value)
         return param
@@ -80,6 +82,14 @@ class Parameter(nnx.Variable[T]):
         self._tag = tag
 
 
+class NonNegativeReal(Parameter[T]):
+    """Parameter that is non-negative."""
+
+    def __init__(self, value: T, tag: ParameterTag = "non_negative", **kwargs):
+        super().__init__(value=value, tag=tag, **kwargs)
+        _safe_assert(_check_is_non_negative, self.value)
+
+
 class PositiveReal(Parameter[T]):
     """Parameter that is strictly positive."""
 
@@ -104,7 +114,7 @@ class SigmoidBounded(Parameter[T]):
         # Only perform validation in non-JIT contexts
         if (
             not isinstance(value, jnp.ndarray)
-            or not getattr(value, "aval", None) is None
+            or getattr(value, "aval", None) is not None
         ):
             _safe_assert(
                 _check_in_bounds,
@@ -133,17 +143,18 @@ class LowerTriangular(Parameter[T]):
         # Only perform validation in non-JIT contexts
         if (
             not isinstance(value, jnp.ndarray)
-            or not getattr(value, "aval", None) is None
+            or getattr(value, "aval", None) is not None
         ):
             _safe_assert(_check_is_square, self.value)
             _safe_assert(_check_is_lower_triangular, self.value)
 
 
 DEFAULT_BIJECTION = {
-    "positive": tfb.Softplus(),
-    "real": tfb.Identity(),
-    "sigmoid": tfb.Sigmoid(low=0.0, high=1.0),
-    "lower_triangular": tfb.FillTriangular(),
+    "positive": npt.SoftplusTransform(),
+    "non_negative": npt.SoftplusTransform(),
+    "real": npt.IdentityTransform(),
+    "sigmoid": npt.SigmoidTransform(),
+    "lower_triangular": FillTriangularTransform(),
 }
 
 
@@ -162,6 +173,13 @@ def _check_is_arraylike(value: T) -> None:
         )
 
 
+@checkify.checkify
+def _check_is_non_negative(value):
+    checkify.check(
+        jnp.all(value >= 0), "value needs to be non-negative, got {value}", value=value
+    )
+
+
 @checkify.checkify
 def _check_is_positive(value):
     checkify.check(

gpjax/variational_families.py

@@ -22,6 +22,7 @@ from cola.linalg.inverse.inv import solve
 from cola.ops.operators import (
     Dense,
     I_like,
+    Identity,
     Triangular,
 )
 from flax import nnx
@@ -296,7 +297,10 @@ class WhitenedVariationalGaussian(VariationalGaussian[L]):
 
         # Compute whitened KL divergence
         qu = GaussianDistribution(loc=jnp.atleast_1d(mu.squeeze()), scale=S)
-        pu = GaussianDistribution(loc=jnp.zeros_like(jnp.atleast_1d(mu.squeeze())))
+        pu_S = Identity(shape=(self.num_inducing, self.num_inducing), dtype=mu.dtype)
+        pu = GaussianDistribution(
+            loc=jnp.zeros_like(jnp.atleast_1d(mu.squeeze())), scale=pu_S
+        )
         return qu.kl_divergence(pu)
 
     def predict(self, test_inputs: Float[Array, "N D"]) -> GaussianDistribution:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.10.2
+Version: 0.11.1
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/JaxGaussianProcesses/GPJax/issues
@@ -24,8 +24,8 @@ Requires-Dist: jax>=0.5.0
 Requires-Dist: jaxlib>=0.5.0
 Requires-Dist: jaxtyping>0.2.10
 Requires-Dist: numpy>=2.0.0
+Requires-Dist: numpyro
 Requires-Dist: optax>0.2.1
-Requires-Dist: tensorflow-probability>=0.24.0
 Requires-Dist: tqdm>4.66.2
 Description-Content-Type: text/markdown
 
@@ -138,65 +138,6 @@ jupytext --to notebook example.py
 jupytext --to py:percent example.ipynb
 ```
 
-# Simple example
-
-Let us import some dependencies and simulate a toy dataset $\mathcal{D}$.
-
-```python
-from jax import config
-
-config.update("jax_enable_x64", True)
-
-import gpjax as gpx
-from jax import grad, jit
-import jax.numpy as jnp
-import jax.random as jr
-import optax as ox
-
-key = jr.key(123)
-
-f = lambda x: 10 * jnp.sin(x)
-
-n = 50
-x = jr.uniform(key=key, minval=-3.0, maxval=3.0, shape=(n,1)).sort()
-y = f(x) + jr.normal(key, shape=(n,1))
-D = gpx.Dataset(X=x, y=y)
-
-# Construct the prior
-meanf = gpx.mean_functions.Zero()
-kernel = gpx.kernels.RBF()
-prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
-
-# Define a likelihood
-likelihood = gpx.likelihoods.Gaussian(num_datapoints = n)
-
-# Construct the posterior
-posterior = prior * likelihood
-
-# Define an optimiser
-optimiser = ox.adam(learning_rate=1e-2)
-
-# Obtain Type 2 MLEs of the hyperparameters
-opt_posterior, history = gpx.fit(
-    model=posterior,
-    objective=lambda p, d: -gpx.objectives.conjugate_mll(p, d),
-    train_data=D,
-    optim=optimiser,
-    num_iters=500,
-    safe=True,
-    key=key,
-)
-
-# Infer the predictive posterior distribution
-xtest = jnp.linspace(-3., 3., 100).reshape(-1, 1)
-latent_dist = opt_posterior(xtest, D)
-predictive_dist = opt_posterior.likelihood(latent_dist)
-
-# Obtain the predictive mean and standard deviation
-pred_mean = predictive_dist.mean()
-pred_std = predictive_dist.stddev()
-```
-
 # Installation
 
 ## Stable version

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

@@ -1,22 +1,23 @@
-gpjax/__init__.py,sha256=F9GVk18tdmvwiDEHZNo_4Wr0TkmPhWIEwl3KzEWQcaQ,1654
+gpjax/__init__.py,sha256=TjAAfeZTCEl_zsibA8pV76M1jcHkeFhNfWk_SllfgHY,1686
 gpjax/citation.py,sha256=f2Hzj5MLyCE7l0hHAzsEQoTORZH5hgV_eis4uoBiWvE,3811
 gpjax/dataset.py,sha256=NsToLKq4lOsHnfLfukrUIRKvhOEuoUk8aHTF0oAqRbU,4079
-gpjax/distributions.py,sha256=X48FJr3reop9maherdMVt7-XZOm2f26T8AJt_IKM_oE,9339
-gpjax/fit.py,sha256=OHv8jUHxa1ndpqMERSDRtYtUDzubk9rMPVIhfCiIH5Q,11551
+gpjax/distributions.py,sha256=8LWmfmRVHOX29Uy8PkKFi2UhcCiunuu-4TMI_5-krHc,9299
+gpjax/fit.py,sha256=7L2veA6aRNiozZD8fWa-MVDoYFUKjGJahmvjz8Wp-P0,15046
 gpjax/gps.py,sha256=97lYGrsmsufQxKEd8qz5wPNvui6FKXTF_Ps-sMFIjnY,31246
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
-gpjax/likelihoods.py,sha256=DOyV1L0ompkpeImMTiOOiWLJfqSqvDX_acOumuFqPEc,9234
+gpjax/likelihoods.py,sha256=99oTZoWld1M7vxgGM0pNY5Hnt2Ajd2lQNqawzrLmwtk,9308
 gpjax/lower_cholesky.py,sha256=3pnHaBrlGckFsrfYJ9Lsbd0pGmO7NIXdyY4aGm48MpY,1952
-gpjax/mean_functions.py,sha256=BpeFkR3Eqa3O_FGp9BtSu9HKNSYZ8M08VtyfPfWbwRg,6479
-gpjax/objectives.py,sha256=XwkPyL_iovTNKpKGVNt0Lt2_OMTJitSPhuyCtUrJpbc,15383
-gpjax/parameters.py,sha256=6VKq6wBzEUtx-GXniC8fEqjTNrTC1YwIOw66QguW6UM,6457
+gpjax/mean_functions.py,sha256=-sVYO1_LWE8f34rllUOuaT5sgGGAdxo99v5kRo2d4oM,6490
+gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
+gpjax/objectives.py,sha256=I_ZqnwTNYIAUAZ9KQNenIl0ish1jDOXb7KaNmjz3Su4,15340
+gpjax/parameters.py,sha256=H-DiXmotdBZCbf-GOjRaJoS_isk3GgFrpKHTq5GpnoA,6998
 gpjax/scan.py,sha256=jStQvwkE9MGttB89frxam1kaeXdWih7cVxkGywyaeHQ,5365
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=s1rk7PtNTjQPabmVu-jBsuJBoqsxAAXwKFZJOEswkNQ,28161
+gpjax/variational_families.py,sha256=Y9J1H91tXPm_hMy3ri_PgjAxqc_3r-BqKV83HRvB_m4,28295
 gpjax/kernels/__init__.py,sha256=WZanH0Tpdkt0f7VfMqnalm_VZAMVwBqeOVaICNj6xQU,1901
 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/approximations/rff.py,sha256=VbitjNuahFE5_IvCj1A0SxHhJXU0O0Qq0FMMVq8xA3E,4125
 gpjax/kernels/computations/__init__.py,sha256=uTVkqvnZVesFLDN92h0ZR0jfR69Eo2WyjOlmSYmCPJ8,1379
 gpjax/kernels/computations/base.py,sha256=zzabLN_-FkTWo6cBYjzjvUGYa7vrYyHxyrhQZxLgHBk,3651
 gpjax/kernels/computations/basis_functions.py,sha256=zY4rUDZDLOYvQPY9xosRmCLPdiXfbzAN5GICjQhFOis,2528
@@ -28,21 +29,21 @@ gpjax/kernels/non_euclidean/__init__.py,sha256=RT7puRPqCTpyxZ16q596EuOQEQi1LK1v3
 gpjax/kernels/non_euclidean/graph.py,sha256=K4WIdX-dx1SsWuNHZnNjHFw8ElKZxGcReUiA3w4aCOI,4204
 gpjax/kernels/non_euclidean/utils.py,sha256=z42aw8ga0zuREzHawemR9okttgrAUPmq-aN5HMt4SuY,1578
 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=7SDMfEcBCqnRn9xyj4iGcYLNvYJZiveN3uLZ_h12p10,3257
+gpjax/kernels/nonstationary/arccosine.py,sha256=2WV6aM0Z3-xXZnoPw-77n2CW62n-AZuJy-7AQ9xrMco,5858
+gpjax/kernels/nonstationary/linear.py,sha256=UIMoCq2hg6dQKr4J5UGiiPqotBleQuYfy00Ia1NaMOo,2571
+gpjax/kernels/nonstationary/polynomial.py,sha256=arP8DK0jnBOaayDWcFvHF0pdu9FVhwzXdqjnHUAL2VI,3293
 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
-gpjax/kernels/stationary/matern32.py,sha256=ZVYbUIQhvKpriC7abH8wV6Pk-mRoxtl3e2YYwH-KM5Y,2000
-gpjax/kernels/stationary/matern52.py,sha256=xfMYbY7MXxgMECtA2qVT5I8HoDGzGxygUvduGT3_Gvs,2053
+gpjax/kernels/stationary/base.py,sha256=25qDqpZP4gNtzbyzDCW-6u7rJfMqkg0dW88XUmTTupU,7078
+gpjax/kernels/stationary/matern12.py,sha256=DGjqw6VveYsyy0TrufyJJvCei7p9slnm2f0TgRGG7_U,1773
+gpjax/kernels/stationary/matern32.py,sha256=laLsJWJozJzpYHBzlkPUq0rWxz1eWEwGC36P2nPJuaQ,1966
+gpjax/kernels/stationary/matern52.py,sha256=VSByD2sb7k-DzRFjaz31P3Rtc4bPPhHvMshrxZNFnns,2019
 gpjax/kernels/stationary/periodic.py,sha256=IAbCxURtJEHGdmYzbdrsqRZ3zJ8F8tGQF9O7sggafZk,3598
 gpjax/kernels/stationary/powered_exponential.py,sha256=8qT91IWKJK7PpEtFcX4MVu1ahWMOFOZierPko4JCjKA,3776
 gpjax/kernels/stationary/rational_quadratic.py,sha256=dYONp3i4rnKj3ET8UyxAKXv6UOl8uOFT3lCutleSvo4,3496
-gpjax/kernels/stationary/rbf.py,sha256=G13gg5phO7ite7D9QgoCy7gB2_y0FM6GZhgFW4RL6Xw,1734
-gpjax/kernels/stationary/utils.py,sha256=Xa9EEnxgFqEi08ZSFAZYYHhJ85_3Ac-ZUyUk18B63M4,2225
+gpjax/kernels/stationary/rbf.py,sha256=euHUs6FdfRICQcabAWE4MX-7GEDr2TxgZWdFQiXr9Bw,1690
+gpjax/kernels/stationary/utils.py,sha256=6BI9EBcCzeeKx-XH-MfW1ORmtU__tPX5zyvfLhpkBsU,2180
 gpjax/kernels/stationary/white.py,sha256=TkdXXZCCjDs7JwR_gj5uvn2s1wyfRbe1vyHhUMJ8jjI,2212
-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,,
+gpjax-0.11.1.dist-info/METADATA,sha256=02crI6D0dsht6XJ8N1ZqNj5ZktmS5NymVfY45pPmEgM,8558
+gpjax-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.11.1.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.11.1.dist-info/RECORD,,