jax.lax.rsqrt#

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

Elementwise reciprocal square root: \(1 \over \sqrt{x}\).

This function lowers directly to the stablehlo.rsqrt operation.

Parameters:

  • x (ArrayLike) – Input array. Must have floating or complex dtype.

  • accuracy – Optional lax.Tolerance or lax.AccuracyMode object that selects the implementation of the op based on the requested accuracy. If the implementation cannot satisfy the requested tolerance, the compiler will return an error. If mode is specified and there are no multiple implementations available, the default implementation will be used.

Returns:

An array of the same shape and dtype as x containing the reciporical square root.

Return type:

Array

See also

jax.lax.pow(): Elementwise power. jax.lax.sqrt(): Elementwise square root. jax.lax.cbrt(): Elementwise cube root.