jax.numpy.vsplit

jax.numpy.vsplit(ary, indices_or_sections)

将数组垂直拆分为子数组。

JAX 实现的 numpy.vsplit().

有关详细信息，请参阅 jax.numpy.split() 的文档；vsplit 等效于 split，其中 axis=0。

示例

1D 数组

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

2D 数组

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

参见

• jax.numpy.split()：沿任何轴拆分数组。

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

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

• jax.numpy.array_split(): 类似于 split，但允许 indices_or_sections 为一个不完全整除数组大小的整数。

参数:

• ary (ArrayLike)

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

返回类型:

list[Array]