jax.numpy.array_split — JAX documentation (original) (raw)

jax.numpy.array_split#

jax.numpy.array_split(ary, indices_or_sections, axis=0)[source]#

Split an array into sub-arrays.

JAX implementation of numpy.array_split().

Refer to the documentation of jax.numpy.split() for details; array_splitis equivalent to split, but allows integer indices_or_sections which does not evenly divide the split axis.

Examples

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]

See also

• jax.numpy.split(): split an array along any axis.
• jax.numpy.vsplit(): split vertically, i.e. along axis=0
• jax.numpy.hsplit(): split horizontally, i.e. along axis=1
• jax.numpy.dsplit(): split depth-wise, i.e. along axis=2

Parameters:

• ary (ArrayLike)
• indices_or_sections (int | Sequence _[_int] | ArrayLike)
• axis (int)

Return type:

list[Array]