CompressΒΆ

Compress an array along a given axis using a boolean condition.

Abstract Signature:

Compress(condition: Tensor, a: Tensor, axis: int | None)

JAX (Core)

API: jax.numpy.compress
Strategy: Direct Mapping

NumPy

API: numpy.compress
Strategy: Direct Mapping

Flax NNX

API: jax.numpy.ndarray.compress
Strategy: Direct Mapping