Toeplitz diag preserves autograd
@@ -89,9 +89,12 @@ class ToeplitzLazyVariable(LazyVariable):
if len(self.J_left) != len(self.J_right):
raise RuntimeError('diag not supported for non-square interpolated Toeplitz matrices.')
WTW_diag = Variable(torch.zeros(len(self.J_right)))
+ n = len(self.J_right)
+ mask = Variable(torch.zeros(n))
for i in range(len(self.J_right)):
- WTW_diag[i] = self[i:i + 1, i:i + 1].evaluate()
-
+ mask.data[i] = 1
+ WTW_diag = WTW_diag + mask * self[i:i + 1, i:i + 1].evaluate().squeeze(1).expand(n)
+ mask.data[i] = 0
return WTW_diag
def evaluate(self):
