ExtractΒΆ

Return the elements of an array that satisfy some condition.

Abstract Signature:

Extract(condition: Tensor, arr: Tensor)

PyTorch

API: torch.masked_select
Strategy: Direct Mapping

JAX (Core)

API: jax.numpy.extract
Strategy: Direct Mapping

NumPy

API: numpy.extract
Strategy: Direct Mapping

TensorFlow

API: tf.boolean_mask
Strategy: Direct Mapping

Flax NNX

API: jax.numpy.extract
Strategy: Direct Mapping

PaxML / Praxis

API: jax.numpy.extract
Strategy: Direct Mapping