gpjax 0.10.2__py3-none-any.whl → 0.11.0__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.2"
+__version__ = "0.11.0"
 
 __all__ = [
     "base",

gpjax/distributions.py

@@ -15,77 +15,76 @@
 
 
 from beartype.typing import (
-    Any,
     Optional,
-    Tuple,
-    TypeVar,
 )
 import cola
+from cola.linalg.decompositions import Cholesky
 from cola.ops import (
-    Identity,
     LinearOperator,
 )
 from jax import vmap
 import jax.numpy as jnp
 import jax.random as jr
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+from numpyro.distributions import constraints
+from numpyro.distributions.distribution import Distribution
+from numpyro.distributions.util import is_prng_key
 
 from gpjax.lower_cholesky import lower_cholesky
 from gpjax.typing import (
     Array,
-    KeyArray,
     ScalarFloat,
 )
 
-tfd = tfp.distributions
-
-from cola.linalg.decompositions import Cholesky
-
 
-class GaussianDistribution(tfd.Distribution):
-    r"""Multivariate Gaussian distribution with a linear operator scale matrix."""
-
-    # TODO: Consider `distrax.transformed.Transformed` object. Can we create a LinearOperator to `distrax.bijector` representation
-    # and modify `distrax.MultivariateNormalFromBijector`?
-    # TODO: Consider natural and expectation parameterisations in future work.
-    # TODO: we don't really need to inherit from `tfd.Distribution` here
+class GaussianDistribution(Distribution):
+    support = constraints.real_vector
 
     def __init__(
         self,
-        loc: Optional[Float[Array, " N"]] = None,
-        scale: Optional[LinearOperator] = None,
-    ) -> None:
-        r"""Initialises the distribution.
-
-        Args:
-            loc: the mean of the distribution as an array of shape (n_points,).
-            scale: the scale matrix of the distribution as a LinearOperator object.
-        """
-        _check_loc_scale(loc, scale)
+        loc: Optional[Float[Array, " N"]],
+        scale: Optional[LinearOperator],
+        validate_args=None,
+    ):
+        self.loc = loc
+        self.scale = cola.PSD(scale)
+        batch_shape = ()
+        event_shape = jnp.shape(self.loc)
+        super().__init__(batch_shape, event_shape, validate_args=validate_args)
 
-        # Find dimensionality of the distribution.
-        if loc is not None:
-            num_dims = loc.shape[-1]
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        # Obtain covariance root.
+        covariance_root = lower_cholesky(self.scale)
 
-        elif scale is not None:
-            num_dims = scale.shape[-1]
+        # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
+        white_noise = jr.normal(
+            key, shape=sample_shape + self.batch_shape + self.event_shape
+        )
 
-        # Set the location to zero vector if unspecified.
-        if loc is None:
-            loc = jnp.zeros((num_dims,))
+        # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
+        def affine_transformation(_x):
+            return self.loc + covariance_root @ _x
 
-        # If not specified, set the scale to the identity matrix.
-        if scale is None:
-            scale = Identity(shape=(num_dims, num_dims), dtype=loc.dtype)
-
-        self.loc = loc
-        self.scale = cola.PSD(scale)
+        return vmap(affine_transformation)(white_noise)
 
+    @property
     def mean(self) -> Float[Array, " N"]:
         r"""Calculates the mean."""
         return self.loc
 
+    @property
+    def variance(self) -> Float[Array, " N"]:
+        r"""Calculates the variance."""
+        return cola.diag(self.scale)
+
+    def entropy(self) -> ScalarFloat:
+        r"""Calculates the entropy of the distribution."""
+        return 0.5 * (
+            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi))
+            + cola.logdet(self.scale, Cholesky(), Cholesky())
+        )
+
     def median(self) -> Float[Array, " N"]:
         r"""Calculates the median."""
         return self.loc
@@ -98,25 +97,19 @@ class GaussianDistribution(tfd.Distribution):
         r"""Calculates the covariance matrix."""
         return self.scale.to_dense()
 
-    def variance(self) -> Float[Array, " N"]:
-        r"""Calculates the variance."""
-        return cola.diag(self.scale)
+    @property
+    def covariance_matrix(self) -> Float[Array, "N N"]:
+        r"""Calculates the covariance matrix."""
+        return self.covariance()
 
     def stddev(self) -> Float[Array, " N"]:
         r"""Calculates the standard deviation."""
         return jnp.sqrt(cola.diag(self.scale))
 
