jax.lax.dot#

jax.lax.dot(lhs, rhs, precision=None, preferred_element_type=None)[source]#

Vector/vector, matrix/vector, and matrix/matrix multiplication.

Wraps XLA’s Dot operator.

For more general contraction, see the dot_general operator.

Parameters:

  • lhs (Array) – an array of dimension 1 or 2.

  • rhs (Array) – an array of dimension 1 or 2.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision]]) – Optional. Either None, which means the default precision for the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two Precision enums indicating precision of lhs` and rhs.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Optional. Either None, which means the default accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Return type:

Array

Returns:

An array containing the product.