jax.numpy.split#

jax.numpy.split(ary, indices_or_sections, axis=0)[原始碼]#

將陣列分割成子陣列。

numpy.split() 的 JAX 實作。

參數:

  • ary (ArrayLike) – 要分割的 N 維類陣列物件

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

    單一整數或索引序列。

    • 如果 indices_or_sections 為整數 N，則 N 必須能整除 ary.shape[axis]，且 ary 將沿著 axis 分割成 N 個大小相等的區塊。

    • 如果 indices_or_sections 是整數序列，則這些整數指定沿著 axis 大小不均的區塊之間的邊界；請參閱以下範例。

  • axis (int) – 要分割的軸；預設為 0。

返回:

陣列列表。如果 indices_or_sections 是整數 N，則列表長度為 N。如果 indices_or_sections 是序列 seq，則列表長度為 len(seq) + 1。

返回類型:

list[Array]

範例

分割一維陣列

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

分割成三個相等的部分

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

依索引分割成多個部分

>>> chunks = jnp.split(x, [2, 7])  # [x[0:2], x[2:7], x[7:]]
>>> print(*chunks)
[1 2] [3 4 5 6 7] [8 9]

沿著軸 1 分割二維陣列

>>> x = jnp.array([[1, 2, 3, 4],
...                [5, 6, 7, 8]])
>>> x1, x2 = jnp.split(x, 2, axis=1)
>>> print(x1)
[[1 2]
 [5 6]]
>>> print(x2)
[[3 4]
 [7 8]]

另請參閱