jax.numpy.dot

jax.numpy.dot(a, b, *, precision=None, preferred_element_type=None)

计算两个数组的点积。

numpy.dot() 的 JAX 实现。

这与 jax.numpy.matmul() 在两个方面有所不同

• 如果 a 或 b 中任何一个是标量，则 dot 的结果等同于 jax.numpy.multiply()，而 matmul 的结果则会报错。

• 如果 a 和 b 具有超过 2 个维度，则批次索引会被堆叠而不是广播。

参数:

• a (Array | ndarray | bool | number | bool | int | float | complex) – 第一个输入数组，形状为 (..., N)。

• b (Array | ndarray | bool | number | bool | int | float | complex) – 第二个输入数组。必须具有形状 (N,) 或 (..., N, M)。在多维情况下，前导维度必须与 a 的前导维度广播兼容。

• precision (None | str | Precision | tuple[str, str] | tuple[Precision, Precision] | DotAlgorithm | DotAlgorithmPreset) – 可以是 None (默认)，表示后端的默认精度，Precision 枚举值 (Precision.DEFAULT, Precision.HIGH 或 Precision.HIGHEST) 或两个此类值的元组，表示 a 和 b 的精度。

• preferred_element_type (str | type[Any] | dtype | SupportsDType | None) – 可以是 None (默认)，表示输入类型的默认累积类型，或一个数据类型，表示将结果累积到该数据类型并返回结果。

返回:

包含输入点积的数组，其中 a 和 b 的批次维度被堆叠而不是广播。

返回类型:

Array

参见

• jax.numpy.matmul()：广播的批次矩阵乘法。

• jax.lax.dot_general()：通用批次矩阵乘法。

示例

对于标量输入，dot 计算元素级乘积

>>> x = jnp.array([1, 2, 3])
>>> jnp.dot(x, 2)
Array([2, 4, 6], dtype=int32)

对于向量或矩阵输入，dot 计算向量或矩阵乘积

>>> 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)

对于更高维度的矩阵乘积，批次维度会被堆叠，而在 matmul() 中，它们会被广播。例如

>>> 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)