Fixes in MultitaskMultivariateNormal
@@ -30,7 +30,7 @@ class MultitaskMultivariateNormal(MultivariateNormal):
if not torch.is_tensor(covariance_matrix) and not isinstance(covariance_matrix, LazyTensor):
raise RuntimeError("The covariance of a MultitaskMultivariateNormal must be a Tensor or LazyTensor")
- if mean.ndimension() not in {2, 3}:
+ if mean.dim() < 2:
raise RuntimeError("mean should be a matrix or a batch matrix (batch mode)")
self._output_shape = mean.shape
@@ -50,12 +50,12 @@ class MultitaskMultivariateNormal(MultivariateNormal):
def from_independent_mvns(cls, mvns):
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")
- if len(mvns[0].batch_shape) > 1:
- raise ValueError("Multiple batch dimensions are not supported in from_independent_mvns.")
mean = torch.stack([mvn.mean for mvn in mvns], -1)
# TODO: To do the following efficiently, we don't want to evaluate the
# covariance matrices. Instead, we want to use the lazies directly in the
@@ -66,10 +66,7 @@ class MultitaskMultivariateNormal(MultivariateNormal):
dim=0,
output_device=mean.device
)
- covar_lazy = BlockDiagLazyTensor(
- covar_blocks_lazy,
- block_dim=0
- )
+ covar_lazy = BlockDiagLazyTensor(covar_blocks_lazy, block_dim=0)
return cls(mean=mean, covariance_matrix=covar_lazy)
def get_base_samples(self, sample_shape=torch.Size()):
