n-order Taylor approximation of matrix exponentiation for PyTorch

expm.py

def expm(x, order=10):
    '''
    nth-order Taylor approximation of matrix exponential
    '''
    if order < 1: 
        raise ValueError("order cannot be smaller than 1")
    
    I = torch.eye(x.shape[-1], dtype=x.dtype, device=x.device)
    result = I
    nom = I
    denom = 1.0
    for i in range(1, order):
        nom = x @ nom
        denom *= i
        result += nom / denom
    
    return result