jax.numpy.vecdot

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

Perform a conjugate multiplication of two batched vectors.

JAX implementation of numpy.vecdot().

Parameters:

• a – left-hand side array.

• b – right-hand side array. Size of b[axis] must match size of a[axis], and remaining dimensions must be broadcast-compatible.

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

• precision (None | str | Precision | tuple[str, str] | tuple[Precision, Precision] | DotAlgorithm | DotAlgorithmPreset) – 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 a and b.

• preferred_element_type (str | type[Any] | dtype | SupportsDType | 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.

• x1 (Array | ndarray | bool | number | bool | int | float | complex)

• x2 (Array | ndarray | bool | number | bool | int | float | complex)

Returns:

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

Return type:

Array

See also

• jax.numpy.vdot(): flattened vector product.

• jax.numpy.vecmat(): vector-matrix product.

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

• jax.lax.dot_general(): general N-dimensional batched dot product.

Examples

Vector conjugate-dot product of two 1D arrays:

>>> a = jnp.array([1j, 2j, 3j])
>>> b = jnp.array([4., 5., 6.])
>>> jnp.linalg.vecdot(a, b)
Array(0.-32.j, dtype=complex64)

Batched vector dot product of two 2D arrays:

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