jax.lax.dot#

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

General dot product/contraction operator.

This operation lowers directly to thestablehlo.dot_general operation.

The semantics ofdot_general are complicated, but most users should not have touse it directly. Instead, you can use higher-level functions likejax.numpy.dot(),jax.numpy.matmul(),jax.numpy.tensordot(),jax.numpy.einsum(),and others which will construct appropriate calls todot_general under the hood.If you really want to understanddot_general itself, we recommend reading XLA’sDotGeneral operator documentation.

Parameters:

• lhs (ArrayLike) – an array

• rhs (ArrayLike) – an array

• dimension_numbers (DotDimensionNumbers |None) – an optional tuple of tuples of sequences of ints of the form((lhs_contracting_dims,rhs_contracting_dims),(lhs_batch_dims,rhs_batch_dims)). This may be left unspecified in the common case ofun-batched matrix-matrix, matrix-vector, or vector-vector dot products, asdetermined by the shape oflhs andrhs.

• precision (PrecisionLike) –

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

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

  • aPrecision enum value or a tuple of twoPrecision enums indicating precision oflhs` andrhs, or

  • aDotAlgorithm or aDotAlgorithmPreset indicating the algorithm thatmust be used to accumulate the dot product.

• preferred_element_type (DTypeLike |None) – Optional. This parameter controls the data typeoutput by the dot product. By default, the output element type of thisoperation will match thelhs andrhs input element types underthe usual type promotion rules. Settingpreferred_element_type to aspecificdtype will mean that the operation returns that element type.Whenprecision is not aDotAlgorithm orDotAlgorithmPreset,preferred_element_type providesa hint to the compiler to accumulate the dot product using this data type.

• out_sharding – an optional sharding specification for the output. If not specified,it will be determined automatically by the compiler.

Returns:

An array whose first dimensions are the (shared) batch dimensions, followedby thelhs non-contracting/non-batch dimensions, and finally therhsnon-contracting/non-batch dimensions.

Return type:

Array