如何在TensorFlow中对复数张量使用tf.reduce_max？

问题原因 #

TensorFlow原生的tf.reduce_max不支持复数类型，因为复数没有统一的"大小"定义；而numpy默认是按**复数的模长（幅值）**来比较并返回对应元素。

解决方案 #

根据不同的比较逻辑，有以下几种实现方式：

1. 按复数模长取最大值（和numpy默认行为一致） #

通过计算复数的模长，找到模长最大的元素索引，再从原张量中提取对应元素：

import tensorflow as tf

h = tf.constant([1.+2.j, 3 + 4.j])
# 计算每个复数的模长
magnitudes = tf.abs(h)
# 获取模长最大的元素索引
max_idx = tf.argmax(magnitudes)
# 根据索引取出对应的复数
max_complex = tf.gather(h, max_idx)

print(max_complex.numpy())  # 输出：(3+4j)

2. 按复数实部取最大值 #

如果需要以实部大小为判断标准：

import tensorflow as tf

h = tf.constant([1.+2.j, 3 + 4.j])
# 提取复数的实部
real_parts = tf.math.real(h)
max_idx = tf.argmax(real_parts)
max_complex = tf.gather(h, max_idx)

print(max_complex.numpy())  # 输出：(3+4j)

3. 按复数虚部取最大值 #

如果需要以虚部大小为判断标准：

import tensorflow as tf

h = tf.constant([1.+2.j, 3 + 4.j])
# 提取复数的虚部
imag_parts = tf.math.imag(h)
max_idx = tf.argmax(imag_parts)
max_complex = tf.gather(h, max_idx)

print(max_complex.numpy())  # 输出：(3+4j)

4. 高维复数张量的处理 #

对于高维张量，只需调整argmax的axis参数，并使用tf.gather_nd处理索引：

import tensorflow as tf

# 二维复数张量
h = tf.constant([[1.+2.j, 3+4.j], [5.+1.j, 2.+6.j]])
magnitudes = tf.abs(h)
# 按行取模长最大的元素索引（axis=1表示行维度）
max_idx = tf.argmax(magnitudes, axis=1)
# 构造高维索引
indices = tf.stack([tf.range(tf.shape(h)[0]), max_idx], axis=1)
max_complex = tf.gather_nd(h, indices)

print(max_complex.numpy())  # 输出：[3.+4.j 2.+6.j]

内容的提问来源于stack exchange，提问作者Royal