jax.experimental.sparse.bcoo_broadcast_in_dim

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

通过复制数据扩展 BCOO 数组的大小和秩（rank）。

BCOO 等效于 jax.lax.broadcast_in_dim。

参数:

• mat (BCOO) — BCOO 格式的数组。

• shape (Shape) – 目标数组的形状。

• broadcast_dimensions (Sequence[int]) – 目标数组形状中的维度，操作数（mat）形状的每个维度与其对应。

返回:

包含目标数组的 BCOO 格式数组。

返回类型:

BCOO