qpytorch 0.1__py3-none-any.whl

qpytorch/kernels/qexponential_symmetrized_kl_kernel.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+
+import torch
+from torch.distributions import Chi2
+from gpytorch.kernels.distributional_input_kernel import DistributionalInputKernel
+
+
+def _symmetrized_kl(dist1, dist2, eps=1e-8, **kwargs):
+    """
+    Symmetrized KL distance between two q-exponential distributions. We assume that
+    the first half of the distribution tensors are the mean, and the second half
+    are the log variances.
+    Args:
+        dist1 (torch.Tensor) has shapes batch x n x dimensions. The first half
+            of the last dimensions are the means, while the second half are the log-variances.
+        dist2 (torch.Tensor) has shapes batch x n x dimensions. The first half
+            of the last dimensions are the means, while the second half are the log-variances.
+        eps (float) jitter term for the noise variance
+    """
+
+    power = kwargs.pop('power', torch.tensor(1.0))
+    exact = kwargs.pop('exact', False)
+    num_dims = int(dist1.shape[-1] / 2)
+
+    dist1_mean = dist1[..., :num_dims].unsqueeze(-3)
+    dist1_logvar = dist1[..., num_dims:].unsqueeze(-3)
+    dist1_var = eps + dist1_logvar.exp()
+
+    dist2_mean = dist2[..., :num_dims].unsqueeze(-2)
+    dist2_logvar = dist2[..., num_dims:].unsqueeze(-2)
+    dist2_var = eps + dist2_logvar.exp()
+
+    var_ratio12 = dist1_var / dist2_var
+    # log_var_ratio12 = var_ratio12.log()
+    # note that the log variance ratio cancels because of the summed KL.
+    loc_sqdiffs = (dist1_mean - dist2_mean).pow(2)
+    kl1 = 0.5 * ((var_ratio12 + loc_sqdiffs / dist2_var).pow(power/2.) - 1)
+    kl2 = 0.5 * ((var_ratio12.reciprocal() + loc_sqdiffs / dist1_var).pow(power/2.) - 1)
+    if power!=2:
+        kl1 += 0.5* (-(1-2./power)*torch.log(2*kl1+1) + (power/2.-1) * (-2./power*Chi2(1).entropy() if exact else 0) )
+        kl2 += 0.5* (-(1-2./power)*torch.log(2*kl2+1) + (power/2.-1) * (-2./power*Chi2(1).entropy() if exact else 0) )
+    symmetrized_kl = kl1 + kl2
+    return symmetrized_kl.sum(-1).transpose(-1, -2)
+
+
+class QExponentialSymmetrizedKLKernel(DistributionalInputKernel):
+    r"""
+    Computes a kernel based on the symmetrized KL divergence, assuming that two q-exponential
+    distributions are inputted. Inputs are assumed to be `batch x N x 2d` tensors where `d` is the
+    dimension of the distribution. The first `d` dimensions are the mean parameters of the
+    `batch x N` distributions, while the second `d` dimensions are the log variances.
+
+    Original citation is Moreno et al, '04
+    (https://papers.nips.cc/paper/2351-a-kullback-leibler-divergence-based-kernel-for-svm-\
+    classification-in-multimedia-applications.pdf) for the symmetrized KL divergence kernel between
+    two Gaussian distributions.
+    """
+
+    def __init__(self, **kwargs):
+        distance_function = lambda dist1, dist2: _symmetrized_kl(dist1, dist2, **kwargs)
+        super(QExponentialSymmetrizedKLKernel, self).__init__(distance_function=distance_function, **kwargs)

