fix noise shaping
@@ -93,7 +93,15 @@ class AddedDiagLazyTensor(SumLazyTensor):
def _init_cache(self):
*batch_shape, n, k = self._piv_chol_self.shape
self._noise = self._diag_tensor.diag().unsqueeze(-1)
- self._constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
+
+ # the check for constant diag needs to be done carefully for batches.
+ if len(self._noise.shape) == 2:
+ noise_first_element = self._noise[0]
+ else:
+ noise_first_element = self._noise[:, 0, ...]
+ noise_first_element = noise_first_element.unsqueeze(-1)
+
+ self._constant_diag = torch.equal(self._noise, noise_first_element * torch.ones_like(self._noise))
eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
if self._constant_diag:
