qpytorch 0.1__py3-none-any.whl

qpytorch/variational/delta_variational_distribution.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+
+from typing import Union
+
+import torch
+
+from ..distributions import Delta, Distribution, MultivariateNormal, MultivariateQExponential
+from ._variational_distribution import _VariationalDistribution
+
+
+class DeltaVariationalDistribution(_VariationalDistribution):
+    """
+    This :obj:`~qpytorch.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 (q-exponential) noise to add to the mean initialization.
+    """
+
+    def __init__(
+        self,
+        num_inducing_points: int,
+        batch_shape: torch.Size = torch.Size([]),
+        mean_init_std: float = 1e-3,
+        **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) -> Distribution:
+        return Delta(self.variational_mean)
+
+    def initialize_variational_distribution(self, prior_dist: Union[MultivariateNormal, MultivariateQExponential]) -> None:
+        self.variational_mean.data.copy_(prior_dist.mean)
+        self.variational_mean.data.add_(torch.randn_like(prior_dist.mean), alpha=self.mean_init_std)

qpytorch/variational/grid_interpolation_variational_strategy.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+
+import torch
+from linear_operator.operators import InterpolatedLinearOperator
+from linear_operator.utils.interpolation import left_interp
+
+from ..distributions import MultivariateNormal, MultivariateQExponential
+from gpytorch.utils.interpolation import Interpolation
+from gpytorch.utils.memoize import cached
+from ._variational_strategy import _VariationalStrategy
+
+
+class GridInterpolationVariationalStrategy(_VariationalStrategy):
+    r"""
+    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 (~qpytorch.models.ApproximateQEP) model: Model this strategy is applied to.
+        Typically passed in when the VariationalStrategy is created in the
+        __init__ method of the user defined model.
+        It should contain power if Q-Exponential distribution is involved in.
+    :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 ~qpytorch.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)
+        # TODO: investigate why smaller than 1e-3 breaks some tests
+        if hasattr(self.model, 'power'):
+            res = MultivariateQExponential(out.mean, out.lazy_covariance_matrix.add_jitter(1e-3), power=self.model.power)
+        else:
+            res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter(1e-3))
+        return res
+
+    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None):
+        if variational_inducing_covar is None:
+            raise RuntimeError(
+                "GridInterpolationVariationalStrategy is only compatible with Gaussian (Q-Exponential) variational "
+                f"distributions. Got ({self.variational_distribution.__class__.__name__}."
+            )
+
+        variational_distribution = self.variational_distribution
+
+        # Get interpolations
+        interp_indices, interp_values = self._compute_grid(x)
+
+        # Compute test mean
+        # Left multiply samples by interpolation matrix
+        predictive_mean = left_interp(interp_indices, interp_values, inducing_values.unsqueeze(-1))
+        predictive_mean = predictive_mean.squeeze(-1)
+
+        # Compute test covar
+        predictive_covar = InterpolatedLinearOperator(
+            variational_distribution.lazy_covariance_matrix,
+            interp_indices,
+            interp_values,
+            interp_indices,
+            interp_values,
+        )
+        if hasattr(self.model, 'power'):
+            output = MultivariateQExponential(predictive_mean, predictive_covar, power=self.model.power)
+        else:
+            output = MultivariateNormal(predictive_mean, predictive_covar)
+        return output

qpytorch/variational/independent_multitask_variational_strategy.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+import warnings
+
+import torch
+from linear_operator.operators import RootLinearOperator
+
+from ..distributions import MultitaskMultivariateNormal, MultivariateNormal
+from ..module import Module
+from ._variational_strategy import _VariationalStrategy
+
+
+class IndependentMultitaskVariationalStrategy(_VariationalStrategy):
+    """
+    IndependentMultitaskVariationalStrategy wraps an existing
+    :obj:`~qpytorch.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 ~qpytorch.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):
+        r"""
+        See :class:`LMCVariationalStrategy`.
+        """
+        function_dist = self.base_variational_strategy(x, prior=prior, **kwargs)
+
+        if task_indices is None:
+            # Every data point will get an output for each task
+            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:
+            # Each data point will get a single output corresponding to a single task
+
+            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
+
+            # Create a mask to choose specific task assignment
+            shape = list(function_dist.batch_shape + function_dist.event_shape)
+            shape[task_dim] = 1
+            task_indices = task_indices.expand(shape).squeeze(task_dim)
+
+            # Create a mask to choose specific task assignment
+            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:`~qpytorch.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 ~qpytorch.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)

