jax_squeeze¶

probly.representation.jax_functions.jax_squeeze(a: JaxArrayLike, /, axis: int | Sequence[int] | None = None) → jax.Array[source]¶

Remove axes of length one, mirroring jax.numpy.squeeze.

Parameters:

a – The array to squeeze.
axis – The axis or axes to remove. If omitted, all length-one axes are removed.

Returns:

The squeezed array.