Efficient iterated discrete dynamics in Flux

I’m trying to implement a Koopman operator layer in Flux. This is the operation of starting with a vector, x, and iteratively multiplying it by a dynamics matrix, K, for n steps, e.g.

x1 = x
x2 = K * x1
x3 = K * x2
etc etc

where n is typically a couple of thousand steps, so you end up with an array with a small number of rows and a lot of columns: [x x1 x2 x3…]

One easy way to do this is to calculate all the matrix powers, e.g. K, K^2, K^3, etc, and then do the matrix-vector multiply and combine everything:

hcat(map(i->(K^i)*x, 1:n)...)

but this is slow as molasses. As best I can determine, the most efficient way to do it in Julia would be to preallocate the array you need and just do a for loop:

xi = Array{Float32}(undef, size(K,1), n)
for i=2:n
    xi[:,i] = K*xi[:, i-1]

but Flux throws a “Mutating arrays is not supported” error when I do this.

What the Julia Way to do this that’s compatible with Flux?