jax.lax.split

jax.lax.split(operand, sizes, axis=0)

沿 axis 分割数组。

参数:

• operand (ArrayLike) – 要分割的数组

• sizes (Sequence[int]) – 分割数组的大小。大小的总和必须等于 operand 的 axis 维度的大小。

• axis (int) – 沿其分割数组的轴。

返回:

一个包含 len(sizes) 个数组的序列。如果 sizes 是 [s1, s2, ...]，则此函数返回沿 axis 取的 s1、s2 大小的块。

返回类型:

Sequence[Array]