jax.numpy.extractΒΆ

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

Return the elements of an array that satisfy some condition.

LAX-backend implementation of 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