jax.lax.rsqrt#
- jax.lax.rsqrt(x, accuracy=None)[source]#
逐元素倒数平方根:\(1 \over \sqrt{x}\)。
此函数直接降低到 stablehlo.rsqrt 操作。
- 参数:
x (ArrayLike) – 输入数组。必须具有浮点型或复数 dtype。
accuracy – 可选的 lax.Tolerance 或 lax.AccuracyMode 对象,用于根据请求的精度选择操作的实现。如果实现无法满足请求的容差,编译器将返回错误。如果指定了模式,并且没有多个可用的实现,则将使用默认实现。
- 返回:
与
x
具有相同形状和 dtype 的数组,包含倒数平方根。- 返回类型:
另请参阅
jax.lax.pow()
: 逐元素幂运算。jax.lax.sqrt()
: 逐元素平方根。jax.lax.cbrt()
: 逐元素立方根。