jax.experimental.sparse.bcoo_reshape

jax.experimental.sparse.bcoo_reshape#

jax.experimental.sparse.bcoo_reshape(mat, *, new_sizes, dimensions=None)[源代码][源代码]#

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

参数:
  • operand – 要重塑的 BCOO 数组。

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

  • dimensions (Sequence[int] | None) – 可选的整数序列,指定输入形状的排列顺序。如果指定,长度必须与 operand.shape 匹配。此外,维度只能在相同类型的维度之间进行排列:批次、稀疏和密集维度不能进行排列。

  • mat (BCOO)

返回:

重塑后的数组。

返回类型:

out