jax.experimental.sparse.bcsr_extract#
- jax.experimental.sparse.bcsr_extract(indices, indptr, mat)[source]#
Extract values from a dense matrix at given BCSR (indices, indptr).
- Parameters:
indices (jax.typing.ArrayLike) – An ndarray; see BCSR indices.
indptr (jax.typing.ArrayLike) – An ndarray; see BCSR indptr.
mat (jax.typing.ArrayLike) – A dense matrix.
- Returns:
An ndarray; see BCSR data.
- Return type: