jax.experimental.sparse.bcoo_extract

jax.experimental.sparse.bcoo_extract#

jax.experimental.sparse.bcoo_extract(sparr, arr, *, assume_unique=None)[source]#

根据稀疏数组的索引从密集数组中提取值。

参数:
  • sparr (BCOO) – BCOO 数组,其索引将用于输出。

  • arr (ArrayLike) – ArrayLike,其形状等于 self.shape

  • assume_unique (bool | None | None) – bool,默认为 sparr.unique_indices 如果为 True,则为每个索引提取值,即使索引包含重复项。如果为 False,则重复索引的值将被求和并返回在第一个索引的位置。

返回值:

一个与自身具有相同稀疏模式的 BCOO 数组。

返回类型:

提取的