gpjax 0.13.3__py3-none-any.whl → 0.13.5__py3-none-any.whl

gpjax/__init__.py

@@ -40,7 +40,7 @@ __license__ = "MIT"
 __description__ = "Gaussian processes in JAX and Flax"
 __url__ = "https://github.com/thomaspinder/GPJax"
 __contributors__ = "https://github.com/thomaspinder/GPJax/graphs/contributors"
-__version__ = "0.13.3"
+__version__ = "0.13.5"
 
 __all__ = [
     "gps",

gpjax/citation.py

@@ -23,6 +23,7 @@ from gpjax.kernels import (
     Matern32,
     Matern52,
 )
+from gpjax.likelihoods import HeteroscedasticGaussian
 
 CitationType = Union[None, str, Dict[str, str]]
 
@@ -149,3 +150,15 @@ def _(tree) -> PaperCitation:
         booktitle="Advances in neural information processing systems",
         citation_type="article",
     )
+
+
+@cite.register(HeteroscedasticGaussian)
+def _(tree) -> PaperCitation:
+    return PaperCitation(
+        citation_key="lazaro2011variational",
+        authors="Lázaro-Gredilla, Miguel and Titsias, Michalis",
+        title="Variational heteroscedastic Gaussian process regression",
+        year="2011",
+        booktitle="Proceedings of the 28th International Conference on Machine Learning (ICML)",
+        citation_type="inproceedings",
+    )

gpjax/gps.py

