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