jax.experimental.sparse.bcoo_dot_general_sampled
jax.experimental.sparse.bcoo_dot_general_sampled#
- jax.experimental.sparse.bcoo_dot_general_sampled(A, B, indices, *, dimension_numbers)[source]#
A contraction operation with output computed at given sparse indices.
- Parameters
lhs – An ndarray.
rhs – An ndarray.
indices (
Array) – BCOO indices.dimension_numbers (
Tuple[Tuple[Sequence[int],Sequence[int]],Tuple[Sequence[int],Sequence[int]]]) – a tuple of tuples of the form ((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims)).A (
Array) –B (
Array) –
- Return type
Array- Returns
BCOO data, an ndarray containing the result.