jax.numpy.split

目录

jax.numpy.split#

jax.numpy.split(ary, indices_or_sections, axis=0)[源代码][源代码]#

将一个数组分割成子数组。

JAX 实现的 numpy.split()。

参数:

  • 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]]

参见

目录