def power(l, a, v=None):
    """
    Compute a^l and the scan sum_i a^i v_i
    Args:
        l (int): power
        a (torch.tensor): (..., N, N)
        v (torch.tensor): (..., N, L)
    """

    i = torch.eye(a.shape[-1]).to(a)  # dtype=a.dtype, device=a.device

    powers = [a]
    ll = 1
    while True:
        if l % 2 == 1:
            i = powers[-1] @ i
        l //= 2
        if l == 0:
            break
        ll *= 2
        powers.append(powers[-1] @ powers[-1])

    if v is None:
        return i

    # Invariants:
    # powers[-1] := A^l
    # l := largest po2 at most L

    # Note that an alternative divide and conquer to compute the reduction is possible
    # and can be embedded into the above loop without caching intermediate powers of A
    # We do this reverse divide-and-conquer for efficiency reasons:
    # 1) it involves fewer padding steps for non-po2 L
    # 2) it involves more contiguous arrays

    # Take care of edge case for non-po2 arrays
    # Note that this initial step is a no-op for the case of power of 2 (l == L)
    k = v.size(-1) - ll
    v_ = powers.pop() @ v[..., ll:]
    v = v[..., :ll]
    v[..., :k] = v[..., :k] + v_

    # Handle reduction for power of 2
    while v.size(-1) > 1:
        v = rearrange(v, '... (z l) -> ... z l', z=2)
        v = v[..., 0, :] + powers.pop() @ v[..., 1, :]
    return i, v.squeeze(-1)
Compute a^l and the scan sum_i a^i v_i
Args:
l (int): power
a (torch.tensor): (..., N, N)
v (torch.tensor): (..., N, L)
towhee/models/vis4mer/utils.py