jax.lax.dot

内容

jax.lax.dot#

jax.lax.dot(lhs, rhs, precision=None, preferred_element_type=None)[source]#

向量/向量、矩阵/向量和矩阵/矩阵乘法。

包装 XLA 的 Dot 运算符。

有关更一般的收缩，请参阅 jax.lax.dot_general() 运算符。

参数：

  • lhs (Array) – 一个维度为 1 或 2 的数组。

  • rhs (Array) – 一个维度为 1 或 2 的数组。

  • precision (PrecisionLike | None) – 可选。可以是 None，表示使用后端默认精度，也可以是 Precision 枚举值（Precision.DEFAULT、Precision.HIGH 或 Precision.HIGHEST），或者是一个包含两个 Precision 枚举值的元组，分别指示 lhs` 和 rhs 的精度。

  • preferred_element_type (DTypeLike | None | None) – 可选。可以是 None，表示使用输入类型的默认累积类型，也可以是数据类型，指示将结果累积到该数据类型并返回具有该数据类型的结果。

返回值:

包含乘积结果的数组。

返回类型:

数组

内容