Address cholesky issue with arrowhead mass trix
@@ -28,12 +28,24 @@ def sqrt(x):
# NB: the complexity is O(N * head_size^2)
# ref: https://en.wikipedia.org/wiki/Schur_complement#Background
Dsqrt = x.bottom_diag.sqrt()
+
+ # On cholesky error, retry with smaller tail part B.
+ num_attempts = 6
+ for i in range(num_attempts):
B_Dsqrt = B / Dsqrt.unsqueeze(-2) # shape: head_size x N
schur_complement = A - B_Dsqrt.matmul(B_Dsqrt.t()) # complexity: head_size^2 x N
# we will decompose schur_complement to U @ U.T (so that the sqrt matrix
# is upper triangular) using some `flip` operators:
# flip(cholesky(flip(schur_complement)))
- top_left = torch.flip(torch.flip(schur_complement, (-2, -1)).cholesky(), (-2, -1))
+ try:
+ top_left = torch.flip(torch.cholesky(torch.flip(schur_complement, (-2, -1))), (-2, -1))
+ break
+ except RuntimeError:
+ B = B / 2
+ continue
+ raise RuntimeError("Singular schur complement in computing Cholesky of the input"
+ " arrowhead matrix")
+
top_right = B_Dsqrt
top = torch.cat([top_left, top_right], -1)
bottom_diag = Dsqrt
