jax.lax.pow

jax.lax.pow(x, y)

Elementwise power: \(x^y\).

This function lowers directly to the stablehlo.pow operation, along with a stablehlo.convert when the argument dtypes do not match.

Parameters:

• x (ArrayLike) – Input array giving the base value. Must have floating or complex type.

• y (ArrayLike) – Input array giving the exponent value. Must have integer, floating, or complex type. Its dtype will be cast to that of x.dtype if necessary. If neither x nor y is a scalar, then x and y must have the same number of dimensions and be broadcast-compatible.

Returns:

An array of the same dtype as x containing the elementwise power.

Return type:

Array

See also

jax.lax.integer_pow(): Elementwise power where y is a static integer.