jax.experimental.sparse.bcoo_slice

Contents

jax.experimental.sparse.bcoo_slice#

jax.experimental.sparse.bcoo_slice(mat, *, start_indices, limit_indices, strides=None)[source]#

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

Parameters:
  • mat (BCOO) – BCOO array to be reshaped.

  • start_indices (Sequence[int]) – sequence of integers of length mat.ndim specifying the starting indices of each slice.

  • limit_indices (Sequence[int]) – sequence of integers of length mat.ndim specifying the ending indices of each slice

  • strides (Sequence[int] | None) – (not implemented) sequence of integers of length mat.ndim specifying the stride for each slice

Returns:

BCOO array containing the slice.

Return type:

out