Migrate RandomVariable to use PyTorch distributions about gpytorch HOT 9 CLOSED

import torch
from torch.distributions import MultivariateNormal as TMultivariateNormal
from torch.distributions.multivariate_normal import _batch_mv

from ..lazy import LazyTensor


class MultivariateNormal(TMultivariateNormal):
    """
    Constructs a multivariate Normal random variable, based on mean and covariance
    Can be multivariate, or a batch of multivariate gaussians

    Passing a vector mean corresponds to a multivariate Gaussian
    Passing a matrix mean corresponds to a batch of multivariate Gaussians

    Args:
        mean (Tensor): vector n or matrix b x n mean of Gaussian distribution
        covar (Tensor): matrix n x n or batch matrix b x n x n covariance of
            Gaussian distribution
    """

    def __init__(self, mean, covariance_matrix, validate_args=False):
        if not isinstance(mean, LazyTensor) and not isinstance(
            covariance_matrix, LazyTensor
        ):
            return TMultivariateNormal(loc=mean, covariance_matrix=covariance_matrix, validate_args=validate_args)
        else:
            return _LazyMultivariateNormal(
                mean=mean, covariance_matrix=covariance_matrix, validate_args=validate_args
            )

    def __add__(self, other):
        if isinstance(other, MultivariateNormal):
            return MultivariateNormal(
                mean=self._mean + other.mean,
                covariance_matrix=self.covariance_matrix + other.covariance_matrix,
            )
        elif isinstance(other, int) or isinstance(other, float):
            return MultivariateNormal(self.mean + other, self.covariance_matrix)
        else:
            raise RuntimeError("Unsupported type for addition w/ MultivariateNormal")

    def __radd__(self, other):
        if other == 0:
            return self
        return self.__add__(other)

    def __div__(self, other):
        return self.__mul__(1. / other)

    def __mul__(self, other):
        if not isinstance(other, int) and not isinstance(other, float):
            raise RuntimeError("Can only multiply by scalars")
        return self.__class__(
            loc=self.mean * other,
            covariance_matrix=self.covariance_matrix * (other ** 2),
        )


class _LazyMultivariateNormal(MultivariateNormal):
    def __init__(self, mean, covariance_matrix, validate_args=False):
        # add validation here...
        self.loc = mean
        self._covar = covariance_matrix

    @property
    def scale_tril(self):
        raise NotImplementedError()

    @property
    def covariance_matrix(self):
        return self._covar

    @property
    def precision_matrix(self):
        raise NotImplementedError()

    def rsample(self, sample_shape=torch.Size()):
        shape = self._extended_shape(sample_shape)
        eps = self.loc.new_empty(shape).normal_()
        # this will fail, rewrite using LazyTensor code
        return self.loc + _batch_mv(self._covar.root_decomposition(), eps)

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        diff = value - self.loc
        # re-write this for LazyTensors
        # M = _batch_mahalanobis(self._unbroadcasted_scale_tril, diff)
        # half_log_det = _batch_diag(self._unbroadcasted_scale_tril).log().sum(-1)
        return -0.5 * (self._event_shape[0] * math.log(2 * math.pi) + M) - half_log_det

    def entropy(self):
        # re-write this for LazyTensors
        # half_log_det = _batch_diag(self._unbroadcasted_scale_tril).log().sum(-1)
        H = 0.5 * self._event_shape[0] * (1.0 + math.log(2 * math.pi)) + half_log_det
        if len(self._batch_shape) == 0:
            return H
        else:
            return H.expand(self._batch_shape)