jax.experimental.sparse.bcoo_squeeze

jax.experimental.sparse.bcoo_squeeze(arr, *, dimensions)

Sparse implementation of jax.lax.squeeze().

Squeeze any number of size 1 dimensions from an array.

Parameters:

• arr (BCOO) – BCOO array to be reshaped.

• dimensions (Sequence[int]) – sequence of integers specifying dimensions to squeeze.

Returns:

reshaped array.

Return type:

out