Comment by iskander

I expected the core of the algorithm to be a parallel prefix scan though (isn't that the point of Mamba?):

    for i in range(l):
            x = deltaA[:, :, i] \* x + deltaB_u[:, :, i]
            y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')
            ys.append(y)