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

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

  • None，表示当前后端的默认精度，

  • 一个 Precision 枚举值或一对 Precision 枚举值，分别表示 lhs 和 rhs 的精度，或者

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

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

返回:

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

返回类型:

Array