TensorFlow带轴支持的自定义归约函数：获取绝对值最大的张量原值

要实现类似reduce_max/reduce_min、按指定轴提取绝对值最大且保留原符号的张量值，可以通过「找绝对值最大的索引→提取原张量对应位置值」的思路实现，以下是完整的自定义函数实现：

import tensorflow as tf

def reduce_maxamplitude(tensor, axis):
    # 计算张量的绝对值
    abs_tensor = tf.abs(tensor)
    # 获取指定轴上绝对值最大的位置索引
    max_abs_indices = tf.argmax(abs_tensor, axis=axis, output_type=tf.int32)
    # 构造用于索引的多维坐标
    grid_dims = [tf.range(d) for d in tensor.shape.as_list()]
    grid_dims[axis] = max_abs_indices
    # 转换为堆叠的坐标张量
    indices = tf.stack(tf.meshgrid(*grid_dims, indexing='ij'), axis=-1)
    # 根据索引提取原张量的值
    result = tf.gather_nd(tensor, indices)
    return result

用你提供的测试张量验证效果：

tensor = tf.constant(
  [
    [[ 1,  5, -3],
     [ 2, -3,  1],
     [ 3, -6,  2]],

    [[-2,  3, -5],
     [-1,  4,  2],
     [ 4, -1,  0]]
   ]
)

# 测试axis=0
print(reduce_maxamplitude(tensor, axis=0))
# 输出：
# tf.Tensor(
# [[-2  5 -5]
#  [ 2  4  2]
#  [ 4 -6  2]], shape=(3, 3), dtype=int32)

# 测试axis=1
print(reduce_maxamplitude(tensor, axis=1))
# 输出：
# tf.Tensor(
# [[ 3 -6 -3]
#  [ 4  4 -5]], shape=(2, 3), dtype=int32)

# 测试axis=2
print(reduce_maxamplitude(tensor, axis=2))
# 输出：
# tf.Tensor(
# [[ 5 -3 -6]
#  [-5  4  4]], shape=(2, 3), dtype=int32)

实现逻辑说明 #

• 绝对值计算：先对原张量取绝对值，后续通过argmax定位绝对值最大的位置；
• 索引构造：利用tf.meshgrid生成除目标轴外的所有维度网格坐标，替换为argmax得到的目标轴索引，组合成tf.gather_nd可识别的多维坐标；
• 值提取：通过tf.gather_nd从原张量中提取对应位置的原值，完整保留符号信息。

该方法支持任意维度的张量与合法axis参数，使用逻辑完全对齐reduce_max。

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