qpytorch/kernels/rbf_kernel_grad.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+
+import torch
+from linear_operator.operators import KroneckerProductLinearOperator
+
+from gpytorch.kernels.rbf_kernel import postprocess_rbf, RBFKernel
+
+
+class RBFKernelGrad(RBFKernel):
+    r"""
+    Computes a covariance matrix of the RBF kernel that models the covariance
+    between the values and partial derivatives for inputs :math:`\mathbf{x_1}`
+    and :math:`\mathbf{x_2}`.
+
+    See :class:`qpytorch.kernels.Kernel` for descriptions of the lengthscale options.
+
+    .. note::
+
+        This kernel does not have an `outputscale` parameter. To add a scaling parameter,
+        decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.
+
+    :param ard_num_dims: Set this if you want a separate lengthscale for each input
+        dimension. It should be `d` if x1 is a `n x d` matrix. (Default: `None`.)
+    :param batch_shape: Set this if you want a separate lengthscale for each batch of input
+        data. It should be :math:`B_1 \times \ldots \times B_k` if :math:`\mathbf x1` is
+        a :math:`B_1 \times \ldots \times B_k \times N \times D` tensor.
+    :param active_dims: Set this if you want to compute the covariance of only
+        a few input dimensions. The ints corresponds to the indices of the
+        dimensions. (Default: `None`.)
+    :param lengthscale_prior: Set this if you want to apply a prior to the
+        lengthscale parameter. (Default: `None`)
+    :param lengthscale_constraint: Set this if you want to apply a constraint
+        to the lengthscale parameter. (Default: `Positive`.)
+    :param eps: The minimum value that the lengthscale can take (prevents
+        divide by zero errors). (Default: `1e-6`.)
+
+    :ivar torch.Tensor lengthscale: The lengthscale parameter. Size/shape of parameter depends on the
+        ard_num_dims and batch_shape arguments.
+
+    Example:
+        >>> x = torch.randn(10, 5)
+        >>> # Non-batch: Simple option
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.RBFKernelGrad())
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (60 x 60), where 60 = n * (d + 1)
+        >>>
+        >>> batch_x = torch.randn(2, 10, 5)
+        >>> # Batch: Simple option
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.RBFKernelGrad())
+        >>> # Batch: different lengthscale for each batch
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.RBFKernelGrad(batch_shape=torch.Size([2])))
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (2 x 60 x 60)
+    """
+
+    def __init__(self, **kwargs):
+        super(RBFKernelGrad, self).__init__(**kwargs)
+        self._interleaved = kwargs.pop('interleaved', True)
+
+    def forward(self, x1, x2, diag=False, **params):
+        batch_shape = x1.shape[:-2]
+        n_batch_dims = len(batch_shape)
+        n1, d = x1.shape[-2:]
+        n2 = x2.shape[-2]
+
+        if not diag:
+            K = torch.zeros(*batch_shape, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)
+
+            # Scale the inputs by the lengthscale (for stability)
+            x1_ = x1.div(self.lengthscale)
+            x2_ = x2.div(self.lengthscale)
+
+            # Form all possible rank-1 products for the gradient and Hessian blocks
+            outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
+            outer = outer / self.lengthscale.unsqueeze(-2)
+            outer = torch.transpose(outer, -1, -2).contiguous()
+
+            # 1) Kernel block
+            diff = self.covar_dist(x1_, x2_, square_dist=True, **params)
+            K_11 = postprocess_rbf(diff)
+            K[..., :n1, :n2] = K_11
+
+            # 2) First gradient block
+            outer1 = outer.view(*batch_shape, n1, n2 * d)
+            K[..., :n1, n2:] = outer1 * K_11.repeat([*([1] * (n_batch_dims + 1)), d])
+
+            # 3) Second gradient block
+            outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
+            outer2 = outer2.transpose(-1, -2)
+            K[..., n1:, :n2] = -outer2 * K_11.repeat([*([1] * n_batch_dims), d, 1])
+
+            # 4) Hessian block
+            outer3 = outer1.repeat([*([1] * n_batch_dims), d, 1]) * outer2.repeat([*([1] * (n_batch_dims + 1)), d])
+            kp = KroneckerProductLinearOperator(
+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            )
+            chain_rule = kp.to_dense() - outer3
+            K[..., n1:, n2:] = chain_rule * K_11.repeat([*([1] * n_batch_dims), d, d])
+
+            # Symmetrize for stability
+            if n1 == n2 and torch.eq(x1, x2).all():
+                K = 0.5 * (K.transpose(-1, -2) + K)
+
+            # Apply a perfect shuffle permutation to match the MutiTask ordering
+            if self._interleaved:
+                pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().reshape((n1 * (d + 1)))
+                pi2 = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().reshape((n2 * (d + 1)))
+                K = K[..., pi1, :][..., :, pi2]
+
+            return K
+
+        else:
+            if not (n1 == n2 and torch.eq(x1, x2).all()):
+                raise RuntimeError("diag=True only works when x1 == x2")
+
+            kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
+            grad_diag = torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow(2)
+            grad_diag = grad_diag.transpose(-1, -2).reshape(*batch_shape, n2 * d)
+            k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
+            if self._interleaved:
+                pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().reshape((n2 * (d + 1)))
+                k_diag = k_diag[..., pi]
+            return k_diag
+
+    def num_outputs_per_input(self, x1, x2):
+        return x1.size(-1) + 1

