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