jax.experimental.sparse.bcoo_broadcast_in_dim

jax.experimental.sparse.bcoo_broadcast_in_dim(mat, *, shape, broadcast_dimensions, sharding=None)

透過複製資料來擴展 BCOO 陣列的大小和秩。

BCOO 等效於 jax.lax.broadcast_in_dim。

參數:

• mat (BCOO) – 一個 BCOO 格式的陣列。

• shape (Shape) – 目標陣列的形狀。

• broadcast_dimensions (Sequence[int]) – 目標陣列形狀中的維度，運算元 (mat) 形狀的每個維度對應到此維度。

回傳:

一個包含目標陣列的 BCOO 格式陣列。

回傳類型:

BCOO