jax.numpy.compressΒΆ

jax.numpy.compress(condition, a, axis=None, out=None)[source]ΒΆ

Return selected slices of an array along given axis.

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

When working along a given axis, a slice along that axis is returned in output for each index where condition evaluates to True. When working on a 1-D array, compress is equivalent to extract.

Parameters
  • condition (1-D array of bools) – Array that selects which entries to return. If len(condition) is less than the size of a along the given axis, then output is truncated to the length of the condition array.

  • a (array_like) – Array from which to extract a part.

  • axis (int, optional) – Axis along which to take slices. If None (default), work on the flattened array.

  • out (ndarray, optional) – Output array. Its type is preserved and it must be of the right shape to hold the output.

Returns

compressed_array – A copy of a without the slices along axis for which condition is false.

Return type

ndarray

See also

take(), choose(), diag(), diagonal(), select()

ndarray.compress()

Equivalent method in ndarray

np.extract()

Equivalent method when working on 1-D arrays

ufuncs-output-type()

Examples

>>> a = np.array([[1, 2], [3, 4], [5, 6]])
>>> a
array([[1, 2],
       [3, 4],
       [5, 6]])
>>> np.compress([0, 1], a, axis=0)
array([[3, 4]])
>>> np.compress([False, True, True], a, axis=0)
array([[3, 4],
       [5, 6]])
>>> np.compress([False, True], a, axis=1)
array([[2],
       [4],
       [6]])

Working on the flattened array does not return slices along an axis but selects elements.

>>> np.compress([False, True], a)
array([2])