jax.numpy.matvec

jax.numpy.matvec(x1, x2, /)

Batched matrix-vector product.

JAX implementation of numpy.matvec().

Parameters:

• x1 (Array | ndarray | bool | number | bool | int | float | complex) – array of shape (..., M, N)

• x2 (Array | ndarray | bool | number | bool | int | float | complex) – array of shape (..., N). Leading dimensions must be broadcast-compatible with leading dimensions of x1.

Returns:

An array of shape (..., M) containing the batched matrix-vector product.

Return type:

Array

See also

• jax.numpy.linalg.vecdot(): batched vector product.

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

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

Examples

Simple matrix-vector product:

>>> x1 = jnp.array([[1, 2, 3],
...                 [4, 5, 6]])
>>> x2 = jnp.array([7, 8, 9])
>>> jnp.matvec(x1, x2)
Array([ 50, 122], dtype=int32)

Batched matrix-vector product:

>>> x2 = jnp.array([[7, 8, 9],
...                 [5, 6, 7]])
>>> jnp.matvec(x1, x2)
Array([[ 50, 122],
       [ 38,  92]], dtype=int32)