batched sqrt inverses using newton schulz

batched_sqrt_inv.py

import torch

def compute_H_inv_cubic(A, num_iters=10):
    X = torch.eye(A.shape[-1], dtype=A.dtype, device=A.device).expand(A.shape)
    for _ in range(num_iters):
        X_update = torch.einsum('...ij,...jk,...kl,...lm->...im', X, A, X, X)
        X = 1.5 * X - 0.5 * X_update
    return X

def compute_H_inv_quintic(A, num_iters=5):
    I = torch.eye(A.shape[-1], dtype=A.dtype, device=A.device).expand(A.shape)
    X = I.clone()
    for _ in range(num_iters):
        Y = torch.einsum('...ij,...jk,...kl->...il', A, X, X)
        X = X @ (1.875 * I - 1.25 * Y + 0.375 * (Y @ Y))
    return X