optax.contrib.split_real_and_imaginary

optax.contrib.split_real_and_imaginary(inner: optax.GradientTransformation) → optax.GradientTransformation

Splits the real and imaginary components of complex updates into two.

The inner transformation processes real parameters and updates, and the pairs of transformed real updates are merged into complex updates.

Parameters and updates that are real before splitting are passed through unmodified.

Parameters:

inner – The inner transformation.

Returns:

An optax.GradientTransformation.