jax.experimental.sparse.bcoo_reshape

jax.experimental.sparse.bcoo_reshape(mat, *, new_sizes, dimensions=None, sharding=None)

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

参数:

• operand – 要重塑的 BCOO 数组。

• new_sizes (Sequence[int]) – 指定结果形状的整数序列。最终数组的大小必须与输入的大小匹配。这必须指定为批处理维度、稀疏维度和密集维度不混合。

• dimensions (Sequence[int] | None | None) – 可选的整数序列，指定输入形状的排列顺序。如果指定，则长度必须与 operand.shape 匹配。此外，维度只能在 mat 的相似维度之间排列：批处理维度、稀疏维度和密集维度不能排列。

• mat (BCOO)

返回值:

重塑的数组。

返回类型:

out