@@ -32,8 +32,10 @@ from gpjax.distributions import GaussianDistribution
 from gpjax.kernels import RFF
 from gpjax.kernels.base import AbstractKernel
 from gpjax.likelihoods import (
+    AbstractHeteroscedasticLikelihood,
     AbstractLikelihood,
     Gaussian,
+    HeteroscedasticGaussian,
     NonGaussian,
 )
 from gpjax.linalg import (
@@ -62,6 +64,7 @@ M = tp.TypeVar("M", bound=AbstractMeanFunction)
 L = tp.TypeVar("L", bound=AbstractLikelihood)
 NGL = tp.TypeVar("NGL", bound=NonGaussian)
 GL = tp.TypeVar("GL", bound=Gaussian)
+HL = tp.TypeVar("HL", bound=AbstractHeteroscedasticLikelihood)
 
 
 class AbstractPrior(nnx.Module, tp.Generic[M, K]):
@@ -476,6 +479,22 @@ class AbstractPosterior(nnx.Module, tp.Generic[P, L]):
         raise NotImplementedError
 
 
+class LatentPosterior(AbstractPosterior[P, L]):
+    r"""A posterior shell used to expose prior structure without inference."""
+
+    def predict(
+        self,
+        test_inputs: Num[Array, "N D"],
+        train_data: Dataset,
+        *,
+        return_covariance_type: Literal["dense", "diagonal"] = "dense",
+    ) -> GaussianDistribution:
+        raise NotImplementedError(
+            "LatentPosteriors are a lightweight wrapper for priors and do not "
+            "implement predictive distributions. Use a variational family for inference."
+        )
+
+
 class ConjugatePosterior(AbstractPosterior[P, GL]):
     r"""A Conjuate Gaussian process posterior object.
 
@@ -839,6 +858,40 @@ class NonConjugatePosterior(AbstractPosterior[P, NGL]):
         return GaussianDistribution(jnp.atleast_1d(mean.squeeze()), cov)
 
 
+class HeteroscedasticPosterior(LatentPosterior[P, HL]):
+    r"""Posterior shell for heteroscedastic likelihoods.
+
+    The posterior retains both the signal and noise priors; inference is delegated
+    to variational families and specialised objectives.
+    """
+
+    def __init__(
+        self,
+        prior: AbstractPrior[M, K],
+        likelihood: HL,
+        jitter: float = 1e-6,
+    ):
+        if likelihood.noise_prior is None:
+            raise ValueError("Heteroscedastic likelihoods require a noise_prior.")
+        super().__init__(prior=prior, likelihood=likelihood, jitter=jitter)
+        self.noise_prior = likelihood.noise_prior
+        self.noise_posterior = LatentPosterior(
+            prior=self.noise_prior, likelihood=likelihood, jitter=jitter
+        )
+
+
+class ChainedPosterior(HeteroscedasticPosterior[P, HL]):
+    r"""Posterior routed for heteroscedastic likelihoods using chained bounds."""
+
+    def __init__(
+        self,
+        prior: AbstractPrior[M, K],
+        likelihood: HL,
+        jitter: float = 1e-6,
+    ):
+        super().__init__(prior=prior, likelihood=likelihood, jitter=jitter)
+
+
 #######################
 # Utils
 #######################
@@ -854,6 +907,18 @@ def construct_posterior(  # noqa: F811
 ) -> NonConjugatePosterior[P, NGL]: ...
 
 
+@tp.overload
+def construct_posterior(  # noqa: F811
+    prior: P, likelihood: HeteroscedasticGaussian
+) -> HeteroscedasticPosterior[P, HeteroscedasticGaussian]: ...
+
+
+@tp.overload
+def construct_posterior(  # noqa: F811
+    prior: P, likelihood: AbstractHeteroscedasticLikelihood
+) -> ChainedPosterior[P, AbstractHeteroscedasticLikelihood]: ...
+
+
 def construct_posterior(prior, likelihood):  # noqa: F811
     r"""Utility function for constructing a posterior object from a prior and
     likelihood. The function will automatically select the correct posterior
@@ -873,6 +938,15 @@ def construct_posterior(prior, likelihood):  # noqa: F811
     if isinstance(likelihood, Gaussian):
         return ConjugatePosterior(prior=prior, likelihood=likelihood)
 
+    if (
+        isinstance(likelihood, HeteroscedasticGaussian)
+        and likelihood.supports_tight_bound()
+    ):
+        return HeteroscedasticPosterior(prior=prior, likelihood=likelihood)
+
+    if isinstance(likelihood, AbstractHeteroscedasticLikelihood):
+        return ChainedPosterior(prior=prior, likelihood=likelihood)
+
     return NonConjugatePosterior(prior=prior, likelihood=likelihood)
 
 
@@ -911,7 +985,10 @@ __all__ = [
     "AbstractPrior",
     "Prior",
     "AbstractPosterior",
+    "LatentPosterior",
     "ConjugatePosterior",
     "NonConjugatePosterior",
+    "HeteroscedasticPosterior",
+    "ChainedPosterior",
     "construct_posterior",
 ]

gpjax/likelihoods.py

@@ -10,15 +10,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from __future__ import annotations
 
 import abc
+from dataclasses import dataclass
 
 import beartype.typing as tp
 from flax import nnx
+import jax
 from jax import vmap
+import jax.nn as jnn
 import jax.numpy as jnp
 import jax.scipy as jsp
 from jaxtyping import Float
+import numpy as np
 import numpyro.distributions as npd
 
 from gpjax.distributions import GaussianDistribution
@@ -36,6 +41,20 @@ from gpjax.typing import (
 )
 
 
+@dataclass
+class NoiseMoments:
+    log_variance: Array
+    inv_variance: Array
+    variance: Array
+
+
+jax.tree_util.register_pytree_node(
+    NoiseMoments,
+    lambda x: ((x.log_variance, x.inv_variance, x.variance), None),
+    lambda _, x: NoiseMoments(*x),
+)
+
+
 class AbstractLikelihood(nnx.Module):
     r"""Abstract base class for likelihoods.
 
@@ -103,6 +122,9 @@ class AbstractLikelihood(nnx.Module):
         y: Float[Array, "N D"],
         mean: Float[Array, "N D"],
         variance: Float[Array, "N D"],
+        mean_g: tp.Optional[Float[Array, "N D"]] = None,
+        variance_g: tp.Optional[Float[Array, "N D"]] = None,
+        **_: tp.Any,
     ) -> Float[Array, " N"]:
         r"""Compute the expected log likelihood.
 
@@ -116,6 +138,12 @@ class AbstractLikelihood(nnx.Module):
             y (Float[Array, 'N D']): The observed response variable.
             mean (Float[Array, 'N D']): The variational mean.
             variance (Float[Array, 'N D']): The variational variance.
+            mean_g (Float[Array, 'N D']): Optional moments of the latent noise
+                process for heteroscedastic likelihoods.
+            variance_g (Float[Array, 'N D']): Optional moments of the latent noise
+                process for heteroscedastic likelihoods.
+            **_: Unused extra arguments for compatibility with specialised
+                likelihoods.
 
         Returns:
             ScalarFloat: The expected log likelihood.
@@ -126,6 +154,143 @@ class AbstractLikelihood(nnx.Module):
         )
 
 
