Custom sum_batch for DiagLazyTensor
@@ -64,6 +64,14 @@ class DiagLazyTensor(LazyTensor):
else:
return super(DiagLazyTensor, self).evaluate()
+ def sum_batch(self, sum_batch_size=None):
+ if sum_batch_size is None:
+ diag = self._diag.view(-1, self._diag.size(-1))
+ else:
+ diag = self._diag.view(-1, sum_batch_size, self._diag.size(-1))
+
+ return self.__class__(diag.sum(-2))
+
def zero_mean_mvn_samples(self, num_samples):
if self.ndimension() == 3:
base_samples = torch.randn(
