jax.experimental.sparse.bcoo_squeeze

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

{func}`jax.lax.squeeze` 的稀疏实现。

从数组中挤压任意数量的大小为 1 的维度。

参数:

• arr (BCOO) – 要重塑的 BCOO 数组。

• dimensions (Sequence[int]) – 指定要挤压的维度的整数序列。

返回:

重塑后的数组。

返回类型:

out