Python 实现取多维数组第n维的前几位

print("# original", input.shape)
input_transpose = input.transpose((1, 0, 2, 3))
print("# transpose", input_transpose.shape)
input_slice = input_transpose[0:3]
print("# slice", input_slice.shape)
output = input_slice.transpose((1, 0, 2, 3))
print("# output", output.shape)

print("# original", input.shape)
print("# output", input[:, [0, 1, 2]].shape)