Nice, I’m glad you got it figured out.
You need to add the initial state yourself, as you discovered. You could just pt.concatenate the result with A_sim[:, :, 0] though; no need to create a new array and fill it. Also you don’t need to evaluate at.shape(A_array), it’s fine to give symbolic tensors as shapes.