import sys
_module = sys.modules[__name__]
del sys
conf = _module
densenet = _module
LBFGS = _module
gpytorch = _module
beta_features = _module
constraints = _module
constraints = _module
distributions = _module
delta = _module
distribution = _module
multitask_multivariate_normal = _module
multivariate_normal = _module
functions = _module
_log_normal_cdf = _module
matern_covariance = _module
rbf_covariance = _module
kernels = _module
additive_structure_kernel = _module
arc_kernel = _module
cosine_kernel = _module
cylindrical_kernel = _module
distributional_input_kernel = _module
gaussian_symmetrized_kl_kernel = _module
grid_interpolation_kernel = _module
grid_kernel = _module
index_kernel = _module
inducing_point_kernel = _module
keops = _module
keops_kernel = _module
matern_kernel = _module
rbf_kernel = _module
kernel = _module
lcm_kernel = _module
linear_kernel = _module
matern_kernel = _module
multi_device_kernel = _module
multitask_kernel = _module
newton_girard_additive_kernel = _module
periodic_kernel = _module
piecewise_polynomial_kernel = _module
polynomial_kernel = _module
polynomial_kernel_grad = _module
product_structure_kernel = _module
rbf_kernel = _module
rbf_kernel_grad = _module
rff_kernel = _module
rq_kernel = _module
scale_kernel = _module
spectral_delta_kernel = _module
spectral_mixture_kernel = _module
lazy = _module
lazy_evaluated_kernel_tensor = _module
lazy_tensor = _module
non_lazy_tensor = _module
likelihoods = _module
bernoulli_likelihood = _module
beta_likelihood = _module
gaussian_likelihood = _module
laplace_likelihood = _module
likelihood = _module
likelihood_list = _module
multitask_gaussian_likelihood = _module
noise_models = _module
softmax_likelihood = _module
student_t_likelihood = _module
means = _module
constant_mean = _module
constant_mean_grad = _module
linear_mean = _module
mean = _module
multitask_mean = _module
zero_mean = _module
metrics = _module
metrics = _module
mlls = _module
_approximate_mll = _module
added_loss_term = _module
deep_approximate_mll = _module
deep_predictive_log_likelihood = _module
exact_marginal_log_likelihood = _module
gamma_robust_variational_elbo = _module
inducing_point_kernel_added_loss_term = _module
kl_gaussian_added_loss_term = _module
leave_one_out_pseudo_likelihood = _module
marginal_log_likelihood = _module
noise_model_added_loss_term = _module
predictive_log_likelihood = _module
sum_marginal_log_likelihood = _module
variational_elbo = _module
models = _module
approximate_gp = _module
deep_gps = _module
deep_gp = _module
dspp = _module
exact_gp = _module
exact_prediction_strategies = _module
gp = _module
gplvm = _module
bayesian_gplvm = _module
latent_variable = _module
model_list = _module
pyro = _module
_pyro_mixin = _module
pyro_gp = _module
module = _module
optim = _module
ngd = _module
priors = _module
horseshoe_prior = _module
lkj_prior = _module
prior = _module
smoothed_box_prior = _module
torch_priors = _module
utils = _module
wishart_prior = _module
settings = _module
test = _module
base_kernel_test_case = _module
base_likelihood_test_case = _module
base_mean_test_case = _module
base_test_case = _module
model_test_case = _module
utils = _module
variational_test_case = _module
broadcasting = _module
cholesky = _module
deprecation = _module
errors = _module
getitem = _module
grid = _module
interpolation = _module
lanczos = _module
memoize = _module
nearest_neighbors = _module
permutation = _module
quadrature = _module
sparse = _module
toeplitz = _module
transforms = _module
warnings = _module
variational = _module
_variational_distribution = _module
_variational_strategy = _module
additive_grid_interpolation_variational_strategy = _module
batch_decoupled_variational_strategy = _module
cholesky_variational_distribution = _module
ciq_variational_strategy = _module
delta_variational_distribution = _module
grid_interpolation_variational_strategy = _module
independent_multitask_variational_strategy = _module
lmc_variational_strategy = _module
mean_field_variational_distribution = _module
natural_variational_distribution = _module
nearest_neighbor_variational_strategy = _module
orthogonally_decoupled_variational_strategy = _module
tril_natural_variational_distribution = _module
unwhitened_variational_strategy = _module
variational_strategy = _module
setup = _module
test_constraints = _module
test_delta = _module
test_multitask_multivariate_normal = _module
test_multivariate_normal = _module
examples = _module
test_batch_decoupled_ppgpr_regression = _module
test_batch_gp_regression = _module
test_batch_multitask_gp_regression = _module
test_batch_svgp_gp_regression = _module
test_decoupled_svgp_regression = _module
test_dspp_regression = _module
test_fixed_noise_fanatasy_updates = _module
test_grid_gp_regression = _module
test_hadamard_multitask_gp_regression = _module
test_independent_multitask_gp_regression = _module
test_kissgp_additive_classification = _module
test_kissgp_additive_regression = _module
test_kissgp_dkl_regression = _module
test_kissgp_gp_classification = _module
test_kissgp_gp_regression = _module
test_kissgp_kronecker_product_classification = _module
test_kissgp_kronecker_product_regression = _module
test_kissgp_multiplicative_regression = _module
test_kissgp_variational_regression = _module
test_kissgp_white_noise_regression = _module
test_kronecker_multitask_gp_regression = _module
test_kronecker_multitask_sgpr_regression = _module
test_kronecker_multitask_ski_gp_regression = _module
test_lcm_kernel_regression = _module
test_lmc_svgp_regression = _module
test_model_list_gp_regression = _module
test_pyro_integration = _module
test_rff_gp_regression = _module
test_sgpr_regression = _module
test_simple_gp_classification = _module
test_simple_gp_regression = _module
test_spectral_mixture_gp_regression = _module
test_svgp_gp_classification = _module
test_svgp_gp_regression = _module
test_unwhitened_svgp_regression = _module
test_white_noise_regression = _module
test_log_normal_cdf = _module
test_matern_covariance = _module
test_rbf_covariance = _module
test_matern_kernel = _module
test_rbf_kernel = _module
test_additive_and_product_kernels = _module
test_arc_kernel = _module
test_cosine_kernel = _module
test_cylindrical_kernel = _module
test_gaussian_symmetrized_kl_kernel = _module
test_grid_interpolation_kernel = _module
test_grid_kernel = _module
test_index_kernel = _module
test_inducing_point_kernel = _module
test_linear_kernel = _module
test_matern_kernel = _module
test_newton_girard_additive_kernel = _module
test_periodic_kernel = _module
test_piecewise_polynomial_kernel = _module
test_polynomial_kernel = _module
test_polynomial_kernel_grad = _module
test_rbf_kernel = _module
test_rbf_kernel_grad = _module
test_rff_kernel = _module
test_rq_kernel = _module
test_scale_kernel = _module
test_spectral_delta_kernel = _module
test_spectral_mixture_kernel = _module
test_lazy_evaluated_kernel_tensor = _module
test_lazy_tensor_deprecation = _module
test_bernoulli_likelihood = _module
test_beta_likelihood = _module
test_gaussian_likelihood = _module
test_general_multitask_gaussian_likelihood = _module
test_laplace_likelihood = _module
test_multitask_gaussian_likelihood = _module
test_softmax_likelihood = _module
test_student_t_likelihood = _module
test_constant_mean = _module
test_constant_mean_grad = _module
test_linear_mean = _module
test_multitask_mean = _module
test_zero_mean = _module
test_metrics = _module
test_exact_marginal_log_likelihood = _module
test_inducing_point_kernel_added_loss_term = _module
test_leave_one_out_pseudo_likelihood = _module
test_exact_gp = _module
test_model_list = _module
test_variational_gp = _module
test_ngd = _module
test_gamma_prior = _module
test_half_cauchy_prior = _module
test_half_normal_prior = _module
test_horseshoe_prior = _module
test_lkj_prior = _module
test_multivariate_normal_prior = _module
test_normal_prior = _module
test_smoothed_box_prior = _module
test_module = _module
test_settings = _module
test_grid = _module
test_interpolation = _module
test_nearest_neighbors = _module
test_quadrature = _module
test_batch_decoupled_variational_strategy = _module
test_ciq_variational_strategy = _module
test_grid_interpolation_variational_strategy = _module
test_independent_multitask_variational_strategy = _module
test_lmc_variational_strategy = _module
test_natural_variational_distribution = _module
test_nearest_neighbor_variational_strategy = _module
test_orthogonally_decoupled_variational_strategy = _module
test_unwhitened_variational_strategy = _module
test_variational_strategy = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import re


import warnings


from typing import ForwardRef


import torch


import torch.nn as nn


import torch.nn.functional as F


from collections import OrderedDict


import numpy as np


import matplotlib.pyplot as plt


from functools import reduce


from copy import deepcopy


from torch.optim import Optimizer


from typing import Optional


from typing import Tuple


from typing import Union


from torch import Tensor


import math


from torch import sigmoid


from torch.nn import Module


import numbers


from torch.distributions import constraints


from torch.distributions.kl import register_kl


from torch.distributions import Distribution as TDistribution


from numbers import Number


from torch.distributions import MultivariateNormal as TMultivariateNormal


from torch.distributions.utils import _standard_normal


from torch.distributions.utils import lazy_property


from typing import Any


from torch.autograd import Function


from torch.distributions import Normal


from math import pi


from typing import Callable


from typing import List


import copy


from abc import abstractmethod


from typing import Dict


from typing import Iterable


from torch.nn import ModuleList


from torch.nn.parallel import DataParallel


import logging


import functools


from functools import wraps


from abc import ABC


from torch.nn import Parameter


from torch.distributions import kl_divergence


import string


from abc import abstractproperty


import inspect


import itertools


from torch import nn


from torch.distributions import Distribution


from torch.distributions import HalfCauchy


from torch.nn import Module as TModule


from torch.distributions import LKJCholesky


from torch.distributions.utils import broadcast_all


from torch.distributions import Gamma


from torch.distributions import HalfNormal


from torch.distributions import LogNormal


from torch.distributions import MultivariateNormal


from torch.distributions import Uniform


import random


from typing import Generator


from torch.distributions.kl import kl_divergence


import abc


from torch.nn.functional import softplus


from torch import optim


from math import exp


from torch.utils.data import TensorDataset


from torch.utils.data import DataLoader


import time


class _DenseLayer(nn.Sequential):

    def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):
        super(_DenseLayer, self).__init__()
        self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
        self.add_module('relu1', nn.ReLU(inplace=True)),
        self.add_module('conv1', nn.Conv2d(num_input_features, bn_size * growth_rate, kernel_size=1, stride=1, bias=False)),
        self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
        self.add_module('relu2', nn.ReLU(inplace=True)),
        self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate, kernel_size=3, stride=1, padding=1, bias=False)),
        self.drop_rate = drop_rate

    def forward(self, x):
        new_features = super(_DenseLayer, self).forward(x)
        if self.drop_rate > 0:
            new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
        return torch.cat([x, new_features], 1)


class _Transition(nn.Sequential):

    def __init__(self, num_input_features, num_output_features):
        super(_Transition, self).__init__()
        self.add_module('norm', nn.BatchNorm2d(num_input_features))
        self.add_module('relu', nn.ReLU(inplace=True))
        self.add_module('conv', nn.Conv2d(num_input_features, num_output_features, kernel_size=1, stride=1, bias=False))
        self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))


class _DenseBlock(nn.Sequential):

    def __init__(self, num_layers, num_input_features, bn_size, growth_rate, drop_rate):
        super(_DenseBlock, self).__init__()
        for i in range(num_layers):
            layer = _DenseLayer(num_input_features + i * growth_rate, growth_rate, bn_size, drop_rate)
            self.add_module('denselayer%d' % (i + 1), layer)


class DenseNet(nn.Module):
    """Densenet-BC model class, based on
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`
    Args:
        growth_rate (int) - how many filters to add each layer (`k` in paper)
        block_config (list of 3 or 4 ints) - how many layers in each pooling block
        num_init_features (int) - the number of filters to learn in the first convolution layer
        bn_size (int) - multiplicative factor for number of bottle neck layers
            (i.e. bn_size * k features in the bottleneck layer)
        drop_rate (float) - dropout rate after each dense layer
        num_classes (int) - number of classification classes
    """

    def __init__(self, growth_rate=12, block_config=(16, 16, 16), compression=0.5, num_init_features=24, bn_size=4, drop_rate=0, avgpool_size=8, num_classes=10):
        super(DenseNet, self).__init__()
        assert 0 < compression <= 1, 'compression of densenet should be between 0 and 1'
        self.avgpool_size = avgpool_size
        self.features = nn.Sequential(OrderedDict([('conv0', nn.Conv2d(3, num_init_features, kernel_size=3, stride=1, padding=1, bias=False))]))
        num_features = num_init_features
        for i, num_layers in enumerate(block_config):
            block = _DenseBlock(num_layers=num_layers, num_input_features=num_features, bn_size=bn_size, growth_rate=growth_rate, drop_rate=drop_rate)
            self.features.add_module('denseblock%d' % (i + 1), block)
            num_features = num_features + num_layers * growth_rate
            if i != len(block_config) - 1:
                trans = _Transition(num_input_features=num_features, num_output_features=int(num_features * compression))
                self.features.add_module('transition%d' % (i + 1), trans)
                num_features = int(num_features * compression)
        self.features.add_module('norm_final', nn.BatchNorm2d(num_features))
        self.classifier = nn.Linear(num_features, num_classes)

    def forward(self, x):
        features = self.features(x)
        out = F.relu(features, inplace=True)
        out = F.avg_pool2d(out, kernel_size=self.avgpool_size).view(features.size(0), -1)
        out = self.classifier(out)
        return out


def inv_sigmoid(x):
    return torch.log(x) - torch.log(1 - x)


def inv_softplus(x):
    return x + torch.log(-torch.expm1(-x))


TRANSFORM_REGISTRY = {torch.exp: torch.log, torch.nn.functional.softplus: inv_softplus, torch.sigmoid: inv_sigmoid}


def _get_inv_param_transform(param_transform, inv_param_transform=None):
    reg_inv_tf = TRANSFORM_REGISTRY.get(param_transform, None)
    if reg_inv_tf is None:
        if inv_param_transform is None:
            raise RuntimeError('Must specify inv_param_transform for custom param_transforms')
        return inv_param_transform
    elif inv_param_transform is not None and reg_inv_tf != inv_param_transform:
        raise RuntimeError('TODO')
    return reg_inv_tf


class Interval(Module):

    def __init__(self, lower_bound, upper_bound, transform=sigmoid, inv_transform=inv_sigmoid, initial_value=None):
        """
        Defines an interval constraint for GP model parameters, specified by a lower bound and upper bound. For usage
        details, see the documentation for :meth:`~gpytorch.module.Module.register_constraint`.

        Args:
            lower_bound (float or torch.Tensor): The lower bound on the parameter.
            upper_bound (float or torch.Tensor): The upper bound on the parameter.
        """
        lower_bound = torch.as_tensor(lower_bound).float()
        upper_bound = torch.as_tensor(upper_bound).float()
        if torch.any(torch.ge(lower_bound, upper_bound)):
            raise RuntimeError('Got parameter bounds with empty intervals.')
        super().__init__()
        self.register_buffer('lower_bound', lower_bound)
        self.register_buffer('upper_bound', upper_bound)
        self._transform = transform
        self._inv_transform = inv_transform
        if transform is not None and inv_transform is None:
            self._inv_transform = _get_inv_param_transform(transform)
        if initial_value is not None:
            if not isinstance(initial_value, torch.Tensor):
                initial_value = torch.tensor(initial_value)
            self._initial_value = self.inverse_transform(initial_value)
        else:
            self._initial_value = None

    def _apply(self, fn):
        self.lower_bound = fn(self.lower_bound)
        self.upper_bound = fn(self.upper_bound)
        return super()._apply(fn)

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        result = super()._load_from_state_dict(state_dict=state_dict, prefix=prefix, local_metadata=local_metadata, strict=False, missing_keys=missing_keys, unexpected_keys=unexpected_keys, error_msgs=error_msgs)
        return result

    @property
    def enforced(self):
        return self._transform is not None

    def check(self, tensor):
        return bool(torch.all(tensor <= self.upper_bound) and torch.all(tensor >= self.lower_bound))

    def check_raw(self, tensor):
        return bool(torch.all(self.transform(tensor) <= self.upper_bound) and torch.all(self.transform(tensor) >= self.lower_bound))

    def intersect(self, other):
        """
        Returns a new Interval constraint that is the intersection of this one and another specified one.

        Args:
            other (Interval): Interval constraint to intersect with

        Returns:
            Interval: intersection if this interval with the other one.
        """
        if self.transform != other.transform:
            raise RuntimeError('Cant intersect Interval constraints with conflicting transforms!')
        lower_bound = torch.max(self.lower_bound, other.lower_bound)
        upper_bound = torch.min(self.upper_bound, other.upper_bound)
        return Interval(lower_bound, upper_bound)

    def transform(self, tensor):
        """
        Transforms a tensor to satisfy the specified bounds.

        If upper_bound is finite, we assume that `self.transform` saturates at 1 as tensor -> infinity. Similarly,
        if lower_bound is finite, we assume that `self.transform` saturates at 0 as tensor -> -infinity.

        Example transforms for one of the bounds being finite include torch.exp and torch.nn.functional.softplus.
        An example transform for the case where both are finite is torch.nn.functional.sigmoid.
        """
        if not self.enforced:
            return tensor
        if settings.debug.on():
            max_bound = torch.max(self.upper_bound)
            min_bound = torch.min(self.lower_bound)
            if max_bound == math.inf or min_bound == -math.inf:
                raise RuntimeError('Cannot make an Interval directly with non-finite bounds. Use a derived class like GreaterThan or LessThan instead.')
        transformed_tensor = self._transform(tensor) * (self.upper_bound - self.lower_bound) + self.lower_bound
        return transformed_tensor

    def inverse_transform(self, transformed_tensor):
        """
        Applies the inverse transformation.
        """
        if not self.enforced:
            return transformed_tensor
        if settings.debug.on():
            max_bound = torch.max(self.upper_bound)
            min_bound = torch.min(self.lower_bound)
            if max_bound == math.inf or min_bound == -math.inf:
                raise RuntimeError('Cannot make an Interval directly with non-finite bounds. Use a derived class like GreaterThan or LessThan instead.')
        tensor = self._inv_transform((transformed_tensor - self.lower_bound) / (self.upper_bound - self.lower_bound))
        return tensor

    @property
    def initial_value(self):
        """
        The initial parameter value (if specified, None otherwise)
        """
        return self._initial_value

    def __repr__(self):
        if self.lower_bound.numel() == 1 and self.upper_bound.numel() == 1:
            return self._get_name() + f'({self.lower_bound:.3E}, {self.upper_bound:.3E})'
        else:
            return super().__repr__()

    def __iter__(self):
        yield self.lower_bound
        yield self.upper_bound


class GreaterThan(Interval):

    def __init__(self, lower_bound, transform=softplus, inv_transform=inv_softplus, initial_value=None):
        super().__init__(lower_bound=lower_bound, upper_bound=math.inf, transform=transform, inv_transform=inv_transform, initial_value=initial_value)

    def __repr__(self):
        if self.lower_bound.numel() == 1:
            return self._get_name() + f'({self.lower_bound:.3E})'
        else:
            return super().__repr__()

    def transform(self, tensor):
        transformed_tensor = self._transform(tensor) + self.lower_bound if self.enforced else tensor
        return transformed_tensor

    def inverse_transform(self, transformed_tensor):
        tensor = self._inv_transform(transformed_tensor - self.lower_bound) if self.enforced else transformed_tensor
        return tensor


class Positive(GreaterThan):

    def __init__(self, transform=softplus, inv_transform=inv_softplus, initial_value=None):
        super().__init__(lower_bound=0.0, transform=transform, inv_transform=inv_transform, initial_value=initial_value)

    def __repr__(self):
        return self._get_name() + '()'

    def transform(self, tensor):
        transformed_tensor = self._transform(tensor) if self.enforced else tensor
        return transformed_tensor

    def inverse_transform(self, transformed_tensor):
        tensor = self._inv_transform(transformed_tensor) if self.enforced else transformed_tensor
        return tensor


class LessThan(Interval):

    def __init__(self, upper_bound, transform=softplus, inv_transform=inv_softplus, initial_value=None):
        super().__init__(lower_bound=-math.inf, upper_bound=upper_bound, transform=transform, inv_transform=inv_transform, initial_value=initial_value)

    def transform(self, tensor):
        transformed_tensor = -self._transform(-tensor) + self.upper_bound if self.enforced else tensor
        return transformed_tensor

    def inverse_transform(self, transformed_tensor):
        tensor = -self._inv_transform(-(transformed_tensor - self.upper_bound)) if self.enforced else transformed_tensor
        return tensor

    def __repr__(self):
        return self._get_name() + f'({self.upper_bound:.3E})'


def sq_dist(x1, x2, x1_eq_x2=False):
    adjustment = x1.mean(-2, keepdim=True)
    x1 = x1 - adjustment
    x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
    x1_pad = torch.ones_like(x1_norm)
    if x1_eq_x2 and not x1.requires_grad and not x2.requires_grad:
        x2, x2_norm, x2_pad = x1, x1_norm, x1_pad
    else:
        x2 = x2 - adjustment
        x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
        x2_pad = torch.ones_like(x2_norm)
    x1_ = torch.cat([-2.0 * x1, x1_norm, x1_pad], dim=-1)
    x2_ = torch.cat([x2, x2_pad, x2_norm], dim=-1)
    res = x1_.matmul(x2_.transpose(-2, -1))
    if x1_eq_x2 and not x1.requires_grad and not x2.requires_grad:
        res.diagonal(dim1=-2, dim2=-1).fill_(0)
    return res.clamp_min_(0)


def dist(x1, x2, x1_eq_x2=False):
    res = sq_dist(x1, x2, x1_eq_x2=x1_eq_x2)
    return res.clamp_min_(1e-30).sqrt_()


class Distance(torch.nn.Module):

    def __init__(self, postprocess: Optional[Callable]=None):
        super().__init__()
        if postprocess is not None:
            warnings.warn('The `postprocess` argument is deprecated. See https://github.com/cornellius-gp/gpytorch/pull/2205 for details.', DeprecationWarning)
        self._postprocess = postprocess

    def _sq_dist(self, x1, x2, x1_eq_x2=False, postprocess=False):
        res = sq_dist(x1, x2, x1_eq_x2=x1_eq_x2)
        return self._postprocess(res) if postprocess else res

    def _dist(self, x1, x2, x1_eq_x2=False, postprocess=False):
        res = dist(x1, x2, x1_eq_x2=x1_eq_x2)
        return self._postprocess(res) if postprocess else res


class Noise(Module):
    pass


def _add_to_cache(obj, name, val, *args, kwargs_pkl):
    """Add a result to the cache of an object (honoring calling args)."""
    if not hasattr(obj, '_memoize_cache'):
        obj._memoize_cache = {}
    obj._memoize_cache[name, args, kwargs_pkl] = val
    return val


class CachingError(RuntimeError):
    pass


def _get_from_cache(obj, name, *args, kwargs_pkl):
    """Get an item from the cache (honoring calling args)."""
    try:
        return obj._memoize_cache[name, args, kwargs_pkl]
    except (AttributeError, KeyError):
        raise CachingError('Object does not have item {} stored in cache.'.format(name))


def _is_in_cache(obj, name, *args, kwargs_pkl):
    return hasattr(obj, '_memoize_cache') and (name, args, kwargs_pkl) in obj._memoize_cache


