jax.lax.complex

jax.lax.complex(x, y)

Elementwise make complex number: \(x + jy\).

This function lowers directly to the stablehlo.complex operation.

Parameters:

• x (ArrayLike) – input arrays. Must have matching floating-point dtypes. If neither is a scalar, the two arrays must have the same number of dimensions and be broadcast-compatible.

• y (ArrayLike) – input arrays. Must have matching floating-point dtypes. If neither is a scalar, the two arrays must have the same number of dimensions and be broadcast-compatible.

Returns:

The complex array with the real part given by x, and the imaginary part given by y. For inputs of dtype float32 or float64, the result will have dtype complex64 or complex128 respectively.

Return type:

Array

See also

• jax.lax.real(): elementwise extract real part.

• jax.lax.imag(): elementwise extract imaginary part.

• jax.lax.conj(): elementwise complex conjugate.