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 jax.lax.dot_general() operator.

Parameters:

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

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

  • precision (PrecisionLike | None) –

    Optional. This parameter controls the numerics of the computation, and it can be one of the following:

    • None, which means the default precision for the current backend,

    • a Precision enum value or a tuple of two Precision enums indicating precision of lhs` and rhs, or

    • a DotAlgorithm or a DotAlgorithmPreset indicating the algorithm that must be used to accumulate the dot product.

  • preferred_element_type (DTypeLike | None | None) – Optional. This parameter controls the data type output by the dot product. By default, the output element type of this operation will match the lhs and rhs input element types under the usual type promotion rules. Setting preferred_element_type to a specific dtype will mean that the operation returns that element type. When precision is not a DotAlgorithm or DotAlgorithmPreset, preferred_element_type provides a hint to the compiler to accumulate the dot product using this data type.

Returns:

An array containing the product.

Return type:

Array