-    @property
-    def event_shape(self) -> Tuple:
-        r"""Returns the event shape."""
-        return self.loc.shape[-1:]
-
-    def entropy(self) -> ScalarFloat:
-        r"""Calculates the entropy of the distribution."""
-        return 0.5 * (
-            self.event_shape[0] * (1.0 + jnp.log(2.0 * jnp.pi))
-            + cola.logdet(self.scale, Cholesky(), Cholesky())
-        )
+    #     @property
+    #     def event_shape(self) -> Tuple:
+    #         r"""Returns the event shape."""
+    #         return self.loc.shape[-1:]
 
     def log_prob(self, y: Float[Array, " N"]) -> ScalarFloat:
         r"""Calculates the log pdf of the multivariate Gaussian.
@@ -141,42 +134,39 @@ class GaussianDistribution(tfd.Distribution):
             + diff.T @ cola.solve(sigma, diff, Cholesky())
         )
 
-    def _sample_n(self, key: KeyArray, n: int) -> Float[Array, "n N"]:
-        r"""Samples from the distribution.
+    #     def _sample_n(self, key: KeyArray, n: int) -> Float[Array, "n N"]:
+    #         r"""Samples from the distribution.
 
-        Args:
-            key (KeyArray): The key to use for sampling.
+    #         Args:
+    #             key (KeyArray): The key to use for sampling.
 
-        Returns:
-            The samples as an array of shape (n_samples, n_points).
-        """
-        # Obtain covariance root.
-        sqrt = lower_cholesky(self.scale)
+    #         Returns:
+    #             The samples as an array of shape (n_samples, n_points).
+    #         """
+    #         # Obtain covariance root.
+    #         sqrt = lower_cholesky(self.scale)
 
-        # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
-        Z = jr.normal(key, shape=(n, *self.event_shape))
+    #         # Gather n samples from standard normal distribution Z = [z₁, ..., zₙ]ᵀ.
+    #         Z = jr.normal(key, shape=(n, *self.event_shape))
 
-        # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
-        def affine_transformation(x):
-            return self.loc + sqrt @ x
+    #         # xᵢ ~ N(loc, cov) <=> xᵢ = loc + sqrt zᵢ, where zᵢ ~ N(0, I).
+    #         def affine_transformation(x):
+    #             return self.loc + sqrt @ x
 
-        return vmap(affine_transformation)(Z)
+    #         return vmap(affine_transformation)(Z)
 
-    def sample(
-        self, seed: KeyArray, sample_shape: Tuple[int, ...]
-    ):  # pylint: disable=useless-super-delegation
-        r"""See `Distribution.sample`."""
-        return self._sample_n(
-            seed, sample_shape[0]
-        )  # TODO this looks weird, why ignore the second entry?
+    #     def sample(
+    #         self, seed: KeyArray, sample_shape: Tuple[int, ...]
+    #     ):  # pylint: disable=useless-super-delegation
+    #         r"""See `Distribution.sample`."""
+    #         return self._sample_n(
+    #             seed, sample_shape[0]
+    #         )  # TODO this looks weird, why ignore the second entry?
 
     def kl_divergence(self, other: "GaussianDistribution") -> ScalarFloat:
         return _kl_divergence(self, other)
 
 
-DistrT = TypeVar("DistrT", bound=tfd.Distribution)
-
-
 def _check_and_return_dimension(
     q: GaussianDistribution, p: GaussianDistribution
 ) -> int:
@@ -245,37 +235,37 @@ def _kl_divergence(q: GaussianDistribution, p: GaussianDistribution) -> ScalarFl
     ) / 2.0
 
 
