jax.numpy.extract

jax.numpy.extract(condition, arr)

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