jax.numpy.linalg.vecdot

jax.numpy.linalg.vecdot(x1, x2, /, *, axis=-1, precision=None, preferred_element_type=None)

Compute the (batched) vector conjugate dot product of two arrays.

JAX implementation of numpy.linalg.vecdot().

Parameters:

• x1 (ArrayLike) – left-hand side array.

• x2 (ArrayLike) – right-hand side array. Size of x2[axis] must match size of x1[axis], and remaining dimensions must be broadcast-compatible.

• axis (int) – axis along which to compute the dot product (default: -1)

• precision (PrecisionLike | None) – either None (default), which means the default precision for the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two such values indicating precision of x1 and x2.

• preferred_element_type (DTypeLike | None | None) – either None (default), 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.

Returns:

array containing the conjugate dot product of x1 and x2 along axis. The non-contracted dimensions are broadcast together.

Return type:

Array

See also

• jax.numpy.vecdot(): similar API in the jax.numpy namespace.

• jax.numpy.linalg.matmul(): matrix multiplication.

• jax.numpy.linalg.tensordot(): general tensor dot product.

Examples

Vector dot product of two 1D arrays:

>>> x1 = jnp.array([1, 2, 3])
>>> x2 = jnp.array([4, 5, 6])
>>> jnp.linalg.vecdot(x1, x2)
Array(32, dtype=int32)

Batched vector dot product of two 2D arrays:

>>> x1 = jnp.array([[1, 2, 3],
...                 [4, 5, 6]])
>>> x2 = jnp.array([[2, 3, 4]])
>>> jnp.linalg.vecdot(x1, x2, axis=-1)
Array([20, 47], dtype=int32)