Perform a matrix multiplication.
JAX implementation of numpy.matmul().
a (Array | ndarray | bool | number | bool | int | float | complex) – first input array, of shape (N,) or (..., K, 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.
out_sharding (NamedSharding | PartitionSpec | None)
array containing the matrix product of the inputs. Shape is a.shape[:-1] if b.ndim == 1, otherwise the shape is (..., K, M), where leading dimensions of a and b are broadcast together.
Examples
Vector dot products:
>>> a = jnp.array([1, 2, 3]) >>> b = jnp.array([4, 5, 6]) >>> jnp.matmul(a, b) Array(32, dtype=int32)
Matrix dot product:
>>> a = jnp.array([[1, 2, 3],
... [4, 5, 6]])
>>> b = jnp.array([[1, 2],
... [3, 4],
... [5, 6]])
>>> jnp.matmul(a, b)
Array([[22, 28],
[49, 64]], dtype=int32)
For convenience, in all cases you can do the same computation using the @ operator:
>>> a @ b
Array([[22, 28],
[49, 64]], dtype=int32)
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