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 (ArrayLike)

• indices_or_sections (int | Sequence[int] | ArrayLike)

• axis (int)

返回類型:

list[Array]