def _cached(method=None, name=None):
    """A decorator allowing for specifying the name of a cache, allowing it to be modified elsewhere.
    This variant honors the calling args to the decorated function.
    """
    if method is None:
        return functools.partial(_cached, name=name)

    @functools.wraps(method)
    def g(self, *args, **kwargs):
        cache_name = name if name is not None else method
        kwargs_pkl = pickle.dumps(kwargs)
        if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
            return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
        return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)
    return g


def _add_to_cache_ignore_args(obj, name, val):
    """Add a result to the cache of an object (ignoring calling args)."""
    if not hasattr(obj, '_memoize_cache'):
        obj._memoize_cache = {}
    obj._memoize_cache[name] = val
    return val


def _get_from_cache_ignore_args(obj, name):
    """Get an item from the cache (ignoring calling args)."""
    try:
        return obj._memoize_cache[name]
    except (AttributeError, KeyError):
        raise CachingError('Object does not have item {} stored in cache.'.format(name))


def _is_in_cache_ignore_args(obj, name):
    return hasattr(obj, '_memoize_cache') and name in obj._memoize_cache


def _cached_ignore_args(method=None, name=None):
    """A decorator allowing for specifying the name of a cache, allowing it to be modified elsewhere.
    This variant ignores the calling args to the decorated function.
    """
    if method is None:
        return functools.partial(_cached_ignore_args, name=name)

    @functools.wraps(method)
    def g(self, *args, **kwargs):
        cache_name = name if name is not None else method
        if not _is_in_cache_ignore_args(self, cache_name):
            return _add_to_cache_ignore_args(self, cache_name, method(self, *args, **kwargs))
        return _get_from_cache_ignore_args(self, cache_name)
    return g


def cached(method=None, name=None, ignore_args=False):
    """A decorator allowing for specifying the name of a cache, allowing it to be modified elsewhere."""
    if ignore_args:
        return _cached_ignore_args(method=method, name=name)
    else:
        return _cached(method=method, name=name)


def recall_grad_state(method: Callable) ->Callable:
    """Decorator for LazyEvaluatedKernelTensor's methods to put their execution
    inside the same grad state as during the instantiation of the lazy object.
    This makes the lazy tensor object behave in the same way as a regular tensor
    with respect to the grad state.
    """

    @functools.wraps(method)
    def wrapped(self, *args, **kwargs):
        with torch.set_grad_enabled(self._is_grad_enabled):
            output = method(self, *args, **kwargs)
        return output
    return wrapped


class Prior(Distribution, Module, ABC):
    """
    Base class for Priors in GPyTorch.
    In GPyTorch, a parameter can be assigned a prior by passing it as the `prior` argument to
    :func:`~gpytorch.module.register_parameter`. GPyTorch performs internal bookkeeping of priors,
    and for each parameter with a registered prior includes the log probability of the parameter under its
    respective prior in computing the Marginal Log-Likelihood.
    """

    def transform(self, x):
        return self._transform(x) if self._transform is not None else x

    def log_prob(self, x):
        """
        :return: log-probability of the parameter value under the prior
        :rtype: torch.Tensor
        """
        return super(Prior, self).log_prob(self.transform(x))


class MaternCovariance(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x1, x2, lengthscale, nu, dist_func):
        if any(ctx.needs_input_grad[:2]):
            raise RuntimeError('MaternCovariance cannot compute gradients with respect to x1 and x2')
        if lengthscale.size(-1) > 1:
            raise ValueError('MaternCovariance cannot handle multiple lengthscales')
        needs_grad = any(ctx.needs_input_grad)
        mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
        x1_ = (x1 - mean).div(lengthscale)
        x2_ = (x2 - mean).div(lengthscale)
        scaled_unitless_dist = dist_func(x1_, x2_).mul_(math.sqrt(2 * nu))
        if nu == 0.5:
            scaled_unitless_dist_ = scaled_unitless_dist.clone() if needs_grad else scaled_unitless_dist
            exp_component = scaled_unitless_dist_.neg_().exp_()
            covar_mat = exp_component
            if needs_grad:
                d_output_d_input = scaled_unitless_dist.div_(lengthscale).mul_(exp_component)
        elif nu == 1.5:
            if needs_grad:
                scaled_unitless_dist_ = scaled_unitless_dist.clone()
            linear_term = scaled_unitless_dist.clone().add_(1)
            exp_component = scaled_unitless_dist.neg_().exp_()
            covar_mat = linear_term.mul_(exp_component)
            if needs_grad:
                d_output_d_input = scaled_unitless_dist_.pow_(2).div_(lengthscale).mul_(exp_component)
        elif nu == 2.5:
            linear_term = scaled_unitless_dist.clone().add_(1)
            quadratic_term = scaled_unitless_dist.clone().pow_(2).div_(3)
            exp_component = scaled_unitless_dist.neg_().exp_()
            if needs_grad:
                covar_mat = (linear_term + quadratic_term).mul_(exp_component)
                d_output_d_input = linear_term.mul_(quadratic_term).mul_(exp_component).div_(lengthscale)
            else:
                covar_mat = exp_component.mul_(linear_term.add_(quadratic_term))
        if needs_grad:
            ctx.save_for_backward(d_output_d_input)
        return covar_mat

    @staticmethod
    def backward(ctx, grad_output):
        d_output_d_input = ctx.saved_tensors[0]
        lengthscale_grad = grad_output * d_output_d_input
        return None, None, lengthscale_grad, None, None


class _feature_flag:
    """Base class for feature flag settings with global scope.
    The default is set via the `_default` class attribute.
    """
    _default = False
    _state = None

    @classmethod
    def is_default(cls):
        return cls._state is None

    @classmethod
    def on(cls):
        if cls.is_default():
            return cls._default
        return cls._state

    @classmethod
    def off(cls):
        return not cls.on()

    @classmethod
    def _set_state(cls, state):
        cls._state = state

    def __init__(self, state=True):
        self.prev = self.__class__._state
        self.state = state

    def __enter__(self):
        self.__class__._set_state(self.state)

    def __exit__(self, *args):
        self.__class__._set_state(self.prev)
        return False


class trace_mode(_feature_flag):
    """
    If set to True, we will generally try to avoid calling our built in PyTorch functions, because these cannot
    be run through torch.jit.trace.

    Note that this will sometimes involve explicitly evaluating lazy tensors and various other slowdowns and
    inefficiencies. As a result, you really shouldn't use this feature context unless you are calling torch.jit.trace
    on a GPyTorch model.

    Our hope is that this flag will not be necessary long term, once https://github.com/pytorch/pytorch/issues/22329
    is fixed.

    (Default: False)
    """
    _default = False


def _fmax(r: Tensor, j: int, q: int) ->Tensor:
    return torch.max(torch.tensor(0.0, dtype=r.dtype, device=r.device), 1 - r).pow(j + q)


def _get_cov(r: Tensor, j: int, q: int) ->Tensor:
    if q == 0:
        return 1
    if q == 1:
        return (j + 1) * r + 1
    if q == 2:
        return 1 + (j + 2) * r + (j + 4 * j + 3) / 3.0 * r ** 2
    if q == 3:
        return 1 + (j + 3) * r + (6 * j ** 2 + 36 * j + 45) / 15.0 * r.square() + (j ** 3 + 9 * j ** 2 + 23 * j + 15) / 15.0 * r ** 3


class RBFCovariance(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x1, x2, lengthscale, sq_dist_func):
        if any(ctx.needs_input_grad[:2]):
            raise RuntimeError('RBFCovariance cannot compute gradients with respect to x1 and x2')
        if lengthscale.size(-1) > 1:
            raise ValueError('RBFCovariance cannot handle multiple lengthscales')
        needs_grad = any(ctx.needs_input_grad)
        x1_ = x1.div(lengthscale)
        x2_ = x2.div(lengthscale)
        unitless_sq_dist = sq_dist_func(x1_, x2_)
        unitless_sq_dist_ = unitless_sq_dist.clone() if needs_grad else unitless_sq_dist
        covar_mat = unitless_sq_dist_.div_(-2.0).exp_()
        if needs_grad:
            d_output_d_input = unitless_sq_dist.mul_(covar_mat).div_(lengthscale)
            ctx.save_for_backward(d_output_d_input)
        return covar_mat

    @staticmethod
    def backward(ctx, grad_output):
        d_output_d_input = ctx.saved_tensors[0]
        lengthscale_grad = grad_output * d_output_d_input
        return None, None, lengthscale_grad, None


def postprocess_rbf(dist_mat):
    return dist_mat.div_(-2).exp_()


logger = logging.getLogger()


class _Likelihood(Module, ABC):

    def __init__(self, max_plate_nesting=1):
        super().__init__()
        self.max_plate_nesting = max_plate_nesting

    def _draw_likelihood_samples(self, function_dist, *args, sample_shape=None, **kwargs):
        if sample_shape is None:
            sample_shape = torch.Size([settings.num_likelihood_samples.value()] + [1] * (self.max_plate_nesting - len(function_dist.batch_shape) - 1))
        else:
            sample_shape = sample_shape[:-len(function_dist.batch_shape) - 1]
        if self.training:
            num_event_dims = len(function_dist.event_shape)
            function_dist = base_distributions.Normal(function_dist.mean, function_dist.variance.sqrt())
            function_dist = base_distributions.Independent(function_dist, num_event_dims - 1)
        function_samples = function_dist.rsample(sample_shape)
        return self.forward(function_samples, *args, **kwargs)

    def expected_log_prob(self, observations, function_dist, *args, **kwargs):
        likelihood_samples = self._draw_likelihood_samples(function_dist, *args, **kwargs)
        res = likelihood_samples.log_prob(observations).mean(dim=0)
        return res

    @abstractmethod
    def forward(self, function_samples, *args, **kwargs):
        raise NotImplementedError

    def get_fantasy_likelihood(self, **kwargs):
        return deepcopy(self)

    def log_marginal(self, observations, function_dist, *args, **kwargs):
        likelihood_samples = self._draw_likelihood_samples(function_dist, *args, **kwargs)
        log_probs = likelihood_samples.log_prob(observations)
        res = log_probs.sub(math.log(log_probs.size(0))).logsumexp(dim=0)
        return res

    def marginal(self, function_dist, *args, **kwargs):
        res = self._draw_likelihood_samples(function_dist, *args, **kwargs)
        return res

    def __call__(self, input, *args, **kwargs):
        if torch.is_tensor(input):
            return super().__call__(input, *args, **kwargs)
        elif isinstance(input, MultivariateNormal):
            return self.marginal(input, *args, **kwargs)
        else:
            raise RuntimeError('Likelihoods expects a MultivariateNormal input to make marginal predictions, or a torch.Tensor for conditional predictions. Got a {}'.format(input.__class__.__name__))


class GP(Module):
    pass


class MarginalLogLikelihood(Module):
    """
    These are modules to compute (or approximate/bound) the marginal log likelihood
    (MLL) of the GP model when applied to data.  I.e., given a GP :math:`f \\sim
    \\mathcal{GP}(\\mu, K)`, and data :math:`\\mathbf X, \\mathbf y`, these modules
    compute/approximate

    .. math::

       \\begin{equation*}
          \\mathcal{L} = p_f(\\mathbf y \\! \\mid \\! \\mathbf X)
          = \\int p \\left( \\mathbf y \\! \\mid \\! f(\\mathbf X) \\right) \\: p(f(\\mathbf X) \\! \\mid \\! \\mathbf X) \\: d f
       \\end{equation*}

    This is computed exactly when the GP inference is computed exactly (e.g. regression w/ a Gaussian likelihood).
    It is approximated/bounded for GP models that use approximate inference.

    These models are typically used as the "loss" functions for GP models (though note that the output of
    these functions must be negated for optimization).
    """

    def __init__(self, likelihood, model):
        super(MarginalLogLikelihood, self).__init__()
        if not isinstance(model, GP):
            raise RuntimeError('All MarginalLogLikelihood objects must be given a GP object as a model. If you are using a more complicated model involving a GP, pass the underlying GP object as the model, not a full PyTorch module.')
        self.likelihood = likelihood
        self.model = model

    def forward(self, output, target, **kwargs):
        """
        Computes the MLL given :math:`p(\\mathbf f)` and `\\mathbf y`

        :param ~gpytorch.distributions.MultivariateNormal output: the outputs of the latent function
            (the :obj:`~gpytorch.models.GP`)
        :param torch.Tensor target: :math:`\\mathbf y` The target values
        :param dict kwargs: Additional arguments to pass to the likelihood's forward function.
        """
        raise NotImplementedError


class ExactMarginalLogLikelihood(MarginalLogLikelihood):
    """
    The exact marginal log likelihood (MLL) for an exact Gaussian process with a
    Gaussian likelihood.

    .. note::
        This module will not work with anything other than a :obj:`~gpytorch.likelihoods.GaussianLikelihood`
        and a :obj:`~gpytorch.models.ExactGP`. It also cannot be used in conjunction with
        stochastic optimization.

    :param ~gpytorch.likelihoods.GaussianLikelihood likelihood: The Gaussian likelihood for the model
    :param ~gpytorch.models.ExactGP model: The exact GP model

    Example:
        >>> # model is a gpytorch.models.ExactGP
        >>> # likelihood is a gpytorch.likelihoods.Likelihood
        >>> mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
        >>>
        >>> output = model(train_x)
        >>> loss = -mll(output, train_y)
        >>> loss.backward()
    """

    def __init__(self, likelihood, model):
        if not isinstance(likelihood, _GaussianLikelihoodBase):
            raise RuntimeError('Likelihood must be Gaussian for exact inference')
        super(ExactMarginalLogLikelihood, self).__init__(likelihood, model)

    def _add_other_terms(self, res, params):
        for added_loss_term in self.model.added_loss_terms():
            res = res.add(added_loss_term.loss(*params))
        res_ndim = res.ndim
        for name, module, prior, closure, _ in self.named_priors():
            prior_term = prior.log_prob(closure(module))
            res.add_(prior_term.view(*prior_term.shape[:res_ndim], -1).sum(dim=-1))
        return res

    def forward(self, function_dist, target, *params):
        """
        Computes the MLL given :math:`p(\\mathbf f)` and :math:`\\mathbf y`.

        :param ~gpytorch.distributions.MultivariateNormal function_dist: :math:`p(\\mathbf f)`
            the outputs of the latent function (the :obj:`gpytorch.models.ExactGP`)
        :param torch.Tensor target: :math:`\\mathbf y` The target values
        :rtype: torch.Tensor
        :return: Exact MLL. Output shape corresponds to batch shape of the model/input data.
        """
        if not isinstance(function_dist, MultivariateNormal):
            raise RuntimeError('ExactMarginalLogLikelihood can only operate on Gaussian random variables')
        output = self.likelihood(function_dist, *params)
        res = output.log_prob(target)
        res = self._add_other_terms(res, params)
        num_data = function_dist.event_shape.numel()
        return res.div_(num_data)


class SumMarginalLogLikelihood(MarginalLogLikelihood):
    """Sum of marginal log likelihoods, to be used with Multi-Output models.

    Args:
        likelihood: A MultiOutputLikelihood
        model: A MultiOutputModel
        mll_cls: The Marginal Log Likelihood class (default: ExactMarginalLogLikelihood)

    In case the model outputs are independent, this provives the MLL of the multi-output model.

    """

    def __init__(self, likelihood, model, mll_cls=ExactMarginalLogLikelihood):
        super().__init__(model.likelihood, model)
        self.mlls = ModuleList([mll_cls(mdl.likelihood, mdl) for mdl in model.models])

    def forward(self, outputs, targets, *params):
        """
        Args:
            outputs: (Iterable[MultivariateNormal]) - the outputs of the latent function
            targets: (Iterable[Tensor]) - the target values
            params: (Iterable[Iterable[Tensor]]) - the arguments to be passed through
                (e.g. parameters in case of heteroskedastic likelihoods)
        """
        if len(params) == 0:
            sum_mll = sum(mll(output, target) for mll, output, target in zip(self.mlls, outputs, targets))
        else:
            sum_mll = sum(mll(output, target, *iparams) for mll, output, target, iparams in zip(self.mlls, outputs, targets, params))
        return sum_mll.div_(len(self.mlls))


class LatentVariable(Module):
    """
    This super class is used to describe the type of inference
    used for the latent variable :math:`\\mathbf X` in GPLVM models.

    :param int n: Size of the latent space.
    :param int latent_dim: Dimensionality of latent space.
    """

    def __init__(self, n, dim):
        super().__init__()
        self.n = n
        self.latent_dim = dim

    def forward(self, x):
        raise NotImplementedError


class PointLatentVariable(LatentVariable):
    """
    This class is used for GPLVM models to recover a MLE estimate of
    the latent variable :math:`\\mathbf X`.

    :param int n: Size of the latent space.
    :param int latent_dim: Dimensionality of latent space.
    :param torch.Tensor X_init: initialization for the point estimate of :math:`\\mathbf X`
    """

    def __init__(self, n, latent_dim, X_init):
        super().__init__(n, latent_dim)
        self.register_parameter('X', X_init)

    def forward(self):
        return self.X


class MAPLatentVariable(LatentVariable):
    """
    This class is used for GPLVM models to recover a MAP estimate of
    the latent variable :math:`\\mathbf X`, based on some supplied prior.

    :param int n: Size of the latent space.
    :param int latent_dim: Dimensionality of latent space.
    :param torch.Tensor X_init: initialization for the point estimate of :math:`\\mathbf X`
    :param ~gpytorch.priors.Prior prior_x: prior for :math:`\\mathbf X`
    """

    def __init__(self, n, latent_dim, X_init, prior_x):
        super().__init__(n, latent_dim)
        self.prior_x = prior_x
        self.register_parameter('X', X_init)
        self.register_prior('prior_x', prior_x, 'X')

    def forward(self):
        return self.X


class VariationalLatentVariable(LatentVariable):
    """
    This class is used for GPLVM models to recover a variational approximation of
    the latent variable :math:`\\mathbf X`. The variational approximation will be
    an isotropic Gaussian distribution.

    :param int n: Size of the latent space.
    :param int data_dim: Dimensionality of the :math:`\\mathbf Y` values.
    :param int latent_dim: Dimensionality of latent space.
    :param torch.Tensor X_init: initialization for the point estimate of :math:`\\mathbf X`
    :param ~gpytorch.priors.Prior prior_x: prior for :math:`\\mathbf X`
    """

    def __init__(self, n, data_dim, latent_dim, X_init, prior_x):
        super().__init__(n, latent_dim)
        self.data_dim = data_dim
        self.prior_x = prior_x
        self.q_mu = torch.nn.Parameter(X_init)
        self.q_log_sigma = torch.nn.Parameter(torch.randn(n, latent_dim))
        self.register_added_loss_term('x_kl')

    def forward(self):
        q_x = torch.distributions.Normal(self.q_mu, torch.nn.functional.softplus(self.q_log_sigma))
        x_kl = KLGaussianAddedLossTerm(q_x, self.prior_x, self.n, self.data_dim)
        self.update_added_loss_term('x_kl', x_kl)
        return q_x.rsample()


class RandomModuleMixin(object):

    def initialize(self, **kwargs):
        """
        Set a value for a parameter

        kwargs: (param_name, value) - parameter to initialize.
        Can also initialize recursively by passing in the full name of a
        parameter. For example if model has attribute model.likelihood,
        we can initialize the noise with either
        `model.initialize(**{'likelihood.noise': 0.1})`
        or
        `model.likelihood.initialize(noise=0.1)`.
        The former method would allow users to more easily store the
        initialization values as one object.

        Value can take the form of a tensor, a float, or an int
        """
        for name, value in kwargs.items():
            if not torch.is_tensor(value):
                raise RuntimeError('Initialize in RandomModules can only be done with tensor values.')
            names = name.rsplit('.')
            if len(names) > 1:
                mod_name, param_name = names
                mod = operator.attrgetter(mod_name)(self)
            else:
                mod, param_name = self, name
            old_param = getattr(mod, param_name)
            is_property = hasattr(type(self), name) and isinstance(getattr(type(self), name), property)
            if not isinstance(old_param, torch.nn.Parameter) or is_property:
                setattr(mod, param_name, value.expand(old_param.shape))
            else:
                delattr(mod, param_name)
                setattr(mod, param_name, value.expand(old_param.shape))
        return self


def _extract_named_added_loss_terms(module, memo=None, prefix=''):
    if memo is None:
        memo = set()
    if hasattr(module, '_added_loss_terms'):
        for name, strategy in module._added_loss_terms.items():
            if strategy is not None and strategy not in memo:
                memo.add(strategy)
                yield prefix + ('.' if prefix else '') + name, strategy
    for mname, module_ in module.named_children():
        submodule_prefix = prefix + ('.' if prefix else '') + mname
        for name, strategy in _extract_named_added_loss_terms(module=module_, memo=memo, prefix=submodule_prefix):
            yield name, strategy


def _extract_named_constraints(module, memo=None, prefix=''):
    if memo is None:
        memo = set()
    if hasattr(module, '_constraints'):
        for name, constraint in module._constraints.items():
            if constraint is not None and constraint not in memo:
                memo.add(constraint)
                full_name = ('.' if prefix else '').join([prefix, name])
                yield full_name, constraint
    for mname, module_ in module.named_children():
        submodule_prefix = prefix + ('.' if prefix else '') + mname
        for name, constraint in _extract_named_constraints(module_, memo=memo, prefix=submodule_prefix):
            yield name, constraint


def _extract_named_priors(module, memo=None, prefix=''):
    if memo is None:
        memo = set()
    if hasattr(module, '_priors'):
        for name, (prior, closure, inv_closure) in module._priors.items():
            if prior is not None and prior not in memo:
                memo.add(prior)
                full_name = ('.' if prefix else '').join([prefix, name])
                yield full_name, module, prior, closure, inv_closure
    for mname, module_ in module.named_children():
        submodule_prefix = prefix + ('.' if prefix else '') + mname
        for name, parent_module, prior, closure, inv_closure in _extract_named_priors(module_, memo=memo, prefix=submodule_prefix):
            yield name, parent_module, prior, closure, inv_closure


def _pyro_load_from_samples(module, samples_dict, memo=None, prefix=''):
    if memo is None:
        memo = set()
    if hasattr(module, '_priors'):
        module.local_load_samples(samples_dict, memo, prefix)
    for mname, module_ in module.named_children():
        submodule_prefix = prefix + ('.' if prefix else '') + mname
        _pyro_load_from_samples(module_, samples_dict, memo=memo, prefix=submodule_prefix)


def _set_strict(module, value, memo=None):
    if memo is None:
        memo = set()
    if hasattr(module, '_strict_init'):
        module._strict_init = value
    for mname, module_ in module.named_children():
        _set_strict(module_, value)


def _validate_module_outputs(outputs):
    if isinstance(outputs, tuple):
        if not all(torch.is_tensor(output) or isinstance(output, Distribution) or isinstance(output, LinearOperator) for output in outputs):
            raise RuntimeError('All outputs must be a Distribution, torch.Tensor, or LinearOperator. Got {}'.format([output.__class__.__name__ for output in outputs]))
        if len(outputs) == 1:
            outputs = outputs[0]
        return outputs
    elif torch.is_tensor(outputs) or isinstance(outputs, Distribution) or isinstance(outputs, LinearOperator):
        return outputs
    else:
        raise RuntimeError('Output must be a Distribution, torch.Tensor, or LinearOperator. Got {}'.format(outputs.__class__.__name__))


