jax.experimental.sparse.bcoo_dynamic_slice

jax.experimental.sparse.bcoo_dynamic_slice#

jax.experimental.sparse.bcoo_dynamic_slice(mat, start_indices, slice_sizes)[source]#

{func}`jax.lax.dynamic_slice` 的稀疏实现。

参数:

  • mat (BCOO) – 要切片的 BCOO 数组。

  • start_indices (Sequence[Any]) – 每个维度一个标量索引的列表。这些值可能是动态的。

  • slice_sizes (Sequence[int]) – 切片的尺寸。必须是长度等于 ndim(operand) 的非负整数序列。在 JIT 编译函数内部，仅支持静态值（JIT 内部的所有 JAX 数组都必须具有静态已知的尺寸）。

返回值:

包含切片的 BCOO 数组。

返回类型:

out

目录