jax.numpy.vsplit

jax.numpy.vsplit(ary, indices_or_sections)

Split an array into sub-arrays vertically.

JAX implementation of numpy.vsplit().

Refer to the documentation of jax.numpy.split() for details; vsplit is equivalent to split with axis=0.

Examples

1D array:

>>> 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 array:

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

See also

• jax.numpy.split(): split an array along any axis.

• jax.numpy.hsplit(): split horizontally, i.e. along axis=1

• jax.numpy.dsplit(): split depth-wise, i.e. along axis=2

• jax.numpy.array_split(): like split, but allows indices_or_sections to be an integer that does not evenly divide the size of the array.

Parameters:

• ary (ArrayLike)

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

Return type:

list[Array]