From here, I believe the correct way to do it is to make the struct mutable and “swap out the arrays” after each application of the layer. I based that on how RNNCell was treating h, but that is not an array. On second read through, perhaps the right thing is to make VariationalAffine a more “cell”-like layer (like LSTMCell and RNNCell) and wrap that in a Recur to keep track of the state.
This is a common early misconception. The key to performance is type-stability
.