qpytorch 0.1__py3-none-any.whl

qpytorch/likelihoods/noise_models.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+
+import warnings
+from typing import Any, Optional, Union
+
+import torch
+from linear_operator.operators import ConstantDiagLinearOperator, DiagLinearOperator, LinearOperator, ZeroLinearOperator
+from torch import Tensor
+from torch.nn import Parameter
+
+from .. import settings
+from ..constraints import GreaterThan
+from ..distributions import MultivariateNormal, MultivariateQExponential
+from ..module import Module
+from ..utils.warnings import NumericalWarning
+
+
+class Noise(Module):
+    def __call__(
+        self, *params: Any, shape: Optional[torch.Size] = None, **kwargs: Any
+    ) -> Union[Tensor, LinearOperator]:
+        # For corredct typing
+        return super().__call__(*params, shape=shape, **kwargs)
+
+
+class _HomoskedasticNoiseBase(Noise):
+    def __init__(self, noise_prior=None, noise_constraint=None, batch_shape=torch.Size(), num_tasks=1):
+        super().__init__()
+        if noise_constraint is None:
+            noise_constraint = GreaterThan(1e-4)
+
+        self.register_parameter(name="raw_noise", parameter=Parameter(torch.zeros(*batch_shape, num_tasks)))
+        if noise_prior is not None:
+            self.register_prior("noise_prior", noise_prior, self._noise_param, self._noise_closure)
+
+        self.register_constraint("raw_noise", noise_constraint)
+
+    def _noise_param(self, m):
+        return m.noise
+
+    def _noise_closure(self, m, v):
+        return m._set_noise(v)
+
+    @property
+    def noise(self):
+        return self.raw_noise_constraint.transform(self.raw_noise)
+
+    @noise.setter
+    def noise(self, value: Tensor) -> None:
+        self._set_noise(value)
+
+    def _set_noise(self, value: Tensor) -> None:
+        if not torch.is_tensor(value):
+            value = torch.as_tensor(value).to(self.raw_noise)
+        self.initialize(raw_noise=self.raw_noise_constraint.inverse_transform(value))
+
+    def forward(self, *params: Any, shape: Optional[torch.Size] = None, **kwargs: Any) -> DiagLinearOperator:
+        """In the homoskedastic case, the parameters are only used to infer the required shape.
+        Here are the possible scenarios:
+        - non-batched noise, non-batched input, non-MT -> noise_diag shape is `n`
+        - non-batched noise, non-batched input, MT -> noise_diag shape is `nt`
+        - non-batched noise, batched input, non-MT -> noise_diag shape is `b x n` with b' the broadcasted batch shape
+        - non-batched noise, batched input, MT -> noise_diag shape is `b x nt`
+        - batched noise, non-batched input, non-MT -> noise_diag shape is `b x n`
+        - batched noise, non-batched input, MT -> noise_diag shape is `b x nt`
+        - batched noise, batched input, non-MT -> noise_diag shape is `b' x n`
+        - batched noise, batched input, MT -> noise_diag shape is `b' x nt`
+        where `n` is the number of evaluation points and `t` is the number of tasks (i.e. `num_tasks` of self.noise).
+        So bascially the shape is always `b' x nt`, with `b'` appropriately broadcast from the noise parameter and
+        input batch shapes. `n` and the input batch shape are determined either from the shape arg or from the params
+        input. For this it is sufficient to take in a single `shape` arg, with the convention that shape[:-1] is the
+        batch shape of the input, and shape[-1] is `n`.
+
+        If a "noise" kwarg (a Tensor) is provided, this noise is used directly.
+        """
+        if "noise" in kwargs:
+            return DiagLinearOperator(kwargs.get("noise"))
+        if shape is None:
+            p = params[0] if torch.is_tensor(params[0]) else params[0][0]
+            shape = p.shape if len(p.shape) == 1 else p.shape[:-1]
+        noise = self.noise
+        *batch_shape, n = shape
+        noise_batch_shape = noise.shape[:-1] if noise.dim() > 1 else torch.Size()
+        num_tasks = noise.shape[-1]
+        batch_shape = torch.broadcast_shapes(noise_batch_shape, batch_shape)
+        noise = noise.unsqueeze(-2)
+        noise_diag = noise.expand(*batch_shape, 1, num_tasks).contiguous()
+        if num_tasks == 1:
+            noise_diag = noise_diag.view(*batch_shape, 1)
+        if noise_diag.shape[-1] != 1:
+            noise_diag = noise_diag.unsqueeze(-1)
+        return ConstantDiagLinearOperator(noise_diag, diag_shape=n)
+
+
+class HomoskedasticNoise(_HomoskedasticNoiseBase):
+    def __init__(self, noise_prior=None, noise_constraint=None, batch_shape=torch.Size()):
+        super().__init__(
+            noise_prior=noise_prior, noise_constraint=noise_constraint, batch_shape=batch_shape, num_tasks=1
+        )
+
+
+class MultitaskHomoskedasticNoise(_HomoskedasticNoiseBase):
+    def __init__(self, num_tasks, noise_prior=None, noise_constraint=None, batch_shape=torch.Size()):
+        super().__init__(
+            noise_prior=noise_prior, noise_constraint=noise_constraint, batch_shape=batch_shape, num_tasks=num_tasks
+        )
+
+
+class HeteroskedasticNoise(Noise):
+    def __init__(self, noise_model, noise_indices=None, noise_constraint=None):
+        if noise_constraint is None:
+            noise_constraint = GreaterThan(1e-4)
+        super().__init__()
+        self.noise_model = noise_model
+        self._noise_constraint = noise_constraint
+        self._noise_indices = noise_indices
+
+    def forward(
+        self,
+        *params: Any,
+        batch_shape: Optional[torch.Size] = None,
+        shape: Optional[torch.Size] = None,
+        noise: Optional[Tensor] = None,
+    ) -> DiagLinearOperator:
+        if noise is not None:
+            return DiagLinearOperator(noise)
+        training = self.noise_model.training  # keep track of mode
+        try:
+            self.noise_model.eval()  # we want the posterior prediction of the noise model
+            with settings.detach_test_caches(False), settings.debug(False):
+                if len(params) == 1 and not torch.is_tensor(params[0]):
+                    output = self.noise_model(*params[0])
+                else:
+                    output = self.noise_model(*params)
+        finally:
+            self.noise_model.train(training)
+        if not isinstance(output, (MultivariateNormal, MultivariateQExponential)):
+            raise NotImplementedError("Currently only noise models that return a MultivariateNormal or MultivariateQExponential are supported")
+        # note: this also works with MultitaskMultivariateNormal (MultitaskMultivariateQExponential), where this
+        # will return a batched DiagLinearOperators of size n x num_tasks x num_tasks
+        noise_diag = output.mean if self._noise_indices is None else output.mean[..., self._noise_indices]
+        return DiagLinearOperator(self._noise_constraint.transform(noise_diag))
+
+
+class FixedNoise(Module):
+    def __init__(self, noise: Tensor) -> None:
+        super().__init__()
+        min_noise = settings.min_fixed_noise.value(noise.dtype)
+        if noise.lt(min_noise).any():
+            warnings.warn(
+                "Very small noise values detected. This will likely "
+                "lead to numerical instabilities. Rounding small noise "
+                f"values up to {min_noise}.",
+                NumericalWarning,
+            )
+            noise = noise.clamp_min(min_noise)
+        self.noise = noise
+
+    def forward(
+        self, *params: Any, shape: Optional[torch.Size] = None, noise: Optional[Tensor] = None, **kwargs: Any
+    ) -> DiagLinearOperator:
+        if shape is None:
+            p = params[0] if torch.is_tensor(params[0]) else params[0][0]
+            shape = p.shape if len(p.shape) == 1 else p.shape[:-1]
+
+        if noise is not None:
+            return DiagLinearOperator(noise)
+        elif shape[-1] == self.noise.shape[-1]:
+            return DiagLinearOperator(self.noise)
+        else:
+            return ZeroLinearOperator()
+
+    def _apply(self, fn):
+        self.noise = fn(self.noise)
+        return super(FixedNoise, self)._apply(fn)
+
+    def __call__(
+        self, *params: Any, shape: Optional[torch.Size] = None, **kwargs: Any
+    ) -> Union[Tensor, LinearOperator]:
+        # For corredct typing
+        return super().__call__(*params, shape=shape, **kwargs)
+
+class FixedGaussianNoise(FixedNoise): # alias to a more general name
+    pass