class Module(nn.Module):

    def __init__(self):
        super().__init__()
        self._added_loss_terms = OrderedDict()
        self._priors = OrderedDict()
        self._constraints = OrderedDict()
        self._strict_init = True
        self._load_strict_shapes = True
        self._register_load_state_dict_pre_hook(self._load_state_hook_ignore_shapes)

    def __call__(self, *inputs, **kwargs):
        outputs = self.forward(*inputs, **kwargs)
        if isinstance(outputs, list):
            return [_validate_module_outputs(output) for output in outputs]
        return _validate_module_outputs(outputs)

    def _clear_cache(self):
        """
        Clear any precomputed caches.
        Should be implemented by any module that caches any computation at test time.
        """
        pass

    def _get_module_and_name(self, parameter_name):
        """Get module and name from full parameter name."""
        module, name = parameter_name.split('.', 1)
        if module in self._modules:
            return self.__getattr__(module), name
        else:
            raise AttributeError('Invalid parameter name {}. {} has no module {}'.format(parameter_name, type(self).__name__, module))

    def _strict(self, value):
        _set_strict(self, value)

    def added_loss_terms(self):
        for _, strategy in self.named_added_loss_terms():
            yield strategy

    def forward(self, *inputs, **kwargs):
        raise NotImplementedError

    def constraints(self):
        for _, constraint in self.named_constraints():
            yield constraint

    def hyperparameters(self):
        for _, param in self.named_hyperparameters():
            yield param

    def initialize(self, **kwargs):
        """
        Set a value for a parameter

        kwargs: (param_name, value) - parameter to initialize.
        Can also initialize recursively by passing in the full name of a
        parameter. For example if model has attribute model.likelihood,
        we can initialize the noise with either
        `model.initialize(**{'likelihood.noise': 0.1})`
        or
        `model.likelihood.initialize(noise=0.1)`.
        The former method would allow users to more easily store the
        initialization values as one object.

        Value can take the form of a tensor, a float, or an int
        """
        for name, val in kwargs.items():
            if isinstance(val, int):
                val = float(val)
            if '.' in name:
                module, name = self._get_module_and_name(name)
                if isinstance(module, nn.ModuleList):
                    idx, name = name.split('.', 1)
                    module[int(idx)].initialize(**{name: val})
                else:
                    module.initialize(**{name: val})
            elif not hasattr(self, name):
                raise AttributeError('Unknown parameter {p} for {c}'.format(p=name, c=self.__class__.__name__))
            elif name not in self._parameters and name not in self._buffers:
                setattr(self, name, val)
            elif torch.is_tensor(val):
                constraint = self.constraint_for_parameter_name(name)
                if constraint is not None and constraint.enforced and not constraint.check_raw(val):
                    raise RuntimeError(f'Attempting to manually set a parameter value that is out of bounds of its current constraints, {constraint}. Most likely, you want to do the following:\n likelihood = GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(better_lower_bound))')
                try:
                    self.__getattr__(name).data.copy_(val.expand_as(self.__getattr__(name)))
                except RuntimeError:
                    if not self._strict_init:
                        self.__getattr__(name).data = val
                    else:
                        self.__getattr__(name).data.copy_(val.view_as(self.__getattr__(name)))
            elif isinstance(val, float):
                constraint = self.constraint_for_parameter_name(name)
                if constraint is not None and not constraint.check_raw(val):
                    raise RuntimeError(f'Attempting to manually set a parameter value that is out of bounds of its current constraints, {constraint}. Most likely, you want to do the following:\n likelihood = GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(better_lower_bound))')
                self.__getattr__(name).data.fill_(val)
            else:
                raise AttributeError('Type {t} not valid for initializing parameter {p}'.format(t=type(val), p=name))
            prior_name = '_'.join([name, 'prior'])
            if prior_name in self._priors:
                prior, closure, _ = self._priors[prior_name]
                try:
                    prior._validate_sample(closure(self))
                except ValueError as e:
                    raise ValueError('Invalid input value for prior {}. Error:\n{}'.format(prior_name, e))
        return self

    def named_added_loss_terms(self):
        """Returns an iterator over module variational strategies, yielding both
        the name of the variational strategy as well as the strategy itself.

        Yields:
            (string, VariationalStrategy): Tuple containing the name of the
                strategy and the strategy

        """
        return _extract_named_added_loss_terms(module=self, memo=None, prefix='')

    def named_hyperparameters(self):
        for module_prefix, module in self.named_modules():
            if not isinstance(module, _VariationalDistribution):
                for elem in module.named_parameters(prefix=module_prefix, recurse=False):
                    yield elem

    def named_priors(self, memo=None, prefix=''):
        """Returns an iterator over the module's priors, yielding the name of the prior,
        the prior, the associated parameter names, and the transformation callable.

        Yields:
            (string, Module, Prior, tuple((Parameter, callable)), callable): Tuple containing:
                - the name of the prior
                - the parent module of the prior
                - the prior
                - a tuple of tuples (param, transform), one for each of the parameters associated with the prior
                - the prior's transform to be called on the parameters
        """
        return _extract_named_priors(module=self, memo=None, prefix='')

    def named_constraints(self, memo=None, prefix=''):
        return _extract_named_constraints(module=self, memo=None, prefix='')

    def named_variational_parameters(self):
        for module_prefix, module in self.named_modules():
            if isinstance(module, _VariationalDistribution):
                for elem in module.named_parameters(prefix=module_prefix, recurse=False):
                    yield elem

    def register_added_loss_term(self, name):
        self._added_loss_terms[name] = None

    def register_parameter(self, name, parameter):
        """
        Adds a parameter to the module. The parameter can be accessed as an attribute using the given name.

        Args:
            name (str):
                The name of the parameter
            parameter (torch.nn.Parameter):
                The parameter
        """
        if '_parameters' not in self.__dict__:
            raise AttributeError('Cannot assign parameter before Module.__init__() call')
        super().register_parameter(name, parameter)

    def register_prior(self, name, prior, param_or_closure, setting_closure=None):
        """
        Adds a prior to the module. The prior can be accessed as an attribute using the given name.

        Args:
            name (str):
                The name of the prior
            prior (Prior):
                The prior to be registered`
            param_or_closure (string or callable):
                Either the name of the parameter, or a closure (which upon calling evalutes a function on
                the module instance and one or more parameters):
                single parameter without a transform: `.register_prior("foo_prior", foo_prior, "foo_param")`
                transform a single parameter (e.g. put a log-Normal prior on it):
                `.register_prior("foo_prior", NormalPrior(0, 1), lambda module: torch.log(module.foo_param))`
                function of multiple parameters:
                `.register_prior("foo2_prior", foo2_prior, lambda module: f(module.param1, module.param2)))`
            setting_closure (callable, optional):
                A function taking in the module instance and a tensor in (transformed) parameter space,
                initializing the internal parameter representation to the proper value by applying the
                inverse transform. Enables setting parametres directly in the transformed space, as well
                as sampling parameter values from priors (see `sample_from_prior`)

        """
        if isinstance(param_or_closure, str):
            if param_or_closure not in self._parameters and not hasattr(self, param_or_closure):
                raise AttributeError('Unknown parameter {name} for {module}'.format(name=param_or_closure, module=self.__class__.__name__) + ' Make sure the parameter is registered before registering a prior.')

            def closure(module):
                return getattr(module, param_or_closure)
            if setting_closure is not None:
                raise RuntimeError('Must specify a closure instead of a parameter name when providing setting_closure')

            def setting_closure(module, val):
                return module.initialize(**{param_or_closure: val})
        else:
            if len(inspect.signature(param_or_closure).parameters) == 0:
                raise ValueError("""As of version 1.4, `param_or_closure` must operate on a module instance. For example:

                    likelihood.noise_covar.register_prior(
                        "noise_std_prior",
                        gpytorch.priors.NormalPrior(0, 1),
                        lambda module: module.noise.sqrt()
                    )
                    """)
            if inspect.isfunction(setting_closure) and len(inspect.signature(setting_closure).parameters) < 2:
                raise ValueError("""As of version 1.4, `setting_closure` must operate on a module instance and a tensor. For example:

                    kernel.register_prior(
                        "radius_prior",
                        gpytorch.priors.LogNormalPrior(0, 1),
                        lambda module: module.radius,
                        lambda module, value: m._set_radius(value),
                    )
                    """)
            closure = param_or_closure
        self.add_module(name, prior)
        self._priors[name] = prior, closure, setting_closure

    def register_constraint(self, param_name, constraint, replace=True):
        if param_name not in self._parameters:
            raise RuntimeError('Attempting to register constraint for nonexistent parameter.')
        constraint_name = param_name + '_constraint'
        if constraint_name in self._constraints:
            current_constraint = self._constraints[constraint_name]
        else:
            current_constraint = None
        if isinstance(current_constraint, Interval) and not replace:
            new_constraint = constraint.intersect(current_constraint)
        else:
            new_constraint = constraint
        self.add_module(constraint_name, new_constraint)
        self._constraints[constraint_name] = new_constraint
        if new_constraint.initial_value is not None:
            self.initialize(**{param_name: new_constraint.initial_value})

    def train(self, mode=True):
        if self.training and not mode or mode:
            self._clear_cache()
        return super().train(mode=mode)

    def constraint_for_parameter_name(self, param_name):
        base_module = self
        base_name = param_name
        while '.' in base_name:
            components = base_name.split('.')
            submodule_name = components[0]
            submodule = getattr(base_module, submodule_name)
            base_module = submodule
            base_name = '.'.join(components[1:])
        try:
            constraint_name = base_name + '_constraint'
            return base_module._constraints.get(constraint_name)
        except AttributeError:
            return None

    def _load_state_hook_ignore_shapes(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        if not self._load_strict_shapes:
            local_name_params = itertools.chain(self._parameters.items(), self._buffers.items())
            local_state = {k: v for k, v in local_name_params if v is not None}
            for name, param in local_state.items():
                key = prefix + name
                if key in state_dict:
                    param.data = state_dict[key].data

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
        self._clear_cache()
        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)

    def load_strict_shapes(self, value):

        def apply_fn(module):
            module._load_strict_shapes = value
        self.apply(apply_fn)

    def named_parameters_and_constraints(self):
        for name, param in self.named_parameters():
            yield name, param, self.constraint_for_parameter_name(name)

    def sample_from_prior(self, prior_name):
        """Sample parameter values from prior. Modifies the module's parameters in-place."""
        if prior_name not in self._priors:
            raise RuntimeError("Unknown prior name '{}'".format(prior_name))
        prior, _, setting_closure = self._priors[prior_name]
        if setting_closure is None:
            raise RuntimeError('Must provide inverse transform to be able to sample from prior.')
        setting_closure(self, prior.sample())

    def to_pyro_random_module(self):
        return self.to_random_module()

    def to_random_module(self):
        random_module_cls = type('_Random' + self.__class__.__name__, (RandomModuleMixin, self.__class__), {})
        if not isinstance(self, random_module_cls):
            new_module = copy.deepcopy(self)
            new_module.__class__ = random_module_cls
        else:
            new_module = copy.deepcopy(self)
        for mname, child in new_module.named_children():
            if isinstance(child, Module):
                setattr(new_module, mname, child.to_random_module())
        return new_module

    def pyro_sample_from_prior(self):
        """
        For each parameter in this Module and submodule that have defined priors, sample a value for that parameter
        from its corresponding prior with a pyro.sample primitive and load the resulting value in to the parameter.

        This method can be used in a Pyro model to conveniently define pyro sample sites for all
        parameters of the model that have GPyTorch priors registered to them.
        """
        new_module = self.to_pyro_random_module()
        return _pyro_sample_from_prior(module=new_module, memo=None, prefix='')

    def local_load_samples(self, samples_dict, memo, prefix):
        """
        Defines local behavior of this Module when loading parameters from a samples_dict generated by a Pyro
        sampling mechanism.

        The default behavior here should almost always be called from any overriding class. However, a class may
        want to add additional functionality, such as reshaping things to account for the fact that parameters will
        acquire an extra batch dimension corresponding to the number of samples drawn.
        """
        self._strict(False)
        for name, (prior, closure, setting_closure) in self._priors.items():
            if prior is not None and prior not in memo:
                memo.add(prior)
                setting_closure(self, samples_dict[prefix + ('.' if prefix else '') + name])
        self._strict(True)

    def pyro_load_from_samples(self, samples_dict):
        """
        Convert this Module in to a batch Module by loading parameters from the given `samples_dict`. `samples_dict`
        is typically produced by a Pyro sampling mechanism.

        Note that the keys of the samples_dict should correspond to prior names (covar_module.outputscale_prior) rather
        than parameter names (covar_module.raw_outputscale), because we will use the setting_closure associated with
        the prior to properly set the unconstrained parameter.

        Args:
            samples_dict (dict): Dictionary mapping *prior names* to sample values.
        """
        return _pyro_load_from_samples(module=self, samples_dict=samples_dict, memo=None, prefix='')

    def update_added_loss_term(self, name, added_loss_term):
        if not isinstance(added_loss_term, AddedLossTerm):
            raise RuntimeError('added_loss_term must be a AddedLossTerm')
        if name not in self._added_loss_terms.keys():
            raise RuntimeError('added_loss_term {} not registered'.format(name))
        self._added_loss_terms[name] = added_loss_term

    def variational_parameters(self):
        for _, param in self.named_variational_parameters():
            yield param

    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError as e:
            try:
                return super().__getattribute__(name)
            except AttributeError:
                raise e


def _bufferize_attributes(module, attributes):
    attr_clones = {attr: getattr(module, attr).clone() for attr in attributes}
    for attr, value in attr_clones.items():
        delattr(module, attr)
        module.register_buffer(attr, value)


class NormalPrior(Prior, Normal):
    """
    Normal (Gaussian) Prior

    pdf(x) = (2 * pi * sigma^2)^-0.5 * exp(-(x - mu)^2 / (2 * sigma^2))

    where mu is the mean and sigma^2 is the variance.
    """

    def __init__(self, loc, scale, validate_args=False, transform=None):
        TModule.__init__(self)
        Normal.__init__(self, loc=loc, scale=scale, validate_args=validate_args)
        _bufferize_attributes(self, ('loc', 'scale'))
        self._transform = transform

    def expand(self, batch_shape):
        batch_shape = torch.Size(batch_shape)
        return NormalPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape))


class SmoothedBoxPrior(Prior):
    """A smoothed approximation of a uniform prior.

    Has full support on the reals and is differentiable everywhere.

    .. math::

        \\begin{equation*}
            B = {x: a_i <= x_i <= b_i}
            d(x, B) = min_{x' in B} |x - x'|
            pdf(x) \\sim exp(- d(x, B)**2 / sqrt(2 * sigma^2))
        \\end{equation*}

    """
    arg_constraints = {'sigma': constraints.positive, 'a': constraints.real, 'b': constraints.real}
    support = constraints.real
    has_rsample = True
    _validate_args = True

    def __init__(self, a, b, sigma=0.01, validate_args=False, transform=None):
        TModule.__init__(self)
        _a = torch.tensor(float(a)) if isinstance(a, Number) else a
        _a = _a.view(-1) if _a.dim() < 1 else _a
        _a, _b, _sigma = broadcast_all(_a, b, sigma)
        if not torch.all(constraints.less_than(_b).check(_a)):
            raise ValueError('must have that a < b (element-wise)')
        batch_shape, event_shape = _a.shape[:-1], _a.shape[-1:]
        self.a, self.b, self.sigma = _a, _b, _sigma
        super(SmoothedBoxPrior, self).__init__(batch_shape, event_shape, validate_args=validate_args)
        del self.a, self.b, self.sigma
        self.register_buffer('a', _a)
        self.register_buffer('b', _b)
        self.register_buffer('sigma', _sigma.clone())
        self.tails = NormalPrior(torch.zeros_like(_a), _sigma, validate_args=validate_args)
        self._transform = transform

    @property
    def _c(self):
        return (self.a + self.b) / 2

    @property
    def _r(self):
        return (self.b - self.a) / 2

    @property
    def _M(self):
        return torch.log(1 + (self.b - self.a) / (math.sqrt(2 * math.pi) * self.sigma))

    def log_prob(self, x):
        return self._log_prob(self.transform(x))

    def _log_prob(self, x):
        X = ((x - self._c).abs_() - self._r).clamp(min=0)
        return (self.tails.log_prob(X) - self._M).sum(-1)

    def rsample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        gauss_max = 1 / (math.sqrt(2 * math.pi) * self.sigma)
        gauss_weight = 1 / (gauss_max * (self.b - self.a) + 1)
        picks = torch.bernoulli(gauss_weight.expand(shape))
        uniform_eps = torch.rand(shape, dtype=self.a.dtype, device=self.a.device)
        uniform_samples = self.a + uniform_eps * (self.b - self.a)
        gaussian_eps = self.tails.rsample(sample_shape)
        gaussian_samples = gaussian_eps + torch.where(gaussian_eps < 0.0, self.a, self.b)
        return torch.where(picks > 0, gaussian_samples, uniform_samples)


class HalfNormalPrior(Prior, HalfNormal):
    """
    Half-Normal prior.
    pdf(x) = 2 * (2 * pi * scale^2)^-0.5 * exp(-x^2 / (2 * scale^2)) for x >= 0; 0 for x < 0
    where scale^2 is the variance.
    """

    def __init__(self, scale, validate_args=None, transform=None):
        TModule.__init__(self)
        HalfNormal.__init__(self, scale=scale, validate_args=validate_args)
        self._transform = transform

    def expand(self, batch_shape):
        return HalfNormal(self.scale.expand(batch_shape))


class LogNormalPrior(Prior, LogNormal):
    """
    Log Normal prior.
    """

    def __init__(self, loc, scale, validate_args=None, transform=None):
        TModule.__init__(self)
        LogNormal.__init__(self, loc=loc, scale=scale, validate_args=validate_args)
        self._transform = transform

    def expand(self, batch_shape):
        batch_shape = torch.Size(batch_shape)
        return LogNormalPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape))


class UniformPrior(Prior, Uniform):
    """
    Uniform prior.
    """

    def __init__(self, a, b, validate_args=None, transform=None):
        TModule.__init__(self)
        Uniform.__init__(self, a, b, validate_args=validate_args)
        self._transform = transform

    def expand(self, batch_shape):
        batch_shape = torch.Size(batch_shape)
        return UniformPrior(self.low.expand(batch_shape), self.high.expand(batch_shape))


class HalfCauchyPrior(Prior, HalfCauchy):
    """
    Half-Cauchy prior.
    """

    def __init__(self, scale, validate_args=None, transform=None):
        TModule.__init__(self)
        HalfCauchy.__init__(self, scale=scale, validate_args=validate_args)
        self._transform = transform

    def expand(self, batch_shape):
        return HalfCauchyPrior(self.scale.expand(batch_shape))


class GammaPrior(Prior, Gamma):
    """Gamma Prior parameterized by concentration and rate

    pdf(x) = beta^alpha / Gamma(alpha) * x^(alpha - 1) * exp(-beta * x)

    were alpha > 0 and beta > 0 are the concentration and rate parameters, respectively.
    """

    def __init__(self, concentration, rate, validate_args=False, transform=None):
        TModule.__init__(self)
        Gamma.__init__(self, concentration=concentration, rate=rate, validate_args=validate_args)
        _bufferize_attributes(self, ('concentration', 'rate'))
        self._transform = transform

    def expand(self, batch_shape):
        batch_shape = torch.Size(batch_shape)
        return GammaPrior(self.concentration.expand(batch_shape), self.rate.expand(batch_shape))

    def __call__(self, *args, **kwargs):
        return super(Gamma, self).__call__(*args, **kwargs)


MVN_LAZY_PROPERTIES = 'covariance_matrix', 'scale_tril', 'precision_matrix'


def _del_attributes(module, attributes, raise_on_error=False):
    for attr in attributes:
        try:
            delattr(module, attr)
        except AttributeError as e:
            if raise_on_error:
                raise e
    return module


class MultivariateNormalPrior(Prior, MultivariateNormal):
    """Multivariate Normal prior

    pdf(x) = det(2 * pi * Sigma)^-0.5 * exp(-0.5 * (x - mu)' Sigma^-1 (x - mu))

    where mu is the mean and Sigma > 0 is the covariance matrix.
    """

    def __init__(self, loc, covariance_matrix=None, precision_matrix=None, scale_tril=None, validate_args=False, transform=None):
        TModule.__init__(self)
        MultivariateNormal.__init__(self, loc=loc, covariance_matrix=covariance_matrix, precision_matrix=precision_matrix, scale_tril=scale_tril, validate_args=validate_args)
        _bufferize_attributes(self, ('loc', '_unbroadcasted_scale_tril'))
        self._transform = transform

    def cuda(self, device=None):
        """Applies module-level cuda() call and resets all lazy properties"""
        module = self._apply(lambda t: t)
        _del_attributes(module, MVN_LAZY_PROPERTIES)
        return module

    def cpu(self):
        """Applies module-level cpu() call and resets all lazy properties"""
        module = self._apply(lambda t: t.cpu())
        _del_attributes(module, MVN_LAZY_PROPERTIES)
        return module

    def expand(self, batch_shape):
        batch_shape = torch.Size(batch_shape)
        cov_shape = batch_shape + self.event_shape
        new_loc = self.loc.expand(batch_shape)
        new_scale_tril = self.scale_tril.expand(cov_shape)
        return MultivariateNormalPrior(loc=new_loc, scale_tril=new_scale_tril)


class WishartPrior(Prior):
    """Wishart prior over n x n positive definite matrices

    pdf(Sigma) \\sim |Sigma|^(nu - n - 1)/2 * exp(-0.5 * Trace(K^-1 Sigma))

    where nu > n - 1 are the degrees of freedom and K > 0 is the p x p scale matrix

    Reference: A. Shah, A. G. Wilson, and Z. Ghahramani. Student-t Processes as
        Alternatives to Gaussian Processes. ArXiv e-prints, Feb. 2014.
    """
    arg_constraints = {'K_inv': constraints.positive_definite, 'nu': constraints.positive}
    support = constraints.positive_definite
    _validate_args = True

    def __init__(self, nu, K, validate_args=False):
        TModule.__init__(self)
        if K.dim() < 2:
            raise ValueError('K must be at least 2-dimensional')
        n = K.shape[-1]
        if K.shape[-2] != K.shape[-1]:
            raise ValueError('K must be square')
        if isinstance(nu, Number):
            nu = torch.tensor(float(nu))
        if torch.any(nu <= n):
            raise ValueError('Must have nu > n - 1')
        self.n = torch.tensor(n, dtype=torch.long, device=nu.device)
        batch_shape = nu.shape
        event_shape = torch.Size([n, n])
        logdetK = torch.logdet(K)
        C = -(nu / 2) * (logdetK + n * math.log(2)) - torch.mvlgamma(nu / 2, n)
        K_inv = torch.inverse(K)
        self.nu = nu
        self.K_inv = K_inv
        self.C = C
        super(WishartPrior, self).__init__(batch_shape, event_shape, validate_args=validate_args)
        del self.nu, self.K_inv, self.C
        self.register_buffer('nu', nu)
        self.register_buffer('K_inv', K_inv)
        self.register_buffer('C', C)

    def log_prob(self, X):
        logdetp = torch.logdet(X)
        Kinvp = torch.matmul(self.K_inv, X)
        trKinvp = torch.diagonal(Kinvp, dim1=-2, dim2=-1).sum(-1)
        return self.C + 0.5 * (self.nu - self.n - 1) * logdetp - trKinvp


class InverseWishartPrior(Prior):
    """Inverse Wishart prior over n x n positive definite matrices

    pdf(Sigma) \\sim |Sigma|^-(nu + 2 * n)/2 * exp(-0.5 * Trace(K Sigma^-1))

    where nu > 0 are the degrees of freedom and K > 0 is the p x p scale matrix

    Reference: A. Shah, A. G. Wilson, and Z. Ghahramani. Student-t Processes as
        Alternatives to Gaussian Processes. ArXiv e-prints, Feb. 2014.
    """
    arg_constraints = {'K': constraints.positive_definite, 'nu': constraints.positive}
    support = constraints.positive_definite
    _validate_args = True

    def __init__(self, nu, K, validate_args=False):
        TModule.__init__(self)
        if K.dim() < 2:
            raise ValueError('K must be at least 2-dimensional')
        n = K.shape[-1]
        if isinstance(nu, Number):
            nu = torch.tensor(float(nu))
        if torch.any(nu <= 0):
            raise ValueError('Must have nu > 0')
        self.n = torch.tensor(n, dtype=torch.long, device=nu.device)
        batch_shape = nu.shape
        event_shape = torch.Size([n, n])
        c = (nu + n - 1) / 2
        logdetK = torch.logdet(K)
        C = c * (logdetK - n * math.log(2)) - torch.mvlgamma(c, n)
        self.nu = nu
        self.K = K
        self.C = C
        super(InverseWishartPrior, self).__init__(batch_shape, event_shape, validate_args=validate_args)
        del self.nu, self.K, self.C
        self.register_buffer('nu', nu)
        self.register_buffer('K', K)
        self.register_buffer('C', C)

    def log_prob(self, X):
        logdetp = torch.logdet(X)
        pinvK = torch.linalg.solve(X, self.K)
        trpinvK = torch.diagonal(pinvK, dim1=-2, dim2=-1).sum(-1)
        return self.C - 0.5 * ((self.nu + 2 * self.n) * logdetp + trpinvK)