qpytorch/kernels/rbf_kernel_gradgrad.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+
+import torch
+from linear_operator.operators import KroneckerProductLinearOperator
+
+from gpytorch.kernels.rbf_kernel import postprocess_rbf, RBFKernel
+
+
+class RBFKernelGradGrad(RBFKernel):
+    r"""
+    Computes a covariance matrix of the RBF kernel that models the covariance
+    between the values and first and second (non-mixed) partial derivatives for inputs :math:`\mathbf{x_1}`
+    and :math:`\mathbf{x_2}`.
+
+    See :class:`qpytorch.kernels.Kernel` for descriptions of the lengthscale options.
+
+    .. note::
+
+        This kernel does not have an `outputscale` parameter. To add a scaling parameter,
+        decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.
+
+    :param ard_num_dims: Set this if you want a separate lengthscale for each input
+        dimension. It should be `d` if x1 is a `n x d` matrix. (Default: `None`.)
+    :param batch_shape: Set this if you want a separate lengthscale for each batch of input
+        data. It should be :math:`B_1 \times \ldots \times B_k` if :math:`\mathbf x1` is
+        a :math:`B_1 \times \ldots \times B_k \times N \times D` tensor.
+    :param active_dims: Set this if you want to compute the covariance of only
+        a few input dimensions. The ints corresponds to the indices of the
+        dimensions. (Default: `None`.)
+    :param lengthscale_prior: Set this if you want to apply a prior to the
+        lengthscale parameter. (Default: `None`)
+    :param lengthscale_constraint: Set this if you want to apply a constraint
+        to the lengthscale parameter. (Default: `Positive`.)
+    :param eps: The minimum value that the lengthscale can take (prevents
+        divide by zero errors). (Default: `1e-6`.)
+
+    :ivar torch.Tensor lengthscale: The lengthscale parameter. Size/shape of parameter depends on the
+        ard_num_dims and batch_shape arguments.
+
+    Example:
+        >>> x = torch.randn(10, 5)
+        >>> # Non-batch: Simple option
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.RBFKernelGradGrad())
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (110 x 110), where 110 = n * (2*d + 1)
+        >>>
+        >>> batch_x = torch.randn(2, 10, 5)
+        >>> # Batch: Simple option
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.RBFKernelGradGrad())
+        >>> # Batch: different lengthscale for each batch
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.RBFKernelGradGrad(batch_shape=torch.Size([2])))
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (2 x 110 x 110)
+    """
+
+    def __init__(self, **kwargs):
+        super(RBFKernelGradGrad, self).__init__(**kwargs)
+        self._interleaved = kwargs.pop('interleaved', True)
+
+    def forward(self, x1, x2, diag=False, **params):
+        batch_shape = x1.shape[:-2]
+        n_batch_dims = len(batch_shape)
+        n1, d = x1.shape[-2:]
+        n2 = x2.shape[-2]
+
+        if not diag:
+            K = torch.zeros(*batch_shape, n1 * (2 * d + 1), n2 * (2 * d + 1), device=x1.device, dtype=x1.dtype)
+
+            # Scale the inputs by the lengthscale (for stability)
+            x1_ = x1.div(self.lengthscale)
+            x2_ = x2.div(self.lengthscale)
+
+            # Form all possible rank-1 products for the gradient and Hessian blocks
+            outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
+            outer = outer / self.lengthscale.unsqueeze(-2)
+            outer = torch.transpose(outer, -1, -2).contiguous()
+
+            # 1) Kernel block
+            diff = self.covar_dist(x1_, x2_, square_dist=True, **params)
+            K_11 = postprocess_rbf(diff)
+            K[..., :n1, :n2] = K_11
+
+            # 2) First gradient block
+            outer1 = outer.view(*batch_shape, n1, n2 * d)
+            K[..., :n1, n2 : (n2 * (d + 1))] = outer1 * K_11.repeat([*([1] * (n_batch_dims + 1)), d])
+
+            # 3) Second gradient block
+            outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
+            outer2 = outer2.transpose(-1, -2)
+            K[..., n1 : (n1 * (d + 1)), :n2] = -outer2 * K_11.repeat([*([1] * n_batch_dims), d, 1])
+
+            # 4) Hessian block
+            outer3 = outer1.repeat([*([1] * n_batch_dims), d, 1]) * outer2.repeat([*([1] * (n_batch_dims + 1)), d])
+            kp = KroneckerProductLinearOperator(
+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            )
+            chain_rule = kp.to_dense() - outer3
+            K[..., n1 : (n1 * (d + 1)), n2 : (n2 * (d + 1))] = chain_rule * K_11.repeat([*([1] * n_batch_dims), d, d])
+
+            # 5) 1-3 block
+            douter1dx2 = KroneckerProductLinearOperator(
+                torch.ones(1, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            ).to_dense()
+
+            K_13 = (-douter1dx2 + outer1 * outer1) * K_11.repeat(
+                [*([1] * (n_batch_dims + 1)), d]
+            )  # verified for n1=n2=1 case
+            K[..., :n1, (n2 * (d + 1)) :] = K_13
+
+            if d>1:
+                douter1dx2 = KroneckerProductLinearOperator(
+                    (torch.ones(1, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2)).transpose(-1, -2),
+                    torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+                ).to_dense()
+            K_31 = (-douter1dx2 + outer2 * outer2) * K_11.repeat(
+                [*([1] * n_batch_dims), d, 1]
+            )  # verified for n1=n2=1 case
+            K[..., (n1 * (d + 1)) :, :n2] = K_31
+
+            # rest of the blocks are all of size (n1*d,n2*d)
+            outer1 = outer1.repeat([*([1] * n_batch_dims), d, 1])
+            outer2 = outer2.repeat([*([1] * (n_batch_dims + 1)), d])
+            # II = (torch.eye(d,d,device=x1.device,dtype=x1.dtype)/lengthscale.pow(2)).repeat(*batch_shape,n1,n2)
+            kp2 = KroneckerProductLinearOperator(
+                torch.ones(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            ).to_dense()
+
+            # II may not be the correct thing to use. It might be more appropriate to use kp instead??
+            II = kp.to_dense()
+            K_11dd = K_11.repeat([*([1] * (n_batch_dims)), d, d])
+
+            K_23 = ((-kp2 + outer1 * outer1) * (-outer2) + 2.0 * II * outer1) * K_11dd  # verified for n1=n2=1 case
+
+            K[..., n1 : (n1 * (d + 1)), (n2 * (d + 1)) :] = K_23
+
+            if d>1:
+                kp2t = KroneckerProductLinearOperator(
+                    (torch.ones(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2)).transpose(-1, -2),
+                    torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+                ).to_dense()
+            K_32 = (
+                (-(kp2t if d>1 else kp2) + outer2 * outer2) * outer1 - 2.0 * II * outer2
+            ) * K_11dd  # verified for n1=n2=1 case
+
+            K[..., (n1 * (d + 1)) :, n2 : (n2 * (d + 1))] = K_32
+
+            K_33 = (
+                (-(kp2t if d>1 else kp2) + outer2 * outer2) * (-kp2) - 2.0 * II * outer2 * outer1 + 2.0 * (II) ** 2
+            ) * K_11dd + (
+                (-(kp2t if d>1 else kp2) + outer2 * outer2) * outer1 - 2.0 * II * outer2
+            ) * outer1 * K_11dd  # verified for n1=n2=1 case
+
+            K[..., (n1 * (d + 1)) :, (n2 * (d + 1)) :] = K_33
+
+            # Symmetrize for stability
+            if n1 == n2 and torch.eq(x1, x2).all():
+                K = 0.5 * (K.transpose(-1, -2) + K)
+
+            # Apply a perfect shuffle permutation to match the MutiTask ordering
+            if self._interleaved:
+                pi1 = torch.arange(n1 * (2 * d + 1)).view(2 * d + 1, n1).t().reshape((n1 * (2 * d + 1)))
+                pi2 = torch.arange(n2 * (2 * d + 1)).view(2 * d + 1, n2).t().reshape((n2 * (2 * d + 1)))
+                K = K[..., pi1, :][..., :, pi2]
+
+            return K
+
+        else:
+            if not (n1 == n2 and torch.eq(x1, x2).all()):
+                raise RuntimeError("diag=True only works when x1 == x2")
+
+            kernel_diag = super(RBFKernelGradGrad, self).forward(x1, x2, diag=True)
+            grad_diag = torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow(2)
+            grad_diag = grad_diag.transpose(-1, -2).reshape(*batch_shape, n2 * d)
+            gradgrad_diag = (
+                3 * torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow(4)
+            )
+            gradgrad_diag = gradgrad_diag.transpose(-1, -2).reshape(*batch_shape, n2 * d)
+            k_diag = torch.cat((kernel_diag, grad_diag, gradgrad_diag), dim=-1)
+            if self._interleaved:
+                pi = torch.arange(n2 * (2 * d + 1)).view(2 * d + 1, n2).t().reshape((n2 * (2 * d + 1)))
+                k_diag = k_diag[..., pi]
+            return k_diag
+
+    def num_outputs_per_input(self, x1, x2):
+        return x1.size(-1) * 2 + 1

qpytorch/kernels/rff_kernel.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+
+import math
+from typing import Optional
+
+import torch
+from linear_operator.operators import LowRankRootLinearOperator, MatmulLinearOperator, RootLinearOperator
+from torch import Tensor
+
+from ..models import exact_prediction_strategies
+from .kernel import Kernel
+
+
+class RFFKernel(Kernel):
+    r"""
+    Computes a covariance matrix based on Random Fourier Features with the RBFKernel.
+
+    Random Fourier features was originally proposed in
+    'Random Features for Large-Scale Kernel Machines' by Rahimi and Recht (2008).
+    Instead of the shifted cosine features from Rahimi and Recht (2008), we use
+    the sine and cosine features which is a lower-variance estimator --- see
+    'On the Error of Random Fourier Features' by Sutherland and Schneider (2015).
+
+    By Bochner's theorem, any continuous kernel :math:`k` is positive definite
+    if and only if it is the Fourier transform of a non-negative measure :math:`p(\omega)`, i.e.
+
+    .. math::
+        \begin{equation}
+            k(x, x') = k(x - x') = \int p(\omega) e^{i(\omega^\top (x - x'))} d\omega.
+        \end{equation}
+
+    where :math:`p(\omega)` is a normalized probability measure if :math:`k(0)=1`.
+
+    For the RBF kernel,
+
+    .. math::
+        \begin{equation}
+        k(\Delta) = \exp{(-\frac{\Delta^2}{2\sigma^2})} \text{ and } p(\omega) = \exp{(-\frac{\sigma^2\omega^2}{2})}
+        \end{equation}
+
+    where :math:`\Delta = x - x'`.
+
+    Given datapoint :math:`x\in \mathbb{R}^d`, we can construct its random Fourier features
+    :math:`z(x) \in \mathbb{R}^{2D}` by
+
+    .. math::
+        \begin{equation}
+        z(x) = \sqrt{\frac{1}{D}}
+        \begin{bmatrix}
+            \cos(\omega_1^\top x)\\
+            \sin(\omega_1^\top x)\\
+            \cdots \\
+            \cos(\omega_D^\top x)\\
+            \sin(\omega_D^\top x)
+        \end{bmatrix}, \omega_1, \ldots, \omega_D \sim p(\omega)
+        \end{equation}
+
+    such that we have an unbiased Monte Carlo estimator
+
+    .. math::
+        \begin{equation}
+            k(x, x') = k(x - x') \approx z(x)^\top z(x') = \frac{1}{D}\sum_{i=1}^D \cos(\omega_i^\top (x - x')).
+        \end{equation}
+
+    .. note::
+        When this kernel is used in batch mode, the random frequencies are drawn
+        independently across the batch dimension as well by default.
+
+    :param num_samples: Number of random frequencies to draw. This is :math:`D` in the above
+        papers. This will produce :math:`D` sine features and :math:`D` cosine
+        features for a total of :math:`2D` random Fourier features.
+    :type num_samples: int
+    :param num_dims: (Default `None`.) Dimensionality of the data space.
+        This is :math:`d` in the above papers. Note that if you want an
+        independent lengthscale for each dimension, set `ard_num_dims` equal to
+        `num_dims`. If unspecified, it will be inferred the first time `forward`
+        is called.
+    :type num_dims: int, optional
+
+    :var torch.Tensor randn_weights: The random frequencies that are drawn once and then fixed.
+
+    Example:
+
+        >>> # This will infer `num_dims` automatically
+        >>> kernel= qpytorch.kernels.RFFKernel(num_samples=5)
+        >>> x = torch.randn(10, 3)
+        >>> kxx = kernel(x, x).to_dense()
+        >>> print(kxx.randn_weights.size())
+        torch.Size([3, 5])
+
+    """
+
+    has_lengthscale = True
+
+    def __init__(self, num_samples: int, num_dims: Optional[int] = None, **kwargs):
+        super().__init__(**kwargs)
+        self.num_samples = num_samples
+        if num_dims is not None:
+            self._init_weights(num_dims, num_samples)
+
+    def _init_weights(
+        self, num_dims: Optional[int] = None, num_samples: Optional[int] = None, randn_weights: Optional[Tensor] = None
+    ):
+        if num_dims is not None and num_samples is not None:
+            d = num_dims
+            D = num_samples
+        if randn_weights is None:
+            randn_shape = torch.Size([*self._batch_shape, d, D])
+            randn_weights = torch.randn(
+                randn_shape, dtype=self.raw_lengthscale.dtype, device=self.raw_lengthscale.device
+            )
+        self.register_buffer("randn_weights", randn_weights)
+
+    def forward(self, x1: Tensor, x2: Tensor, diag: bool = False, last_dim_is_batch: bool = False, **kwargs) -> Tensor:
+        if last_dim_is_batch:
+            x1 = x1.transpose(-1, -2).unsqueeze(-1)
+            x2 = x2.transpose(-1, -2).unsqueeze(-1)
+        num_dims = x1.size(-1)
+        if not hasattr(self, "randn_weights"):
+            self._init_weights(num_dims, self.num_samples)
+        x1_eq_x2 = torch.equal(x1, x2)
+        z1 = self._featurize(x1, normalize=False)
+        if not x1_eq_x2:
+            z2 = self._featurize(x2, normalize=False)
+        else:
+            z2 = z1
+        D = float(self.num_samples)
+        if diag:
+            return (z1 * z2).sum(-1) / D
+        if x1_eq_x2:
+            # Exploit low rank structure, if there are fewer features than data points
+            if z1.size(-1) < z2.size(-2):
+                return LowRankRootLinearOperator(z1 / math.sqrt(D))
+            else:
+                return RootLinearOperator(z1 / math.sqrt(D))
+        else:
+            return MatmulLinearOperator(z1 / D, z2.transpose(-1, -2))
+
+    def _featurize(self, x: Tensor, normalize: bool = False) -> Tensor:
+        # Recompute division each time to allow backprop through lengthscale
+        # Transpose lengthscale to allow for ARD
+        x = x.matmul(self.randn_weights / self.lengthscale.transpose(-1, -2))
+        z = torch.cat([torch.cos(x), torch.sin(x)], dim=-1)
+        if normalize:
+            D = self.num_samples
+            z = z / math.sqrt(D)
+        return z
+
+    def prediction_strategy(self, train_inputs, train_prior_dist, train_labels, likelihood):
+        # Allow for fast sampling
+        return exact_prediction_strategies.RFFPredictionStrategy(
+            train_inputs, train_prior_dist, train_labels, likelihood
+        )

qpytorch/lazy/__init__.py

@@ -0,0 +1,9 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+from gpytorch.lazy.lazy_evaluated_kernel_tensor import LazyEvaluatedKernelTensor
+
+__all__ = [
+    "LazyEvaluatedKernelTensor",
+]

qpytorch/likelihoods/__init__.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+from .bernoulli_likelihood import BernoulliLikelihood
+from .beta_likelihood import BetaLikelihood
+from .gaussian_likelihood import (
+    _GaussianLikelihoodBase,
+    GaussianLikelihood,
+    GaussianLikelihoodWithMissingObs,
+    FixedNoiseGaussianLikelihood,
+    DirichletClassificationLikelihood,
+)
+from .qexponential_likelihood import (
+    _QExponentialLikelihoodBase,
+    QExponentialLikelihood,
+    QExponentialLikelihoodWithMissingObs,
+    FixedNoiseQExponentialLikelihood,
+    QExponentialDirichletClassificationLikelihood,
+)
+from .laplace_likelihood import LaplaceLikelihood
+from .likelihood import _OneDimensionalLikelihood, Likelihood
+from .likelihood_list import LikelihoodList
+from .multitask_gaussian_likelihood import (
+    _MultitaskGaussianLikelihoodBase, 
+    MultitaskGaussianLikelihood,
+    MultitaskFixedNoiseGaussianLikelihood,
+    MultitaskDirichletClassificationLikelihood,
+)
+from .multitask_qexponential_likelihood import (
+    _MultitaskQExponentialLikelihoodBase, 
+    MultitaskQExponentialLikelihood,
+    MultitaskFixedNoiseQExponentialLikelihood,
+    MultitaskQExponentialDirichletClassificationLikelihood,
+)
+from .noise_models import HeteroskedasticNoise
+from .softmax_likelihood import SoftmaxLikelihood
+from .student_t_likelihood import StudentTLikelihood
+
+__all__ = [
+    "_GaussianLikelihoodBase",
+    "_QExponentialLikelihoodBase",
+    "_OneDimensionalLikelihood",
+    "_MultitaskGaussianLikelihoodBase",
+    "_MultitaskQExponentialLikelihoodBase",
+    "BernoulliLikelihood",
+    "BetaLikelihood",
+    "DirichletClassificationLikelihood",
+    "QExponentialDirichletClassificationLikelihood",
+    "FixedNoiseGaussianLikelihood",
+    "FixedNoiseQExponentialLikelihood",
+    "GaussianLikelihood",
+    "QExponentialLikelihood",
+    "GaussianLikelihoodWithMissingObs",
+    "QExponentialLikelihoodWithMissingObs",
+    "HeteroskedasticNoise",
+    "LaplaceLikelihood",
+    "Likelihood",
+    "LikelihoodList",
+    "MultitaskGaussianLikelihood",
+    "MultitaskFixedNoiseGaussianLikelihood",
+    "MultitaskDirichletClassificationLikelihood",
+    "MultitaskQExponentialLikelihood",
+    "MultitaskFixedNoiseQExponentialLikelihood",
+    "MultitaskQExponentialDirichletClassificationLikelihood",
+    "SoftmaxLikelihood",
+    "StudentTLikelihood",
+]

qpytorch/likelihoods/bernoulli_likelihood.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+
+import warnings
+from typing import Any, Union
+
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli
+
+from ..distributions import base_distributions, MultivariateNormal, MultivariateQExponential
+from ..functions import log_normal_cdf
+from .likelihood import _OneDimensionalLikelihood
+
+
+class BernoulliLikelihood(_OneDimensionalLikelihood):
+    r"""
+    Implements the Bernoulli likelihood used for GP/QEP classification, using
+    Probit regression (i.e., the latent function is warped to be in [0,1]
+    using the standard Normal CDF :math:`\Phi(x)`). Given the identity
+    :math:`\Phi(-x) = 1-\Phi(x)`, we can write the likelihood compactly as:
+
+    .. math::
+        \begin{equation*}
+            p(Y=y|f)=\Phi((2y - 1)f)
+        \end{equation*}
+
+    .. note::
+        BernoulliLikelihood has an analytic marginal distribution.
+
+    .. note::
+        The labels should take values in {0, 1}.
+    """
+
+    has_analytic_marginal: bool = True
+
+    def __init__(self) -> None:
+        return super().__init__()
+
+    def forward(self, function_samples: Tensor, *args: Any, **kwargs: Any) -> Bernoulli:
+        output_probs = base_distributions.Normal(0, 1).cdf(function_samples)
+        return base_distributions.Bernoulli(probs=output_probs)
+
+    def log_marginal(
+        self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any
+    ) -> Tensor:
+        marginal = self.marginal(function_dist, *args, **kwargs)
+        return marginal.log_prob(observations)
+
+    def marginal(self, function_dist: Union[MultivariateNormal, MultivariateQExponential], *args: Any, **kwargs: Any) -> Bernoulli:
+        r"""
+        :return: Analytic marginal :math:`p(\mathbf y)`.
+        """
+        mean = function_dist.mean
+        var = function_dist.variance
+        link = mean.div(torch.sqrt(1 + var))
+        output_probs = base_distributions.Normal(0, 1).cdf(link)
+        return base_distributions.Bernoulli(probs=output_probs)
+
+    def expected_log_prob(
+        self, observations: Tensor, function_dist: Union[MultivariateNormal, MultivariateQExponential], *params: Any, **kwargs: Any
+    ) -> Tensor:
+        if torch.any(observations.eq(-1)):
+            # Remove after 1.0
+            warnings.warn(
+                "BernoulliLikelihood.expected_log_prob expects observations with labels in {0, 1}. "
+                "Observations with labels in {-1, 1} are deprecated.",
+                DeprecationWarning,
+            )
+        else:
+            observations = observations.mul(2).sub(1)
+        # Custom function here so we can use log_normal_cdf rather than Normal.cdf
+        # This is going to be less prone to overflow errors
+        log_prob_lambda = lambda function_samples: log_normal_cdf(function_samples.mul(observations))
+        log_prob = self.quadrature(log_prob_lambda, function_dist)
+        return log_prob