np.split()

Разбивает массив на несколько подмассивов вдоль оси.

Сигнатураnp.split(a, indices_or_sections, axis=0)

Делит массив на равные части или по заданным точкам разреза. Возвращает список массивов.

import numpy as np
a = np.arange(6)
print(np.split(a, 3))   # [array([0, 1]), array([2, 3]), array([4, 5])]