import torch


def woodbury_logdet(diag_a, mat_u, rank):
    """
    Apply the matrix determinant lemma of Equation (7) in the paper.
    """
    mat_a = torch.diag_embed(diag_a, dim1=1, dim2=2)
    mat_a_inv = torch.diag_embed(1 / diag_a, dim1=1, dim2=2)
    mat_ut = mat_u.transpose(1, 2)
    mat_i = torch.eye(rank, rank, device=diag_a.device).unsqueeze(0).repeat(diag_a.size(0), 1, 1)
    out1 = torch.logdet(mat_i + mat_ut.matmul(mat_a_inv).matmul(mat_u))
    return out1 + torch.logdet(mat_a)