I want to do something like
x[i, :, :] = (rhs[i, :, :]-diag[i] * x[i+1, :, :])/diag[i]
where x and rhs are 3D numpy arrays of size (T,L,S). diag is a 1D array of size T.
This will broadcast properly.
But now I'd like to write a similar function to work on 2D arrays or some other number of dimensions. How can I write this generically so that it will work on any array that has first dimension of size T. I don't want to duplicate code with just a different number of colons since there are a lot of these kinds of lines in the function.