qpytorch/likelihoods/qexponential_likelihood.py

@@ -0,0 +1,494 @@
+#!/usr/bin/env python3
+import math
+import warnings
+from copy import deepcopy
+from typing import Any, Optional, Tuple, Union
+
+import torch
+from linear_operator.operators import LinearOperator, MaskedLinearOperator, ZeroLinearOperator, DiagLinearOperator
+from torch import Tensor
+from torch.distributions import Distribution
+
+from .. import settings
+from ..constraints import Interval
+from ..distributions import base_distributions, QExponential, MultivariateQExponential
+from ..priors import Prior
+from ..utils.warnings import QEPInputWarning
+from .likelihood import Likelihood
+from .noise_models import FixedNoise, HomoskedasticNoise, Noise
+
+
+class _QExponentialLikelihoodBase(Likelihood):
+    """Base class for QExponential Likelihoods, supporting general heteroskedastic noise models."""
+
+    has_analytic_marginal = True
+
+    def __init__(self, noise_covar: Union[Noise, FixedNoise], **kwargs: Any) -> None:
+        super().__init__()
+        param_transform = kwargs.get("param_transform")
+        if param_transform is not None:
+            warnings.warn(
+                "The 'param_transform' argument is now deprecated. If you want to use a different "
+                "transformaton, specify a different 'noise_constraint' instead.",
+                DeprecationWarning,
+            )
+
+        self.noise_covar = noise_covar
+        self.power = kwargs.pop('power', torch.tensor(2.0))
+        self.miu = kwargs.pop('miu', False) # marginally identical but uncorrelated
+
+    def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any) -> Union[Tensor, LinearOperator]:
+        return self.noise_covar(*params, shape=base_shape, **kwargs)
+
+    def expected_log_prob(self, target: Tensor, input: MultivariateQExponential, *params: Any, **kwargs: Any) -> Tensor:
+
+        noise = self._shaped_noise_covar(input.mean.shape, *params, **kwargs).diagonal(dim1=-1, dim2=-2)
+        # Potentially reshape the noise to deal with the multitask case
+        noise = noise.view(*noise.shape[:-1], *input.event_shape)
+
+        # Handle NaN values if enabled
+        nan_policy = settings.observation_nan_policy.value()
+        if nan_policy == "mask":
+            observed = settings.observation_nan_policy._get_observed(target, input.event_shape)
+            input = MultivariateQExponential(
+                mean=input.mean[..., observed],
+                covariance_matrix=MaskedLinearOperator(
+                    input.lazy_covariance_matrix, observed.reshape(-1), observed.reshape(-1)
+                ),
+                power=input.power
+            )
+            noise = noise[..., observed]
+            target = target[..., observed]
+        elif nan_policy == "fill":
+            missing = torch.isnan(target)
+            target = settings.observation_nan_policy._fill_tensor(target)
+
+        mean, variance, power = input.mean, input.variance, self.power
+        trace_plus_inv_quad_form = ((target - mean).square() + variance) / noise # <r>
+        res = noise.log() + math.log(2 * math.pi) #+ trace_plus_inv_quad_form**(power/2.)
+        # res = res.mul(-0.5)
+        # if power!=2: res += 0.5 * (power/2.-1) * trace_plus_inv_quad_form.log() + torch.log(power/2.) # exact value is intractable; only lower bound is provided.
+
+        if nan_policy == "fill":
+            # res = res * ~missing
+            trace_plus_inv_quad_form = trace_plus_inv_quad_form * ~missing
+
+        # Do appropriate summation for multitask QExponential likelihoods
+        num_event_dim = len(input.event_shape)
+        if num_event_dim > 1 and self.miu:
+            res = res.sum(list(range(-1, -num_event_dim, -1)))
+            trace_plus_inv_quad_form = trace_plus_inv_quad_form.sum(list(range(-1, -num_event_dim, -1)))
+        
+        res = res + trace_plus_inv_quad_form**(power/2.)
+        res = res.mul(-0.5)
+        if power!=2: res += 0.5 * (power/2.-1) * trace_plus_inv_quad_form.log() * (torch.tensor(input.event_shape[-1:-num_event_dim:-1]).prod() if self.miu else 1.0) + torch.log(power/2.) # exact value is intractable; only lower bound is provided.
+
+        if num_event_dim > 1 and not self.miu:
+            res = res.sum(list(range(-1, -num_event_dim, -1)))
+
+        if 'reduction' in kwargs:
+            res = eval('res.'+kwargs.pop('reduction'))
+
+        return res
+
+    def forward(self, function_samples: Tensor, *params: Any, **kwargs: Any) -> QExponential:
+        noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diagonal(dim1=-1, dim2=-2)
+        return QExponential(function_samples, noise.sqrt(), self.power)
+
+    def log_marginal(
+        self, observations: Tensor, function_dist: MultivariateQExponential, *params: Any, **kwargs: Any
+    ) -> Tensor:
+        marginal = self.marginal(function_dist, *params, **kwargs)
+
+        # Handle NaN values if enabled
+        nan_policy = settings.observation_nan_policy.value()
+        if nan_policy == "mask":
+            observed = settings.observation_nan_policy._get_observed(observations, marginal.event_shape)
+            marginal = MultivariateQExponential(
+                mean=marginal.mean[..., observed],
+                covariance_matrix=MaskedLinearOperator(
+                    marginal.lazy_covariance_matrix, observed.reshape(-1), observed.reshape(-1)
+                ),
+                power=marginal.power
+            )
+            observations = observations[..., observed]
+        elif nan_policy == "fill":
+            missing = torch.isnan(observations)
+            observations = settings.observation_nan_policy._fill_tensor(observations)
+
+        if self.miu:
+            if type(marginal) is MultivariateQExponential:
+                marginal = marginal.to_data_uncorrelated_dist()
+            else:
+                marginal.lazy_covariance_matrix = DiagLinearOperator(marginal.lazy_covariance_matrix.diagonal(dim1=-1, dim2=-2))
+            res = marginal.log_prob(observations)/marginal.event_shape[0]
+        else:
+            # We're making everything conditionally independent
+            indep_dist = QExponential(marginal.mean, marginal.variance.clamp_min(1e-8).sqrt(), marginal.power)
+            res = indep_dist.log_prob(observations)
+
+        if nan_policy == "fill":
+            res = res * ~missing
+
+        # Do appropriate summation for multitask QExponential likelihoods
+        num_event_dim = len(marginal.event_shape)
+        if num_event_dim > 1 and not self.miu:
+            res = res.sum(list(range(-1, -num_event_dim, -1)))
+        return res
+
+    def marginal(self, function_dist: MultivariateQExponential, *params: Any, **kwargs: Any) -> MultivariateQExponential:
+        mean, covar = function_dist.mean, function_dist.lazy_covariance_matrix
+        noise_covar = self._shaped_noise_covar(mean.shape, *params, **kwargs)
+        full_covar = covar + noise_covar
+        return function_dist.__class__(mean, full_covar, self.power)
+
+
+class QExponentialLikelihood(_QExponentialLikelihoodBase):
+    r"""
+    The standard likelihood for regression.
+    Assumes a standard homoskedastic noise model:
+
+    .. math::
+        p(y \mid f) = f + \epsilon, \quad \epsilon \sim \mathcal Q-EP (0, \sigma^2)
+
+    where :math:`\sigma^2` is a noise parameter.
+
+    .. note::
+        This likelihood can be used for exact or approximate inference.
+
+    .. note::
+        QExponentialLikelihood has an analytic marginal distribution.
+
+    :param noise_prior: Prior for noise parameter :math:`\sigma^2`.
+    :param noise_constraint: Constraint for noise parameter :math:`\sigma^2`.
+    :param batch_shape: The batch shape of the learned noise parameter (default: []).
+    :param kwargs: power (default: 2.0), miu (default: False).
+
+    :ivar torch.Tensor noise: :math:`\sigma^2` parameter (noise)
+    """
+
+    def __init__(
+        self,
+        noise_prior: Optional[Prior] = None,
+        noise_constraint: Optional[Interval] = None,
+        batch_shape: torch.Size = torch.Size(),
+        **kwargs: Any,
+    ) -> None:
+        noise_covar = HomoskedasticNoise(
+            noise_prior=noise_prior, noise_constraint=noise_constraint, batch_shape=batch_shape
+        )
+        super().__init__(noise_covar=noise_covar, **kwargs)
+
+    @property
+    def noise(self) -> Tensor:
+        return self.noise_covar.noise
+
+    @noise.setter
+    def noise(self, value: Tensor) -> None:
+        self.noise_covar.initialize(noise=value)
+
+    @property
+    def raw_noise(self) -> Tensor:
+        return self.noise_covar.raw_noise
+
+    @raw_noise.setter
+    def raw_noise(self, value: Tensor) -> None:
+        self.noise_covar.initialize(raw_noise=value)
+
+    def marginal(self, function_dist: MultivariateQExponential, *args: Any, **kwargs: Any) -> MultivariateQExponential:
+        r"""
+        :return: Analytic marginal :math:`p(\mathbf y)`.
+        """
+        return super().marginal(function_dist, *args, **kwargs)
+
+
+class QExponentialLikelihoodWithMissingObs(QExponentialLikelihood):
+    r"""
+    The standard likelihood for regression with support for missing values.
+    Assumes a standard homoskedastic noise model:
+
+    .. math::
+        p(y \mid f) = f + \epsilon, \quad \epsilon \sim \mathcal Q-EP (0, \sigma^2)
+
+    where :math:`\sigma^2` is a noise parameter. Values of y that are nan do
+    not impact the likelihood calculation.
+
+    .. note::
+        This likelihood can be used for exact or approximate inference.
+
+    .. warning::
+        This likelihood is deprecated in favor of :class:`gpytorch.settings.observation_nan_policy`.
+
+    :param noise_prior: Prior for noise parameter :math:`\sigma^2`.
+    :type noise_prior: ~gpytorch.priors.Prior, optional
+    :param noise_constraint: Constraint for noise parameter :math:`\sigma^2`.
+    :type noise_constraint: ~gpytorch.constraints.Interval, optional
+    :param batch_shape: The batch shape of the learned noise parameter (default: []).
+    :type batch_shape: torch.Size, optional
+    :var torch.Tensor noise: :math:`\sigma^2` parameter (noise)
+
+    .. note::
+        QExponentialLikelihoodWithMissingObs has an analytic marginal distribution.
+    """
+
+    MISSING_VALUE_FILL: float = -999.0
+
+    def __init__(self, **kwargs: Any) -> None:
+        warnings.warn(
+            "QExponentialLikelihoodWithMissingObs is replaced by qpytorch.settings.observation_nan_policy('fill').",
+            DeprecationWarning,
+        )
+        super().__init__(**kwargs)
+
+    def _get_masked_obs(self, x: Tensor) -> Tuple[Tensor, Tensor]:
+        missing_idx = x.isnan()
+        x_masked = x.masked_fill(missing_idx, self.MISSING_VALUE_FILL)
+        return missing_idx, x_masked
+
+    def expected_log_prob(self, target: Tensor, input: MultivariateQExponential, *params: Any, **kwargs: Any) -> Tensor:
+        missing_idx, target = self._get_masked_obs(target)
+        res = super().expected_log_prob(target, input, *params, **kwargs)
+        return res * ~missing_idx
+
+    def log_marginal(
+        self, observations: Tensor, function_dist: MultivariateQExponential, *params: Any, **kwargs: Any
+    ) -> Tensor:
+        missing_idx, observations = self._get_masked_obs(observations)
+        res = super().log_marginal(observations, function_dist, *params, **kwargs)
+        return res * ~missing_idx
+
+    def marginal(self, function_dist: MultivariateQExponential, *args: Any, **kwargs: Any) -> MultivariateQExponential:
+        r"""
+        :return: Analytic marginal :math:`p(\mathbf y)`.
+        """
+        return super().marginal(function_dist, *args, **kwargs)
+
+
+class FixedNoiseQExponentialLikelihood(_QExponentialLikelihoodBase):
+    r"""
+    A Likelihood that assumes fixed heteroscedastic noise. This is useful when you have fixed, known observation
+    noise for each training example.
+
+    Note that this likelihood takes an additional argument when you call it, `noise`, that adds a specified amount
+    of noise to the passed MultivariateQExponential. This allows for adding known observational noise to test data.
+
+    .. note::
+        This likelihood can be used for exact or approximate inference.
+
+    :param noise: Known observation noise (variance) for each training example.
+    :type noise: torch.Tensor (... x N)
+    :param learn_additional_noise: Set to true if you additionally want to
+        learn added diagonal noise, similar to QExponentialLikelihood.
+    :type learn_additional_noise: bool, optional
+    :param batch_shape: The batch shape of the learned noise parameter (default
+        []) if :obj:`learn_additional_noise=True`.
+    :type batch_shape: torch.Size, optional
+
+    :var torch.Tensor noise: :math:`\sigma^2` parameter (noise)
+
+    .. note::
+        FixedNoiseQExponentialLikelihood has an analytic marginal distribution.
+
+    Example:
+        >>> train_x = torch.randn(55, 2)
+        >>> noises = torch.ones(55) * 0.01
+        >>> likelihood = FixedNoiseQExponentialLikelihood(noise=noises, learn_additional_noise=True)
+        >>> pred_y = likelihood(qep_model(train_x))
+        >>>
+        >>> test_x = torch.randn(21, 2)
+        >>> test_noises = torch.ones(21) * 0.02
+        >>> pred_y = likelihood(qep_model(test_x), noise=test_noises)
+    """
+
+    def __init__(
+        self,
+        noise: Tensor,
+        learn_additional_noise: Optional[bool] = False,
+        batch_shape: Optional[torch.Size] = torch.Size(),
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(noise_covar=FixedNoise(noise=noise), **kwargs)
+
+        self.second_noise_covar: Optional[HomoskedasticNoise] = None
+        if learn_additional_noise:
+            noise_prior = kwargs.get("noise_prior", None)
+            noise_constraint = kwargs.get("noise_constraint", None)
+            self.second_noise_covar = HomoskedasticNoise(
+                noise_prior=noise_prior, noise_constraint=noise_constraint, batch_shape=batch_shape
+            )
+
+    @property
+    def noise(self) -> Tensor:
+        return self.noise_covar.noise + self.second_noise
+
+    @noise.setter
+    def noise(self, value: Tensor) -> None:
+        self.noise_covar.initialize(noise=value)
+
+    @property
+    def second_noise(self) -> Union[float, Tensor]:
+        if self.second_noise_covar is None:
+            return 0.0
+        else:
+            return self.second_noise_covar.noise
+
+    @second_noise.setter
+    def second_noise(self, value: Tensor) -> None:
+        if self.second_noise_covar is None:
+            raise RuntimeError(
+                "Attempting to set secondary learned noise for FixedNoiseQExponentialLikelihood, "
+                "but learn_additional_noise must have been False!"
+            )
+        self.second_noise_covar.initialize(noise=value)
+
+    def get_fantasy_likelihood(self, **kwargs: Any) -> "FixedNoiseQExponentialLikelihood":
+        if "noise" not in kwargs:
+            raise RuntimeError("FixedNoiseQExponentialLikelihood.fantasize requires a `noise` kwarg")
+        old_noise_covar = self.noise_covar
+        self.noise_covar = None  # pyre-fixme[8]
+        fantasy_liklihood = deepcopy(self)
+        self.noise_covar = old_noise_covar
+
+        old_noise = old_noise_covar.noise
+        new_noise = kwargs.get("noise")
+        if old_noise.dim() != new_noise.dim():
+            old_noise = old_noise.expand(*new_noise.shape[:-1], old_noise.shape[-1])
+        fantasy_liklihood.noise_covar = FixedNoise(noise=torch.cat([old_noise, new_noise], -1))
+        return fantasy_liklihood
+
+    def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any) -> Union[Tensor, LinearOperator]:
+        if len(params) > 0:
+            # we can infer the shape from the params
+            shape = None
+        else:
+            # here shape[:-1] is the batch shape requested, and shape[-1] is `n`, the number of points
+            shape = base_shape
+
+        res = self.noise_covar(*params, shape=shape, **kwargs)
+
+        if self.second_noise_covar is not None:
+            res = res + self.second_noise_covar(*params, shape=shape, **kwargs)
+        elif isinstance(res, ZeroLinearOperator):
+            warnings.warn(
+                "You have passed data through a FixedNoiseQExponentialLikelihood that did not match the size "
+                "of the fixed noise, *and* you did not specify noise. This is treated as a no-op.",
+                QEPInputWarning,
+            )
+
+        return res
+
+    def marginal(self, function_dist: MultivariateQExponential, *args: Any, **kwargs: Any) -> MultivariateQExponential:
+        r"""
+        :return: Analytic marginal :math:`p(\mathbf y)`.
+        """
+        return super().marginal(function_dist, *args, **kwargs)
+
+
+class QExponentialDirichletClassificationLikelihood(FixedNoiseQExponentialLikelihood):
+    r"""
+    A classification likelihood that treats the labels as regression targets with fixed heteroscedastic noise.
+    From Milios et al, NeurIPS, 2018 [https://arxiv.org/abs/1805.10915].
+
+    .. note::
+        This likelihood can be used for exact or approximate inference.
+
+    :param targets: (... x N) Classification labels.
+    :param alpha_epsilon: Tuning parameter for the scaling of the likeihood targets. We'd suggest 0.01 or setting
+        via cross-validation.
+    :param learn_additional_noise: Set to true if you additionally want to
+        learn added diagonal noise, similar to QExponentialLikelihood.
+    :param batch_shape: The batch shape of the learned noise parameter (default
+        []) if :obj:`learn_additional_noise=True`.
+
+    :ivar torch.Tensor noise: :math:`\sigma^2` parameter (noise)
+
+    .. note::
+        DirichletClassificationLikelihood has an analytic marginal distribution.
+
+    Example:
+        >>> train_x = torch.randn(55, 1)
+        >>> labels = torch.round(train_x).long()
+        >>> likelihood = DirichletClassificationLikelihood(targets=labels, learn_additional_noise=True)
+        >>> pred_y = likelihood(qep_model(train_x))
+        >>>
+        >>> test_x = torch.randn(21, 1)
+        >>> test_labels = torch.round(test_x).long()
+        >>> pred_y = likelihood(qep_model(test_x), targets=labels)
+    """
+
+    def _prepare_targets(
+        self, targets: Tensor, num_classes: Optional = None, alpha_epsilon: float = 0.01, dtype: torch.dtype = torch.float
+    ) -> Tuple[Tensor, Tensor, int]:
+        if num_classes is None: num_classes = int(targets.max() + 1)
+        # set alpha = \alpha_\epsilon
+        alpha = alpha_epsilon * torch.ones(targets.shape[-1], num_classes, device=targets.device, dtype=dtype)
+
+        # alpha[class_labels] = 1 + \alpha_\epsilon
+        alpha[torch.arange(len(targets)), targets] = alpha[torch.arange(len(targets)), targets] + 1.0
+
+        # sigma^2 = log(1 / alpha + 1)
+        sigma2_i = torch.log(alpha.reciprocal() + 1.0)
+
+        # y = log(alpha) - 0.5 * sigma^2
+        transformed_targets = alpha.log() - 0.5 * sigma2_i
+
+        return sigma2_i.transpose(-2, -1).type(dtype), transformed_targets.type(dtype), num_classes
+
+    def __init__(
+        self,
+        targets: Tensor,
+        alpha_epsilon: float = 0.01,
+        learn_additional_noise: Optional[bool] = False,
+        batch_shape: torch.Size = torch.Size(),
+        dtype: torch.dtype = torch.float,
+        **kwargs: Any,
+    ) -> None:
+        sigma2_labels, transformed_targets, num_classes = self._prepare_targets(
+            targets, alpha_epsilon=alpha_epsilon, dtype=dtype
+        )
+        super().__init__(
+            noise=sigma2_labels,
+            learn_additional_noise=learn_additional_noise,
+            batch_shape=torch.Size((num_classes,)),
+            **kwargs,
+        )
+        self.transformed_targets: Tensor = transformed_targets.transpose(-2, -1)
+        self.num_classes: int = num_classes
+        self.targets: Tensor = targets
+        self.alpha_epsilon: float = alpha_epsilon
+
+    def get_fantasy_likelihood(self, **kwargs: Any) -> "DirichletClassificationLikelihood":
+        # we assume that the number of classes does not change.
+
+        if "targets" not in kwargs:
+            raise RuntimeError("FixedNoiseQExponentialLikelihood.fantasize requires a `targets` kwarg")
+
+        old_noise_covar = self.noise_covar
+        self.noise_covar = None  # pyre-fixme[8]
+        fantasy_liklihood = deepcopy(self)
+        self.noise_covar = old_noise_covar
+
+        old_noise = old_noise_covar.noise
+        new_targets = kwargs.get("noise")
+        new_noise, new_targets, _ = fantasy_liklihood._prepare_targets(new_targets, self.alpha_epsilon)
+        fantasy_liklihood.targets = torch.cat([fantasy_liklihood.targets, new_targets], -1)
+
+        if old_noise.dim() != new_noise.dim():
+            old_noise = old_noise.expand(*new_noise.shape[:-1], old_noise.shape[-1])
+
+        fantasy_liklihood.noise_covar = FixedNoise(noise=torch.cat([old_noise, new_noise], -1))
+        return fantasy_liklihood
+
+    def marginal(self, function_dist: MultivariateQExponential, *args: Any, **kwargs: Any) -> MultivariateQExponential:
+        r"""
+        :return: Analytic marginal :math:`p(\mathbf y)`.
+        """
+        return super().marginal(function_dist, *args, **kwargs)
+
+    def __call__(self, input: Union[Tensor, MultivariateQExponential], *args: Any, **kwargs: Any) -> Distribution:
+        if "targets" in kwargs:
+            targets = kwargs.pop("targets")
+            dtype = self.transformed_targets.dtype
+            new_noise, _, _ = self._prepare_targets(targets, dtype=dtype)
+            kwargs["noise"] = new_noise
+        return super().__call__(input, *args, **kwargs)