jax.experimental.sparse.bcoo_dot_general#

jax.experimental.sparse.bcoo_dot_general(lhs, rhs, *, dimension_numbers, precision=None, preferred_element_type=None, out_sharding=None)[source]#

A general contraction operation.

Parameters:

  • lhs (BCOO | Array) – An ndarray or BCOO-format sparse array.

  • rhs (BCOO | Array) – An ndarray or BCOO-format sparse array..

  • dimension_numbers (DotDimensionNumbers) – a tuple of tuples of the form ((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims)).

  • precision (None | None) – unused

  • preferred_element_type (None | None) – unused

Returns:

An ndarray or BCOO-format sparse array containing the result. If both inputs are sparse, the result will be sparse, of type BCOO. If either input is dense, the result will be dense, of type ndarray.

Return type:

BCOO | Array