+class AbstractNoiseTransform(nnx.Module):
+    """Abstract base class for noise transformations."""
+
+    @abc.abstractmethod
+    def __call__(self, x: Float[Array, "..."]) -> Float[Array, "..."]:
+        """Transform the input noise signal."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def moments(
+        self, mean: Float[Array, "..."], variance: Float[Array, "..."]
+    ) -> NoiseMoments:
+        """Compute the moments of the transformed noise signal."""
+        raise NotImplementedError
+
+
+class LogNormalTransform(AbstractNoiseTransform):
+    """Log-normal noise transformation."""
+
+    def __call__(self, x: Float[Array, "..."]) -> Float[Array, "..."]:
+        return jnp.exp(x)
+
+    def moments(
+        self, mean: Float[Array, "..."], variance: Float[Array, "..."]
+    ) -> NoiseMoments:
+        expected_variance = jnp.exp(mean + 0.5 * variance)
+        expected_log_variance = mean
+        expected_inv_variance = jnp.exp(-mean + 0.5 * variance)
+        return NoiseMoments(
+            log_variance=expected_log_variance,
+            inv_variance=expected_inv_variance,
+            variance=expected_variance,
+        )
+
+
+class SoftplusTransform(AbstractNoiseTransform):
+    """Softplus noise transformation."""
+
+    def __init__(self, num_points: int = 20):
+        self.num_points = num_points
+
+    def __call__(self, x: Float[Array, "..."]) -> Float[Array, "..."]:
+        return jnn.softplus(x)
+
+    def moments(
+        self, mean: Float[Array, "..."], variance: Float[Array, "..."]
+    ) -> NoiseMoments:
+        quad_x, quad_w = np.polynomial.hermite.hermgauss(self.num_points)
+        quad_w = jnp.asarray(quad_w / jnp.sqrt(jnp.pi))
+        quad_x = jnp.asarray(quad_x)
+
+        std = jnp.sqrt(variance)
+        samples = mean[..., None] + jnp.sqrt(2.0) * std[..., None] * quad_x
+        sigma2 = self(samples)
+        log_sigma2 = jnp.log(sigma2)
+        inv_sigma2 = 1.0 / sigma2
+
+        expected_variance = jnp.sum(sigma2 * quad_w, axis=-1)
+        expected_log_variance = jnp.sum(log_sigma2 * quad_w, axis=-1)
+        expected_inv_variance = jnp.sum(inv_sigma2 * quad_w, axis=-1)
+
+        return NoiseMoments(
+            log_variance=expected_log_variance,
+            inv_variance=expected_inv_variance,
+            variance=expected_variance,
+        )
+
+
+class AbstractHeteroscedasticLikelihood(AbstractLikelihood):
+    r"""Base class for heteroscedastic likelihoods with latent noise processes."""
+
+    def __init__(
+        self,
+        num_datapoints: int,
+        noise_prior,
+        noise_transform: tp.Union[
+            AbstractNoiseTransform,
+            tp.Callable[[Float[Array, "..."]], Float[Array, "..."]],
+        ] = SoftplusTransform(),
+        integrator: AbstractIntegrator = GHQuadratureIntegrator(),
+    ):
+        self.noise_prior = noise_prior
+
+        if isinstance(noise_transform, AbstractNoiseTransform):
+            self.noise_transform = noise_transform
+        else:
+            transform_name = getattr(noise_transform, "__name__", "")
+            if noise_transform is jnp.exp or transform_name == "exp":
+                self.noise_transform = LogNormalTransform()
+            else:
+                # Default to SoftplusTransform for softplus or unknown callables (legacy behavior used quadrature)
+                # Note: If an unknown callable is passed, we technically use SoftplusTransform which applies softplus.
+                # Users should implement AbstractNoiseTransform for custom transforms.
+                self.noise_transform = SoftplusTransform()
+
+        super().__init__(num_datapoints=num_datapoints, integrator=integrator)
+
+    def __call__(
+        self,
+        dist: tp.Union[npd.MultivariateNormal, GaussianDistribution],
+        noise_dist: tp.Optional[
+            tp.Union[npd.MultivariateNormal, GaussianDistribution]
+        ] = None,
+    ) -> npd.Distribution:
+        return self.predict(dist, noise_dist)
+
+    def supports_tight_bound(self) -> bool:
+        """Return whether the tighter bound from Lázaro-Gredilla & Titsias (2011)
+        is applicable."""
+        return False
+
+    def noise_statistics(
+        self, mean: Float[Array, "N D"], variance: Float[Array, "N D"]
+    ) -> NoiseMoments:
+        r"""Moment matching of the transformed noise process.
+
+        Args:
+            mean: Mean of the latent noise GP.
+            variance: Variance of the latent noise GP.
+
+        Returns:
+            NoiseMoments: Expected log variance, inverse variance, and variance.
+        """
+        return self.noise_transform.moments(mean, variance)
+
+    def expected_log_likelihood(
+        self,
+        y: Float[Array, "N D"],
+        mean: Float[Array, "N D"],
+        variance: Float[Array, "N D"],
+        mean_g: tp.Optional[Float[Array, "N D"]] = None,
+        variance_g: tp.Optional[Float[Array, "N D"]] = None,
+        **kwargs: tp.Any,
+    ) -> Float[Array, " N"]:
+        raise NotImplementedError
+
+
 class Gaussian(AbstractLikelihood):
     r"""Gaussian likelihood object."""
 
@@ -186,6 +351,69 @@ class Gaussian(AbstractLikelihood):
         return npd.MultivariateNormal(dist.mean, noisy_cov)
 
 
+class HeteroscedasticGaussian(AbstractHeteroscedasticLikelihood):
+    def predict(
+        self,
+        dist: tp.Union[npd.MultivariateNormal, GaussianDistribution],
+        noise_dist: tp.Optional[
+            tp.Union[npd.MultivariateNormal, GaussianDistribution]
+        ] = None,
+    ) -> npd.MultivariateNormal:
+        if noise_dist is None:
+            raise ValueError(
+                "noise_dist must be provided for heteroscedastic prediction."
+            )
+
+        n_data = dist.event_shape[0]
+        noise_mean = noise_dist.mean
+        noise_variance = jnp.diag(noise_dist.covariance_matrix)
+        noise_stats = self.noise_statistics(
+            noise_mean[..., None], noise_variance[..., None]
+        )
+
+        cov = dist.covariance_matrix
+        noisy_cov = cov.at[jnp.diag_indices(n_data)].add(noise_stats.variance.squeeze())
+
+        return npd.MultivariateNormal(dist.mean, noisy_cov)
+
+    def link_function(self, f: Float[Array, "..."]) -> npd.Normal:
+        sigma2 = self.noise_transform(jnp.zeros_like(f))
+        return npd.Normal(loc=f, scale=jnp.sqrt(sigma2))
+
+    def expected_log_likelihood(
+        self,
+        y: Float[Array, "N D"],
+        mean: Float[Array, "N D"],
+        variance: Float[Array, "N D"],
+        mean_g: tp.Optional[Float[Array, "N D"]] = None,
+        variance_g: tp.Optional[Float[Array, "N D"]] = None,
+        noise_stats: tp.Optional[NoiseMoments] = None,
+        return_parts: bool = False,
+        **_: tp.Any,
+    ) -> tp.Union[Float[Array, " N"], tuple[Float[Array, " N"], NoiseMoments]]:
+        if mean_g is None or variance_g is None:
+            raise ValueError(
+                "mean_g and variance_g must be provided for heteroscedastic models."
+            )
+
+        if noise_stats is None:
+            noise_stats = self.noise_statistics(mean_g, variance_g)
+        sq_error = jnp.square(y - mean)
+        log2pi = jnp.log(2.0 * jnp.pi)
+        expected = -0.5 * (
+            log2pi
+            + noise_stats.log_variance
+            + (sq_error + variance) * noise_stats.inv_variance
+        )
+        expected_sum = jnp.sum(expected, axis=1)
+        if return_parts:
+            return expected_sum, noise_stats
+        return expected_sum
+
+    def supports_tight_bound(self) -> bool:
+        return True
+
+
 class Bernoulli(AbstractLikelihood):
     def link_function(self, f: Float[Array, "..."]) -> npd.BernoulliProbs:
         r"""The probit link function of the Bernoulli likelihood.
