@fritzo

def test_multivariate_normal_shape(self):
    mean = torch.randn(5, 3, requires_grad=True)
    mean_no_batch = torch.randn(3, requires_grad=True)
    mean_multi_batch = torch.randn(6, 5, 3, requires_grad=True)

    # construct PSD covariance
    tmp = torch.randn(3, 10)
    cov = torch.tensor(torch.matmul(tmp, tmp.t()) / tmp.shape[-1], requires_grad=True)
    prec = torch.tensor(cov.inverse(), requires_grad=True)
    scale_tril = torch.tensor(torch.potrf(cov, upper=False), requires_grad=True)

    # construct batch of PSD covariances
    tmp = torch.randn(6, 5, 3, 10)
    cov_batched = torch.tensor((tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1), requires_grad=True)
    prec_batched = [C.inverse() for C in cov_batched.view((-1, 3, 3))]
    prec_batched = torch.stack(prec_batched).view(cov_batched.shape)
    scale_tril_batched = [torch.potrf(C, upper=False) for C in cov_batched.view((-1, 3, 3))]
    scale_tril_batched = torch.stack(scale_tril_batched).view(cov_batched.shape)

    # ensure that sample, batch, event shapes all handled correctly
    self.assertEqual(MultivariateNormal(mean, cov).sample().size(), (5, 3))
    self.assertEqual(MultivariateNormal(mean_no_batch, cov).sample().size(), (3,))
    self.assertEqual(MultivariateNormal(mean_multi_batch, cov).sample().size(), (6, 5, 3))
    self.assertEqual(MultivariateNormal(mean, cov).sample((2,)).size(), (2, 5, 3))
    self.assertEqual(MultivariateNormal(mean_no_batch, cov).sample((2,)).size(), (2, 3))
    self.assertEqual(MultivariateNormal(mean_multi_batch, cov).sample((2,)).size(), (2, 6, 5, 3))
    self.assertEqual(MultivariateNormal(mean, cov).sample((2, 7)).size(), (2, 7, 5, 3))
    self.assertEqual(MultivariateNormal(mean_no_batch, cov).sample((2, 7)).size(), (2, 7, 3))
    self.assertEqual(MultivariateNormal(mean_multi_batch, cov).sample((2, 7)).size(), (2, 7, 6, 5, 3))
    self.assertEqual(MultivariateNormal(mean, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
    self.assertEqual(MultivariateNormal(mean_no_batch, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
    self.assertEqual(MultivariateNormal(mean_multi_batch, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
    self.assertEqual(MultivariateNormal(mean, precision_matrix=prec).sample((2, 7)).size(), (2, 7, 5, 3))
    self.assertEqual(MultivariateNormal(mean, precision_matrix=prec_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
    self.assertEqual(MultivariateNormal(mean, scale_tril=scale_tril).sample((2, 7)).size(), (2, 7, 5, 3))
    self.assertEqual(MultivariateNormal(mean, scale_tril=scale_tril_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))

    # check gradients

def fail_test(msg):
    if raise_exception:

The gradient seems exactly 5x off.

I think that I have found the bug at #8649

Although logically, this is supposed to be faster, I am interested in knowing what kind of speedup is obtained in the calculation of the Mahalanobis distance and scale_tril. Is it possible for you to get some numbers?

@vishwakftw Here is some comparison https://gist.github.com/fehiepsi/80810668bae0fbaa926e5b6319c1ecd3

Using trtrs is much faster than inverse when the size of matrix is large.

The timings are amazing!! Thanks!!

@fehiepsi For GPU, you should always torch.cuda.synchronize() before measuring time.

Regarding the gradcheck tests, you could modify the _gradcheck_log_prob function to convert the ctor_params to double. Maybe that will help.

@vishwakftw Since the gradients are exactly 5x off, I doubt it is a precision issue.

Ah yes, sorry - missed that detail above.

About the gradcheck which is 5x off, it seems that we have caught the issue at #8649 and @ssnl will take a look for it.

I am wondering about two other issues. The first is if scale_factor and scale_diag are good names. In tensorflow, they used scale_perturb_factor and scale_diag. factor term comes from "factor loading matrix" in factor analysis. It is also called components in PCA. To me, scale_factor is a good name but I am open to other choices.

Another issue is when to do "broadcasting". Different from other distributions, it seems that it is better to avoid using .expand in the first place. For example, computing scale_tril from expanded_covariance_matrix is less efficient than computing scale_tril from covariance_matrix then expand. So I think that depending on "whatever" scale input (cov, tril, precision) is, we will calculate _scale_tril_unexpand in the init method. Other properties will be expanded from the corresponding calculation with this _scale_tril_unexpand. This saves both memory/computation and makes properties of MVN consistency with other distributions.

Regarding "broadcasting", to solve Linv_y, I broadcast L to match the batch shape of y, then use trtrs to solve. This might be inefficient when L and y have different batch shape. Assume that batch_shape of the distribution is 9, y has shape 10 x 9 x 8, and L has shape 3 x 8 x 8, then it is better to convert L to 9 x 8 x 8, y to 9 x 80, and then apply batch_trtrs with L and y, then reshape the result. I am planning to implement that way after the gradcheck issue is fixed but wondering if there is a better workaround. FYI, we don't face this issue if we take the inverse of L first and then multiply because .matmul will automatically do "broadcasting".

@fritzo I think that you will have better ideas than me on these issues. Could you please give some thoughts on them?

@fehiepsi The names seem fine; it would be nice to match tensorflow so that code is more portable.

Re: broadcasting, I've faced similar issues and haven't yet found a good solution. If you find a good pattern here, we may want to copy it elsewhere, e.g. in computing a batch of VonMises normalizing factors when scale is broadcast.

Here's a proposal:

1. in the __init__ method, save the original unbroadcast covariance as self._unbroadcasted_covariance_matrix
2. in the .scale_tril lazy property, use self._unbroadcasted_covariance_matrix and then apply broadcast_all or similar logic after computing the cholesky decomposition, and return this broadcasted result from the lazy property

We may additionally want to store self._unbroadcasted_scale_tril for other computations.

@fritzo Currently, the implementation for .log_prob is fast (O(NM^2)), but the implementation for .sample is slow (O(N^3)) because I just compute scale_tril of W @ W.t() + D, which is a matrix of size N x N. This has been in my mind for a while and I have thought that there is no solution for it.

Today, I checked again the slack's distributions room and found that @stepelu has suggested a very nice idea to reparameterize: take two random variables eps1, eps2 and compute W.eps1 + D^(1/2).eps2. It is a "whoa" moment for me. Sampling this way is pretty much fast! I will make changes corresponding to that.

@fritzo @apaszke In addition, I checked again tensorflow's MultivariateNormalDiagPlusLowRank. It is different from us: their low rank version is for scale, not covariance_matrix (which is scale @ scale.t()). So I wonder if we stick with the names: scale_factor and scale_diag or change it to something like components and diagonal. If we keep scale_factor and scale_diag, then people might be confused with the ones in tensorflow.

@stepelu has suggested a very nice idea to reparameterize

Nice! That sounds very similar to our pyro.contrib.gp.kernels.Coregionalize sampler.

change [param names] to something like components and diagonal

I agree, those look like reasonable names. Alternatively maybe we should use scale_factor and cov_diag? Or cov_factor and cov_diag? Whatever you think is clearest.

@fritzo I have fixed the bug and revert examples to the old version (which failed for this bug). Thanks a lot!

@fehiepsi Why was the .rsample bug not caught in the .mean and .variance tests? Are those tests missing?

Thank you, @fritzo!

@apaszke Could you please help me retest this? Travis-ci has been interupted for some reason.

@pytorchbot retest this please

@ssnl retest this please doesn't apply to Travis :) Go to the Travis UI and just click the "retry" button (I have done so)

@ezyang Good to know. Thanks! :)

@ssnl @ezyang we have a failure at caffe2 build; should we retest or this is ready to merge?

caffe2 biuld failure looks spurious