Compute the dot product of two arrays.
JAX implementation of numpy.dot().
This differs from jax.numpy.matmul() in two respects:
if either a or b is a scalar, the result of dot is equivalent to jax.numpy.multiply(), while the result of matmul is an error.
if a and b have more than 2 dimensions, the batch indices are stacked rather than broadcast.
a (Array | ndarray | bool | number | bool | int | float | complex) – first input array, of shape (..., N).
b (Array | ndarray | bool | number | bool | int | float | complex) – second input array. Must have shape (N,) or (..., N, M). In the multi-dimensional case, leading dimensions must be broadcast-compatible with the leading dimensions of a.
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.
array containing the dot product of the inputs, with batch dimensions of a and b stacked rather than broadcast.
Examples
For scalar inputs, dot computes the element-wise product:
>>> x = jnp.array([1, 2, 3]) >>> jnp.dot(x, 2) Array([2, 4, 6], dtype=int32)
For vector or matrix inputs, dot computes the vector or matrix product:
>>> M = jnp.array([[2, 3, 4],
... [5, 6, 7],
... [8, 9, 0]])
>>> jnp.dot(M, x)
Array([20, 38, 26], dtype=int32)
>>> jnp.dot(M, M)
Array([[ 51, 60, 29],
[ 96, 114, 62],
[ 61, 78, 95]], dtype=int32)
For higher-dimensional matrix products, batch dimensions are stacked, whereas in matmul() they are broadcast. For example:
>>> a = jnp.zeros((3, 2, 4)) >>> b = jnp.zeros((3, 4, 1)) >>> jnp.dot(a, b).shape (3, 2, 3, 1) >>> jnp.matmul(a, b).shape (3, 2, 1)
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4