qpytorch/variational/lmc_variational_strategy.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+
+from typing import Optional, Union
+
+import torch
+from linear_operator.operators import KroneckerProductLinearOperator, RootLinearOperator
+from linear_operator.utils.interpolation import left_interp
+from torch import LongTensor, Tensor
+
+from .. import settings
+from ..distributions import MultitaskMultivariateNormal, MultitaskMultivariateQExponential, MultivariateNormal, MultivariateQExponential
+from ..module import Module
+from ._variational_strategy import _VariationalStrategy
+
+
+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])
+
+    # We will use the left_interp helper to do the indexing
+    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):
+    r"""
+    LMCVariationalStrategy is an implementation of the "Linear Model of Coregionalization"
+    for multitask GPs (QEPs). 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 (QEP).
+    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:`~qpytorch.variational.VariationalStrategy`.
+    The output will either be a :obj:`~gpytorch.distributions.MultitaskMultivariateNormal` (:obj:`~qpytorch.distributions.MultitaskMultivariateQExponential`) distribution
+    (if we wish to evaluate all tasks for each input) or a :obj:`~gpytorch.distributions.MultivariateNormal` (:obj:`~qpytorch.distributions.MultivariateQExponential`)
+    (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 (QEPs), where one
+    of the batch dimensions corresponds to the latent function dimension.
+
+    .. note::
+
+        The batch shape of the base :obj:`~qpytorch.variational.VariationalStrategy` does not
+        necessarily have to correspond to the batch shape of the underlying GP (QEP) objects.
+
+        For example, if the base variational strategy has a batch shape of `[3]` (corresponding
+        to 3 latent functions), the GP (QEP) 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(qpytorch.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 = qpytorch.variational.MeanFieldVariationalDistribution(
+        >>>             inducing_points.size(-1), batch_shape=torch.Size([3]),
+        >>>         )
+        >>>         variational_strategy = qpytorch.variational.LMCVariationalStrategy(
+        >>>             qpytorch.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 = qpytorch.means.ConstantMean(batch_shape=torch.Size([3]))
+        >>>         self.covar_module = qpytorch.kernels.ScaleKernel(
+        >>>             qpytorch.kernels.RBFKernel(batch_shape=torch.Size([3])),
+        >>>             batch_shape=torch.Size([3]),
+        >>>         )
+        >>>
+
+    :param base_variational_strategy: Base variational strategy
+    :param num_tasks: The total number of tasks (output functions)
+    :param 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**
+    :param jitter_val: Amount of diagonal jitter to add for Cholesky factorization numerical stability
+    """
+
+    def __init__(
+        self,
+        base_variational_strategy: _VariationalStrategy,
+        num_tasks: int,
+        num_latents: int = 1,
+        latent_dim: int = -1,
+        jitter_val: Optional[float] = 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
+
+        # Check if no functions
+        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}, "
+                f"expected the function dim {latent_dim} to be {num_latents}."
+            )
+        self.num_latents = num_latents
+        self.latent_dim = latent_dim
+
+        # Make the batch_shape
+        self.batch_shape = list(batch_shape)
+        del self.batch_shape[self.latent_dim]
+        self.batch_shape = torch.Size(self.batch_shape)
+
+        # LCM coefficients
+        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) -> Union[MultivariateNormal, MultivariateQExponential]:
+        return self.base_variational_strategy.prior_distribution
+
+    @property
+    def variational_distribution(self) -> Union[MultivariateNormal, MultivariateQExponential]:
+        return self.base_variational_strategy.variational_distribution
+
+    @property
+    def variational_params_initialized(self) -> bool:
+        return self.base_variational_strategy.variational_params_initialized
+
+    def kl_divergence(self) -> Tensor:
+        return super().kl_divergence().sum(dim=self.latent_dim)
+
+    def __call__(
+        self, x: Tensor, prior: bool = False, task_indices: Optional[LongTensor] = None, **kwargs
+    ) -> Union[MultitaskMultivariateNormal, MultitaskMultivariateQExponential, MultivariateNormal, MultivariateQExponential]:
+        r"""
+        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` (:class:`~qpytorch.distributions.MultitaskMultivariateQExponential`).
+        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` (:class:`~qpytorch.distributions.MultivariateQExponential`).
+
+        :param x: (... x N x D) Input locations to evaluate variational strategy
+        :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` or :class:`~qpytorch.distributions.MultitaskMultivariateQExponential`).
+            If this **is** provided, then the returned distribution evaluates each input only on its assigned task.
+            (returns :class:`~gpytorch.distributions.MultivariateNormal` or :class:`~qpytorch.distributions.MultivariateQExponential`).
+        :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)`.
+        :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 (~qpytorch.distributions.MultitaskMultivariateQExponential) (... x N x num_tasks)
+            or ~gpytorch.distributions.MultivariateNormal (~qpytorch.distributions.MultivariateQExponential) (... 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)
+
+            # Every data point will get an output for each task
+            # Therefore, we will set up the lmc_coefficients shape for a matmul
+            lmc_coefficients = self.lmc_coefficients.expand(*latent_dist.batch_shape, self.lmc_coefficients.size(-1))
+
+            # Mean: ... x N x num_tasks
+            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
+            )
+
+            # Covar: ... x (N x num_tasks) x (N x num_tasks)
+            latent_covar = latent_dist.lazy_covariance_matrix
+            lmc_factor = RootLinearOperator(lmc_coefficients.unsqueeze(-1))
+            covar = KroneckerProductLinearOperator(latent_covar, lmc_factor).sum(latent_dim)
+            # Add a bit of jitter to make the covar PD
+            covar = covar.add_jitter(self.jitter_val)
+
+            # Done!
+            if isinstance(latent_dist, MultivariateNormal):
+                function_dist = MultitaskMultivariateNormal(mean, covar)
+            elif isinstance(latent_dist, MultivariateQExponential):
+                function_dist = MultitaskMultivariateQExponential(mean, covar, power=latent_dist.power)
+
+        else:
+            # Each data point will get a single output corresponding to a single task
+            # Therefore, we will select the appropriate lmc coefficients for each task
+            lmc_coefficients = _select_lmc_coefficients(self.lmc_coefficients, task_indices)
+
+            # Mean: ... x N
+            mean = (latent_dist.mean * lmc_coefficients).sum(latent_dim)
+
+            # Covar: ... x N x N
+            latent_covar = latent_dist.lazy_covariance_matrix
+            lmc_factor = RootLinearOperator(lmc_coefficients.unsqueeze(-1))
+            covar = (latent_covar * lmc_factor).sum(latent_dim)
+            # Add a bit of jitter to make the covar PD
+            covar = covar.add_jitter(self.jitter_val)
+
+            # Done!
+            if isinstance(latent_dist, MultivariateNormal):
+                function_dist = MultivariateNormal(mean, covar)
+            elif isinstance(latent_dist, MultivariateQExponential):
+                function_dist = MultivariateQExponential(mean, covar, power=latent_dist.power)
+
+        return function_dist

qpytorch/variational/mean_field_variational_distribution.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python3
+
+import torch
+from linear_operator.operators import DiagLinearOperator
+
+from ..distributions import MultivariateNormal, MultivariateQExponential
+from ._variational_distribution import _VariationalDistribution
+
+
+class MeanFieldVariationalDistribution(_VariationalDistribution):
+    """
+    A :obj:`~qpytorch.variational._VariationalDistribution` that is defined to be a multivariate normal (q-exponential) distribution
+    with a diagonal covariance matrix. This will not be as flexible/expressive as a
+    :obj:`~qpytorch.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 (q-exponential) noise to add to the mean initialization.
+    """
+
+    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **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))
+        
+        if 'power' in kwargs: self.power = kwargs.pop('power')
+
+    @property
+    def variational_stddev(self):
+        # TODO: if we don't multiply self._variational_stddev by a mask of one, Pyro models fail
+        # not sure where this bug is occuring (in Pyro or PyTorch)
+        # throwing this in as a hotfix for now - we should investigate later
+        mask = torch.ones_like(self._variational_stddev)
+        return self._variational_stddev.mul(mask).abs().clamp_min(1e-8)
+
+    def forward(self):
+        # TODO: if we don't multiply self._variational_stddev by a mask of one, Pyro models fail
+        # not sure where this bug is occuring (in Pyro or PyTorch)
+        # throwing this in as a hotfix for now - we should investigate later
+        mask = torch.ones_like(self._variational_stddev)
+        variational_covar = DiagLinearOperator(self._variational_stddev.mul(mask).pow(2))
+        if not hasattr(self, 'power'):
+            return MultivariateNormal(self.variational_mean, variational_covar)
+        else:
+            return MultivariateQExponential(self.variational_mean, variational_covar, power=self.power)
+
+    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)