jax.lax.dot_general#

jax.lax.dot_general(lhs, rhs, dimension_numbers, precision=None, preferred_element_type=None, *, out_sharding=None)[源代码]#

通用点积/收缩运算符。

包装 XLA 的 DotGeneral 运算符。

dot_general 的语义很复杂，但大多数用户不应该直接使用它。相反，您可以使用更高级别的函数，例如 jax.numpy.dot(), jax.numpy.matmul(), jax.numpy.tensordot(), jax.numpy.einsum() 和其他函数，它们将在底层构建对 dot_general 的适当调用。如果您真的想理解 dot_general 本身，我们建议阅读 XLA 的 DotGeneral 运算符文档。

参数:

  • lhs (ArrayLike) – 一个数组

  • rhs (ArrayLike) – 一个数组

  • dimension_numbers (DotDimensionNumbers) – 一个元组，其中包含整数序列的元组，格式为 ((lhs_contracting_dims, rhs_contracting_dims), (lhs_batch_dims, rhs_batch_dims))

  • precision (PrecisionLike) –

    可选。此参数控制计算的数值，它可以是以下之一

    • None，这意味着当前后端的默认精度，

    • 一个 Precision 枚举值或两个 Precision 枚举的元组，指示 lhs` 和 rhs 的精度，或者

    • 一个 DotAlgorithm 或一个 DotAlgorithmPreset，指示必须用于累积点积的算法。

  • preferred_element_type (DTypeLike | None) – 可选。此参数控制点积输出的数据类型。默认情况下，此操作的输出元素类型将根据通常的类型提升规则匹配 lhs 和 rhs 输入元素类型。将 preferred_element_type 设置为特定的 dtype 将意味着该操作返回该元素类型。当 precision 不是 DotAlgorithm 或 DotAlgorithmPreset 时，preferred_element_type 为编译器提供了一个提示，可以使用此数据类型累积点积。

返回:

一个数组，其第一个维度是（共享的）批次维度，后跟 lhs 非收缩/非批次维度，最后是 rhs 非收缩/非批次维度。

返回类型:

Array