jax.experimental.sparse.bcoo_dynamic_slice#

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

Sparse implementation of {func}`jax.lax.dynamic_slice`.

Parameters:

  • mat (BCOO) – BCOO array to slice.

  • start_indices (Sequence[Any]) – a list of scalar indices, one per dimension. These values may be dynamic.

  • slice_sizes (Sequence[int]) – the size of the slice. Must be a sequence of non-negative integers with length equal to ndim(operand). Inside a JIT compiled function, only static values are supported (all JAX arrays inside JIT must have statically known size).

Returns:

BCOO array containing the slice.

Return type:

out