class ScaleToBounds(torch.nn.Module):
    """
    Scale the input data so that it lies in between the lower and upper bounds.

    In training (`self.train()`), this module adjusts the scaling factor to the minibatch of data.
    During evaluation (`self.eval()`), this module uses the scaling factor from the previous minibatch of data.

    :param float lower_bound: lower bound of scaled data
    :param float upper_bound: upper bound of scaled data

    Example:
        >>> train_x = torch.randn(10, 5)
        >>> module = gpytorch.utils.grid.ScaleToBounds(lower_bound=-1., upper_bound=1.)
        >>>
        >>> module.train()
        >>> scaled_train_x = module(train_x)  # Data should be between -0.95 and 0.95
        >>>
        >>> module.eval()
        >>> test_x = torch.randn(10, 5)
        >>> scaled_test_x = module(test_x)  # Scaling is based on train_x
    """

    def __init__(self, lower_bound, upper_bound):
        super().__init__()
        self.lower_bound = float(lower_bound)
        self.upper_bound = float(upper_bound)
        self.register_buffer('min_val', torch.tensor(lower_bound))
        self.register_buffer('max_val', torch.tensor(upper_bound))

    def forward(self, x):
        if self.training:
            min_val = x.min()
            max_val = x.max()
            self.min_val.data = min_val
            self.max_val.data = max_val
        else:
            min_val = self.min_val
            max_val = self.max_val
            x = x.clamp(min_val, max_val)
        diff = max_val - min_val
        x = (x - min_val) * (0.95 * (self.upper_bound - self.lower_bound) / diff) + 0.95 * self.lower_bound
        return x


class NNUtil(Module):
    """
    Utility for nearest neighbor search. It would first try to use `faiss`_ (requiring separate pacakge installment)
    as the backend for better computational performance. Otherwise, `scikit-learn` would be used as it is pre-installed
    with gpytorch.

    :param int k: number of nearest neighbors
    :param int dim: dimensionality of data
    :param torch.Size batch_shape: batch shape for train data
    :param str preferred_nnlib: currently supports `faiss` and `scikit-learn` (default: faiss).
    :param torch.device device: device that the NN search will be performed on.

    Example:
        >>> train_x = torch.randn(10, 5)
        >>> nn_util = NNUtil(k=3, dim=train_x.size(-1), device=train_x.device)
        >>> nn_util.set_nn_idx(train_x)
        >>> test_x = torch.randn(2, 5)
        >>> test_nn_indices = nn_util.find_nn_idx(test_x) # finding 3 nearest neighbors for test_x
        >>> test_nn_indices = nn_util.find_nn_idx(test_x, k=2) # finding 2 nearest neighbors for test_x
        >>> sequential_nn_idx = nn_util.build_sequential_nn_idx(train_x) # build up sequential nearest neighbor
        >>>     # structure for train_x

    .. _faiss:
        https://github.com/facebookresearch/faiss
    """

    def __init__(self, k, dim, batch_shape=torch.Size([]), preferred_nnlib='faiss', device='cpu'):
        super().__init__()
        assert k > 0, f'k must be greater than 0, but got k = {k}.'
        self.k = k
        self.dim = dim
        if not isinstance(batch_shape, torch.Size):
            raise RuntimeError(f'batch_shape must be an instance of torch.Size, but got {type(batch_shape)}')
        self.batch_shape = batch_shape
        self.train_n = None
        if preferred_nnlib == 'faiss':
            try:
                self.nnlib = 'faiss'
                self.cpu()
            except ImportError:
                warnings.warn('Tried to import faiss, but failed. Falling back to scikit-learn nearest neighbor search.', ImportWarning)
                self.nnlib = 'sklearn'
                self.train_neighbors = None
        else:
            self.nnlib = 'sklearn'
            self.train_neighbors = None
        self

    def cuda(self, device=None):
        super()
        if self.nnlib == 'faiss':
            self.res = StandardGpuResources()
            self.index = [GpuIndexFlatL2(self.res, self.dim) for _ in range(self.batch_shape.numel())]
        return self

    def cpu(self):
        super().cpu()
        if self.nnlib == 'faiss':
            self.res = None
            self.index = [IndexFlatL2(self.dim) for _ in range(self.batch_shape.numel())]
        return self

    def find_nn_idx(self, test_x, k=None):
        """
        Find :math:`k` nearest neighbors for test data `test_x` among the training data stored in this utility

        :param test_x: test data, shape (... x N x D)
        :param int k: number of nearest neighbors. Default is the value used in utility initialization.
        :rtype: torch.LongTensor
        :return: the indices of nearest neighbors in the training data
        """
        assert self.train_n is not None, 'Please initialize with training data first.'
        if k is None:
            k = self.k
        else:
            assert k > 0, f'k must be greater than 0, but got k = {k}.'
        assert k <= self.train_n, f'k should be smaller than number of train data, but got k = {k}, number of train data = {self.train_n}.'
        test_x = self._expand_and_check_shape(test_x)
        test_n = test_x.shape[-2]
        test_x = test_x.view(-1, test_n, self.dim)
        nn_idx = torch.empty(self.batch_shape.numel(), test_n, k, dtype=torch.int64, device=test_x.device)
        with torch.no_grad():
            if self.nnlib == 'sklearn':
                if self.train_neighbors is None:
                    raise RuntimeError('The nearest neighbor set has not been defined. First call `set_nn_idx`')
                for i in range(self.batch_shape.numel()):
                    nn_idx_i = torch.from_numpy(self.train_neighbors[i].kneighbors(test_x[i].cpu().numpy())[1][..., :k])
                    nn_idx[i] = nn_idx_i.long()
            else:
                for i in range(self.batch_shape.numel()):
                    nn_idx[i] = self.index[i].search(test_x[i], k)[1]
        nn_idx = nn_idx.view(*self.batch_shape, test_n, k)
        return nn_idx

    def set_nn_idx(self, train_x):
        """
        Set the indices of training data to facilitate nearest neighbor search.
        This function needs to be called every time that the data changes.

        :param torch.Tensor train_x: training data points (... x N x D)
        """
        train_x = self._expand_and_check_shape(train_x)
        self.train_n = train_x.shape[-2]
        with torch.no_grad():
            if self.nnlib == 'sklearn':
                self.train_neighbors = []
                from sklearn.neighbors import NearestNeighbors
                train_x = train_x.view(-1, self.train_n, self.dim)
                for i in range(self.batch_shape.numel()):
                    x = train_x[i].cpu().numpy()
                    self.train_neighbors.append(NearestNeighbors(n_neighbors=self.k, algorithm='auto').fit(x))
            elif self.nnlib == 'faiss':
                train_x = train_x.view(-1, self.train_n, self.dim)
                for i in range(self.batch_shape.numel()):
                    self.index[i].reset()
                    self.index[i].add(train_x[i])

    def build_sequential_nn_idx(self, x):
        """
        Build the sequential :math:`k` nearest neighbor structure within training data in the following way:
        for the :math:`i`-th data point :math:`x_i`, find its :math:`k` nearest neighbors among preceding
        training data :math:`x_1, \\cdots, x_{i-1}`, for `i=k+1:N` where `N` is the size of training data.

        :param x: training data. Shape `(N, D)`
        :rtype: torch.LongTensor
        :return: indices of nearest neighbors. Shape: `(N-k, k)`
        """
        x = self._expand_and_check_shape(x)
        N = x.shape[-2]
        assert self.k < N, f'k should be smaller than number of data, but got k = {self.k}, number of data = {N}.'
        nn_idx = torch.empty(self.batch_shape.numel(), N - self.k, self.k, dtype=torch.int64)
        x_np = x.view(-1, N, self.dim).data.float().cpu().numpy()
        if self.nnlib == 'faiss':
            index = IndexFlatL2(self.dim)
            with torch.no_grad():
                if self.res is not None:
                    index = index_cpu_to_gpu(self.res, 0, index)
                for bi in range(self.batch_shape.numel()):
                    index.reset()
                    index.add(x_np[bi][:self.k])
                    for i in range(self.k, N):
                        row = x_np[bi][i][None, :]
                        nn_idx[bi][i - self.k].copy_(torch.from_numpy(index.search(row, self.k)[1][..., 0, :]).long())
                        index.add(row)
        else:
            assert self.nnlib == 'sklearn'
            from sklearn.neighbors import NearestNeighbors
            for bi in range(self.batch_shape.numel()):
                for i in range(self.k, N):
                    train_neighbors = NearestNeighbors(n_neighbors=self.k, algorithm='auto').fit(x_np[bi][:i])
                    nn_idx_i = torch.from_numpy(train_neighbors.kneighbors(x_np[bi][i][None,])[1]).squeeze()
                    nn_idx[bi][i - self.k].copy_(nn_idx_i)
        nn_idx = nn_idx.view(*self.batch_shape, N - self.k, self.k)
        return nn_idx

    def to(self, device):
        """
        Put the utility to a cpu or gpu device.

        :param torch.device device: Target device.
        """
        if str(device) == 'cpu':
            return self.cpu()
        elif 'cuda' in str(device):
            return self
        else:
            raise ValueError(f'Unknown device {device}')

    def _expand_and_check_shape(self, x):
        if len(x.shape) == 1:
            x = x.unsqueeze(-1)
        assert x.shape[:-2] == self.batch_shape, f"x's batch shape must be equal to self.batch_shape, but got x's batch shape={x.shape[:-2]}, self.batch_shape={self.batch_shape}."
        assert x.shape[-1] == self.dim, f"x's dim must be equal to self.dim, but got x's dim = {x.shape[-1]}, self.dim = {self.dim}"
        return x


def _pad_with_singletons(obj: torch.Tensor, num_singletons_before: int=0, num_singletons_after: int=0):
    """
    Pad obj with singleton dimensions on the left and right

    Example:
        >>> x = torch.randn(10, 5)
        >>> _pad_width_singletons(x, 2, 3).shape
        >>> # [1, 1, 10, 5, 1, 1, 1]
    """
    new_shape = [1] * num_singletons_before + list(obj.shape) + [1] * num_singletons_after
    return obj.view(*new_shape)


class GaussHermiteQuadrature1D(Module):
    """
    Implements Gauss-Hermite quadrature for integrating a function with respect to several 1D Gaussian distributions
    in batch mode. Within GPyTorch, this is useful primarily for computing expected log likelihoods for variational
    inference.

    This is implemented as a Module because Gauss-Hermite quadrature has a set of locations and weights that it
    should initialize one time, but that should obey parent calls to .cuda(), .double() etc.
    """

    def __init__(self, num_locs=None):
        super().__init__()
        if num_locs is None:
            num_locs = settings.num_gauss_hermite_locs.value()
        self.num_locs = num_locs
        locations, weights = self._locs_and_weights(num_locs)
        self.locations = locations
        self.weights = weights

    def _apply(self, fn):
        self.locations = fn(self.locations)
        self.weights = fn(self.weights)
        return super(GaussHermiteQuadrature1D, self)._apply(fn)

    def _locs_and_weights(self, num_locs):
        """
        Get locations and weights for Gauss-Hermite quadrature. Note that this is **not** intended to be used
        externally, because it directly creates tensors with no knowledge of a device or dtype to cast to.

        Instead, create a GaussHermiteQuadrature1D object and get the locations and weights from buffers.
        """
        locations, weights = np.polynomial.hermite.hermgauss(num_locs)
        locations = torch.Tensor(locations)
        weights = torch.Tensor(weights)
        return locations, weights

    def forward(self, func, gaussian_dists):
        """
        Runs Gauss-Hermite quadrature on the callable func, integrating against the Gaussian distributions specified
        by gaussian_dists.

        Args:
            - func (callable): Function to integrate
            - gaussian_dists (Distribution): Either a MultivariateNormal whose covariance is assumed to be diagonal
                or a :obj:`torch.distributions.Normal`.
        Returns:
            - Result of integrating func against each univariate Gaussian in gaussian_dists.
        """
        means = gaussian_dists.mean
        variances = gaussian_dists.variance
        locations = _pad_with_singletons(self.locations, num_singletons_before=0, num_singletons_after=means.dim())
        shifted_locs = torch.sqrt(2.0 * variances) * locations + means
        log_probs = func(shifted_locs)
        weights = _pad_with_singletons(self.weights, num_singletons_before=0, num_singletons_after=log_probs.dim() - 1)
        res = 1 / math.sqrt(math.pi) * (log_probs * weights)
        res = res.sum(tuple(range(self.locations.dim())))
        return res


class _VariationalDistribution(Module, ABC):
    """
    Abstract base class for all Variational Distributions.

    :ivar torch.dtype dtype: The dtype of the VariationalDistribution parameters
    :ivar torch.dtype device: The device of the VariationalDistribution parameters
    """

    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=0.001):
        super().__init__()
        self.num_inducing_points = num_inducing_points
        self.batch_shape = batch_shape
        self.mean_init_std = mean_init_std

    @property
    def device(self):
        return next(self.parameters()).device

    @property
    def dtype(self):
        return next(self.parameters()).dtype

    def forward(self):
        """
        Constructs and returns the variational distribution

        :rtype: ~gpytorch.distributions.MultivariateNormal
        :return: The distribution :math:`q(\\mathbf u)`
        """
        raise NotImplementedError

    def shape(self) ->torch.Size:
        """
        Event + batch shape of VariationalDistribution object
        :rtype: torch.Size
        """
        return torch.Size([*self.batch_shape, self.num_inducing_points])

    @abstractmethod
    def initialize_variational_distribution(self, prior_dist):
        """
        Method for initializing the variational distribution, based on the prior distribution.

        :param ~gpytorch.distributions.Distribution prior_dist: The prior distribution :math:`p(\\mathbf u)`.
        """
        raise NotImplementedError

    def __call__(self):
        return self.forward()


class GPInputWarning(UserWarning):
    """
    Warning thrown when a GP model receives an unexpected input.
    For example, when an :obj:`~gpytorch.models.ExactGP` in eval mode receives the training data as input.
    """
    pass


def add_to_cache(obj, name, val, *args, **kwargs):
    """Add a result to the cache of an object (honoring calling args)."""
    return _add_to_cache(obj, name, val, *args, kwargs_pkl=pickle.dumps(kwargs))


def clear_cache_hook(module, *args, **kwargs):
    module._memoize_cache = {}


def prediction_strategy(train_inputs, train_prior_dist, train_labels, likelihood):
    train_train_covar = train_prior_dist.lazy_covariance_matrix
    if isinstance(train_train_covar, LazyEvaluatedKernelTensor):
        cls = train_train_covar.kernel.prediction_strategy
    else:
        cls = DefaultPredictionStrategy
    return cls(train_inputs, train_prior_dist, train_labels, likelihood)


class ExactGP(GP):
    """
    The base class for any Gaussian process latent function to be used in conjunction
    with exact inference.

    :param torch.Tensor train_inputs: (size n x d) The training features :math:`\\mathbf X`.
    :param torch.Tensor train_targets: (size n) The training targets :math:`\\mathbf y`.
    :param ~gpytorch.likelihoods.GaussianLikelihood likelihood: The Gaussian likelihood that defines
        the observational distribution. Since we're using exact inference, the likelihood must be Gaussian.

    The :meth:`forward` function should describe how to compute the prior latent distribution
    on a given input. Typically, this will involve a mean and kernel function.
    The result must be a :obj:`~gpytorch.distributions.MultivariateNormal`.

    Calling this model will return the posterior of the latent Gaussian process when conditioned
    on the training data. The output will be a :obj:`~gpytorch.distributions.MultivariateNormal`.

    Example:
        >>> class MyGP(gpytorch.models.ExactGP):
        >>>     def __init__(self, train_x, train_y, likelihood):
        >>>         super().__init__(train_x, train_y, likelihood)
        >>>         self.mean_module = gpytorch.means.ZeroMean()
        >>>         self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
        >>>
        >>>     def forward(self, x):
        >>>         mean = self.mean_module(x)
        >>>         covar = self.covar_module(x)
        >>>         return gpytorch.distributions.MultivariateNormal(mean, covar)
        >>>
        >>> # train_x = ...; train_y = ...
        >>> likelihood = gpytorch.likelihoods.GaussianLikelihood()
        >>> model = MyGP(train_x, train_y, likelihood)
        >>>
        >>> # test_x = ...;
        >>> model(test_x)  # Returns the GP latent function at test_x
        >>> likelihood(model(test_x))  # Returns the (approximate) predictive posterior distribution at test_x
    """

    def __init__(self, train_inputs, train_targets, likelihood):
        if train_inputs is not None and torch.is_tensor(train_inputs):
            train_inputs = train_inputs,
        if train_inputs is not None and not all(torch.is_tensor(train_input) for train_input in train_inputs):
            raise RuntimeError('Train inputs must be a tensor, or a list/tuple of tensors')
        if not isinstance(likelihood, _GaussianLikelihoodBase):
            raise RuntimeError('ExactGP can only handle Gaussian likelihoods')
        super(ExactGP, self).__init__()
        if train_inputs is not None:
            self.train_inputs = tuple(tri.unsqueeze(-1) if tri.ndimension() == 1 else tri for tri in train_inputs)
            self.train_targets = train_targets
        else:
            self.train_inputs = None
            self.train_targets = None
        self.likelihood = likelihood
        self.prediction_strategy = None

    @property
    def train_targets(self):
        return self._train_targets

    @train_targets.setter
    def train_targets(self, value):
        object.__setattr__(self, '_train_targets', value)

    def _apply(self, fn):
        if self.train_inputs is not None:
            self.train_inputs = tuple(fn(train_input) for train_input in self.train_inputs)
            self.train_targets = fn(self.train_targets)
        return super(ExactGP, self)._apply(fn)

    def _clear_cache(self):
        self.prediction_strategy = None

    def local_load_samples(self, samples_dict, memo, prefix):
        """
        Replace the model's learned hyperparameters with samples from a posterior distribution.
        """
        num_samples = next(iter(samples_dict.values())).size(0)
        self.train_inputs = tuple(tri.unsqueeze(0).expand(num_samples, *tri.shape) for tri in self.train_inputs)
        self.train_targets = self.train_targets.unsqueeze(0).expand(num_samples, *self.train_targets.shape)
        super().local_load_samples(samples_dict, memo, prefix)

    def set_train_data(self, inputs=None, targets=None, strict=True):
        """
        Set training data (does not re-fit model hyper-parameters).

        :param torch.Tensor inputs: The new training inputs.
        :param torch.Tensor targets: The new training targets.
        :param bool strict: (default True) If `True`, the new inputs and
            targets must have the same shape, dtype, and device
            as the current inputs and targets. Otherwise, any shape/dtype/device are allowed.
        """
        if inputs is not None:
            if torch.is_tensor(inputs):
                inputs = inputs,
            inputs = tuple(input_.unsqueeze(-1) if input_.ndimension() == 1 else input_ for input_ in inputs)
            if strict:
                for input_, t_input in zip(inputs, self.train_inputs or (None,)):
                    for attr in {'shape', 'dtype', 'device'}:
                        expected_attr = getattr(t_input, attr, None)
                        found_attr = getattr(input_, attr, None)
                        if expected_attr != found_attr:
                            msg = 'Cannot modify {attr} of inputs (expected {e_attr}, found {f_attr}).'
                            msg = msg.format(attr=attr, e_attr=expected_attr, f_attr=found_attr)
                            raise RuntimeError(msg)
            self.train_inputs = inputs
        if targets is not None:
            if strict:
                for attr in {'shape', 'dtype', 'device'}:
                    expected_attr = getattr(self.train_targets, attr, None)
                    found_attr = getattr(targets, attr, None)
                    if expected_attr != found_attr:
                        msg = 'Cannot modify {attr} of targets (expected {e_attr}, found {f_attr}).'
                        msg = msg.format(attr=attr, e_attr=expected_attr, f_attr=found_attr)
                        raise RuntimeError(msg)
            self.train_targets = targets
        self.prediction_strategy = None

    def get_fantasy_model(self, inputs, targets, **kwargs):
        """
        Returns a new GP model that incorporates the specified inputs and targets as new training data.

        Using this method is more efficient than updating with `set_train_data` when the number of inputs is relatively
        small, because any computed test-time caches will be updated in linear time rather than computed from scratch.

        .. note::
            If `targets` is a batch (e.g. `b x m`), then the GP returned from this method will be a batch mode GP.
            If `inputs` is of the same (or lesser) dimension as `targets`, then it is assumed that the fantasy points
            are the same for each target batch.

        :param torch.Tensor inputs: (`b1 x ... x bk x m x d` or `f x b1 x ... x bk x m x d`) Locations of fantasy
            observations.
        :param torch.Tensor targets: (`b1 x ... x bk x m` or `f x b1 x ... x bk x m`) Labels of fantasy observations.
        :return: An `ExactGP` model with `n + m` training examples, where the `m` fantasy examples have been added
            and all test-time caches have been updated.
        :rtype: ~gpytorch.models.ExactGP
        """
        if self.prediction_strategy is None:
            raise RuntimeError('Fantasy observations can only be added after making predictions with a model so that all test independent caches exist. Call the model on some data first!')
        model_batch_shape = self.train_inputs[0].shape[:-2]
        if self.train_targets.dim() > len(model_batch_shape) + 1:
            raise RuntimeError('Cannot yet add fantasy observations to multitask GPs, but this is coming soon!')
        if not isinstance(inputs, list):
            inputs = [inputs]
        inputs = [(i.unsqueeze(-1) if i.ndimension() == 1 else i) for i in inputs]
        target_batch_shape = targets.shape[:-1]
        input_batch_shape = inputs[0].shape[:-2]
        tbdim, ibdim = len(target_batch_shape), len(input_batch_shape)
        if not (tbdim == ibdim + 1 or tbdim == ibdim):
            raise RuntimeError(f'Unsupported batch shapes: The target batch shape ({target_batch_shape}) must have either the same dimension as or one more dimension than the input batch shape ({input_batch_shape})')
        try:
            torch.broadcast_shapes(model_batch_shape, target_batch_shape)
        except RuntimeError:
            raise RuntimeError(f'Model batch shape ({model_batch_shape}) and target batch shape ({target_batch_shape}) are not broadcastable.')
        if len(model_batch_shape) > len(input_batch_shape):
            input_batch_shape = model_batch_shape
        if len(model_batch_shape) > len(target_batch_shape):
            target_batch_shape = model_batch_shape
        train_inputs = [tin.expand(input_batch_shape + tin.shape[-2:]) for tin in self.train_inputs]
        train_targets = self.train_targets.expand(target_batch_shape + self.train_targets.shape[-1:])
        full_inputs = [torch.cat([train_input, input.expand(input_batch_shape + input.shape[-2:])], dim=-2) for train_input, input in zip(train_inputs, inputs)]
        full_targets = torch.cat([train_targets, targets.expand(target_batch_shape + targets.shape[-1:])], dim=-1)
        try:
            fantasy_kwargs = {'noise': kwargs.pop('noise')}
        except KeyError:
            fantasy_kwargs = {}
        full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
        old_pred_strat = self.prediction_strategy
        old_train_inputs = self.train_inputs
        old_train_targets = self.train_targets
        old_likelihood = self.likelihood
        self.prediction_strategy = None
        self.train_inputs = None
        self.train_targets = None
        self.likelihood = None
        new_model = deepcopy(self)
        self.prediction_strategy = old_pred_strat
        self.train_inputs = old_train_inputs
        self.train_targets = old_train_targets
        self.likelihood = old_likelihood
        new_model.likelihood = old_likelihood.get_fantasy_likelihood(**fantasy_kwargs)
        new_model.prediction_strategy = old_pred_strat.get_fantasy_strategy(inputs, targets, full_inputs, full_targets, full_output, **fantasy_kwargs)
        if tbdim == ibdim + 1:
            new_model.train_inputs = [fi.expand(target_batch_shape + fi.shape[-2:]) for fi in full_inputs]
        else:
            new_model.train_inputs = full_inputs
        new_model.train_targets = full_targets
        return new_model

    def __call__(self, *args, **kwargs):
        train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
        inputs = [(i.unsqueeze(-1) if i.ndimension() == 1 else i) for i in args]
        if self.training:
            if self.train_inputs is None:
                raise RuntimeError('train_inputs, train_targets cannot be None in training mode. Call .eval() for prior predictions, or call .set_train_data() to add training data.')
            if settings.debug.on():
                if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
                    raise RuntimeError('You must train on the training inputs!')
            res = super().__call__(*inputs, **kwargs)
            return res
        elif settings.prior_mode.on() or self.train_inputs is None or self.train_targets is None:
            full_inputs = args
            full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
            if settings.debug().on():
                if not isinstance(full_output, MultivariateNormal):
                    raise RuntimeError('ExactGP.forward must return a MultivariateNormal')
            return full_output
        else:
            if settings.debug.on():
                if all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
                    warnings.warn('The input matches the stored training data. Did you forget to call model.train()?', GPInputWarning)
            if self.prediction_strategy is None:
                train_output = super().__call__(*train_inputs, **kwargs)
                self.prediction_strategy = prediction_strategy(train_inputs=train_inputs, train_prior_dist=train_output, train_labels=self.train_targets, likelihood=self.likelihood)
            full_inputs = []
            batch_shape = train_inputs[0].shape[:-2]
            for train_input, input in zip(train_inputs, inputs):
                if batch_shape != train_input.shape[:-2]:
                    batch_shape = torch.broadcast_shapes(batch_shape, train_input.shape[:-2])
                    train_input = train_input.expand(*batch_shape, *train_input.shape[-2:])
                if batch_shape != input.shape[:-2]:
                    batch_shape = torch.broadcast_shapes(batch_shape, input.shape[:-2])
                    train_input = train_input.expand(*batch_shape, *train_input.shape[-2:])
                    input = input.expand(*batch_shape, *input.shape[-2:])
                full_inputs.append(torch.cat([train_input, input], dim=-2))
            full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
            if settings.debug().on():
                if not isinstance(full_output, MultivariateNormal):
                    raise RuntimeError('ExactGP.forward must return a MultivariateNormal')
            full_mean, full_covar = full_output.loc, full_output.lazy_covariance_matrix
            batch_shape = full_output.batch_shape
            joint_shape = full_output.event_shape
            tasks_shape = joint_shape[1:]
            test_shape = torch.Size([joint_shape[0] - self.prediction_strategy.train_shape[0], *tasks_shape])
            with settings.cg_tolerance(settings.eval_cg_tolerance.value()):
                predictive_mean, predictive_covar = self.prediction_strategy.exact_prediction(full_mean, full_covar)
            predictive_mean = predictive_mean.view(*batch_shape, *test_shape).contiguous()
            return full_output.__class__(predictive_mean, predictive_covar)


