python numpy マトリックスを不均等な区間で抽出するプログラムを書きたい

Question

マトリックスを不均等な区間で抽出するプログラムを書きたいと思っています。

numpyを用いて一応以下のようなコードで動くかと思うのですが,generatorの作り方とかが雑な気もします。

def assign_interval(A,B,axis):
if axis == 0:
for i in range(len(B)-1):
yield A[:,B[i]:B[i+1]]
else:
for i in range(len(B)-1):
yield A[B[i]:B[i+1],:]

もっと効率の良い方法を思いつく方いらっしゃれば教えていただけると助かります。

Answer 1

こちらで解決作を提示していただきました。
np.splitでgeneratorの生成までやってくれるようです。