Split an array into sub-arrays.
JAX implementation of numpy.split()
.
ary (ArrayLike) – N-dimensional array-like object to split
indices_or_sections (int | Sequence[int] | ArrayLike) –
either a single integer or a sequence of indices.
if indices_or_sections
is an integer N, then N must evenly divide ary.shape[axis]
and ary
will be divided into N equally-sized chunks along axis
.
if indices_or_sections
is a sequence of integers, then these integers specify the boundary between unevenly-sized chunks along axis
; see examples below.
axis (int) – the axis along which to split; defaults to 0.
A list of arrays. If indices_or_sections
is an integer N, then the list is of length N. If indices_or_sections
is a sequence seq, then the list is is of length len(seq) + 1.
Examples
Splitting a 1-dimensional array:
>>> x = jnp.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
Split into three equal sections:
>>> chunks = jnp.split(x, 3) >>> print(*chunks) [1 2 3] [4 5 6] [7 8 9]
Split into sections by index:
>>> 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]
Splitting a two-dimensional array along axis 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]]
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4