class _BaseExactGP(ExactGP):

    def __init__(self, train_inputs, train_targets, likelihood, mean_module, covar_module):
        super().__init__(train_inputs, train_targets, likelihood)
        self.mean_module = mean_module
        self.covar_module = covar_module

    def forward(self, x):
        mean = self.mean_module(x)
        covar = self.covar_module(x)
        return MultivariateNormal(mean, covar)


def _add_cache_hook(tsr, pred_strat):
    if tsr.grad_fn is not None:
        wrapper = functools.partial(clear_cache_hook, pred_strat)
        functools.update_wrapper(wrapper, clear_cache_hook)
        tsr.grad_fn.register_hook(wrapper)
    return tsr


class _VariationalStrategy(Module, ABC):
    """
    Abstract base class for all Variational Strategies.
    """
    has_fantasy_strategy = False

    def __init__(self, model, inducing_points, variational_distribution, learn_inducing_locations=True, jitter_val=None):
        super().__init__()
        if jitter_val is None:
            self.jitter_val = settings.variational_cholesky_jitter.value(inducing_points.dtype)
        else:
            self.jitter_val = jitter_val
        object.__setattr__(self, 'model', model)
        inducing_points = inducing_points.clone()
        if inducing_points.dim() == 1:
            inducing_points = inducing_points.unsqueeze(-1)
        if learn_inducing_locations:
            self.register_parameter(name='inducing_points', parameter=torch.nn.Parameter(inducing_points))
        else:
            self.register_buffer('inducing_points', inducing_points)
        self._variational_distribution = variational_distribution
        self.register_buffer('variational_params_initialized', torch.tensor(0))

    def _clear_cache(self):
        clear_cache_hook(self)

    def _expand_inputs(self, x, inducing_points):
        """
        Pre-processing step in __call__ to make x the same batch_shape as the inducing points
        """
        batch_shape = torch.broadcast_shapes(inducing_points.shape[:-2], x.shape[:-2])
        inducing_points = inducing_points.expand(*batch_shape, *inducing_points.shape[-2:])
        x = x.expand(*batch_shape, *x.shape[-2:])
        return x, inducing_points

    @abstractproperty
    @cached(name='prior_distribution_memo')
    def prior_distribution(self):
        """
        The :func:`~gpytorch.variational.VariationalStrategy.prior_distribution` method determines how to compute the
        GP prior distribution of the inducing points, e.g. :math:`p(u) \\sim N(\\mu(X_u), K(X_u, X_u))`. Most commonly,
        this is done simply by calling the user defined GP prior on the inducing point data directly.

        :rtype: :obj:`~gpytorch.distributions.MultivariateNormal`
        :return: The distribution :math:`p( \\mathbf u)`
        """
        raise NotImplementedError

    @property
    @cached(name='variational_distribution_memo')
    def variational_distribution(self):
        return self._variational_distribution()

    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None, **kwargs):
        """
        The :func:`~gpytorch.variational.VariationalStrategy.forward` method determines how to marginalize out the
        inducing point function values. Specifically, forward defines how to transform a variational distribution
        over the inducing point values, :math:`q(u)`, in to a variational distribution over the function values at
        specified locations x, :math:`q(f|x)`, by integrating :math:`\\int p(f|x, u)q(u)du`

        :param torch.Tensor x: Locations :math:`\\mathbf X` to get the
            variational posterior of the function values at.
        :param torch.Tensor inducing_points: Locations :math:`\\mathbf Z` of the inducing points
        :param torch.Tensor inducing_values: Samples of the inducing function values :math:`\\mathbf u`
            (or the mean of the distribution :math:`q(\\mathbf u)` if q is a Gaussian.
        :param ~linear_operator.operators.LinearOperator variational_inducing_covar: If
            the distribuiton :math:`q(\\mathbf u)` is
            Gaussian, then this variable is the covariance matrix of that Gaussian.
            Otherwise, it will be None.

        :rtype: :obj:`~gpytorch.distributions.MultivariateNormal`
        :return: The distribution :math:`q( \\mathbf f(\\mathbf X))`
        """
        raise NotImplementedError

    def kl_divergence(self):
        """
        Compute the KL divergence between the variational inducing distribution :math:`q(\\mathbf u)`
        and the prior inducing distribution :math:`p(\\mathbf u)`.

        :rtype: torch.Tensor
        """
        with settings.max_preconditioner_size(0):
            kl_divergence = torch.distributions.kl.kl_divergence(self.variational_distribution, self.prior_distribution)
        return kl_divergence

    @cached(name='amortized_exact_gp')
    def amortized_exact_gp(self, mean_module=None, covar_module=None):
        mean_module = self.model.mean_module if mean_module is None else mean_module
        covar_module = self.model.covar_module if covar_module is None else covar_module
        with torch.no_grad():
            pseudo_target_covar, pseudo_target_mean = self.pseudo_points
            pseudo_inputs = self.inducing_points.detach()
            if pseudo_inputs.ndim < pseudo_target_mean.ndim:
                pseudo_inputs = pseudo_inputs.expand(*pseudo_target_mean.shape[:-2], *pseudo_inputs.shape)
            new_covar_module = deepcopy(covar_module)
            pseudo_target_mean = pseudo_target_mean.squeeze() + mean_module(pseudo_inputs)
            inducing_exact_model = _BaseExactGP(pseudo_inputs, pseudo_target_mean, mean_module=deepcopy(mean_module), covar_module=new_covar_module, likelihood=deepcopy(self.model.likelihood))
            faked_points = torch.randn(*pseudo_target_mean.shape[:-2], 1, pseudo_inputs.shape[-1], device=pseudo_inputs.device, dtype=pseudo_inputs.dtype)
            inducing_exact_model.eval()
            _ = inducing_exact_model(faked_points)
            pred_strat = inducing_exact_model.prediction_strategy
            pred_strat._memoize_cache = {}
            with torch.no_grad():
                updated_lik_train_train_covar = pred_strat.train_prior_dist.lazy_covariance_matrix + pseudo_target_covar
                pred_strat.lik_train_train_covar = updated_lik_train_train_covar
            train_mean = inducing_exact_model.mean_module(*inducing_exact_model.train_inputs)
            train_labels_offset = (inducing_exact_model.prediction_strategy.train_labels - train_mean).unsqueeze(-1)
            mean_cache = updated_lik_train_train_covar.solve(train_labels_offset).squeeze(-1)
            mean_cache = _add_cache_hook(mean_cache, inducing_exact_model.prediction_strategy)
            add_to_cache(pred_strat, 'mean_cache', mean_cache)
            inducing_exact_model.prediction_strategy = pred_strat
        return inducing_exact_model

    def pseudo_points(self):
        raise NotImplementedError('Each variational strategy must implement its own pseudo points method')

    def get_fantasy_model(self, inputs, targets, mean_module=None, covar_module=None, **kwargs):
        """
        Performs the online variational conditioning (OVC) strategy of Maddox et al, '21 to return
        an exact GP model that incorporates the inputs and targets alongside the variational model's inducing
        points and targets.

        Currently, instead of directly updating the variational parameters (and inducing points), we instead
        return an ExactGP model rather than an updated variational GP model. This is done primarily for
        numerical stability.

        Unlike the ExactGP's call for get_fantasy_model, we enable options for mean_module and covar_module
        that allow specification of the mean / covariance. We expect that either the mean and covariance
        modules are attributes of the model itself called mean_module and covar_module respectively OR that you
        pass them into this method explicitly.

        :param torch.Tensor inputs: (`b1 x ... x bk x m x d` or `f x b1 x ... x bk x m x d`) Locations of fantasy
            observations.
        :param torch.Tensor targets: (`b1 x ... x bk x m` or `f x b1 x ... x bk x m`) Labels of fantasy observations.
        :param torch.nn.Module mean_module: torch module describing the mean function of the GP model. Optional if
            `mean_module` is already an attribute of the variational GP.
        :param torch.nn.Module covar_module: torch module describing the covariance function of the GP model. Optional
            if `covar_module` is already an attribute of the variational GP.
        :return: An `ExactGP` model with `k + m` training examples, where the `m` fantasy examples have been added
            and all test-time caches have been updated. We assume that there are `k` inducing points in this variational
            GP. Note that we return an `ExactGP` rather than a variational GP.
        :rtype: ~gpytorch.models.ExactGP

        Reference: "Conditioning Sparse Variational Gaussian Processes for Online Decision-Making,"
            Maddox, Stanton, Wilson, NeurIPS, '21
            https://papers.nips.cc/paper/2021/hash/325eaeac5bef34937cfdc1bd73034d17-Abstract.html
        """
        if not self.has_fantasy_strategy:
            raise NotImplementedError('No fantasy model support for ', self.__name__, '. Only VariationalStrategy and UnwhitenedVariationalStrategy are currently supported.')
        if not isinstance(self.model.likelihood, GaussianLikelihood):
            raise NotImplementedError('No fantasy model support for ', self.model.likelihood, '. Only GaussianLikelihoods are currently supported.')
        if mean_module is None:
            mean_module = getattr(self.model, 'mean_module', None)
            if mean_module is None:
                raise ModuleNotFoundError('Either you must provide a mean_module as input to get_fantasy_model', 'or it must be an attribute of the model called mean_module.')
        if covar_module is None:
            covar_module = getattr(self.model, 'covar_module', None)
            if covar_module is None:
                raise ModuleNotFoundError('Either you must provide a covar_module as input to get_fantasy_model', 'or it must be an attribute of the model called covar_module.')
        inducing_exact_model = self.amortized_exact_gp(mean_module=mean_module, covar_module=covar_module)
        fantasy_model = inducing_exact_model.get_fantasy_model(inputs, targets, **kwargs)
        fant_pred_strat = fantasy_model.prediction_strategy
        train_mean = fantasy_model.mean_module(*fantasy_model.train_inputs)
        train_labels_offset = (fant_pred_strat.train_labels - train_mean).unsqueeze(-1)
        fantasy_lik_train_root_inv = fant_pred_strat.lik_train_train_covar.root_inv_decomposition()
        mean_cache = fantasy_lik_train_root_inv.matmul(train_labels_offset).squeeze(-1)
        mean_cache = _add_cache_hook(mean_cache, fant_pred_strat)
        add_to_cache(fant_pred_strat, 'mean_cache', mean_cache)
        fantasy_model.prediction_strategy = fant_pred_strat
        return fantasy_model

    def __call__(self, x, prior=False, **kwargs):
        if prior:
            return self.model.forward(x, **kwargs)
        if self.training:
            self._clear_cache()
        if not self.variational_params_initialized.item():
            prior_dist = self.prior_distribution
            self._variational_distribution.initialize_variational_distribution(prior_dist)
            self.variational_params_initialized.fill_(1)
        inducing_points = self.inducing_points
        if inducing_points.shape[:-2] != x.shape[:-2]:
            x, inducing_points = self._expand_inputs(x, inducing_points)
        variational_dist_u = self.variational_distribution
        if isinstance(variational_dist_u, MultivariateNormal):
            return super().__call__(x, inducing_points, inducing_values=variational_dist_u.mean, variational_inducing_covar=variational_dist_u.lazy_covariance_matrix, **kwargs)
        elif isinstance(variational_dist_u, Delta):
            return super().__call__(x, inducing_points, inducing_values=variational_dist_u.mean, variational_inducing_covar=None, **kwargs)
        else:
            raise RuntimeError(f'Invalid variational distribuition ({type(variational_dist_u)}). Expected a multivariate normal or a delta distribution.')


class CholeskyVariationalDistribution(_VariationalDistribution):
    """
    A :obj:`~gpytorch.variational._VariationalDistribution` that is defined to be a multivariate normal distribution
    with a full covariance matrix.

    The most common way this distribution is defined is to parameterize it in terms of a mean vector and a covariance
    matrix. In order to ensure that the covariance matrix remains positive definite, we only consider the lower
    triangle.

    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
        should be this size, and the variational covariance matrix should have this many rows and columns.
    :param batch_shape: Specifies an optional batch size
        for the variational parameters. This is useful for example when doing additive variational inference.
    :type batch_shape: :obj:`torch.Size`, optional
    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
    """

    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=0.001, **kwargs):
        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
        mean_init = torch.zeros(num_inducing_points)
        covar_init = torch.eye(num_inducing_points, num_inducing_points)
        mean_init = mean_init.repeat(*batch_shape, 1)
        covar_init = covar_init.repeat(*batch_shape, 1, 1)
        self.register_parameter(name='variational_mean', parameter=torch.nn.Parameter(mean_init))
        self.register_parameter(name='chol_variational_covar', parameter=torch.nn.Parameter(covar_init))

    def forward(self):
        chol_variational_covar = self.chol_variational_covar
        dtype = chol_variational_covar.dtype
        device = chol_variational_covar.device
        lower_mask = torch.ones(self.chol_variational_covar.shape[-2:], dtype=dtype, device=device).tril(0)
        chol_variational_covar = TriangularLinearOperator(chol_variational_covar.mul(lower_mask))
        variational_covar = CholLinearOperator(chol_variational_covar)
        return MultivariateNormal(self.variational_mean, variational_covar)

    def initialize_variational_distribution(self, prior_dist):
        self.variational_mean.data.copy_(prior_dist.mean)
        self.variational_mean.data.add_(torch.randn_like(prior_dist.mean), alpha=self.mean_init_std)
        self.chol_variational_covar.data.copy_(prior_dist.lazy_covariance_matrix.cholesky().to_dense())


def _phi_for_cholesky_(A):
    """Modifies A to be the phi function used in differentiating through Cholesky"""
    A.tril_().diagonal(offset=0, dim1=-2, dim2=-1).mul_(0.5)
    return A


def _cholesky_backward(dout_dL, L, L_inverse):
    A = L.transpose(-1, -2) @ dout_dL
    phi = _phi_for_cholesky_(A)
    grad_input = L_inverse.transpose(-1, -2) @ phi @ L_inverse
    return grad_input.add(grad_input.transpose(-1, -2)).mul_(0.5)


def _triangular_inverse(A, upper=False):
    eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)
    return torch.linalg.solve_triangular(A, eye, upper=upper)


class _NaturalToMuVarSqrt(torch.autograd.Function):

    @staticmethod
    def _forward(nat_mean, nat_covar):
        try:
            L_inv = psd_safe_cholesky(-2.0 * nat_covar, upper=False)
        except RuntimeError as e:
            if str(e).startswith('cholesky'):
                raise RuntimeError('Non-negative-definite natural covariance. You probably updated it using an optimizer other than gpytorch.optim.NGD (such as Adam). This is not supported.')
            else:
                raise e
        L = _triangular_inverse(L_inv, upper=False)
        S = L.transpose(-1, -2) @ L
        mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)
        return mu, psd_safe_cholesky(S, upper=False)

    @staticmethod
    def forward(ctx, nat_mean, nat_covar):
        mu, L = _NaturalToMuVarSqrt._forward(nat_mean, nat_covar)
        ctx.save_for_backward(mu, L)
        return mu, L

    @staticmethod
    def _backward(dout_dmu, dout_dL, mu, L, C):
        """Calculate dout/d(eta1, eta2), which are:
        eta1 = mu
        eta2 = mu*mu^T + LL^T = mu*mu^T + Sigma

        Thus:
        dout/deta1 = dout/dmu + dout/dL dL/deta1
        dout/deta2 = dout/dL dL/deta1

        For L = chol(eta2 - eta1*eta1^T).
        dout/dSigma = _cholesky_backward(dout/dL, L)
        dout/deta2 = dout/dSigma
        dSigma/deta1 = -2* (dout/dSigma) mu
        """
        dout_dSigma = _cholesky_backward(dout_dL, L, C)
        dout_deta1 = dout_dmu - 2 * (dout_dSigma @ mu.unsqueeze(-1)).squeeze(-1)
        return dout_deta1, dout_dSigma

    @staticmethod
    def backward(ctx, dout_dmu, dout_dL):
        """Calculates the natural gradient with respect to nat_mean, nat_covar"""
        mu, L = ctx.saved_tensors
        C = _triangular_inverse(L, upper=False)
        return _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)


class _NaturalVariationalDistribution(_VariationalDistribution, abc.ABC):
    """Any :obj:`~gpytorch.variational._VariationalDistribution` which calculates
    natural gradients with respect to its parameters.
    """
    pass


class NaturalVariationalDistribution(_NaturalVariationalDistribution):
    """A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,
    parameterized by **natural** parameters.

    .. note::
       The :obj:`~gpytorch.variational.NaturalVariationalDistribution` can only
       be used with :obj:`gpytorch.optim.NGD`, or other optimizers
       that follow exactly the gradient direction. Failure to do so will cause
       the natural matrix :math:`\\mathbf \\Theta_\\text{mat}` to stop being
       positive definite, and a :obj:`~RuntimeError` will be raised.

    .. seealso::
        The `natural gradient descent tutorial
        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
        for use instructions.

        The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` for
        a more numerically stable parameterization, at the cost of needing more
        iterations to make variational regression converge.

    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
        should be this size, and the variational covariance matrix should have this many rows and columns.
    :param batch_shape: Specifies an optional batch size
        for the variational parameters. This is useful for example when doing additive variational inference.
    :type batch_shape: :obj:`torch.Size`, optional
    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.

    """

    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=0.001, **kwargs):
        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
        scaled_mean_init = torch.zeros(num_inducing_points)
        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points).mul(-0.5)
        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)
        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)
        self.register_parameter(name='natural_vec', parameter=torch.nn.Parameter(scaled_mean_init))
        self.register_parameter(name='natural_mat', parameter=torch.nn.Parameter(neg_prec_init))

    def forward(self):
        mean, chol_covar = _NaturalToMuVarSqrt.apply(self.natural_vec, self.natural_mat)
        res = MultivariateNormal(mean, CholLinearOperator(TriangularLinearOperator(chol_covar)))
        return res

    def initialize_variational_distribution(self, prior_dist):
        prior_prec = prior_dist.covariance_matrix.inverse()
        prior_mean = prior_dist.mean
        noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)
        self.natural_vec.data.copy_((prior_prec @ prior_mean.unsqueeze(-1)).squeeze(-1).add_(noise))
        self.natural_mat.data.copy_(prior_prec.mul(-0.5))


class _NgdInterpTerms(torch.autograd.Function):
    """
    This function takes in

        - the kernel interpolation term K_ZZ^{-1/2} k_ZX
        - the natural parameters of the variational distribution

    and returns

        - the predictive distribution mean/covariance
        - the inducing KL divergence KL( q(u) || p(u))

    However, the gradients will be with respect to the **cannonical parameters**
    of the variational distribution, rather than the **natural parameters**.
    This corresponds to performing natural gradient descent on the variational distribution.
    """

    @staticmethod
    def forward(ctx, interp_term: torch.Tensor, natural_vec: torch.Tensor, natural_mat: torch.Tensor) ->Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        prec = natural_mat.mul(-2.0)
        diag = prec.diagonal(dim1=-1, dim2=-2).unsqueeze(-1)
        batch_shape = torch.broadcast_shapes(interp_term.shape[:-2], natural_vec.shape[:-1])
        expanded_interp_term = interp_term.expand(*batch_shape, *interp_term.shape[-2:])
        expanded_natural_vec = natural_vec.expand(*batch_shape, natural_vec.size(-1))
        solves = linear_cg(prec.matmul, torch.cat([expanded_natural_vec.unsqueeze(-1), expanded_interp_term], dim=-1), n_tridiag=0, max_iter=settings.max_cg_iterations.value(), tolerance=min(settings.eval_cg_tolerance.value(), settings.cg_tolerance.value()), max_tridiag_iter=settings.max_lanczos_quadrature_iterations.value(), preconditioner=lambda x: x / diag)
        expec_vec = solves[..., 0]
        s_times_interp_term = solves[..., 1:]
        interp_mean = (s_times_interp_term.transpose(-1, -2) @ natural_vec.unsqueeze(-1)).squeeze(-1)
        interp_var = (s_times_interp_term * interp_term).sum(dim=-2)
        kl_div = torch.zeros_like(interp_mean[..., 0])
        ctx.save_for_backward(interp_term, s_times_interp_term, interp_mean, natural_vec, expec_vec, prec)
        return interp_mean, interp_var, kl_div

    @staticmethod
    def backward(ctx, interp_mean_grad: torch.Tensor, interp_var_grad: torch.Tensor, kl_div_grad: torch.Tensor) ->Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        interp_term, s_times_interp_term, interp_mean, natural_vec, expec_vec, prec = ctx.saved_tensors
        interp_mean_grad = interp_mean_grad.unsqueeze(-2)
        interp_var_grad = interp_var_grad.unsqueeze(-2)
        interp_term_grad = (interp_var_grad * s_times_interp_term).mul(2.0) + interp_mean_grad * expec_vec.unsqueeze(-1)
        expec_vec_grad = sum([(interp_var_grad * interp_mean.unsqueeze(-2) * interp_term).sum(dim=-1).mul(-2), (interp_mean_grad * interp_term).sum(dim=-1), kl_div_grad.unsqueeze(-1) * natural_vec])
        eye = torch.eye(expec_vec.size(-1), device=expec_vec.device, dtype=expec_vec.dtype)
        expec_mat_grad = torch.add(interp_var_grad * interp_term @ interp_term.transpose(-1, -2), kl_div_grad.unsqueeze(-1).unsqueeze(-1) * (eye - prec).mul(0.5))
        return interp_term_grad, expec_vec_grad, expec_mat_grad, None


