def linear_system_from_krylov(u, C, D, k):
    """
    Computes the state-space system y = Cx + Du from Krylov matrix K(A, B)

    u: (L, B, ...) ... = H
    C: (..., M, N) ... = H
    D: (..., M)
    k: (..., N, L) Krylov matrix representing b, Ab, A^2b...

    y: (L, B, ..., M)
    """


    # Equivalent ways to perform C @ k, slight speed differences
    k = C @ k # (..., M, L)
    # k = torch.einsum('... m n, ... n l -> ... m l', C, k) # C @ k
    # k = torch.sum(k.unsqueeze(-3) * C.unsqueeze(-1), dim=-2) # (..., M, L) C @ k

    k = rearrange(k, '... m l -> m ... l')
    k = k.to(u) # if training in half precision, need to go back to float32 for the fft
    k = k.unsqueeze(1) # (M, 1, ..., L)

    v = u.unsqueeze(-1).transpose(0, -1) # (1, B, ..., L)
    y = causal_convolution(k, v, fast=True) # (M, B, ..., L)
    y = y.transpose(0, -1) # (L, B, ..., M)
    y = y + u.unsqueeze(-1) * D # (L, B, ..., M)
    return y

    Computes the state-space system y = Cx + Du from Krylov matrix K(A, B)

    u: (L, B, ...) ... = H
    C: (..., M, N) ... = H
    D: (..., M)
    k: (..., N, L) Krylov matrix representing b, Ab, A^2b...

    y: (L, B, ..., M)
    
analyze/convnexts4nd/src/models/sequence/modules/lssl.py