@@ -268,7 +496,13 @@ __all__ = [
     "AbstractLikelihood",
     "NonGaussian",
     "Gaussian",
+    "AbstractHeteroscedasticLikelihood",
+    "HeteroscedasticGaussian",
     "Bernoulli",
     "Poisson",
     "inv_probit",
+    "NoiseMoments",
+    "AbstractNoiseTransform",
+    "LogNormalTransform",
+    "SoftplusTransform",
 ]

gpjax/mean_functions.py

@@ -147,9 +147,9 @@ class Constant(AbstractMeanFunction):
             Float[Array, "1"]: The evaluated mean function.
         """
         if isinstance(self.constant, Parameter):
-            return jnp.ones((x.shape[0], 1)) * self.constant.value
+            return jnp.ones((x.shape[0], 1), dtype=x.dtype) * self.constant.value
         else:
-            return jnp.ones((x.shape[0], 1)) * self.constant
+            return jnp.ones((x.shape[0], 1), dtype=x.dtype) * self.constant
 
 
 class Zero(Constant):

gpjax/objectives.py

@@ -14,6 +14,9 @@ from gpjax.gps import (
     ConjugatePosterior,
     NonConjugatePosterior,
 )
+from gpjax.likelihoods import (
+    AbstractHeteroscedasticLikelihood,
+)
 from gpjax.linalg import (
     Dense,
     lower_cholesky,
@@ -25,9 +28,13 @@ from gpjax.typing import (
     Array,
     ScalarFloat,
 )
-from gpjax.variational_families import AbstractVariationalFamily
+from gpjax.variational_families import (
+    AbstractVariationalFamily,
+    HeteroscedasticVariationalFamily,
+)
 
 VF = TypeVar("VF", bound=AbstractVariationalFamily)
+HVF = TypeVar("HVF", bound=HeteroscedasticVariationalFamily)
 
 
 Objective = tpe.Callable[[nnx.Module, Dataset], ScalarFloat]
@@ -414,3 +421,51 @@ def collapsed_elbo(variational_family: VF, data: Dataset) -> ScalarFloat:
 
     # log N(y; μx, Io² + KxzKzz⁻¹Kzx) - 1/2o² tr(Kxx - KxzKzz⁻¹Kzx)
     return (two_log_prob - two_trace).squeeze() / 2.0
+
+
+def heteroscedastic_elbo_conjugate(
+    variational_family: HVF, data: Dataset
+) -> ScalarFloat:
+    r"""Tight bound from Lázaro-Gredilla & Titsias (2011) for heteroscedastic Gaussian likelihoods."""
+    likelihood = variational_family.posterior.likelihood
+    mean_f, var_f, mean_g, var_g = variational_family.predict(data.X)
+
+    expected_ll, _ = likelihood.expected_log_likelihood(
+        data.y,
+        mean_f,
+        var_f,
+        mean_g=mean_g,
+        variance_g=var_g,
+        return_parts=True,
+    )
+
+    scale = likelihood.num_datapoints / data.n
+    return scale * jnp.sum(expected_ll) - variational_family.prior_kl()
+
+
+def heteroscedastic_elbo_chained(variational_family: HVF, data: Dataset) -> ScalarFloat:
+    r"""Generic chained bound for heteroscedastic likelihoods."""
+    likelihood: AbstractHeteroscedasticLikelihood = (
+        variational_family.posterior.likelihood
+    )
+    mean_f, var_f, mean_g, var_g = variational_family.predict(data.X)
+    noise_stats = likelihood.noise_statistics(mean_g, var_g)
+
+    expected_ll = likelihood.expected_log_likelihood(
+        data.y,
+        mean_f,
+        var_f,
+        mean_g=mean_g,
+        variance_g=var_g,
+        noise_stats=noise_stats,
+    )
+
+    scale = likelihood.num_datapoints / data.n
+    return scale * jnp.sum(expected_ll) - variational_family.prior_kl()
+
+
+def heteroscedastic_elbo(variational_family: HVF, data: Dataset) -> ScalarFloat:
+    likelihood = variational_family.posterior.likelihood
+    if likelihood.supports_tight_bound():
+        return heteroscedastic_elbo_conjugate(variational_family, data)
+    return heteroscedastic_elbo_chained(variational_family, data)

gpjax/parameters.py

@@ -1,6 +1,7 @@
 import typing as tp
 
 from flax import nnx
+import jax
 from jax.experimental import checkify
 import jax.numpy as jnp
 import jax.tree_util as jtu
@@ -77,7 +78,14 @@ class Parameter(nnx.Variable[T]):
         _check_is_arraylike(value)
 
         super().__init__(value=jnp.asarray(value), **kwargs)
-        self.tag = tag
+
+        # nnx.Variable metadata must be set via set_metadata (direct setattr is disallowed).
+        self.set_metadata(tag=tag)
+
+    @property
+    def tag(self) -> ParameterTag:
+        """Return the parameter's constraint tag."""
+        return self.metadata.get("tag", "real")
 
 
 class NonNegativeReal(Parameter[T]):