class CiqVariationalStrategy(_VariationalStrategy):
    """
    Similar to :class:`~gpytorch.variational.VariationalStrategy`,
    except the whitening operation is performed using Contour Integral Quadrature
    rather than Cholesky (see `Pleiss et al. (2020)`_ for more info).
    See the `CIQ-SVGP tutorial`_ for an example.

    Contour Integral Quadrature uses iterative matrix-vector multiplication to approximate
    the :math:`\\mathbf K_{\\mathbf Z \\mathbf Z}^{-1/2}` matrix used for the whitening operation.
    This can be more efficient than the standard variational strategy for large numbers
    of inducing points (e.g. :math:`M > 1000`) or when the inducing points have structure
    (e.g. they lie on an evenly-spaced grid).

    .. note::

        It is recommended that this object is used in conjunction with
        :obj:`~gpytorch.variational.NaturalVariationalDistribution` and
        `natural gradient descent`_.

    :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.
        Typically passed in when the VariationalStrategy is created in the
        __init__ method of the user defined model.
    :param torch.Tensor inducing_points: Tensor containing a set of inducing
        points to use for variational inference.
    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A
        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`
    :param learn_inducing_locations: (Default True): Whether or not
        the inducing point locations :math:`\\mathbf Z` should be learned (i.e. are they
        parameters of the model).
    :type learn_inducing_locations: `bool`, optional

    .. _Pleiss et al. (2020):
        https://arxiv.org/pdf/2006.11267.pdf
    .. _CIQ-SVGP tutorial:
        examples/04_Variational_and_Approximate_GPs/SVGP_CIQ.html
    .. _natural gradient descent:
        examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.html
    """

    def _ngd(self):
        return isinstance(self._variational_distribution, NaturalVariationalDistribution)

    @property
    @cached(name='prior_distribution_memo')
    def prior_distribution(self):
        zeros = torch.zeros(self._variational_distribution.shape(), dtype=self._variational_distribution.dtype, device=self._variational_distribution.device)
        ones = torch.ones_like(zeros)
        res = MultivariateNormal(zeros, DiagLinearOperator(ones))
        return res

    @property
    @cached(name='variational_distribution_memo')
    def variational_distribution(self):
        if self._ngd():
            raise RuntimeError('Variational distribution for NGD-CIQ should be computed during forward calls. This is probably a bug in GPyTorch.')
        return super().variational_distribution

    def forward(self, x: torch.Tensor, inducing_points: torch.Tensor, inducing_values: torch.Tensor, variational_inducing_covar: Optional[MultivariateNormal]=None, **kwargs) ->MultivariateNormal:
        full_inputs = torch.cat([inducing_points, x], dim=-2)
        full_output = self.model.forward(full_inputs)
        full_covar = full_output.lazy_covariance_matrix
        num_induc = inducing_points.size(-2)
        test_mean = full_output.mean[..., num_induc:]
        induc_induc_covar = full_covar[..., :num_induc, :num_induc].evaluate_kernel().add_jitter(self.jitter_val)
        induc_data_covar = full_covar[..., :num_induc, num_induc:].to_dense()
        data_data_covar = full_covar[..., num_induc:, num_induc:].add_jitter(self.jitter_val)
        with settings.max_preconditioner_size(0):
            interp_term = to_linear_operator(induc_induc_covar).sqrt_inv_matmul(induc_data_covar)
        if self._ngd():
            interp_mean, interp_var, kl_div = _NgdInterpTerms.apply(interp_term, self._variational_distribution.natural_vec, self._variational_distribution.natural_mat)
            predictive_var = data_data_covar.diagonal(dim1=-1, dim2=-2) - interp_term.pow(2).sum(dim=-2) + interp_var
            predictive_var = torch.clamp_min(predictive_var, settings.min_variance.value(predictive_var.dtype))
            predictive_covar = DiagLinearOperator(predictive_var)
            if not hasattr(self, '_memoize_cache'):
                self._memoize_cache = dict()
            self._memoize_cache['kl'] = kl_div
        else:
            interp_mean = torch.matmul(interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)).squeeze(-1)
            middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)
            if variational_inducing_covar is not None:
                middle_term = SumLinearOperator(variational_inducing_covar, middle_term)
            predictive_covar = SumLinearOperator(data_data_covar.add_jitter(self.jitter_val), MatmulLinearOperator(interp_term.transpose(-1, -2), middle_term @ interp_term))
        predictive_mean = interp_mean + test_mean
        return MultivariateNormal(predictive_mean, predictive_covar)

    def kl_divergence(self):
        """
        Compute the KL divergence between the variational inducing distribution :math:`q(\\mathbf u)`
        and the prior inducing distribution :math:`p(\\mathbf u)`.

        :rtype: torch.Tensor
        """
        if self._ngd():
            if hasattr(self, '_memoize_cache') and 'kl' in self._memoize_cache:
                return self._memoize_cache['kl']
            else:
                raise RuntimeError('KL divergence for NGD-CIQ should be computed during forward calls.This is probably a bug in GPyTorch.')
        else:
            return super().kl_divergence()

    def __call__(self, x: torch.Tensor, prior: bool=False, **kwargs) ->MultivariateNormal:
        if prior:
            return self.model.forward(x)
        if self.training:
            self._clear_cache()
        if not self.variational_params_initialized.item():
            if self._ngd():
                noise = torch.randn_like(self.prior_distribution.mean).mul_(0.001)
                eye = torch.eye(noise.size(-1), dtype=noise.dtype, device=noise.device).mul(-0.5)
                self._variational_distribution.natural_vec.data.copy_(noise)
                self._variational_distribution.natural_mat.data.copy_(eye)
                self.variational_params_initialized.fill_(1)
            else:
                prior_dist = self.prior_distribution
                self._variational_distribution.initialize_variational_distribution(prior_dist)
                self.variational_params_initialized.fill_(1)
        inducing_points = self.inducing_points
        if inducing_points.shape[:-2] != x.shape[:-2]:
            x, inducing_points = self._expand_inputs(x, inducing_points)
        if self._ngd():
            return Module.__call__(self, x, inducing_points, inducing_values=None, variational_inducing_covar=None, **kwargs)
        else:
            variational_dist_u = self.variational_distribution
            if isinstance(variational_dist_u, MultivariateNormal):
                return Module.__call__(self, x, inducing_points, inducing_values=variational_dist_u.mean, variational_inducing_covar=variational_dist_u.lazy_covariance_matrix, **kwargs)
            elif isinstance(variational_dist_u, Delta):
                return Module.__call__(self, x, inducing_points, inducing_values=variational_dist_u.mean, variational_inducing_covar=None, ngd=False, **kwargs)
            else:
                raise RuntimeError(f'Invalid variational distribuition ({type(variational_dist_u)}). Expected a multivariate normal or a delta distribution.')


class DeltaVariationalDistribution(_VariationalDistribution):
    """
    This :obj:`~gpytorch.variational._VariationalDistribution` object replaces a variational distribution
    with a single particle. It is equivalent to doing MAP inference.

    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
        should be this size, and the variational covariance matrix should have this many rows and columns.
    :param batch_shape: Specifies an optional batch size
        for the variational parameters. This is useful for example when doing additive variational inference.
    :type batch_shape: :obj:`torch.Size`, optional
    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
    """

    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=0.001, **kwargs):
        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
        mean_init = torch.zeros(num_inducing_points)
        mean_init = mean_init.repeat(*batch_shape, 1)
        self.register_parameter(name='variational_mean', parameter=torch.nn.Parameter(mean_init))

    def forward(self):
        return Delta(self.variational_mean)

    def initialize_variational_distribution(self, prior_dist):
        self.variational_mean.data.copy_(prior_dist.mean)
        self.variational_mean.data.add_(torch.randn_like(prior_dist.mean), alpha=self.mean_init_std)


def convert_legacy_grid(grid: torch.Tensor) ->List[torch.Tensor]:
    return [grid[:, i] for i in range(grid.size(-1))]


class Interpolation(object):

    def _cubic_interpolation_kernel(self, scaled_grid_dist):
        """
        Computes the interpolation kernel u() for points X given the scaled
        grid distances:
                                    (X-x_{t})/s
        where s is the distance between neighboring grid points. Note that,
        in this context, the word "kernel" is not used to mean a covariance
        function as in the rest of the package. For more details, see the
        original paper Keys et al., 1989, equation (4).

        scaled_grid_dist should be an n-by-g matrix of distances, where the
        (ij)th element is the distance between the ith data point in X and the
        jth element in the grid.

        Note that, although this method ultimately expects a scaled distance matrix,
        it is only intended to be used on single dimensional data.
        """
        U = scaled_grid_dist.abs()
        res = torch.zeros(U.size(), dtype=U.dtype, device=U.device)
        U_lt_1 = 1 - U.floor().clamp(0, 1)
        res = res + ((1.5 * U - 2.5).mul(U).mul(U) + 1) * U_lt_1
        U_ge_1_le_2 = 1 - U_lt_1
        res = res + (((-0.5 * U + 2.5).mul(U) - 4).mul(U) + 2) * U_ge_1_le_2
        return res

    def interpolate(self, x_grid: List[torch.Tensor], x_target: torch.Tensor, interp_points=range(-2, 2), eps=1e-10):
        if torch.is_tensor(x_grid):
            x_grid = convert_legacy_grid(x_grid)
        num_target_points = x_target.size(0)
        num_dim = x_target.size(-1)
        assert num_dim == len(x_grid)
        grid_sizes = [len(x_grid[i]) for i in range(num_dim)]
        x_target_max = x_target.max(0)[0]
        x_target_min = x_target.min(0)[0]
        grid_mins = torch.stack([x_grid[i].min() for i in range(num_dim)], dim=0)
        grid_maxs = torch.stack([x_grid[i].max() for i in range(num_dim)], dim=0)
        lt_min_mask = (x_target_min - grid_mins).lt(-1e-07)
        gt_max_mask = (x_target_max - grid_maxs).gt(1e-07)
        if lt_min_mask.sum().item():
            first_out_of_range = lt_min_mask.nonzero(as_tuple=False).squeeze(1)[0].item()
            raise RuntimeError('Received data that was out of bounds for the specified grid. Grid bounds were ({:.3f}, {:.3f}), but min = {:.3f}, max = {:.3f}'.format(grid_mins[first_out_of_range].item(), grid_maxs[first_out_of_range].item(), x_target_min[first_out_of_range].item(), x_target_max[first_out_of_range].item()))
        if gt_max_mask.sum().item():
            first_out_of_range = gt_max_mask.nonzero(as_tuple=False).squeeze(1)[0].item()
            raise RuntimeError('Received data that was out of bounds for the specified grid. Grid bounds were ({:.3f}, {:.3f}), but min = {:.3f}, max = {:.3f}'.format(grid_mins[first_out_of_range].item(), grid_maxs[first_out_of_range].item(), x_target_min[first_out_of_range].item(), x_target_max[first_out_of_range].item()))
        interp_points = torch.tensor(interp_points, dtype=x_grid[0].dtype, device=x_grid[0].device)
        interp_points_flip = interp_points.flip(0)
        num_coefficients = len(interp_points)
        interp_values = torch.ones(num_target_points, num_coefficients ** num_dim, dtype=x_grid[0].dtype, device=x_grid[0].device)
        interp_indices = torch.zeros(num_target_points, num_coefficients ** num_dim, dtype=torch.long, device=x_grid[0].device)
        for i in range(num_dim):
            num_grid_points = x_grid[i].size(0)
            grid_delta = (x_grid[i][1] - x_grid[i][0]).clamp_min_(eps)
            lower_grid_pt_idxs = torch.floor((x_target[:, i] - x_grid[i][0]) / grid_delta)
            lower_pt_rel_dists = (x_target[:, i] - x_grid[i][0]) / grid_delta - lower_grid_pt_idxs
            lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()
            lower_grid_pt_idxs.detach_()
            if len(lower_grid_pt_idxs.shape) == 0:
                lower_grid_pt_idxs = lower_grid_pt_idxs.unsqueeze(0)
            scaled_dist = lower_pt_rel_dists.unsqueeze(-1) + interp_points_flip.unsqueeze(-2)
            dim_interp_values = self._cubic_interpolation_kernel(scaled_dist)
            left_boundary_pts = (lower_grid_pt_idxs < 0).nonzero(as_tuple=False)
            num_left = len(left_boundary_pts)
            if num_left > 0:
                left_boundary_pts.squeeze_(1)
                x_grid_first = x_grid[i][:num_coefficients].unsqueeze(1).t().expand(num_left, num_coefficients)
                grid_targets = x_target.select(1, i)[left_boundary_pts].unsqueeze(1).expand(num_left, num_coefficients)
                dists = torch.abs(x_grid_first - grid_targets)
                closest_from_first = torch.min(dists, 1)[1]
                for j in range(num_left):
                    dim_interp_values[left_boundary_pts[j], :] = 0
                    dim_interp_values[left_boundary_pts[j], closest_from_first[j]] = 1
                    lower_grid_pt_idxs[left_boundary_pts[j]] = 0
            right_boundary_pts = (lower_grid_pt_idxs > num_grid_points - num_coefficients).nonzero(as_tuple=False)
            num_right = len(right_boundary_pts)
            if num_right > 0:
                right_boundary_pts.squeeze_(1)
                x_grid_last = x_grid[i][-num_coefficients:].unsqueeze(1).t().expand(num_right, num_coefficients)
                grid_targets = x_target.select(1, i)[right_boundary_pts].unsqueeze(1)
                grid_targets = grid_targets.expand(num_right, num_coefficients)
                dists = torch.abs(x_grid_last - grid_targets)
                closest_from_last = torch.min(dists, 1)[1]
                for j in range(num_right):
                    dim_interp_values[right_boundary_pts[j], :] = 0
                    dim_interp_values[right_boundary_pts[j], closest_from_last[j]] = 1
                    lower_grid_pt_idxs[right_boundary_pts[j]] = num_grid_points - num_coefficients
            offset = (interp_points - interp_points.min()).long().unsqueeze(-2)
            dim_interp_indices = lower_grid_pt_idxs.long().unsqueeze(-1) + offset
            n_inner_repeat = num_coefficients ** i
            n_outer_repeat = num_coefficients ** (num_dim - i - 1)
            index_coeff = reduce(mul, grid_sizes[i + 1:], 1)
            dim_interp_indices = dim_interp_indices.unsqueeze(-1).repeat(1, n_inner_repeat, n_outer_repeat)
            dim_interp_values = dim_interp_values.unsqueeze(-1).repeat(1, n_inner_repeat, n_outer_repeat)
            interp_indices = interp_indices.add(dim_interp_indices.view(num_target_points, -1).mul(index_coeff))
            interp_values = interp_values.mul(dim_interp_values.view(num_target_points, -1))
        return interp_indices, interp_values


def left_interp(interp_indices: torch.LongTensor, interp_values: torch.Tensor, rhs: torch.Tensor) ->torch.Tensor:
    warnings.warn('gpytorch.utils.interpolation.left_interp is deprecated in favor of linear_operator.utils.interpolation.left_interp.', DeprecationWarning)
    return _left_interp(interp_indices=interp_indices, interp_values=interp_values, rhs=rhs)


class GridInterpolationVariationalStrategy(_VariationalStrategy):
    """
    This strategy constrains the inducing points to a grid and applies a deterministic
    relationship between :math:`\\mathbf f` and :math:`\\mathbf u`.
    It was introduced by `Wilson et al. (2016)`_.

    Here, the inducing points are not learned. Instead, the strategy
    automatically creates inducing points based on a set of grid sizes and grid
    bounds.

    .. _Wilson et al. (2016):
        https://arxiv.org/abs/1611.00336

    :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.
        Typically passed in when the VariationalStrategy is created in the
        __init__ method of the user defined model.
    :param int grid_size: Size of the grid
    :param list grid_bounds: Bounds of each dimension of the grid (should be a list of (float, float) tuples)
    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A
        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`
    """

    def __init__(self, model, grid_size, grid_bounds, variational_distribution):
        grid = torch.zeros(grid_size, len(grid_bounds))
        for i in range(len(grid_bounds)):
            grid_diff = float(grid_bounds[i][1] - grid_bounds[i][0]) / (grid_size - 2)
            grid[:, i] = torch.linspace(grid_bounds[i][0] - grid_diff, grid_bounds[i][1] + grid_diff, grid_size)
        inducing_points = torch.zeros(int(pow(grid_size, len(grid_bounds))), len(grid_bounds))
        prev_points = None
        for i in range(len(grid_bounds)):
            for j in range(grid_size):
                inducing_points[j * grid_size ** i:(j + 1) * grid_size ** i, i].fill_(grid[j, i])
                if prev_points is not None:
                    inducing_points[j * grid_size ** i:(j + 1) * grid_size ** i, :i].copy_(prev_points)
            prev_points = inducing_points[:grid_size ** (i + 1), :i + 1]
        super(GridInterpolationVariationalStrategy, self).__init__(model, inducing_points, variational_distribution, learn_inducing_locations=False)
        object.__setattr__(self, 'model', model)
        self.register_buffer('grid', grid)

    def _compute_grid(self, inputs):
        n_data, n_dimensions = inputs.size(-2), inputs.size(-1)
        batch_shape = inputs.shape[:-2]
        inputs = inputs.reshape(-1, n_dimensions)
        interp_indices, interp_values = Interpolation().interpolate(self.grid, inputs)
        interp_indices = interp_indices.view(*batch_shape, n_data, -1)
        interp_values = interp_values.view(*batch_shape, n_data, -1)
        if interp_indices.dim() - 2 != len(self._variational_distribution.batch_shape):
            batch_shape = torch.broadcast_shapes(interp_indices.shape[:-2], self._variational_distribution.batch_shape)
            interp_indices = interp_indices.expand(*batch_shape, *interp_indices.shape[-2:])
            interp_values = interp_values.expand(*batch_shape, *interp_values.shape[-2:])
        return interp_indices, interp_values

    @property
    @cached(name='prior_distribution_memo')
    def prior_distribution(self):
        out = self.model.forward(self.inducing_points)
        res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter(0.001))
        return res

    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None):
        if variational_inducing_covar is None:
            raise RuntimeError(f'GridInterpolationVariationalStrategy is only compatible with Gaussian variational distributions. Got ({self.variational_distribution.__class__.__name__}.')
        variational_distribution = self.variational_distribution
        interp_indices, interp_values = self._compute_grid(x)
        predictive_mean = left_interp(interp_indices, interp_values, inducing_values.unsqueeze(-1))
        predictive_mean = predictive_mean.squeeze(-1)
        predictive_covar = InterpolatedLinearOperator(variational_distribution.lazy_covariance_matrix, interp_indices, interp_values, interp_indices, interp_values)
        output = MultivariateNormal(predictive_mean, predictive_covar)
        return output


