Fix test failures by adding _expand_batch and _sum_batch methods to ConstantDiagLazyTensor
@@ -189,4 +189,11 @@ class ConstantDiagLazyTensor(DiagLazyTensor):
of `n x n` diagonal matrices
"""
super(DiagLazyTensor, self).__init__(diag_values, diag_shape=diag_shape)
+ self.diag_shape = diag_shape
self._diag = diag_values.expand(*diag_values.shape[:-1], diag_shape)
+
+ def _expand_batch(self, batch_shape):
+ return self.__class__(self._diag.expand(*batch_shape, self._diag.size(-1)), diag_shape=self.diag_shape)
+
+ def _sum_batch(self, dim):
+ return self.__class__(self._diag.sum(dim), diag_shape=self.diag_shape)
