jax.lax.squeeze

jax.lax.squeeze(array, dimensions)

Squeeze any number of size 1 dimensions from an array.

Parameters:

• array (ArrayLike)

• dimensions (Sequence[int])

Return type:

Array