jax.lax.sqrt#

jax.lax.sqrt(x, accuracy=None)[source]#

逐元素平方根：\(\sqrt{x}\)。

此函数直接降低到 stablehlo.sqrt 操作。

参数:

  • x (ArrayLike) – 输入数组。必须具有浮点型或复数类型。

  • accuracy – 可选的 lax.Tolerance 或 lax.AccuracyMode 对象，用于根据请求的精度选择操作的实现。如果实现无法满足请求的容差，编译器将返回错误。如果指定了模式并且没有多个实现可用，则将使用默认实现。

返回:

一个与 x 具有相同形状和数据类型的数组，其中包含平方根。

返回类型:

数组 (Array)

另请参阅

jax.lax.pow(): 逐元素幂运算。 jax.lax.cbrt(): 逐元素立方根。 jax.lax.rsqrt(): 逐元素倒数平方根。