| import torch |
| import math |
|
|
| def flat_to_sym(V, N): |
| A = torch.zeros(N, N, dtype=V.dtype, device=V.device) |
| idxs = torch.tril_indices(N, N, device=V.device) |
| A[idxs.unbind()] = V |
| A[idxs[1, :], idxs[0, :]] = V |
| return A |
|
|
| def block_LDL(H, b, check_nan=True): |
| n = H.shape[0] |
| assert (n % b == 0) |
| m = n // b |
| try: |
| L = torch.linalg.cholesky(H) |
| except: |
| return None |
| DL = torch.diagonal(L.reshape(m, b, m, b), dim1=0, dim2=2).permute(2, 0, 1) |
| D = (DL @ DL.permute(0, 2, 1)).cpu() |
| DL = torch.linalg.inv(DL) |
| L = L.view(n, m, b) |
| for i in range(m): |
| L[:, i, :] = L[:, i, :] @ DL[i, :, :] |
|
|
| if check_nan and L.isnan().any(): |
| return None |
| |
| L = L.reshape(n, n) |
| return (L, D.to(DL.device)) |
|
|
| def approx_int_sqrt(n): |
| p = int(math.floor(math.sqrt(n))) |
| while (n % p != 0): |
| p -= 1 |
| return (p, n // p) |
|
|
| def regularize_H(H, n, sigma_reg): |
| H.div_(torch.diag(H).mean()) |
| idx = torch.arange(n) |
| H[idx,idx] += sigma_reg |
| return H |
| |
|
|