@@ -155,7 +163,7 @@ def _check_is_arraylike(value: T) -> None:
     Raises:
         TypeError: If the value is not array-like.
     """
-    if not isinstance(value, (ArrayLike, list)):
+    if not isinstance(value, (jax.Array, ArrayLike, list)):
         raise TypeError(
             f"Expected parameter value to be an array-like type. Got {value}."
         )

gpjax/variational_families.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 
 import abc
+from dataclasses import dataclass
 
 import beartype.typing as tp
 from flax import nnx
@@ -29,9 +30,12 @@ from gpjax.distributions import GaussianDistribution
 from gpjax.gps import (
     AbstractPosterior,
     AbstractPrior,
+    ChainedPosterior,
+    HeteroscedasticPosterior,
 )
 from gpjax.kernels.base import AbstractKernel
 from gpjax.likelihoods import (
+    AbstractHeteroscedasticLikelihood,
     Gaussian,
     NonGaussian,
 )
@@ -59,8 +63,10 @@ M = tp.TypeVar("M", bound=AbstractMeanFunction)
 L = tp.TypeVar("L", Gaussian, NonGaussian)
 NGL = tp.TypeVar("NGL", bound=NonGaussian)
 GL = tp.TypeVar("GL", bound=Gaussian)
+HL = tp.TypeVar("HL", bound=AbstractHeteroscedasticLikelihood)
 P = tp.TypeVar("P", bound=AbstractPrior)
 PP = tp.TypeVar("PP", bound=AbstractPosterior)
+HP = tp.TypeVar("HP", HeteroscedasticPosterior, ChainedPosterior)
 
 
 class AbstractVariationalFamily(nnx.Module, tp.Generic[L]):
@@ -870,6 +876,126 @@ class CollapsedVariationalGaussian(AbstractVariationalGaussian[GL]):
         )
 
 
+@dataclass
+class VariationalGaussianInit:
+    """Initialization parameters for a variational Gaussian distribution."""
+
+    inducing_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]]
+    variational_mean: tp.Union[Float[Array, "N 1"], None] = None
+    variational_root_covariance: tp.Union[Float[Array, "N N"], None] = None
+
+
+class HeteroscedasticPrediction(tp.NamedTuple):
+    """Mean and variance of the signal and noise latent processes."""
+
+    mean_f: Float[Array, "N 1"]
+    variance_f: Float[Array, "N 1"]
+    mean_g: Float[Array, "N 1"]
+    variance_g: Float[Array, "N 1"]
+
+
+class HeteroscedasticVariationalFamily(AbstractVariationalFamily[HL]):
+    r"""Variational family for two independent latent processes f and g."""
+
+    def __init__(
+        self,
+        posterior: HP,
+        inducing_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]] = None,
+        inducing_inputs_g: tp.Union[
+            Int[Array, "M D"], Float[Array, "M D"], None
+        ] = None,
+        variational_mean_f: tp.Union[Float[Array, "N 1"], None] = None,
+        variational_root_covariance_f: tp.Union[Float[Array, "N N"], None] = None,
+        variational_mean_g: tp.Union[Float[Array, "M 1"], None] = None,
+        variational_root_covariance_g: tp.Union[Float[Array, "M M"], None] = None,
+        jitter: ScalarFloat = 1e-6,
+        signal_init: tp.Optional[VariationalGaussianInit] = None,
+        noise_init: tp.Optional[VariationalGaussianInit] = None,
+    ):
+        self.jitter = jitter
+
+        if signal_init is not None:
+            self.signal_variational = VariationalGaussian(
+                posterior=posterior,
+                inducing_inputs=signal_init.inducing_inputs,
+                variational_mean=signal_init.variational_mean,
+                variational_root_covariance=signal_init.variational_root_covariance,
+                jitter=jitter,
+            )
+        elif inducing_inputs is not None:
+            self.signal_variational = VariationalGaussian(
+                posterior=posterior,
+                inducing_inputs=inducing_inputs,
+                variational_mean=variational_mean_f,
+                variational_root_covariance=variational_root_covariance_f,
+                jitter=jitter,
+            )
+        else:
+            raise ValueError("Either signal_init or inducing_inputs must be provided.")
+
+        if noise_init is not None:
+            self.noise_variational = VariationalGaussian(
+                posterior=posterior.noise_posterior,
+                inducing_inputs=noise_init.inducing_inputs,
+                variational_mean=noise_init.variational_mean,
+                variational_root_covariance=noise_init.variational_root_covariance,
+                jitter=jitter,
+            )
+        else:
+            noise_inducing = (
+                inducing_inputs if inducing_inputs_g is None else inducing_inputs_g
+            )
+            if noise_inducing is None and signal_init is not None:
+                noise_inducing = signal_init.inducing_inputs
+
+            if noise_inducing is None:
+                raise ValueError(
+                    "Could not determine inducing inputs for noise process."
+                )
+
+            self.noise_variational = VariationalGaussian(
+                posterior=posterior.noise_posterior,
+                inducing_inputs=noise_inducing,
+                variational_mean=variational_mean_g,
+                variational_root_covariance=variational_root_covariance_g,
+                jitter=jitter,
+            )
+        super().__init__(posterior)
+
+    def prior_kl(self) -> ScalarFloat:
+        return self.signal_variational.prior_kl() + self.noise_variational.prior_kl()
+
+    def predict(
+        self, test_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]]
+    ) -> HeteroscedasticPrediction:
+        dist_f = self.signal_variational.predict(test_inputs)
+        dist_g = self.noise_variational.predict(test_inputs)
+
+        mean_f = dist_f.mean[:, None] if dist_f.mean.ndim == 1 else dist_f.mean
+        var_f = (
+            dist_f.variance[:, None] if dist_f.variance.ndim == 1 else dist_f.variance
+        )
+        mean_g = dist_g.mean[:, None] if dist_g.mean.ndim == 1 else dist_g.mean
+        var_g = (
+            dist_g.variance[:, None] if dist_g.variance.ndim == 1 else dist_g.variance
+        )
+
+        return HeteroscedasticPrediction(
+            mean_f=mean_f,
+            variance_f=var_f,
+            mean_g=mean_g,
+            variance_g=var_g,
+        )
+
+    def predict_latents(
+        self, test_inputs: tp.Union[Int[Array, "N D"], Float[Array, "N D"]]
+    ) -> tuple[GaussianDistribution, GaussianDistribution]:
+        return (
+            self.signal_variational.predict(test_inputs),
+            self.noise_variational.predict(test_inputs),
+        )
+
+
 __all__ = [
     "AbstractVariationalFamily",
     "AbstractVariationalGaussian",
@@ -879,4 +1005,7 @@ __all__ = [
     "NaturalVariationalGaussian",
     "ExpectationVariationalGaussian",
     "CollapsedVariationalGaussian",
+    "HeteroscedasticVariationalFamily",
+    "VariationalGaussianInit",
+    "HeteroscedasticPrediction",
 ]

{gpjax-0.13.3.dist-info → gpjax-0.13.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpjax
-Version: 0.13.3
+Version: 0.13.5
 Summary: Gaussian processes in JAX.
 Project-URL: Documentation, https://docs.jaxgaussianprocesses.com/
 Project-URL: Issues, https://github.com/thomaspinder/GPJax/issues
@@ -25,6 +25,7 @@ Requires-Dist: jaxtyping>0.2.10
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: numpyro
 Requires-Dist: optax>0.2.1
+Requires-Dist: tensorstore!=0.1.76; sys_platform == 'darwin'
 Requires-Dist: tqdm>4.66.2
 Provides-Extra: dev
 Requires-Dist: absolufy-imports>=0.3.1; extra == 'dev'
@@ -59,7 +60,7 @@ Requires-Dist: mkdocs-jupyter>=0.24.3; extra == 'docs'
 Requires-Dist: mkdocs-literate-nav>=0.6.0; extra == 'docs'
 Requires-Dist: mkdocs-material>=9.5.12; extra == 'docs'
 Requires-Dist: mkdocs>=1.5.3; extra == 'docs'
-Requires-Dist: mkdocstrings[python]<0.31.0; extra == 'docs'
+Requires-Dist: mkdocstrings[python]<1.1.0; extra == 'docs'
 Requires-Dist: nbconvert>=7.16.2; extra == 'docs'
 Requires-Dist: networkx>=3.0; extra == 'docs'
 Requires-Dist: pandas>=1.5.3; extra == 'docs'
@@ -141,7 +142,7 @@ GPJax into the package it is today.
 > - [**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/)
+> - [**Heteroscedastic Inference**](https://docs.jaxgaussianprocesses.com/_examples/heteroscedastic_inference/)
 > - [**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/)

{gpjax-0.13.3.dist-info → gpjax-0.13.5.dist-info}/RECORD

@@ -1,18 +1,18 @@
-gpjax/__init__.py,sha256=8EULPS_vtq4TeN6MjdLVHLlTVVyQADOoHtuVoRN8z5Y,1625
-gpjax/citation.py,sha256=pwFS8h1J-LE5ieRS0zDyuwhmQHNxkFHYE7iSMlVNmQc,3928
+gpjax/__init__.py,sha256=8guaWVHiNXlZ9m_XvjPo8nherGpCpGWCCxB5fenq_Gw,1625
+gpjax/citation.py,sha256=9L-OVIDCrwjM3TeWinC5UHS3Lxn_9sVIsj_H9i08BT8,4428
 gpjax/dataset.py,sha256=Ef5JGrl4jJS1mQmL3JdO0fdqbVmflT_Cu5VrlpYdJY4,4071
 gpjax/distributions.py,sha256=iKmeQ_NN2CIjRiuOeJlwEGASzGROi4ZCerVi1uY7zRM,7758
 gpjax/fit.py,sha256=tOXmM3l5-N3Jlnq8MOVkSpRj-0fRWOy2t9GxQRyUqxY,15318
-gpjax/gps.py,sha256=NcPXkkx0kXrSBTUje4QR6PS0DGrD7p5c-DyiOATwUz8,35338
+gpjax/gps.py,sha256=83rki591ma3B_yr1Bq1OUlAH75sThpg-x00ZZXeqwN0,37906
 gpjax/integrators.py,sha256=eyJPqWNPKj6pKP5da0fEj4HW7BVyevqeGrurEuy_XPw,5694
-gpjax/likelihoods.py,sha256=xwnSQpn6Aa-FPpEoDn_3xpBdPQAmHP97jP-9iJmT4G8,9087
-gpjax/mean_functions.py,sha256=Aq09I5h5DZe1WokRxLa0Mpj8U0kJxpQ8CmdJJyCjNOc,6541
+gpjax/likelihoods.py,sha256=SzGh3uAkULVMfAVA3QpKQwWp-XAcx46B7QhhgKymn6A,17381
+gpjax/mean_functions.py,sha256=cDv3X5E4-to8pgzHpbKhoMKDLdWg3WhEq1W21kvbDDg,6571
 gpjax/numpyro_extras.py,sha256=-vWJ7SpZVNhSdCjjrlxIkovMFrM1IzpsMJK3B4LioGE,3411
-gpjax/objectives.py,sha256=GvKbDIPqYjsc9FpiTccmZwRdHr2lCykgfxI9BX9I_GA,15362
-gpjax/parameters.py,sha256=hnyIKr6uIzd7Kb3KZC9WowR88ruQwUvdcto3cx2ZDv4,6756
+gpjax/objectives.py,sha256=QScq5e77-jugtt2wKZNLgk3VtkqSDddJyYBs-hQ7dZc,17206
+gpjax/parameters.py,sha256=iGab7IMAANuhFm8ojcK5x5fJNzUOdflr2h2pBaVLe9Q,7035
 gpjax/scan.py,sha256=Z_V1yd76dL4B7_rnJnr_fohom6xzN_WRYxlTAAdqfa0,5357
 gpjax/typing.py,sha256=M3CvWsYtZ3PFUvBvvbRNjpwerNII0w4yGuP0I-sLeYI,1705
-gpjax/variational_families.py,sha256=x4VnUh9GiW77ijnDEwrYzH0aWFOW6NLPQXEp--I9R-g,31566
+gpjax/variational_families.py,sha256=J7L1lh0McYSgr1BKq5vsVjtfu2XeyOCn-I_AHQym-xA,36723
 gpjax/kernels/__init__.py,sha256=GFiku1s8KPPzvfwDYzEjENPQReywvJ29M9Kl2JjkcoU,1893
 gpjax/kernels/base.py,sha256=gov60BgXOQK0PxbKQRVgjU55kCVZYAkDZrAInODvugo,11549
 gpjax/kernels/approximations/__init__.py,sha256=bK9HlGd-PZeGrqtG5RpXxUTXNUrZTgfjH1dP626yNMA,68
@@ -46,7 +46,7 @@ gpjax/linalg/__init__.py,sha256=F8mxk_9Zc2nFd7Q-unjJ50_6rXEKzZj572WsU_jUKqI,547
 gpjax/linalg/operations.py,sha256=xvhOy5P4FmUCPWjIVNdg1yDXaoFQ48anFUfR-Tnfr6k,6480
 gpjax/linalg/operators.py,sha256=arxRGwcoAy_RqUYqBpZ3XG6OXbjShUl7m8sTpg85npE,11608
 gpjax/linalg/utils.py,sha256=fKV8G_iKZVhNkNvN20D_dQEi93-8xosGbXBP-v7UEyo,2020
-gpjax-0.13.3.dist-info/METADATA,sha256=1Au0xXIXWvS64SPLP8R_mPQ6nuUMayx2qdNLD-HdhW8,10296
-gpjax-0.13.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-gpjax-0.13.3.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
-gpjax-0.13.3.dist-info/RECORD,,
+gpjax-0.13.5.dist-info/METADATA,sha256=qJ8ZeG1TGMX6a_2WQnNSwTl2JBXVp8107QqY9f1XWyE,10382
+gpjax-0.13.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+gpjax-0.13.5.dist-info/licenses/LICENSE.txt,sha256=3umwi0h8wmKXOZO8XwRBwSl3vJt2hpWKEqSrSXLR7-I,1084
+gpjax-0.13.5.dist-info/RECORD,,

{gpjax-0.13.3.dist-info → gpjax-0.13.5.dist-info}/WHEEL

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