def krylov(l, a, b, c=None, return_power=False):
    """
    Compute the Krylov matrix (b, ab, a^2b, ...) using the squaring trick.
    If return_power=True, return A^{L-1} as well
    Args:
        a (torch.tensor): torch.tensor
        b (torch.tensor): torch.tensor
        c (torch.tensor): torch.tensor
        return_power (bool): return_power
    """
    # TODO There is an edge case if l=1 where output doesn't get broadcasted, which might be an issue
    #  if caller is expecting broadcasting semantics... can deal with it if it arises

    x = b.unsqueeze(-1)  # (..., N, 1)
    a_ = a

    al = None
    if return_power:
        al = torch.eye(a.shape[-1], dtype=a.dtype, device=a.device)
        l_ = l-1

    done = l == 1
    # loop invariant: l_ represents how many indices left to compute
    while not done:
        if return_power:
            if l_ % 2 == 1:
                al = a_ @ al
            l_ //= 2

        # Save memory on last iteration
        ll = x.shape[-1]
        if l - ll <= ll:
            done = True
            x_ = x[..., :l-ll]
        else:
            x_ = x

        x_ = a_ @ x_
        x = torch.cat([x, x_], dim=-1)  # there might be a more efficient way of ordering axes
        if not done:
            a_ = a_ @ a_

    assert x.shape[-1] == l

    if c is not None:
        x = torch.einsum('...nl, ...n -> ...l', x, c)
    x = x.contiguous()  # WOW!!
    if return_power:
        return x, al
    else:
        return x
Compute the Krylov matrix (b, ab, a^2b, ...) using the squaring trick.
If return_power=True, return A^{L-1} as well
Args:
a (torch.tensor): torch.tensor
b (torch.tensor): torch.tensor
c (torch.tensor): torch.tensor
return_power (bool): return_power
towhee/models/vis4mer/utils.py