我想做点什么
x[i, :, :] = (rhs[i, :, :]-diag[i] * x[i+1, :, :])/diag[i]
其中x和rhs是尺寸为（T，L，S）的3D numpy数组。 diag是一个大小为T的一维数组。
这将正确播放。
但是现在我想编写一个类似的函数来处理2D数组或其他一些维度。我如何一般地编写它，以便它可以在任何大小为T的第一个维度的数组上工作。我不想只用不同数量的冒号重复代码，因为函数中有很多这样的行。
答案 0 :(得分：2)
x[i] = (rhs[i] - diag[i] * x[i+1])/diag[i]
那些冒号完全没必要。