python - 任意の軸でボリュームを反復処理するより Pythonic な方法は?

Question

3D numpy 配列 (ボリュームと呼びます) を取り、それを 2D スライスのリストに変換する関数があります。ユーザーがスライスする軸を指定できるようにしたい。以下のコードでこれを管理していますが、トリプル if ステートメントは、これを行う最もエレガントな方法とは思えません。これがより良い方法で達成できるかどうかについて、人々の考えをいただければ幸いです。

axis = 0 # Can be set to 0, 1, or 2 

volume = np.ones((100, 100, 100))

n_slices = volume.shape[axis]

slices = []

for i in range(n_slices):

    if axis == 0:
        my_slice = volume[i, :, :]
    elif axis == 1:
        my_slice = volume[:, i, :]
    elif axis == 2:
        my_slice = volume[:, :, i]

    slices.append(my_slice)