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]