class MultitaskMultivariateNormal(MultivariateNormal):
    """
    Constructs a multi-output multivariate Normal random variable, based on mean and covariance
    Can be multi-output multivariate, or a batch of multi-output multivariate Normal

    Passing a matrix mean corresponds to a multi-output multivariate Normal
    Passing a matrix mean corresponds to a batch of multivariate Normals

    :param torch.Tensor mean:  An `n x t` or batch `b x n x t` matrix of means for the MVN distribution.
    :param ~linear_operator.operators.LinearOperator covar: An `... x NT x NT` (batch) matrix.
        covariance matrix of MVN distribution.
    :param bool validate_args: (default=False) If True, validate `mean` anad `covariance_matrix` arguments.
    :param bool interleaved: (default=True) If True, covariance matrix is interpreted as block-diagonal w.r.t.
        inter-task covariances for each observation. If False, it is interpreted as block-diagonal
        w.r.t. inter-observation covariance for each task.
    """

    def __init__(self, mean, covariance_matrix, validate_args=False, interleaved=True):
        if not torch.is_tensor(mean) and not isinstance(mean, LinearOperator):
            raise RuntimeError('The mean of a MultitaskMultivariateNormal must be a Tensor or LinearOperator')
        if not torch.is_tensor(covariance_matrix) and not isinstance(covariance_matrix, LinearOperator):
            raise RuntimeError('The covariance of a MultitaskMultivariateNormal must be a Tensor or LinearOperator')
        if mean.dim() < 2:
            raise RuntimeError('mean should be a matrix or a batch matrix (batch mode)')
        batch_shape = torch.broadcast_shapes(mean.shape[:-2], covariance_matrix.shape[:-2])
        if mean.shape[-2:].numel() != covariance_matrix.size(-1):
            if covariance_matrix.size(-1) % mean.shape[-2:].numel():
                raise RuntimeError(f'mean shape {mean.shape} is incompatible with covariance shape {covariance_matrix.shape}')
            elif mean.size(-2) == 1:
                mean = mean.expand(*batch_shape, covariance_matrix.size(-1) // mean.size(-1), mean.size(-1))
            elif mean.size(-1) == 1:
                mean = mean.expand(*batch_shape, mean.size(-2), covariance_matrix.size(-2) // mean.size(-2))
            else:
                raise RuntimeError(f'mean shape {mean.shape} is incompatible with covariance shape {covariance_matrix.shape}')
        else:
            mean = mean.expand(*batch_shape, *mean.shape[-2:])
        self._output_shape = mean.shape
        self._interleaved = interleaved
        if self._interleaved:
            mean_mvn = mean.reshape(*mean.shape[:-2], -1)
        else:
            mean_mvn = mean.transpose(-1, -2).reshape(*mean.shape[:-2], -1)
        super().__init__(mean=mean_mvn, covariance_matrix=covariance_matrix, validate_args=validate_args)

    @property
    def base_sample_shape(self):
        """
        Returns the shape of a base sample (without batching) that is used to
        generate a single sample.
        """
        base_sample_shape = self.event_shape
        return base_sample_shape

    @property
    def event_shape(self):
        return self._output_shape[-2:]

    @classmethod
    def from_batch_mvn(cls, batch_mvn, task_dim=-1):
        """
        Reinterprate a batch of multivariate normal distributions as an (independent) multitask multivariate normal
        distribution.

        :param ~gpytorch.distributions.MultivariateNormal batch_mvn: The base MVN distribution.
            (This distribution should have at least one batch dimension).
        :param int task_dim: Which batch dimension should be interpreted as the dimension for the independent tasks.
        :returns: the independent multitask distribution
        :rtype: gpytorch.distributions.MultitaskMultivariateNormal

        Example:
            >>> # model is a gpytorch.models.VariationalGP
            >>> # likelihood is a gpytorch.likelihoods.Likelihood
            >>> mean = torch.randn(4, 2, 3)
            >>> covar_factor = torch.randn(4, 2, 3, 3)
            >>> covar = covar_factor @ covar_factor.transpose(-1, -2)
            >>> mvn = gpytorch.distributions.MultivariateNormal(mean, covar)
            >>> print(mvn.event_shape, mvn.batch_shape)
            >>> # torch.Size([3]), torch.Size([4, 2])
            >>>
            >>> mmvn = MultitaskMultivariateNormal.from_batch_mvn(mvn, task_dim=-1)
            >>> print(mmvn.event_shape, mmvn.batch_shape)
            >>> # torch.Size([3, 2]), torch.Size([4])
        """
        orig_task_dim = task_dim
        task_dim = task_dim if task_dim >= 0 else len(batch_mvn.batch_shape) + task_dim
        if task_dim < 0 or task_dim > len(batch_mvn.batch_shape):
            raise ValueError(f'task_dim of {orig_task_dim} is incompatible with MVN batch shape of {batch_mvn.batch_shape}')
        num_dim = batch_mvn.mean.dim()
        res = cls(mean=batch_mvn.mean.permute(*range(0, task_dim), *range(task_dim + 1, num_dim), task_dim), covariance_matrix=BlockInterleavedLinearOperator(batch_mvn.lazy_covariance_matrix, block_dim=task_dim))
        return res

    @classmethod
    def from_independent_mvns(cls, mvns):
        """
        Convert an iterable of MVNs into a :obj:`~gpytorch.distributions.MultitaskMultivariateNormal`.
        The resulting distribution will have ``len(mvns)`` tasks, and the tasks will be independent.

        :param ~gpytorch.distributions.MultitaskNormal mvn: The base MVN distributions.
        :returns: the independent multitask distribution
        :rtype: gpytorch.distributions.MultitaskMultivariateNormal

        Example:
            >>> # model is a gpytorch.models.VariationalGP
            >>> # likelihood is a gpytorch.likelihoods.Likelihood
            >>> mean = torch.randn(4, 3)
            >>> covar_factor = torch.randn(4, 3, 3)
            >>> covar = covar_factor @ covar_factor.transpose(-1, -2)
            >>> mvn1 = gpytorch.distributions.MultivariateNormal(mean, covar)
            >>>
            >>> mean = torch.randn(4, 3)
            >>> covar_factor = torch.randn(4, 3, 3)
            >>> covar = covar_factor @ covar_factor.transpose(-1, -2)
            >>> mvn2 = gpytorch.distributions.MultivariateNormal(mean, covar)
            >>>
            >>> mmvn = MultitaskMultivariateNormal.from_independent_mvns([mvn1, mvn2])
            >>> print(mmvn.event_shape, mmvn.batch_shape)
            >>> # torch.Size([3, 2]), torch.Size([4])
        """
        if len(mvns) < 2:
            raise ValueError('Must provide at least 2 MVNs to form a MultitaskMultivariateNormal')
        if any(isinstance(mvn, MultitaskMultivariateNormal) for mvn in mvns):
            raise ValueError('Cannot accept MultitaskMultivariateNormals')
        if not all(m.batch_shape == mvns[0].batch_shape for m in mvns[1:]):
            raise ValueError('All MultivariateNormals must have the same batch shape')
        if not all(m.event_shape == mvns[0].event_shape for m in mvns[1:]):
            raise ValueError('All MultivariateNormals must have the same event shape')
        mean = torch.stack([mvn.mean for mvn in mvns], -1)
        covar_blocks_lazy = CatLinearOperator(*[mvn.lazy_covariance_matrix.unsqueeze(0) for mvn in mvns], dim=0, output_device=mean.device)
        covar_lazy = BlockDiagLinearOperator(covar_blocks_lazy, block_dim=0)
        return cls(mean=mean, covariance_matrix=covar_lazy, interleaved=False)

    @classmethod
    def from_repeated_mvn(cls, mvn, num_tasks):
        """
        Convert a single MVN into a :obj:`~gpytorch.distributions.MultitaskMultivariateNormal`,
        where each task shares the same mean and covariance.

        :param ~gpytorch.distributions.MultitaskNormal mvn: The base MVN distribution.
        :param int num_tasks: How many tasks to create.
        :returns: the independent multitask distribution
        :rtype: gpytorch.distributions.MultitaskMultivariateNormal

        Example:
            >>> # model is a gpytorch.models.VariationalGP
            >>> # likelihood is a gpytorch.likelihoods.Likelihood
            >>> mean = torch.randn(4, 3)
            >>> covar_factor = torch.randn(4, 3, 3)
            >>> covar = covar_factor @ covar_factor.transpose(-1, -2)
            >>> mvn = gpytorch.distributions.MultivariateNormal(mean, covar)
            >>> print(mvn.event_shape, mvn.batch_shape)
            >>> # torch.Size([3]), torch.Size([4])
            >>>
            >>> mmvn = MultitaskMultivariateNormal.from_repeated_mvn(mvn, num_tasks=2)
            >>> print(mmvn.event_shape, mmvn.batch_shape)
            >>> # torch.Size([3, 2]), torch.Size([4])
        """
        return cls.from_batch_mvn(mvn.expand(torch.Size([num_tasks]) + mvn.batch_shape), task_dim=0)

    def expand(self, batch_size):
        new_mean = self.mean.expand(torch.Size(batch_size) + self.mean.shape[-2:])
        new_covar = self._covar.expand(torch.Size(batch_size) + self._covar.shape[-2:])
        res = self.__class__(new_mean, new_covar, interleaved=self._interleaved)
        return res

    def get_base_samples(self, sample_shape=torch.Size()):
        base_samples = super().get_base_samples(sample_shape)
        if not self._interleaved:
            new_shape = sample_shape + self._output_shape[:-2] + self._output_shape[:-3:-1]
            return base_samples.view(new_shape).transpose(-1, -2).contiguous()
        return base_samples.view(*sample_shape, *self._output_shape)

    def log_prob(self, value):
        if not self._interleaved:
            new_shape = value.shape[:-2] + value.shape[:-3:-1]
            value = value.view(new_shape).transpose(-1, -2).contiguous()
        return super().log_prob(value.reshape(*value.shape[:-2], -1))

    @property
    def mean(self):
        mean = super().mean
        if not self._interleaved:
            new_shape = self._output_shape[:-2] + self._output_shape[:-3:-1]
            return mean.view(new_shape).transpose(-1, -2).contiguous()
        return mean.view(self._output_shape)

    @property
    def num_tasks(self):
        return self._output_shape[-1]

    def rsample(self, sample_shape=torch.Size(), base_samples=None):
        if base_samples is not None:
            mean_shape = self.mean.shape
            base_sample_shape = base_samples.shape[-self.mean.ndimension():]
            if mean_shape != base_sample_shape:
                raise RuntimeError('The shape of base_samples (minus sample shape dimensions) should agree with the shape of self.mean. Expected ...{} but got {}'.format(mean_shape, base_sample_shape))
            sample_shape = base_samples.shape[:-self.mean.ndimension()]
            base_samples = base_samples.view(*sample_shape, *self.loc.shape)
        samples = super().rsample(sample_shape=sample_shape, base_samples=base_samples)
        if not self._interleaved:
            new_shape = sample_shape + self._output_shape[:-2] + self._output_shape[:-3:-1]
            return samples.view(new_shape).transpose(-1, -2).contiguous()
        return samples.view(sample_shape + self._output_shape)

    def to_data_independent_dist(self, jitter_val=0.0001):
        """
        Convert a multitask MVN into a batched (non-multitask) MVNs
        The result retains the intertask covariances, but gets rid of the inter-data covariances.
        The resulting distribution will have ``len(mvns)`` tasks, and the tasks will be independent.

        :returns: the bached data-independent MVN
        :rtype: gpytorch.distributions.MultivariateNormal
        """
        full_covar = self.lazy_covariance_matrix
        num_data, num_tasks = self.mean.shape[-2:]
        if self._interleaved:
            data_indices = torch.arange(0, num_data * num_tasks, num_tasks, device=full_covar.device).view(-1, 1, 1)
            task_indices = torch.arange(num_tasks, device=full_covar.device)
        else:
            data_indices = torch.arange(num_data, device=full_covar.device).view(-1, 1, 1)
            task_indices = torch.arange(0, num_data * num_tasks, num_data, device=full_covar.device)
        task_covars = full_covar[..., data_indices + task_indices.unsqueeze(-2), data_indices + task_indices.unsqueeze(-1)]
        return MultivariateNormal(self.mean, to_linear_operator(task_covars).add_jitter(jitter_val=jitter_val))

    @property
    def variance(self):
        var = super().variance
        if not self._interleaved:
            new_shape = self._output_shape[:-2] + self._output_shape[:-3:-1]
            return var.view(new_shape).transpose(-1, -2).contiguous()
        return var.view(self._output_shape)


class IndependentMultitaskVariationalStrategy(_VariationalStrategy):
    """
    IndependentMultitaskVariationalStrategy wraps an existing
    :obj:`~gpytorch.variational.VariationalStrategy` to produce vector-valued (multi-task)
    output distributions. Each task will be independent of one another.

    The output will either be a :obj:`~gpytorch.distributions.MultitaskMultivariateNormal` distribution
    (if we wish to evaluate all tasks for each input) or a :obj:`~gpytorch.distributions.MultivariateNormal`
    (if we wish to evaluate a single task for each input).

    The base variational strategy is assumed to operate on a batch of GPs. One of the batch
    dimensions corresponds to the multiple tasks.

    :param ~gpytorch.variational.VariationalStrategy base_variational_strategy: Base variational strategy
    :param int num_tasks: Number of tasks. Should correspond to the batch size of task_dim.
    :param int task_dim: (Default: -1) Which batch dimension is the task dimension
    """

    def __init__(self, base_variational_strategy, num_tasks, task_dim=-1):
        Module.__init__(self)
        self.base_variational_strategy = base_variational_strategy
        self.task_dim = task_dim
        self.num_tasks = num_tasks

    @property
    def prior_distribution(self):
        return self.base_variational_strategy.prior_distribution

    @property
    def variational_distribution(self):
        return self.base_variational_strategy.variational_distribution

    @property
    def variational_params_initialized(self):
        return self.base_variational_strategy.variational_params_initialized

    def kl_divergence(self):
        return super().kl_divergence().sum(dim=-1)

    def __call__(self, x, task_indices=None, prior=False, **kwargs):
        """
        See :class:`LMCVariationalStrategy`.
        """
        function_dist = self.base_variational_strategy(x, prior=prior, **kwargs)
        if task_indices is None:
            if self.task_dim > 0 and self.task_dim > len(function_dist.batch_shape) or self.task_dim < 0 and self.task_dim + len(function_dist.batch_shape) < 0:
                return MultitaskMultivariateNormal.from_repeated_mvn(function_dist, num_tasks=self.num_tasks)
            else:
                function_dist = MultitaskMultivariateNormal.from_batch_mvn(function_dist, task_dim=self.task_dim)
                assert function_dist.event_shape[-1] == self.num_tasks
                return function_dist
        else:
            if self.task_dim > 0:
                raise RuntimeError(f'task_dim must be a negative indexed batch dimension: got {self.task_dim}.')
            num_batch = len(function_dist.batch_shape)
            task_dim = num_batch + self.task_dim
            shape = list(function_dist.batch_shape + function_dist.event_shape)
            shape[task_dim] = 1
            task_indices = task_indices.expand(shape).squeeze(task_dim)
            task_mask = torch.nn.functional.one_hot(task_indices, num_classes=self.num_tasks)
            task_mask = task_mask.permute(*range(0, task_dim), *range(task_dim + 1, num_batch + 1), task_dim)
            mean = (function_dist.mean * task_mask).sum(task_dim)
            covar = (function_dist.lazy_covariance_matrix * RootLinearOperator(task_mask[..., None])).sum(task_dim)
            return MultivariateNormal(mean, covar)


class MultitaskVariationalStrategy(IndependentMultitaskVariationalStrategy):
    """
    IndependentMultitaskVariationalStrategy wraps an existing
    :obj:`~gpytorch.variational.VariationalStrategy`
    to produce a :obj:`~gpytorch.variational.MultitaskMultivariateNormal` distribution.
    All outputs will be independent of one another.

    The base variational strategy is assumed to operate on a batch of GPs. One of the batch
    dimensions corresponds to the multiple tasks.

    :param ~gpytorch.variational.VariationalStrategy base_variational_strategy: Base variational strategy
    :param int num_tasks: Number of tasks. Should correspond to the batch size of task_dim.
    :param int task_dim: (Default: -1) Which batch dimension is the task dimension
    """

    def __init__(self, base_variational_strategy, num_tasks, task_dim=-1):
        warnings.warn('MultitaskVariationalStrategy has been renamed to IndependentMultitaskVariationalStrategy', DeprecationWarning)
        super().__init__(base_variational_strategy, num_tasks, task_dim=-1)


def _select_lmc_coefficients(lmc_coefficients: torch.Tensor, indices: torch.LongTensor) ->torch.Tensor:
    """
    Given a list of indices for ... x N datapoints,
      select the row from lmc_coefficient that corresponds to each datapoint

    lmc_coefficients: torch.Tensor ... x num_latents x ... x num_tasks
    indices: torch.Tesnor ... x N
    """
    batch_shape = torch.broadcast_shapes(lmc_coefficients.shape[:-1], indices.shape[:-1])
    lmc_coefficients = lmc_coefficients.expand(*batch_shape, lmc_coefficients.shape[-1])[..., None]
    indices = indices.expand(*batch_shape, indices.shape[-1])[..., None]
    res = left_interp(indices, torch.ones(indices.shape, dtype=torch.long, device=indices.device), lmc_coefficients).squeeze(-1)
    return res


class LMCVariationalStrategy(_VariationalStrategy):
    """
    LMCVariationalStrategy is an implementation of the "Linear Model of Coregionalization"
    for multitask GPs. This model assumes that there are :math:`Q` latent functions
    :math:`\\mathbf g(\\cdot) = [g^{(1)}(\\cdot), \\ldots, g^{(q)}(\\cdot)]`,
    each of which is modelled by a GP.
    The output functions (tasks) are linear combination of the latent functions:

    .. math::

        f_{\\text{task } i}( \\mathbf x) = \\sum_{q=1}^Q a_i^{(q)} g^{(q)} ( \\mathbf x )

    LMCVariationalStrategy wraps an existing :obj:`~gpytorch.variational.VariationalStrategy`.
    The output will either be a :obj:`~gpytorch.distributions.MultitaskMultivariateNormal` distribution
    (if we wish to evaluate all tasks for each input) or a :obj:`~gpytorch.distributions.MultivariateNormal`
    (if we wish to evaluate a single task for each input).

    The base variational strategy is assumed to operate on a multi-batch of GPs, where one
    of the batch dimensions corresponds to the latent function dimension.

    .. note::

        The batch shape of the base :obj:`~gpytorch.variational.VariationalStrategy` does not
        necessarily have to correspond to the batch shape of the underlying GP objects.

        For example, if the base variational strategy has a batch shape of `[3]` (corresponding
        to 3 latent functions), the GP kernel object could have a batch shape of `[3]` or no
        batch shape. This would correspond to each of the latent functions having different kernels
        or the same kernel, respectivly.

    Example:
        >>> class LMCMultitaskGP(gpytorch.models.ApproximateGP):
        >>>     '''
        >>>     3 latent functions
        >>>     5 output dimensions (tasks)
        >>>     '''
        >>>     def __init__(self):
        >>>         # Each latent function shares the same inducing points
        >>>         # We'll have 32 inducing points, and let's assume the input dimensionality is 2
        >>>         inducing_points = torch.randn(32, 2)
        >>>
        >>>         # The variational parameters have a batch_shape of [3] - for 3 latent functions
        >>>         variational_distribution = gpytorch.variational.MeanFieldVariationalDistribution(
        >>>             inducing_points.size(-1), batch_shape=torch.Size([3]),
        >>>         )
        >>>         variational_strategy = gpytorch.variational.LMCVariationalStrategy(
        >>>             gpytorch.variational.VariationalStrategy(
        >>>                 self, inducing_points, variational_distribution, learn_inducing_locations=True,
        >>>             ),
        >>>             num_tasks=5,
        >>>             num_latents=3,
        >>>             latent_dim=-1,
        >>>         )
        >>>
        >>>         # Each latent function has its own mean/kernel function
        >>>         super().__init__(variational_strategy)
        >>>         self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([3]))
        >>>         self.covar_module = gpytorch.kernels.ScaleKernel(
        >>>             gpytorch.kernels.RBFKernel(batch_shape=torch.Size([3])),
        >>>             batch_shape=torch.Size([3]),
        >>>         )
        >>>

    :param ~gpytorch.variational.VariationalStrategy base_variational_strategy: Base variational strategy
    :param int num_tasks: The total number of tasks (output functions)
    :param int num_latents: The total number of latent functions in each group
    :param latent_dim: (Default: -1) Which batch dimension corresponds to the latent function batch.
        **Must be negative indexed**
    :type latent_dim: `int` < 0
    """

    def __init__(self, base_variational_strategy, num_tasks, num_latents=1, latent_dim=-1, jitter_val=None):
        Module.__init__(self)
        self.base_variational_strategy = base_variational_strategy
        self.num_tasks = num_tasks
        batch_shape = self.base_variational_strategy._variational_distribution.batch_shape
        if latent_dim >= 0:
            raise RuntimeError(f'latent_dim must be a negative indexed batch dimension: got {latent_dim}.')
        if not (batch_shape[latent_dim] == num_latents or batch_shape[latent_dim] == 1):
            raise RuntimeError(f'Mismatch in num_latents: got a variational distribution of batch shape {batch_shape}, expected the function dim {latent_dim} to be {num_latents}.')
        self.num_latents = num_latents
        self.latent_dim = latent_dim
        self.batch_shape = list(batch_shape)
        del self.batch_shape[self.latent_dim]
        self.batch_shape = torch.Size(self.batch_shape)
        lmc_coefficients = torch.randn(*batch_shape, self.num_tasks)
        self.register_parameter('lmc_coefficients', torch.nn.Parameter(lmc_coefficients))
        if jitter_val is None:
            self.jitter_val = settings.variational_cholesky_jitter.value(self.base_variational_strategy.inducing_points.dtype)
        else:
            self.jitter_val = jitter_val

    @property
    def prior_distribution(self):
        return self.base_variational_strategy.prior_distribution

    @property
    def variational_distribution(self):
        return self.base_variational_strategy.variational_distribution

    @property
    def variational_params_initialized(self):
        return self.base_variational_strategy.variational_params_initialized

    def kl_divergence(self):
        return super().kl_divergence().sum(dim=self.latent_dim)

    def __call__(self, x, task_indices=None, prior=False, **kwargs):
        """
        Computes the variational (or prior) distribution
        :math:`q( \\mathbf f \\mid \\mathbf X)` (or :math:`p( \\mathbf f \\mid \\mathbf X)`).
        There are two modes:

        1.  Compute **all tasks** for all inputs.
            If this is the case, the task_indices attribute should be None.
            The return type will be a (... x N x num_tasks)
            :class:`~gpytorch.distributions.MultitaskMultivariateNormal`.
        2.  Compute **one task** per inputs.
            If this is the case, the (... x N) task_indices tensor should contain
            the indices of each input's assigned task.
            The return type will be a (... x N)
            :class:`~gpytorch.distributions.MultivariateNormal`.

        :param x: Input locations to evaluate variational strategy
        :type x: torch.Tensor (... x N x D)
        :param task_indices: (Default: None) Task index associated with each input.
            If this **is not** provided, then the returned distribution evaluates every input on every task
            (returns :class:`~gpytorch.distributions.MultitaskMultivariateNormal`).
            If this **is** provided, then the returned distribution evaluates each input only on its assigned task.
            (returns :class:`~gpytorch.distributions.MultivariateNormal`).
        :type task_indices: torch.Tensor (... x N), optional
        :param prior: (Default: False) If False, returns the variational distribution
            :math:`q( \\mathbf f \\mid \\mathbf X)`.
            If True, returns the prior distribution
            :math:`p( \\mathbf f \\mid \\mathbf X)`.
        :type prior: bool
        :return: :math:`q( \\mathbf f \\mid \\mathbf X)` (or the prior),
            either for all tasks (if `task_indices == None`)
            or for a specific task (if `task_indices != None`).
        :rtype: ~gpytorch.distributions.MultitaskMultivariateNormal (... x N x num_tasks)
            or ~gpytorch.distributions.MultivariateNormal (... x N)
        """
        latent_dist = self.base_variational_strategy(x, prior=prior, **kwargs)
        num_batch = len(latent_dist.batch_shape)
        latent_dim = num_batch + self.latent_dim
        if task_indices is None:
            num_dim = num_batch + len(latent_dist.event_shape)
            lmc_coefficients = self.lmc_coefficients.expand(*latent_dist.batch_shape, self.lmc_coefficients.size(-1))
            latent_mean = latent_dist.mean.permute(*range(0, latent_dim), *range(latent_dim + 1, num_dim), latent_dim)
            mean = latent_mean @ lmc_coefficients.permute(*range(0, latent_dim), *range(latent_dim + 1, num_dim - 1), latent_dim, -1)
            latent_covar = latent_dist.lazy_covariance_matrix
            lmc_factor = RootLinearOperator(lmc_coefficients.unsqueeze(-1))
            covar = KroneckerProductLinearOperator(latent_covar, lmc_factor).sum(latent_dim)
            covar = covar.add_jitter(self.jitter_val)
            function_dist = MultitaskMultivariateNormal(mean, covar)
        else:
            lmc_coefficients = _select_lmc_coefficients(self.lmc_coefficients, task_indices)
            mean = (latent_dist.mean * lmc_coefficients).sum(latent_dim)
            latent_covar = latent_dist.lazy_covariance_matrix
            lmc_factor = RootLinearOperator(lmc_coefficients.unsqueeze(-1))
            covar = (latent_covar * lmc_factor).sum(latent_dim)
            covar = covar.add_jitter(self.jitter_val)
            function_dist = MultivariateNormal(mean, covar)
        return function_dist


class MeanFieldVariationalDistribution(_VariationalDistribution):
    """
    A :obj:`~gpytorch.variational._VariationalDistribution` that is defined to be a multivariate normal distribution
    with a diagonal covariance matrix. This will not be as flexible/expressive as a
    :obj:`~gpytorch.variational.CholeskyVariationalDistribution`.

    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
        should be this size, and the variational covariance matrix should have this many rows and columns.
    :param batch_shape: Specifies an optional batch size
        for the variational parameters. This is useful for example when doing additive variational inference.
    :type batch_shape: :obj:`torch.Size`, optional
    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
    """

    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=0.001, **kwargs):
        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
        mean_init = torch.zeros(num_inducing_points)
        covar_init = torch.ones(num_inducing_points)
        mean_init = mean_init.repeat(*batch_shape, 1)
        covar_init = covar_init.repeat(*batch_shape, 1)
        self.register_parameter(name='variational_mean', parameter=torch.nn.Parameter(mean_init))
        self.register_parameter(name='_variational_stddev', parameter=torch.nn.Parameter(covar_init))

    @property
    def variational_stddev(self):
        mask = torch.ones_like(self._variational_stddev)
        return self._variational_stddev.mul(mask).abs().clamp_min(1e-08)

    def forward(self):
        mask = torch.ones_like(self._variational_stddev)
        variational_covar = DiagLinearOperator(self._variational_stddev.mul(mask).pow(2))
        return MultivariateNormal(self.variational_mean, variational_covar)

    def initialize_variational_distribution(self, prior_dist):
        self.variational_mean.data.copy_(prior_dist.mean)
        self.variational_mean.data.add_(torch.randn_like(prior_dist.mean), alpha=self.mean_init_std)
        self._variational_stddev.data.copy_(prior_dist.stddev)


class OrthogonallyDecoupledVariationalStrategy(_VariationalStrategy):
    """
    Implements orthogonally decoupled VGPs as defined in `Salimbeni et al. (2018)`_.
    This variational strategy uses a different set of inducing points for the mean and covariance functions.
    The idea is to use more inducing points for the (computationally efficient) mean and fewer inducing points for the
    (computationally expensive) covaraince.

    This variational strategy defines the inducing points/:obj:`~gpytorch.variational._VariationalDistribution`
    for the mean function.
    It then wraps a different :obj:`~gpytorch.variational._VariationalStrategy` which
    defines the covariance inducing points.

    :param ~gpytorch.variational._VariationalStrategy covar_variational_strategy:
        The variational strategy for the covariance term.
    :param torch.Tensor inducing_points: Tensor containing a set of inducing
        points to use for variational inference.
    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A
        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`

    Example:
        >>> mean_inducing_points = torch.randn(1000, train_x.size(-1), dtype=train_x.dtype, device=train_x.device)
        >>> covar_inducing_points = torch.randn(100, train_x.size(-1), dtype=train_x.dtype, device=train_x.device)
        >>>
        >>> covar_variational_strategy = gpytorch.variational.VariationalStrategy(
        >>>     model, covar_inducing_points,
        >>>     gpytorch.variational.CholeskyVariationalDistribution(covar_inducing_points.size(-2)),
        >>>     learn_inducing_locations=True
        >>> )
        >>>
        >>> variational_strategy = gpytorch.variational.OrthogonallyDecoupledVariationalStrategy(
        >>>     covar_variational_strategy, mean_inducing_points,
        >>>     gpytorch.variational.DeltaVariationalDistribution(mean_inducing_points.size(-2)),
        >>> )

    .. _Salimbeni et al. (2018):
        https://arxiv.org/abs/1809.08820
    """

    def __init__(self, covar_variational_strategy, inducing_points, variational_distribution):
        if not isinstance(variational_distribution, DeltaVariationalDistribution):
            raise NotImplementedError('OrthogonallyDecoupledVariationalStrategy currently works with DeltaVariationalDistribution')
        super().__init__(covar_variational_strategy, inducing_points, variational_distribution, learn_inducing_locations=True)
        self.base_variational_strategy = covar_variational_strategy

    @property
    @cached(name='prior_distribution_memo')
    def prior_distribution(self):
        out = self.model(self.inducing_points)
        res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter(self.jitter_val))
        return res

    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None, **kwargs):
        if variational_inducing_covar is not None:
            raise NotImplementedError('OrthogonallyDecoupledVariationalStrategy currently works with DeltaVariationalDistribution')
        num_data = x.size(-2)
        full_output = self.model(torch.cat([x, inducing_points], dim=-2), **kwargs)
        full_mean = full_output.mean
        full_covar = full_output.lazy_covariance_matrix
        if self.training:
            induc_mean = full_mean[..., num_data:]
            induc_induc_covar = full_covar[..., num_data:, num_data:]
            prior_dist = MultivariateNormal(induc_mean, induc_induc_covar)
            add_to_cache(self, 'prior_distribution_memo', prior_dist)
        test_mean = full_mean[..., :num_data]
        data_induc_covar = full_covar[..., :num_data, num_data:]
        predictive_mean = (data_induc_covar @ inducing_values.unsqueeze(-1)).squeeze(-1).add(test_mean)
        predictive_covar = full_covar[..., :num_data, :num_data]
        return MultivariateNormal(predictive_mean, predictive_covar)

    def kl_divergence(self):
        mean = self.variational_distribution.mean
        induc_induc_covar = self.prior_distribution.lazy_covariance_matrix
        kl = self.model.kl_divergence() + ((induc_induc_covar @ mean.unsqueeze(-1)).squeeze(-1) * mean).sum(-1).mul(0.5)
        return kl


