qpytorch 0.1__py3-none-any.whl

qpytorch/kernels/matern32_kernel_grad.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+
+import math
+
+import torch
+from linear_operator.operators import KroneckerProductLinearOperator
+
+from gpytorch.kernels.matern_kernel import MaternKernel
+
+sqrt3 = math.sqrt(3)
+
+
+class Matern32KernelGrad(MaternKernel):
+    r"""
+    Computes a covariance matrix of the Matern32 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.Matern32KernelGrad())
+        >>> 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.Matern32KernelGrad())
+        >>> # Batch: different lengthscale for each batch
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.Matern32KernelGrad(batch_shape=torch.Size([2]))) # noqa: E501
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (2 x 60 x 60)
+    """
+
+    def __init__(self, **kwargs):
+
+        # remove nu in case it was set
+        kwargs.pop("nu", None)
+        super(Matern32KernelGrad, self).__init__(nu=1.5, **kwargs)
+        self._interleaved = kwargs.pop('interleaved', True)
+
+    def forward(self, x1, x2, diag=False, **params):
+
+        lengthscale = self.lengthscale
+
+        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)
+
+            distance_matrix = self.covar_dist(x1.div(lengthscale), x2.div(lengthscale), diag=diag, **params)
+            exp_neg_sqrt3r = torch.exp(-sqrt3 * distance_matrix)
+
+            # differences matrix in each dimension to be used for derivatives
+            # shape of n1 x n2 x d
+            outer = x1.view(*batch_shape, n1, 1, d) - x2.view(*batch_shape, 1, n2, d)
+            outer = outer / lengthscale.unsqueeze(-2) ** 2
+            # shape of n1 x d x n2
+            outer = torch.transpose(outer, -1, -2).contiguous()
+
+            # 1) Kernel block, cov(f^m, f^n)
+            # shape is n1 x n2
+            # exp_component = torch.exp(-sqrt3 * distance_matrix)
+            constant_component = (sqrt3 * distance_matrix).add(1)
+
+            K[..., :n1, :n2] = constant_component * exp_neg_sqrt3r #exp_component
+
+            # 2) First gradient block, cov(f^m, omega^n_i)
+            outer1 = outer.view(*batch_shape, n1, n2 * d)
+            # the - signs on -outer1 and -five_thirds cancel out
+            K[..., :n1, n2:] = 3 * outer1 * exp_neg_sqrt3r.repeat(
+                [*([1] * (n_batch_dims + 1)), d]
+            )
+
+            # 3) Second gradient block, cov(omega^m_j, f^n)
+            outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
+            outer2 = outer2.transpose(-1, -2)
+            K[..., n1:, :n2] = -3 * outer2 * exp_neg_sqrt3r.repeat(
+                [*([1] * n_batch_dims), d, 1]
+            )
+
+            # 4) Hessian block, cov(omega^m_j, omega^n_i)
+            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) / lengthscale**2,
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            )
+
+            # part1 = -3 * exp_neg_sqrt3r
+            # part2 = sqrt3 * invrdd * outer3
+            invrdd = (distance_matrix+self.eps).pow(-1)
+            # invrdd[torch.arange(min(n1,n2)),torch.arange(min(n1,n2))] = distance_matrix.diagonal()
+            invrdd = invrdd.repeat([*([1] * (n_batch_dims)), d, d])
+            # invrdd = distance_matrix.pow(-1).fill_diagonal_(0).repeat([*([1] * (n_batch_dims)), d, d]).fill_diagonal_(1)
+
+            K[..., n1:, n2:] = -3 * exp_neg_sqrt3r.repeat([*([1] * n_batch_dims), d, d]).mul_(
+                (sqrt3*invrdd * outer3).sub_(kp.to_dense())
+            )
+
+            # 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")
+
+            # nu is set to 2.5
+            kernel_diag = super(Matern32KernelGrad, self).forward(x1, x2, diag=True)
+            grad_diag = (
+                3 * torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype)
+            ) / lengthscale**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/matern52_kernel_grad.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+
+import math
+
+import torch
+from linear_operator.operators import KroneckerProductLinearOperator
+
+from gpytorch.kernels.matern_kernel import MaternKernel
+
+sqrt5 = math.sqrt(5)
+five_thirds = 5.0 / 3.0
+
+
+class Matern52KernelGrad(MaternKernel):
+    r"""
+    Computes a covariance matrix of the Matern52 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`.
+
+    .. note::
+
+        A perfect shuffle permutation is applied after the calculation of the matrix blocks
+        in order to match the MutiTask ordering.
+
+    The Matern52 kernel is defined as
+
+    .. math::
+
+        k(r) = (1 + \sqrt{5}r + \frac{5}{3} r^2) \exp(- \sqrt{5}r)
+
+    where :math:`r` is defined as
+
+    .. math::
+
+        r(\mathbf{x}^m , \mathbf{x}^n) = \sqrt{\sum_d{\frac{(x^m_d - x^n_d)^2}{l^2_d}}}
+
+    The first gradient block containing :math:`\frac{\partial k}{\partial x^n_i}` is defined as
+
+    .. math::
+
+        \frac{\partial k}{\partial x^n_i} = \frac{5}{3} \left( 1 + \sqrt{5}r \right) \exp(- \sqrt{5}r) \left(\frac{x^m_i - x^n_i}{l^2_i} \right)
+
+    The second gradient block containing :math:`\frac{\partial k}{\partial x^m_j}` is defined as
+
+    .. math::
+
+        \frac{\partial k}{\partial x^m_j} = - \frac{5}{3} \left( 1 + \sqrt{5}r \right) \exp(- \sqrt{5}r) \left(\frac{x^m_j - x^n_j}{l^2_j} \right)
+
+    The Hessian block containing :math:`\frac{\partial^2 k}{\partial x^m_j \partial x^n_i}` is defined as
+
+    .. math::
+
+        \frac{\partial^2 k}{\partial x^m_j \partial x^n_i} = - \frac{5}{3} \exp(- \sqrt{5}r) \left[ 5\left(\frac{x^m_i - x^n_i}{l^2_i} \right) \left( \frac{x^m_j - x^n_j}{l^2_j} \right) - \frac{\delta_{ij}}{l^2_i} \left( 1 + \sqrt{5}r \right) \right]
+
+    The derivations can be found `here <https://github.com/cornellius-gp/gpytorch/pull/2512>`__.
+
+    :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.Matern52KernelGrad())
+        >>> 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.Matern52KernelGrad())
+        >>> # Batch: different lengthscale for each batch
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.Matern52KernelGrad(batch_shape=torch.Size([2]))) # noqa: E501
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (2 x 60 x 60)
+    """
+
+    def __init__(self, **kwargs):
+
+        # remove nu in case it was set
+        kwargs.pop("nu", None)
+        super(Matern52KernelGrad, self).__init__(nu=2.5, **kwargs)
+        self._interleaved = kwargs.pop('interleaved', True)
+
+    def forward(self, x1, x2, diag=False, **params):
+
+        lengthscale = self.lengthscale
+
+        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)
+
+            distance_matrix = self.covar_dist(x1.div(lengthscale), x2.div(lengthscale), diag=diag, **params)
+            exp_neg_sqrt5r = torch.exp(-sqrt5 * distance_matrix)
+            one_plus_sqrt5r = 1 + sqrt5 * distance_matrix
+
+            # differences matrix in each dimension to be used for derivatives
+            # shape of n1 x n2 x d
+            outer = x1.view(*batch_shape, n1, 1, d) - x2.view(*batch_shape, 1, n2, d)
+            outer = outer / lengthscale.unsqueeze(-2) ** 2
+            # shape of n1 x d x n2
+            outer = torch.transpose(outer, -1, -2).contiguous()
+
+            # 1) Kernel block, cov(f^m, f^n)
+            # shape is n1 x n2
+            # exp_component = torch.exp(-sqrt5 * distance_matrix)
+            constant_component = one_plus_sqrt5r.add(five_thirds * distance_matrix**2)
+
+            K[..., :n1, :n2] = constant_component * exp_neg_sqrt5r #exp_component
+
+            # 2) First gradient block, cov(f^m, omega^n_i)
+            outer1 = outer.view(*batch_shape, n1, n2 * d)
+            # the - signs on -outer1 and -five_thirds cancel out
+            K[..., :n1, n2:] = five_thirds * outer1 * (one_plus_sqrt5r * exp_neg_sqrt5r).repeat(
+                [*([1] * (n_batch_dims + 1)), d]
+            )
+
+            # 3) Second gradient block, cov(omega^m_j, f^n)
+            outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
+            outer2 = outer2.transpose(-1, -2)
+            K[..., n1:, :n2] = -five_thirds * outer2 * (one_plus_sqrt5r * exp_neg_sqrt5r).repeat(
+                [*([1] * n_batch_dims), d, 1]
+            )
+
+            # 4) Hessian block, cov(omega^m_j, omega^n_i)
+            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) / lengthscale**2,
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            )
+
+            # part1 = -five_thirds * exp_neg_sqrt5r
+            # part2 = 5 * outer3
+            # part3 = 1 + sqrt5 * distance_matrix
+
+            K[..., n1:, n2:] = -five_thirds * exp_neg_sqrt5r.repeat([*([1] * n_batch_dims), d, d]).mul_(
+                # need to use kp.to_dense().mul instead of kp.to_dense().mul_
+                # because otherwise a RuntimeError is raised due to how autograd works with
+                # view + inplace operations in the case of 1-dimensional input
+                (5 * outer3).sub_(kp.to_dense().mul(one_plus_sqrt5r.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")
+
+            # nu is set to 2.5
+            kernel_diag = super(Matern52KernelGrad, self).forward(x1, x2, diag=True)
+            grad_diag = (
+                five_thirds * torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype)
+            ) / lengthscale**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/matern52_kernel_gradgrad.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+
+import math
+
+import torch
+from linear_operator.operators import KroneckerProductLinearOperator
+
+from gpytorch.kernels.matern_kernel import MaternKernel
+
+sqrt5 = math.sqrt(5)
+five_thirds = 5.0 / 3.0
+
+
+class Matern52KernelGradGrad(MaternKernel):
+    r"""
+    Computes a covariance matrix of the Matern52 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.Matern52KernelGradGrad())
+        >>> 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.Matern52KernelGradGrad())
+        >>> # Batch: different lengthscale for each batch
+        >>> covar_module = qpytorch.kernels.ScaleKernel(qpytorch.kernels.Matern52KernelGradGrad(batch_shape=torch.Size([2]))) # noqa: E501
+        >>> covar = covar_module(x)  # Output: LinearOperator of size (2 x 110 x 110)
+    """
+
+    def __init__(self, **kwargs):
+
+        # remove nu in case it was set
+        kwargs.pop("nu", None)
+        super(Matern52KernelGradGrad, self).__init__(nu=2.5, **kwargs)
+        self._interleaved = kwargs.pop('interleaved', True)
+
+    def forward(self, x1, x2, diag=False, **params):
+
+        lengthscale = self.lengthscale
+
+        batch_shape = x1.shape[:-2]
+        n_batch_dims = len(batch_shape)
+        n1, d = x1.shape[-2:]
+        n2 = x2.shape[-2]
+        
+        mask_idx1 = params.pop('mask_idx1', None) # mask off-diagonal covariance
+        if mask_idx1 is not None:
+            mask_idx2 = params.pop('mask_idx2', None)
+            if mask_idx2 is None:
+                mask_idx2 = mask_idx1
+            else:
+                assert mask_idx1.shape[:-1] == mask_idx2.shape[:-1], 'Batch shapes of mask indices do not match!'
+
+        if not diag:
+
+            K = torch.zeros(*batch_shape, n1 * (2 * d + 1), n2 * (2 * d + 1), device=x1.device, dtype=x1.dtype)
+
+            distance_matrix = self.covar_dist(x1.div(lengthscale), x2.div(lengthscale), diag=diag, **params)
+            exp_neg_sqrt5r = torch.exp(-sqrt5 * distance_matrix)
+            one_plus_sqrt5r = 1 + sqrt5 * distance_matrix
+
+            # differences matrix in each dimension to be used for derivatives
+            # shape of n1 x n2 x d
+            outer = x1.view(*batch_shape, n1, 1, d) - x2.view(*batch_shape, 1, n2, d)
+            outer = outer / lengthscale.unsqueeze(-2) ** 2
+            # shape of n1 x d x n2
+            outer = torch.transpose(outer, -1, -2).contiguous()
+
+            # 1) Kernel block, cov(f^m, f^n)
+            # shape is n1 x n2
+            # exp_component = torch.exp(-sqrt5 * distance_matrix)
+            constant_component = one_plus_sqrt5r.add(five_thirds * distance_matrix**2)
+
+            K[..., :n1, :n2] = constant_component * exp_neg_sqrt5r #exp_component
+
+            # 2) First gradient block, cov(f^m, omega^n_i)
+            outer1 = outer.view(*batch_shape, n1, n2 * d)
+            # the - signs on -outer1 and -five_thirds cancel out
+            K[..., :n1, n2: (n2 * (d + 1))] = five_thirds * outer1 * (one_plus_sqrt5r * exp_neg_sqrt5r).repeat(
+                [*([1] * (n_batch_dims + 1)), d]
+            )
+
+            # 3) Second gradient block, cov(omega^m_j, f^n)
+            outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
+            outer2 = outer2.transpose(-1, -2)
+            K[..., n1: (n1 * (d + 1)), :n2] = -five_thirds * outer2 * (one_plus_sqrt5r * exp_neg_sqrt5r).repeat(
+                [*([1] * n_batch_dims), d, 1]
+            )
+
+            # 4) Hessian block, cov(omega^m_j, omega^n_i)
+            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) / lengthscale**2,
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1),
+            )
+
+            # part1 = -five_thirds * exp_neg_sqrt5r
+            # part2 = 5 * outer3
+            # part3 = 1 + sqrt5 * distance_matrix
+            exp_neg_sqrt5rdd = exp_neg_sqrt5r.repeat([*([1] * (n_batch_dims)), d, d])
+
+            K[..., n1: (n1 * (d + 1)), n2: (n2 * (d + 1))] = -five_thirds * exp_neg_sqrt5rdd.mul(
+                # need to use kp.to_dense().mul instead of kp.to_dense().mul_
+                # because otherwise a RuntimeError is raised due to how autograd works with
+                # view + inplace operations in the case of 1-dimensional input
+                (5 * outer3).sub(kp.to_dense().mul(one_plus_sqrt5r.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 = five_thirds * (-douter1dx2 * one_plus_sqrt5r.repeat([*([1] * (n_batch_dims + 1)), d]) + 5* outer1 * outer1) * exp_neg_sqrt5r.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 = five_thirds * (-douter1dx2 * one_plus_sqrt5r.repeat([*([1] * n_batch_dims), d, 1]) + 5* outer2 * outer2) * exp_neg_sqrt5r.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()
+            # exp_neg_sqrt5rdd = exp_neg_sqrt5r.repeat([*([1] * (n_batch_dims)), d, d])
+            invrdd = (distance_matrix+self.eps).pow(-1)
+            # invrdd[torch.arange(min(n1,n2)),torch.arange(min(n1,n2))] = distance_matrix.diagonal()
+            invrdd = invrdd.repeat([*([1] * (n_batch_dims)), d, d])
+            # invrdd = distance_matrix.pow(-1).fill_diagonal_(0).repeat([*([1] * (n_batch_dims)), d, d]).fill_diagonal_(1)
+
+            K_23 = five_thirds * 5* ((kp2 - sqrt5*invrdd* outer1 * outer1) * outer2 + 2.0 * II * outer1) * exp_neg_sqrt5rdd  # 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 = five_thirds * 5* (
+                (-(kp2t if d>1 else kp2) + sqrt5*invrdd* outer2 * outer2) * outer1 - 2.0 * II * outer2
+            ) * exp_neg_sqrt5rdd  # verified for n1=n2=1 case
+
+            K[..., (n1 * (d + 1)) :, n2 : (n2 * (d + 1))] = K_32
+
+            # K_33 = five_thirds * 5*(
+            #     ((-(kp2t if d>1 else kp2) + sqrt5*invrdd*outer2 * outer2) * (-kp2) - 2.0 *sqrt5*invrdd * II * outer2 * outer1 + 2.0 * (II) ** 2
+            # )  + (
+            #     (-(kp2t if d>1 else kp2)*sqrt5*invrdd + (5+sqrt5*invrdd)*invrdd**2*outer2 * outer2) * outer1 - 2.0 *sqrt5*invrdd * II * outer2
+            # ) * outer1) * exp_neg_sqrt5rdd  # verified for n1=n2=1 case
+            K_33 = five_thirds * 5*(
+                (kp2 - sqrt5*invrdd*outer1 * outer1) * ((kp2t if d>1 else kp2)-sqrt5*invrdd* outer2 * outer2) 
+                + sqrt5*invrdd*(invrdd**2*outer3-4*(II))*outer3 + 2*(II)**2
+            ) * exp_neg_sqrt5rdd
+
+            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]
+
+            if mask_idx1 is not None:
+                if mask_idx1.ndim==1:
+                    diag2keep = K[...,mask_idx1,mask_idx2]
+                    K[...,mask_idx1,:] = 0; K[...,mask_idx2] = 0
+                    K[...,mask_idx1,mask_idx2] = diag2keep * self.eps
+                elif mask_idx1.ndim==2:
+                    for b in range(mask_idx1.shape[0]):
+                        diag2keep = K[b,...,mask_idx1[b],mask_idx2[b]]
+                        K[b,...,mask_idx1[b],:] = 0; K[b,...,mask_idx2[b]] = 0
+                        K[b,...,mask_idx1[b],mask_idx2[b]] = diag2keep * self.eps
+                else:
+                    raise NotImplementedError('Mask indices of batch dimension bigger than 1 not implemented!')
+
+            return K
+        else:
+            if not (n1 == n2 and torch.eq(x1, x2).all()):
+                raise RuntimeError("diag=True only works when x1 == x2")
+
+            # nu is set to 2.5
+            kernel_diag = super(Matern52KernelGradGrad, self).forward(x1, x2, diag=True)
+            grad_diag = (
+                five_thirds * torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype)
+            ) / lengthscale**2
+            grad_diag = grad_diag.transpose(-1, -2).reshape(*batch_shape, n2 * d)
+            gradgrad_diag = (
+                5**2 * torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype) / 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/polynomial_kernel_grad.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+
+from typing import Optional
+
+import torch
+
+from gpytorch.kernels.polynomial_kernel import PolynomialKernel
+
+
+class PolynomialKernelGrad(PolynomialKernel):
+
+    def __init__(self, **kwargs):
+        super(PolynomialKernelGrad, self).__init__(**kwargs)
+        self._interleaved = kwargs.pop('interleaved', True)
+
+    def forward(
+        self,
+        x1: torch.Tensor,
+        x2: torch.Tensor,
+        diag: Optional[bool] = False,
+        last_dim_is_batch: Optional[bool] = False,
+        **params,
+    ) -> torch.Tensor:
+        offset = self.offset.view(*self.batch_shape, 1, 1)
+
+        batch_shape = x1.shape[:-2]
+        n1, d = x1.shape[-2:]
+        n2 = x2.shape[-2]
+
+        if diag:
+            base_diag = (x1 * x2).sum(dim=-1) + self.offset
+            K11_diag = base_diag.pow(self.power)
+
+            all_outers_diag = (x1 * x2).transpose(-2, -1).reshape(*batch_shape, -1)
+            K22_base_diag = self.power * (self.power - 1) * base_diag.pow(self.power - 2)
+            K12_base_diag = self.power * base_diag.pow(self.power - 1)
+
+            K22_diag = torch.add(
+                all_outers_diag * K22_base_diag.repeat(*([1] * (K22_base_diag.dim() - 1)), d),
+                K12_base_diag.repeat(*([1] * (K12_base_diag.dim() - 1)), d),
+            )
+
+            K_diag = torch.cat([K11_diag, K22_diag], dim=-1)
+            # Apply perfect shuffle
+            if self._interleaved:
+                pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().reshape((n1 * (d + 1)))
+                K_diag = K_diag[..., pi1]
+            return K_diag
+
+        else:
+            base_inner_prod = torch.matmul(x1, x2.transpose(-2, -1)) + offset
+            K11 = base_inner_prod.pow(self.power)
+
+            K12_base = self.power * base_inner_prod.pow(self.power - 1)
+            K12 = torch.zeros(*batch_shape, n1, n2 * d, dtype=x1.dtype, device=x1.device)
+
+            ones_ = torch.ones(*batch_shape, d, 1, n2, dtype=x1.dtype, device=x1.device)
+            K12_outer_prods = torch.matmul(x1.transpose(-2, -1).unsqueeze(-1), ones_)
+            K12 = (K12_base.unsqueeze(-3) * K12_outer_prods).transpose(-3, -2).reshape(*batch_shape, n1, d * n2)
+
+            ones_ = torch.ones(*batch_shape, d, n1, 1, dtype=x1.dtype, device=x1.device)
+            K21_outer_prods = torch.matmul(ones_, x2.transpose(-2, -1).unsqueeze(-2))
+            K21 = (K12_base.unsqueeze(-3) * K21_outer_prods).view(*batch_shape, d * n1, n2)
+
+            K22_base = self.power * (self.power - 1) * base_inner_prod.pow(self.power - 2)
+            K22 = torch.zeros(*batch_shape, n1 * d, n2 * d, dtype=x1.dtype, device=x1.device)
+            all_outers = x1.unsqueeze(-2).unsqueeze(-2).transpose(-2, -1).matmul(x2.unsqueeze(-3).unsqueeze(-2))
+            all_outers = all_outers.transpose(-4, -2).transpose(-3, -1)
+            K22 = K22_base.unsqueeze(-3).unsqueeze(-3) * all_outers  # d x d x n1 x n2
+
+            # Can't avoid this for loop without unnecessary memory duplication, which is worse.
+            for i in range(d):
+                K22[..., i, i, :, :] = K22[..., i, i, :, :] + K12_base
+
+            K22 = K22.transpose(-4, -3).transpose(-3, -2).reshape(*batch_shape, n1 * d, n2 * d)
+
+            K = torch.cat([torch.cat([K11, K12], dim=-1), torch.cat([K21, K22], dim=-1)], dim=-2)
+
+            # Apply perfect shuffle
+            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
+
+    def num_outputs_per_input(self, x1, x2):
+        return x1.size(-1) + 1