jax.default_matmul_precision#

jax.default_matmul_precision = <jax._src.config.State object>#

用于 jax_default_matmul_precision 配置选项的上下文管理器。

控制 32 位输入的默认矩阵乘法和卷积精度。

某些平台(例如 TPU)为矩阵乘法和卷积计算提供可配置的精度级别,以牺牲速度来换取准确性。精度可以为每个操作控制;例如,请参阅 jax.lax.conv_general_dilated()jax.lax.dot() 的文档字符串。但控制当操作未指定特定精度时获得的默认行为很有用。

此选项可用于控制 32 位输入上涉及矩阵乘法和卷积的计算的默认精度级别。这些级别大致描述了计算标量积的精度。“bfloat16”选项速度最快、精度最低;“float32”类似于完整的 float32 精度;“tensorfloat32”介于两者之间。

此参数还可以用于为执行矩阵乘法的函数(例如 jax.lax.dot())指定累加“算法”。要指定算法,请将此选项设置为 DotAlgorithmPreset 的名称。

参数:

new_val (任意类型)