Using PyTensor.scan to loop over multidimensional array

@jessegrabowski:
The issue I seem to have with both version is that A_new[i+1] is informed by the previous step A_new[i]. In my example above I used at.set_subtensor(), but how does that work together with scan?