Fix toeplitz getitem flaky-ness
@@ -216,24 +216,31 @@ class ToeplitzLazyVariable(LazyVariable):
def __getitem__(self, i):
if isinstance(i, tuple):
+ first_index = i[0]
+ if not isinstance(first_index, slice):
+ first_index = slice(first_index, first_index + 1, None)
+ second_index = i[1]
+ if not isinstance(second_index, slice):
+ second_index = slice(second_index, second_index + 1, None)
+
if self.J_left is None:
# Pretend that the matrix is WTW, where W is an identity matrix, with appropriate slices
- # J[i[0], :], C[i[0], :]
- J_left_new = self.c.data.new(range(len(self.c))[i[0]]).unsqueeze(1)
+ # J[first_index, :], C[first_index, :]
+ J_left_new = self.c.data.new(range(len(self.c))[first_index]).unsqueeze(1)
C_left_new = self.c.data.new().resize_as_(J_left_new).fill_(1)
J_left_new = J_left_new.long()
- # J[i[1], :] C[i[1], :]
- J_right_new = self.c.data.new(range(len(self.c))[i[1]]).unsqueeze(1)
+ # J[second_index, :] C[second_index, :]
+ J_right_new = self.c.data.new(range(len(self.c))[second_index]).unsqueeze(1)
C_right_new = self.c.data.new().resize_as_(J_right_new).fill_(1)
J_right_new = J_right_new.long()
else:
- # J[i[0], :], C[i[0], :]
- J_left_new = self.J_left[i[0]]
- C_left_new = self.C_left[i[0]]
+ # J[first_index, :], C[first_index, :]
+ J_left_new = self.J_left[first_index]
+ C_left_new = self.C_left[first_index]
- # J[i[1], :] C[i[1], :]
- J_right_new = self.J_right[i[1]]
- C_right_new = self.C_right[i[1]]
+ # J[second_index, :] C[second_index, :]
+ J_right_new = self.J_right[second_index]
+ C_right_new = self.C_right[second_index]
if self.added_diag is not None:
if len(J_left_new) != len(J_right_new):
@@ -241,7 +248,7 @@ class ToeplitzLazyVariable(LazyVariable):
diagonal component to make it non-square is probably not intended.\
It is ambiguous which diagonal elements to choose')
- diag_new = self.added_diag[i[0]]
+ diag_new = self.added_diag[first_index]
else:
diag_new = None
@@ -249,6 +256,9 @@ class ToeplitzLazyVariable(LazyVariable):
J_right_new, C_right_new, diag_new)
else:
+ if not isinstance(i, slice):
+ i = slice(i, i + 1, None)
+
if self.J_left is not None:
J_left_new = self.J_left[i]
C_left_new = self.C_left[i]