-def _check_loc_scale(loc: Optional[Any], scale: Optional[Any]) -> None:
-    r"""Checks that the inputs are correct."""
-    if loc is None and scale is None:
-        raise ValueError("At least one of `loc` or `scale` must be specified.")
-
-    if loc is not None and loc.ndim < 1:
-        raise ValueError("The parameter `loc` must have at least one dimension.")
-
-    if scale is not None and len(scale.shape) < 2:  # scale.ndim < 2:
-        raise ValueError(
-            "The `scale` must have at least two dimensions, but "
-            f"`scale.shape = {scale.shape}`."
-        )
-
-    if scale is not None and not isinstance(scale, LinearOperator):
-        raise ValueError(
-            f"The `scale` must be a CoLA LinearOperator but got {type(scale)}"
-        )
-
-    if scale is not None and (scale.shape[-1] != scale.shape[-2]):
-        raise ValueError(
-            f"The `scale` must be a square matrix, but `scale.shape = {scale.shape}`."
-        )
-
-    if loc is not None:
-        num_dims = loc.shape[-1]
-        if scale is not None and (scale.shape[-1] != num_dims):
-            raise ValueError(
-                f"Shapes are not compatible: `loc.shape = {loc.shape}` and "
-                f"`scale.shape = {scale.shape}`."
-            )
+# def _check_loc_scale(loc: Optional[Any], scale: Optional[Any]) -> None:
+#     r"""Checks that the inputs are correct."""
+#     if loc is None and scale is None:
+#         raise ValueError("At least one of `loc` or `scale` must be specified.")
+
+#     if loc is not None and loc.ndim < 1:
+#         raise ValueError("The parameter `loc` must have at least one dimension.")
+
+#     if scale is not None and len(scale.shape) < 2:  # scale.ndim < 2:
+#         raise ValueError(
+#             "The `scale` must have at least two dimensions, but "
+#             f"`scale.shape = {scale.shape}`."
+#         )
+
+#     if scale is not None and not isinstance(scale, LinearOperator):
+#         raise ValueError(
+#             f"The `scale` must be a CoLA LinearOperator but got {type(scale)}"
+#         )
+
+#     if scale is not None and (scale.shape[-1] != scale.shape[-2]):
+#         raise ValueError(
+#             f"The `scale` must be a square matrix, but `scale.shape = {scale.shape}`."
+#         )
+
+#     if loc is not None:
+#         num_dims = loc.shape[-1]
+#         if scale is not None and (scale.shape[-1] != num_dims):
+#             raise ValueError(
+#                 f"Shapes are not compatible: `loc.shape = {loc.shape}` and "
+#                 f"`scale.shape = {scale.shape}`."
+#             )
 
 
 __all__ = [

gpjax/fit.py

@@ -20,9 +20,9 @@ import jax
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random as jr
+from numpyro.distributions.transforms import Transform
 import optax as ox
 from scipy.optimize import minimize
-from tensorflow_probability.substrates.jax.bijectors import Bijector
 
 from gpjax.dataset import Dataset
 from gpjax.objectives import Objective
@@ -47,7 +47,7 @@ def fit(  # noqa: PLR0913
     objective: Objective,
     train_data: Dataset,
     optim: ox.GradientTransformation,
-    params_bijection: tp.Union[dict[Parameter, Bijector], None] = DEFAULT_BIJECTION,
+    params_bijection: tp.Union[dict[Parameter, Transform], None] = DEFAULT_BIJECTION,
     key: KeyArray = jr.PRNGKey(42),
     num_iters: int = 100,
     batch_size: int = -1,

gpjax/kernels/approximations/rff.py

@@ -68,7 +68,7 @@ class RFF(AbstractKernel):
 
             self.frequencies = Static(
                 self.base_kernel.spectral_density.sample(
-                    seed=key, sample_shape=(self.num_basis_fns, n_dims)
+                    key=key, sample_shape=(self.num_basis_fns, n_dims)
                 )
             )
         self.name = f"{self.base_kernel.name} (RFF)"

gpjax/kernels/stationary/base.py

@@ -18,7 +18,7 @@ import beartype.typing as tp
 from flax import nnx
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.base import AbstractKernel
 from gpjax.kernels.computations import (
@@ -92,7 +92,7 @@ class StationaryKernel(AbstractKernel):
                 self.variance = tp.cast(PositiveReal[ScalarFloat], self.variance)
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.Normal | npd.StudentT:
         r"""The spectral density of the kernel.
 
         Returns:

gpjax/kernels/stationary/matern12.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import (
@@ -48,5 +48,5 @@ class Matern12(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.StudentT:
         return build_student_t_distribution(nu=1)

gpjax/kernels/stationary/matern32.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import (
@@ -54,5 +54,5 @@ class Matern32(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.StudentT:
         return build_student_t_distribution(nu=3)

gpjax/kernels/stationary/matern52.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax.distributions as tfd
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import (
@@ -53,5 +53,5 @@ class Matern52(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfd.Distribution:
+    def spectral_density(self) -> npd.StudentT:
         return build_student_t_distribution(nu=5)

gpjax/kernels/stationary/rbf.py

@@ -15,7 +15,7 @@
 
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
 
 from gpjax.kernels.stationary.base import StationaryKernel
 from gpjax.kernels.stationary.utils import squared_distance
@@ -44,5 +44,5 @@ class RBF(StationaryKernel):
         return K.squeeze()
 
     @property
-    def spectral_density(self) -> tfp.distributions.Normal:
-        return tfp.distributions.Normal(0.0, 1.0)
+    def spectral_density(self) -> npd.Normal:
+        return npd.Normal(0.0, 1.0)

gpjax/kernels/stationary/utils.py

@@ -14,17 +14,15 @@
 # ==============================================================================
 import jax.numpy as jnp
 from jaxtyping import Float
-import tensorflow_probability.substrates.jax as tfp
+import numpyro.distributions as npd
 
 from gpjax.typing import (
     Array,
     ScalarFloat,
 )
 
-tfd = tfp.distributions
 
-
-def build_student_t_distribution(nu: int) -> tfd.Distribution:
+def build_student_t_distribution(nu: int) -> npd.StudentT:
     r"""Build a Student's t distribution with a fixed smoothness parameter.
 
     For a fixed half-integer smoothness parameter, compute the spectral density of a
@@ -37,7 +35,7 @@ def build_student_t_distribution(nu: int) -> tfd.Distribution:
     -------
         tfp.Distribution: A Student's t distribution with the same smoothness parameter.
     """
-    dist = tfd.StudentT(df=nu, loc=0.0, scale=1.0)
+    dist = npd.StudentT(df=nu, loc=0.0, scale=1.0)
     return dist
 
 

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 (
@@ -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
@@ -157,20 +154,20 @@ class Gaussian(AbstractLikelihood):
 
         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

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
@@ -104,7 +106,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 +135,17 @@ 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(),
+    "real": npt.IdentityTransform(),
+    "sigmoid": npt.SigmoidTransform(),
+    "lower_triangular": FillTriangularTransform(),
 }
 
 

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.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.10.2
+Version: 0.11.0
 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.0.dist-info}/RECORD

@@ -1,22 +1,23 @@
-gpjax/__init__.py,sha256=F9GVk18tdmvwiDEHZNo_4Wr0TkmPhWIEwl3KzEWQcaQ,1654
+gpjax/__init__.py,sha256=wXJtQa_3W7wZEw_t1Dk0uHUzNQQDv8QzsVbnwXCMXcQ,1654
 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=STwpeqSuu2pgT6uZU7xd7koPZbAjPDzhcZ8nHfozR7Q,11538
 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=VcCibgihaskmvNJT4kuPa7ehgjlnR9LgMz_2KJJvHY0,9296
 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=gIPz7exEhish3yeJQxZp5Q_jlf2-gCE-KVAnL2Rumkc,6489
+gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
+gpjax/objectives.py,sha256=I_ZqnwTNYIAUAZ9KQNenIl0ish1jDOXb7KaNmjz3Su4,15340
+gpjax/parameters.py,sha256=Vj1xzrziSLxfBSqyc-BacyKBwkbE9Sjq4b1HV5HZiOg,6507
 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
@@ -32,17 +33,17 @@ gpjax/kernels/nonstationary/arccosine.py,sha256=UCTVJEhTZFQjARGFsYMImLnTDyTyxobI
 gpjax/kernels/nonstationary/linear.py,sha256=UKDHFCQzKWDMYo76qcb5-ujjnP2_iL-1tcN017xjK48,2562
 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
-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=FlsXMsXyZ5cI80jbsIo8Jv-H6gsV3C7v6plIhyCl-GI,7042
+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.0.dist-info/METADATA,sha256=eSWVc5y9WNrUmKpaOVq1CcHjrKjMwlmSvwovN9h9aCk,8558
+gpjax-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+gpjax-0.11.0.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.11.0.dist-info/RECORD,,