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:
- Returns:
BCOO data, an ndarray containing the result.