jax.numpy.array_split

jax.numpy.array_split(ary, indices_or_sections, axis=0)

将数组分割成子数组。

numpy.array_split() 的 JAX 实现。

有关详细信息，请参阅 jax.numpy.split() 的文档；array_split 与 split 等效，但允许使用不能均匀分割分割轴的整数 indices_or_sections。

示例

>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> chunks = jnp.array_split(x, 4)
>>> print(*chunks)
[1 2 3] [4 5] [6 7] [8 9]

另请参阅

• jax.numpy.split()：沿任意轴分割数组。

• jax.numpy.vsplit()：垂直分割，即沿 axis=0 分割

• jax.numpy.hsplit()：水平分割，即沿 axis=1 分割

• jax.numpy.dsplit()：深度分割，即沿 axis=2 分割

参数:

• ary (类似数组)

• indices_or_sections (int | Sequence[int] | 类似数组)

• axis (int)

返回类型:

list[Array]