jax.numpy.extract#

jax.numpy.extract(condition, arr)[source]#

Return the elements of an array that satisfy some condition.

LAX-backend implementation of numpy.extract().

Original docstring below.

This is equivalent to np.compress(ravel(condition), ravel(arr)). If condition is boolean np.extract is equivalent to arr[condition].

Note that place does the exact opposite of extract.

Parameters
  • condition (array_like) – An array whose nonzero or True entries indicate the elements of arr to extract.

  • arr (array_like) – Input array of the same size as condition.

Returns

extract – Rank 1 array of values from arr where condition is True.

Return type

ndarray