jax.experimental.sparse.bcoo_动态切片

jax.experimental.sparse.bcoo_动态切片#

jax.experimental.sparse.bcoo_dynamic_slice(mat, start_indices, slice_sizes)[源代码][源代码]#

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

参数:
  • mat (BCOO) – BCOO 数组到切片。

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

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

返回:

包含切片信息的 BCOO 数组。

返回类型:

out