jax.experimental.sparse.bcoo_dynamic_slice

jax.experimental.sparse.bcoo_dynamic_slice(mat, start_indices, slice_sizes)

jax.lax.dynamic_slice() 的稀疏实现。

参数:

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

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

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

返回:

包含切片的 BCOO 数组。

返回类型:

out