jax.numpy.linalg.vecdot — JAX documentation (original) (raw)

jax.numpy.linalg.vecdot#

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

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 (lax.PrecisionLike) – 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) – 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)