Introduction

Yes, you can train a TensorFlow model on complex-valued data, but you need to be precise about what “complex input or output” means in practice. TensorFlow supports complex tensors at the framework level, yet many high-level Keras layers and losses are still designed around real-valued tensors, so the common production approach is to represent the real and imaginary parts explicitly.

What TensorFlow Supports Directly

TensorFlow has tf.complex, tf.math.real, and tf.math.imag, so the underlying tensor system understands complex numbers. That means you can store and manipulate data such as Fourier coefficients or IQ samples without inventing your own serialization format.

1import tensorflow as tf
2
3real = tf.constant([1.0, 2.0, 3.0], dtype=tf.float32)
4imag = tf.constant([0.5, -0.5, 1.5], dtype=tf.float32)
5z = tf.complex(real, imag)
6
7print(z)
8print(tf.math.real(z))
9print(tf.math.imag(z))

The limitation appears one level higher. Many standard Keras layers, initializers, metrics, and optimizers assume floating-point tensors such as float32, not complex64.

The Usual Practical Design

For most supervised learning problems, split the complex value into two real-valued channels. One channel stores the real part, and one stores the imaginary part. The model then works with ordinary dense or convolution layers.

Suppose each sample contains 128 complex measurements. You can encode that as a tensor of shape (128, 2).

1import numpy as np
2import tensorflow as tf
3from tensorflow import keras
4
5num_samples = 1000
6x_complex = np.random.randn(num_samples, 128) + 1j * np.random.randn(num_samples, 128)
7y_complex = 2.0 * x_complex
8
9x = np.stack([x_complex.real, x_complex.imag], axis=-1).astype("float32")
10y = np.stack([y_complex.real, y_complex.imag], axis=-1).astype("float32")
11
12model = keras.Sequential([
13    keras.layers.Input(shape=(128, 2)),
14    keras.layers.Flatten(),
15    keras.layers.Dense(256, activation="relu"),
16    keras.layers.Dense(128 * 2),
17    keras.layers.Reshape((128, 2)),
18])
19
20model.compile(optimizer="adam", loss="mse")
21model.fit(x, y, epochs=3, batch_size=32, verbose=0)

This model does not know anything special about complex arithmetic, but it can learn relationships between the two channels.

When You Need True Complex Arithmetic

Some domains, especially signal processing and physics, really do care about complex multiplication, phase, and conjugation. In those cases, simply splitting into two channels may lose useful inductive structure.

You can still use TensorFlow, but you often need custom layers that implement the real and imaginary parts of the math manually. For example, a complex dense layer can be written from real-valued weights:

1class ComplexDense(tf.keras.layers.Layer):
2    def __init__(self, units):
3        super().__init__()
4        self.units = units
5
6    def build(self, input_shape):
7        features = input_shape[-1]
8        self.w_real = self.add_weight(shape=(features, self.units), initializer="glorot_uniform")
9        self.w_imag = self.add_weight(shape=(features, self.units), initializer="glorot_uniform")
10
11    def call(self, inputs):
12        x_real, x_imag = inputs[..., 0], inputs[..., 1]
13        y_real = tf.matmul(x_real, self.w_real) - tf.matmul(x_imag, self.w_imag)
14        y_imag = tf.matmul(x_real, self.w_imag) + tf.matmul(x_imag, self.w_real)
15        return tf.stack([y_real, y_imag], axis=-1)

This keeps the model in real-valued tensors while preserving complex-number algebra.

Handling Complex Outputs

Complex outputs are easier to manage than they first appear. In most cases, you train the network to predict two real numbers per logical output: one for the real part and one for the imaginary part.

At inference time, rebuild the complex number:

pred = model.predict(x[:2], verbose=0)
pred_complex = pred[..., 0] + 1j * pred[..., 1]
print(pred_complex.shape)

This is usually enough for regression tasks. If your final evaluation uses amplitude or phase, compute those derived quantities after prediction.

Choosing the Right Loss

A common and reliable choice is mean squared error over the stacked real and imaginary channels. That corresponds to minimizing squared Euclidean distance in the two-dimensional representation.

If the task is especially sensitive to magnitude and phase, you may define a custom loss. For example, you could penalize magnitude error and phase error separately. The main requirement is consistency: the model output shape, label shape, and loss function must all agree.

Common Pitfalls

The biggest mistake is feeding complex64 tensors directly into ordinary Keras layers and assuming the full stack supports them. Some operations do, but many common layers do not.

Another problem is mixing representations. If your inputs are split into real and imaginary channels, your labels and loss must follow the same convention. Otherwise shape mismatches or silent semantic bugs appear quickly.

A third issue is evaluating the model only on channel-wise error when the actual domain metric is phase, magnitude, or spectral distance. Train in a representation Keras can support, but validate with domain-appropriate metrics.

Summary

• TensorFlow understands complex tensors, but many high-level Keras layers still expect real-valued inputs.
• The most practical solution is to split complex numbers into real and imaginary channels.
• For true complex algebra, implement custom layers using paired real-valued weights and operations.
• Complex outputs are typically modeled as two real outputs and reconstructed after prediction.
• Keep the data representation, model output, and loss function aligned throughout training.