jax.numpy.hsplitÂ¶

jax.numpy.hsplit(ary, indices_or_sections)Â¶

Split an array into multiple sub-arrays horizontally (column-wise).

LAX-backend implementation of hsplit(). Original docstring below.

Please refer to the split documentation. hsplit is equivalent to split with axis=1, the array is always split along the second axis regardless of the array dimension.

split : Split an array into multiple sub-arrays of equal size.

>>> x = np.arange(16.0).reshape(4, 4)
>>> x
array([[ 0.,   1.,   2.,   3.],
[ 4.,   5.,   6.,   7.],
[ 8.,   9.,  10.,  11.],
[12.,  13.,  14.,  15.]])
>>> np.hsplit(x, 2)
[array([[  0.,   1.],
[  4.,   5.],
[  8.,   9.],
[12.,  13.]]),
array([[  2.,   3.],
[  6.,   7.],
[10.,  11.],
[14.,  15.]])]
>>> np.hsplit(x, np.array([3, 6]))
[array([[ 0.,   1.,   2.],
[ 4.,   5.,   6.],
[ 8.,   9.,  10.],
[12.,  13.,  14.]]),
array([[ 3.],
[ 7.],
[11.],
[15.]]),
array([], shape=(4, 0), dtype=float64)]


With a higher dimensional array the split is still along the second axis.

>>> x = np.arange(8.0).reshape(2, 2, 2)
>>> x
array([[[0.,  1.],
[2.,  3.]],
[[4.,  5.],
[6.,  7.]]])
>>> np.hsplit(x, 2)
[array([[[0.,  1.]],
[[4.,  5.]]]),
array([[[2.,  3.]],
[[6.,  7.]]])]