force add to zero sign in stable_qr
@@ -26,6 +26,9 @@ def stable_qr(mat):
if torch.any(zeroish):
# can't use in-place operation here b/c it would mess up backward pass
# haven't found a more elegant way to add a jitter diagonal yet...
- jitter_diag = 1e-6 * torch.sign(Rdiag) * zeroish.to(Rdiag)
+ Rdiag_sign = torch.sign(Rdiag)
+ # force zero diagonals to have jitter added to them.
+ Rdiag_sign[Rdiag_sign == 0] = 1.0
+ jitter_diag = 1e-4 * Rdiag_sign * zeroish.to(Rdiag)
R = R + torch.diag_embed(jitter_diag)
return Q, R