class _TrilNaturalToMuVarSqrt(torch.autograd.Function):

    @staticmethod
    def _forward(nat_mean, tril_nat_covar):
        L = _triangular_inverse(tril_nat_covar, upper=False)
        mu = L @ (L.transpose(-1, -2) @ nat_mean.unsqueeze(-1))
        return mu.squeeze(-1), L

    @staticmethod
    def forward(ctx, nat_mean, tril_nat_covar):
        mu, L = _TrilNaturalToMuVarSqrt._forward(nat_mean, tril_nat_covar)
        ctx.save_for_backward(mu, L, tril_nat_covar)
        return mu, L

    @staticmethod
    def backward(ctx, dout_dmu, dout_dL):
        mu, L, C = ctx.saved_tensors
        dout_dnat1, dout_dnat2 = _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)
        """
        Now we need to do the Jacobian-Vector Product for the transformation:
        L = inv(chol(inv(-2 theta_cov)))

        C^T C = -2 theta_cov

        so we need to do forward differentiation, starting with sensitivity (sensitivities marked with .dots.)
        .theta_cov. = dout_dnat2

        and ending with sensitivity .C.

        if B = inv(-2 theta_cov) then:

        .B.  =  d inv(-2 theta_cov)/dtheta_cov * .theta_cov.  =  -B (-2 .theta_cov.) B

        if L = chol(B), B = LL^T then (https://homepages.inf.ed.ac.uk/imurray2/pub/16choldiff/choldiff.pdf):

        .L. = L phi(L^{-1} .B. (L^{-1})^T) = L phi(2 L^T .theta_cov. L)

        Then C = inv(L), so

        .C. = -C .L. C = phi(-2 L^T .theta_cov. L)C
        """
        A = L.transpose(-2, -1) @ dout_dnat2 @ L
        phi = _phi_for_cholesky_(A.mul_(-2))
        dout_dtril = phi @ C
        return dout_dnat1, dout_dtril


class TrilNaturalVariationalDistribution(_NaturalVariationalDistribution):
    """A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,
    parameterized by the natural vector, and a triangular decomposition of the
    natural matrix (which is not the Cholesky).

    .. note::
       The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` should only
       be used with :obj:`gpytorch.optim.NGD`, or other optimizers
       that follow exactly the gradient direction.

    .. seealso::
        The `natural gradient descent tutorial
        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
        for use instructions.

        The :obj:`~gpytorch.variational.NaturalVariationalDistribution`, which
        needs less iterations to make variational regression converge, at the
        cost of introducing numerical instability.

    .. note::
        The relationship of the parameter :math:`\\mathbf \\Theta_\\text{tril_mat}`
        to the natural parameter :math:`\\mathbf \\Theta_\\text{mat}` from
        :obj:`~gpytorch.variational.NaturalVariationalDistribution` is
        :math:`\\mathbf \\Theta_\\text{mat} = -1/2 {\\mathbf \\Theta_\\text{tril_mat}}^T {\\mathbf \\Theta_\\text{tril_mat}}`.
        Note that this is not the form of the Cholesky decomposition of :math:`\\boldsymbol \\Theta_\\text{mat}`.

    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
        should be this size, and the variational covariance matrix should have this many rows and columns.
    :param batch_shape: Specifies an optional batch size
        for the variational parameters. This is useful for example when doing additive variational inference.
    :type batch_shape: :obj:`torch.Size`, optional
    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
    """

    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=0.001, **kwargs):
        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
        scaled_mean_init = torch.zeros(num_inducing_points)
        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points)
        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)
        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)
        self.register_parameter(name='natural_vec', parameter=torch.nn.Parameter(scaled_mean_init))
        self.register_parameter(name='natural_tril_mat', parameter=torch.nn.Parameter(neg_prec_init))

    def forward(self):
        mean, chol_covar = _TrilNaturalToMuVarSqrt.apply(self.natural_vec, self.natural_tril_mat)
        return MultivariateNormal(mean, CholLinearOperator(TriangularLinearOperator(chol_covar)))

    def initialize_variational_distribution(self, prior_dist):
        prior_cov = prior_dist.lazy_covariance_matrix
        chol = prior_cov.cholesky().to_dense()
        tril_mat = _triangular_inverse(chol, upper=False)
        natural_vec = prior_cov.solve(prior_dist.mean.unsqueeze(-1)).squeeze(-1)
        noise = torch.randn_like(natural_vec).mul_(self.mean_init_std)
        self.natural_vec.data.copy_(natural_vec.add_(noise))
        self.natural_tril_mat.data.copy_(tril_mat)


class UnwhitenedVariationalStrategy(_VariationalStrategy):
    """
    Similar to :obj:`~gpytorch.variational.VariationalStrategy`, but does not perform the
    whitening operation. In almost all cases :obj:`~gpytorch.variational.VariationalStrategy`
    is preferable, with a few exceptions:

    - When the inducing points are exactly equal to the training points (i.e. :math:`\\mathbf Z = \\mathbf X`).
      Unwhitened models are faster in this case.

    - When the number of inducing points is very large (e.g. >2000). Unwhitened models can use CG for faster
      computation.

    :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.
        Typically passed in when the VariationalStrategy is created in the
        __init__ method of the user defined model.
    :param torch.Tensor inducing_points: Tensor containing a set of inducing
        points to use for variational inference.
    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A
        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`
    :param bool learn_inducing_points: (optional, default True): Whether or not
        the inducing point locations :math:`\\mathbf Z` should be learned (i.e. are they
        parameters of the model).
    """
    has_fantasy_strategy = True

    @cached(name='cholesky_factor', ignore_args=True)
    def _cholesky_factor(self, induc_induc_covar):
        L = psd_safe_cholesky(to_dense(induc_induc_covar))
        return TriangularLinearOperator(L)

    @property
    @cached(name='prior_distribution_memo')
    def prior_distribution(self):
        out = self.model.forward(self.inducing_points)
        res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter())
        return res

    @property
    @cached(name='pseudo_points_memo')
    def pseudo_points(self):
        if not isinstance(self._variational_distribution, CholeskyVariationalDistribution):
            raise NotImplementedError('Only CholeskyVariationalDistribution has pseudo-point support currently, ', 'but your _variational_distribution is a ', self._variational_distribution.__name__)
        var_cov_root = TriangularLinearOperator(self._variational_distribution.chol_variational_covar)
        var_cov = CholLinearOperator(var_cov_root)
        var_mean = self.variational_distribution.mean
        if var_mean.shape[-1] != 1:
            var_mean = var_mean.unsqueeze(-1)
        Kmm = self.model.covar_module(self.inducing_points)
        res = Kmm - var_cov
        cov_diff = res
        eval_lhs = var_cov.to_dense()
        eval_rhs = cov_diff.transpose(-1, -2).matmul(eval_lhs)
        inner_term = cov_diff.matmul(cov_diff.transpose(-1, -2))
        inner_solve = inner_term.add_jitter(self.jitter_val).solve(eval_rhs, eval_lhs.transpose(-1, -2))
        inducing_covar = var_cov + inner_solve
        rhs = cov_diff.transpose(-1, -2).matmul(var_mean)
        inner_rhs_mean_solve = inner_term.add_jitter(self.jitter_val).solve(rhs)
        pseudo_target_mean = var_mean + var_cov.matmul(inner_rhs_mean_solve)
        try:
            pseudo_target_covar = CholLinearOperator(inducing_covar.add_jitter(self.jitter_val).cholesky()).to_dense()
        except NotPSDError:
            evals, evecs = torch.linalg.eigh(inducing_covar)
            pseudo_target_covar = evecs.matmul(DiagLinearOperator(evals + self.jitter_val)).matmul(evecs.transpose(-1, -2)).to_dense()
        return pseudo_target_covar, pseudo_target_mean

    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None):
        if torch.equal(x, inducing_points):
            if variational_inducing_covar is None:
                raise RuntimeError
            else:
                return MultivariateNormal(inducing_values, variational_inducing_covar)
        num_induc = inducing_points.size(-2)
        full_inputs = torch.cat([inducing_points, x], dim=-2)
        full_output = self.model.forward(full_inputs)
        full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
        test_mean = full_mean[..., num_induc:]
        induc_mean = full_mean[..., :num_induc]
        mean_diff = (inducing_values - induc_mean).unsqueeze(-1)
        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter(self.jitter_val)
        induc_data_covar = full_covar[..., :num_induc, num_induc:].to_dense()
        data_data_covar = full_covar[..., num_induc:, num_induc:]
        if settings.fast_computations.log_prob.off() or num_induc <= settings.max_cholesky_size.value():
            induc_induc_covar = CholLinearOperator(self._cholesky_factor(induc_induc_covar))
        if not self.training and settings.skip_posterior_variances.on():
            self._mean_cache = induc_induc_covar.solve(mean_diff).detach()
            predictive_mean = torch.add(test_mean, induc_data_covar.transpose(-2, -1).matmul(self._mean_cache).squeeze(-1))
            predictive_covar = ZeroLinearOperator(test_mean.size(-1), test_mean.size(-1))
            return MultivariateNormal(predictive_mean, predictive_covar)
        shapes = [mean_diff.shape[:-1], induc_data_covar.shape[:-1], induc_induc_covar.shape[:-1]]
        if variational_inducing_covar is not None:
            root_variational_covar = variational_inducing_covar.root_decomposition().root.to_dense()
            shapes.append(root_variational_covar.shape[:-1])
        shape = torch.broadcast_shapes(*shapes)
        mean_diff = mean_diff.expand(*shape, mean_diff.size(-1))
        induc_data_covar = induc_data_covar.expand(*shape, induc_data_covar.size(-1))
        induc_induc_covar = induc_induc_covar.expand(*shape, induc_induc_covar.size(-1))
        if variational_inducing_covar is not None:
            root_variational_covar = root_variational_covar.expand(*shape, root_variational_covar.size(-1))
        if self.training:
            prior_dist = MultivariateNormal(induc_mean, induc_induc_covar)
            add_to_cache(self, 'prior_distribution_memo', prior_dist)
        if variational_inducing_covar is None:
            left_tensors = mean_diff
        else:
            left_tensors = torch.cat([mean_diff, root_variational_covar], -1)
        inv_products = induc_induc_covar.solve(induc_data_covar, left_tensors.transpose(-1, -2))
        predictive_mean = torch.add(test_mean, inv_products[..., 0, :])
        if self.training:
            interp_data_data_var, _ = induc_induc_covar.inv_quad_logdet(induc_data_covar, logdet=False, reduce_inv_quad=False)
            data_covariance = DiagLinearOperator((data_data_covar.diagonal(dim1=-1, dim2=-2) - interp_data_data_var).clamp(0, math.inf))
        else:
            neg_induc_data_data_covar = torch.matmul(induc_data_covar.transpose(-1, -2).mul(-1), induc_induc_covar.solve(induc_data_covar))
            data_covariance = data_data_covar + neg_induc_data_data_covar
        predictive_covar = PsdSumLinearOperator(RootLinearOperator(inv_products[..., 1:, :].transpose(-1, -2)), data_covariance)
        return MultivariateNormal(predictive_mean, predictive_covar)


class OldVersionWarning(UserWarning):
    """
    Warning thrown when loading a saved model from an outdated version of GPyTorch.
    """
    pass


def _ensure_updated_strategy_flag_set(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
    device = state_dict[list(state_dict.keys())[0]].device
    if prefix + 'updated_strategy' not in state_dict:
        state_dict[prefix + 'updated_strategy'] = torch.tensor(False, device=device)
        warnings.warn('You have loaded a variational GP model (using `VariationalStrategy`) from a previous version of GPyTorch. We have updated the parameters of your model to work with the new version of `VariationalStrategy` that uses whitened parameters.\nYour model will work as expected, but we recommend that you re-save your model.', OldVersionWarning)


def pop_from_cache_ignore_args(obj, name):
    """Pop an item from the cache (honoring calling args)."""
    try:
        return obj._memoize_cache.pop(name)
    except (KeyError, AttributeError):
        raise CachingError('Object does not have item {} stored in cache.'.format(name))


class VariationalStrategy(_VariationalStrategy):
    """
    The standard variational strategy, as defined by `Hensman et al. (2015)`_.
    This strategy takes a set of :math:`m \\ll n` inducing points :math:`\\mathbf Z`
    and applies an approximate distribution :math:`q( \\mathbf u)` over their function values.
    (Here, we use the common notation :math:`\\mathbf u = f(\\mathbf Z)`.
    The approximate function distribution for any abitrary input :math:`\\mathbf X` is given by:

    .. math::

        q( f(\\mathbf X) ) = \\int p( f(\\mathbf X) \\mid \\mathbf u) q(\\mathbf u) \\: d\\mathbf u

    This variational strategy uses "whitening" to accelerate the optimization of the variational
    parameters. See `Matthews (2017)`_ for more info.

    :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.
        Typically passed in when the VariationalStrategy is created in the
        __init__ method of the user defined model.
    :param torch.Tensor inducing_points: Tensor containing a set of inducing
        points to use for variational inference.
    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A
        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`
    :param learn_inducing_locations: (Default True): Whether or not
        the inducing point locations :math:`\\mathbf Z` should be learned (i.e. are they
        parameters of the model).
    :type learn_inducing_locations: `bool`, optional

    .. _Hensman et al. (2015):
        http://proceedings.mlr.press/v38/hensman15.pdf
    .. _Matthews (2017):
        https://www.repository.cam.ac.uk/handle/1810/278022
    """

    def __init__(self, model, inducing_points, variational_distribution, learn_inducing_locations=True, jitter_val=None):
        super().__init__(model, inducing_points, variational_distribution, learn_inducing_locations, jitter_val=jitter_val)
        self.register_buffer('updated_strategy', torch.tensor(True))
        self._register_load_state_dict_pre_hook(_ensure_updated_strategy_flag_set)
        self.has_fantasy_strategy = True

    @cached(name='cholesky_factor', ignore_args=True)
    def _cholesky_factor(self, induc_induc_covar):
        L = psd_safe_cholesky(to_dense(induc_induc_covar).type(_linalg_dtype_cholesky.value()))
        return TriangularLinearOperator(L)

    @property
    @cached(name='prior_distribution_memo')
    def prior_distribution(self):
        zeros = torch.zeros(self._variational_distribution.shape(), dtype=self._variational_distribution.dtype, device=self._variational_distribution.device)
        ones = torch.ones_like(zeros)
        res = MultivariateNormal(zeros, DiagLinearOperator(ones))
        return res

    @property
    @cached(name='pseudo_points_memo')
    def pseudo_points(self):
        if not isinstance(self._variational_distribution, CholeskyVariationalDistribution):
            raise NotImplementedError('Only CholeskyVariationalDistribution has pseudo-point support currently, ', 'but your _variational_distribution is a ', self._variational_distribution.__name__)
        var_cov_root = TriangularLinearOperator(self._variational_distribution.chol_variational_covar)
        var_cov = CholLinearOperator(var_cov_root)
        var_mean = self.variational_distribution.mean
        if var_mean.shape[-1] != 1:
            var_mean = var_mean.unsqueeze(-1)
        cov_diff = var_cov.add_jitter(-1.0)
        cov_diff = -1.0 * cov_diff
        Kmm = self.model.covar_module(self.inducing_points)
        Kmm_root = Kmm.cholesky()
        eval_var_cov = var_cov.to_dense()
        eval_rhs = cov_diff.transpose(-1, -2).matmul(eval_var_cov)
        inner_term = cov_diff.matmul(cov_diff.transpose(-1, -2))
        inner_solve = inner_term.add_jitter(self.jitter_val).solve(eval_rhs, eval_var_cov.transpose(-1, -2))
        inducing_covar = var_cov + inner_solve
        inducing_covar = Kmm_root.matmul(inducing_covar).matmul(Kmm_root.transpose(-1, -2))
        rhs = cov_diff.transpose(-1, -2).matmul(var_mean)
        inner_rhs_mean_solve = inner_term.add_jitter(self.jitter_val).solve(rhs)
        pseudo_target_mean = Kmm_root.matmul(inner_rhs_mean_solve)
        try:
            pseudo_target_covar = CholLinearOperator(inducing_covar.add_jitter(self.jitter_val).cholesky()).to_dense()
        except NotPSDError:
            evals, evecs = torch.linalg.eigh(inducing_covar)
            pseudo_target_covar = evecs.matmul(DiagLinearOperator(evals + self.jitter_val)).matmul(evecs.transpose(-1, -2)).to_dense()
        return pseudo_target_covar, pseudo_target_mean

    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None, **kwargs):
        full_inputs = torch.cat([inducing_points, x], dim=-2)
        full_output = self.model.forward(full_inputs, **kwargs)
        full_covar = full_output.lazy_covariance_matrix
        num_induc = inducing_points.size(-2)
        test_mean = full_output.mean[..., num_induc:]
        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter(self.jitter_val)
        induc_data_covar = full_covar[..., :num_induc, num_induc:].to_dense()
        data_data_covar = full_covar[..., num_induc:, num_induc:]
        L = self._cholesky_factor(induc_induc_covar)
        if L.shape != induc_induc_covar.shape:
            try:
                pop_from_cache_ignore_args(self, 'cholesky_factor')
            except CachingError:
                pass
            L = self._cholesky_factor(induc_induc_covar)
        interp_term = L.solve(induc_data_covar.type(_linalg_dtype_cholesky.value()))
        predictive_mean = (interp_term.transpose(-1, -2) @ inducing_values.unsqueeze(-1)).squeeze(-1) + test_mean
        middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)
        if variational_inducing_covar is not None:
            middle_term = SumLinearOperator(variational_inducing_covar, middle_term)
        if trace_mode.on():
            predictive_covar = data_data_covar.add_jitter(self.jitter_val).to_dense() + interp_term.transpose(-1, -2) @ middle_term.to_dense() @ interp_term
        else:
            predictive_covar = SumLinearOperator(data_data_covar.add_jitter(self.jitter_val), MatmulLinearOperator(interp_term.transpose(-1, -2), middle_term @ interp_term))
        return MultivariateNormal(predictive_mean, predictive_covar)

    def __call__(self, x, prior=False, **kwargs):
        if not self.updated_strategy.item() and not prior:
            with torch.no_grad():
                prior_function_dist = self(self.inducing_points, prior=True)
                prior_mean = prior_function_dist.loc
                L = self._cholesky_factor(prior_function_dist.lazy_covariance_matrix.add_jitter(self.jitter_val))
                orig_mean_init_std = self._variational_distribution.mean_init_std
                self._variational_distribution.mean_init_std = 0.0
                variational_dist = self.variational_distribution
                mean_diff = (variational_dist.loc - prior_mean).unsqueeze(-1).type(_linalg_dtype_cholesky.value())
                whitened_mean = L.solve(mean_diff).squeeze(-1)
                covar_root = variational_dist.lazy_covariance_matrix.root_decomposition().root.to_dense()
                covar_root = covar_root.type(_linalg_dtype_cholesky.value())
                whitened_covar = RootLinearOperator(L.solve(covar_root))
                whitened_variational_distribution = variational_dist.__class__(whitened_mean, whitened_covar)
                self._variational_distribution.initialize_variational_distribution(whitened_variational_distribution)
                self._variational_distribution.mean_init_std = orig_mean_init_std
                clear_cache_hook(self)
                self.updated_strategy.fill_(True)
        return super().__call__(x, prior=prior, **kwargs)


data_dim = 1


class SmallFeatureExtractor(nn.Sequential):

    def __init__(self):
        super(SmallFeatureExtractor, self).__init__()
        self.add_module('linear1', nn.Linear(data_dim, 10))
        self.add_module('relu3', nn.ReLU())
        self.add_module('linear4', nn.Linear(10, 1))


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (ScaleToBounds,
     lambda: ([], {'lower_bound': 4, 'upper_bound': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (_DenseBlock,
     lambda: ([], {'num_layers': 1, 'num_input_features': 4, 'bn_size': 4, 'growth_rate': 4, 'drop_rate': 0.5}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (_DenseLayer,
     lambda: ([], {'num_input_features': 4, 'growth_rate': 4, 'bn_size': 4, 'drop_rate': 0.5}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (_Transition,
     lambda: ([], {'num_input_features': 4, 'num_output_features': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
]

class Test_cornellius_gp_gpytorch(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

