Repository: cornellius-gp/gpytorch
Pull request #965 (2019-11-23 12:33:18)

Indexing MultivariateNormals
if `dist` is a `MultivariateNormal` object, now it is possible to write `dist[index]` and create a new distribution with the appropriate marginalized mean and covariance.

The test provides examples of usage.

This is helpful in my research and it might also be helpful to upstream it?


Patch:
diff --git a/gpytorch/distributions/multivariate_normal.py b/gpytorch/distributions/multivariate_normal.py
index 2dcf9cb6..92e673cd 100644
--- a/gpytorch/distributions/multivariate_normal.py
+++ b/gpytorch/distributions/multivariate_normal.py
@@ -8,7 +8,7 @@ from torch.distributions.kl import register_kl
 from torch.distributions.utils import _standard_normal, lazy_property
 
 from .. import settings
-from ..lazy import LazyTensor, delazify, lazify
+from ..lazy import DiagLazyTensor, LazyTensor, delazify, lazify
 from ..utils.broadcasting import _mul_broadcast_shape
 from .distribution import Distribution
 
@@ -211,6 +211,30 @@ class MultivariateNormal(TMultivariateNormal, Distribution):
     def __truediv__(self, other):
         return self.__mul__(1.0 / other)
 
+    def __getitem__(self, idx):
+        if not isinstance(idx, tuple):
+            idx = (idx,)
+        rest_idx = idx[:-1]
+        last_idx = idx[-1]
+        new_mean = self.mean[idx]
+
+        if len(idx) <= self.mean.dim() - 1 and (Ellipsis not in rest_idx):
+            new_cov = self.lazy_covariance_matrix[idx]
+        elif len(idx) > self.mean.dim():
+            raise IndexError(f"Index {idx} has too many dimensions")
+        else:
+            # In this case we know last_idx corresponds to the last dimension
+            # of mean and the last two dimensions of lazy_covariance_matrix
+            if isinstance(last_idx, int):
+                new_cov = DiagLazyTensor(self.lazy_covariance_matrix.diag()[(*rest_idx, last_idx)])
+            elif isinstance(last_idx, slice):
+                new_cov = self.lazy_covariance_matrix[(*rest_idx, last_idx, last_idx)]
+            elif last_idx is (...):
+                new_cov = self.lazy_covariance_matrix[rest_idx]
+            else:
+                new_cov = self.lazy_covariance_matrix[(*rest_idx, last_idx, slice(None, None, None))][..., last_idx]
+        return self.__class__(mean=new_mean, covariance_matrix=new_cov)
+
 
 @register_kl(MultivariateNormal, MultivariateNormal)
 def kl_mvn_mvn(p_dist, q_dist):


Test patch:
diff --git a/test/distributions/test_multivariate_normal.py b/test/distributions/test_multivariate_normal.py
index 9566de15..6fb88761 100644
--- a/test/distributions/test_multivariate_normal.py
+++ b/test/distributions/test_multivariate_normal.py
@@ -267,6 +267,35 @@ class TestMultivariateNormal(BaseTestCase, unittest.TestCase):
             with least_used_cuda_device():
                 self.test_kl_divergence(cuda=True)
 
+    def test_getitem(self):
+        shape = (2, 4, 3, 2)
+        cov = torch.randn(*shape, shape[1])
+        cov = cov @ cov.transpose(-1, -2)
+        mean = torch.randn(*shape)
+        dist = MultivariateNormal(mean, cov)
+        dist_cov = dist.covariance_matrix
+
+        d = dist[1]
+        assert torch.equal(d.mean, dist.mean[1])
+        self.assertAllClose(d.covariance_matrix, dist_cov[1])
+
+        d = dist[..., 1]
+        assert torch.equal(d.mean, dist.mean[..., 1])
+        cov = dist_cov[..., 1, 1]
+        self.assertAllClose(d.covariance_matrix, cov.unsqueeze(-1) * torch.eye(shape[-2]))
+
+        d = dist[:, [2, 3], :, 1:]
+        assert torch.equal(d.mean, dist.mean[:, [2, 3], :, 1:])
+        self.assertAllClose(d.covariance_matrix, dist_cov[:, [2, 3], :, 1:, 1:])
+
+        d = dist[:, :, ..., [0, 1, 1, 0]]
+        assert torch.equal(d.mean, dist.mean[..., [0, 1, 1, 0]])
+        self.assertAllClose(d.covariance_matrix, dist_cov[..., [0, 1, 1, 0], :][..., [0, 1, 1, 0]])
+
+        d = dist[1, 2, 2, ...]
+        assert torch.equal(d.mean, dist.mean[1, 2, 2, :])
+        self.assertAllClose(d.covariance_matrix, dist_cov[1, 2, 2, :, :])
+
 
 if __name__ == "__main__":
     unittest.main()