TensorFlow忽略零值的均值与标准差池化报错问题咨询

需要对批次内每个向量的各个维度执行忽略零值的平均池化、标准差池化操作，最小可复现代码如下：

import tensorflow as tf

av_pool = tf.keras.layers.Lambda(
    lambda z: tf.math.reduce_mean(z, axis=1)
)

sd_pool = tf.keras.layers.Lambda(
    lambda z: tf.math.reduce_std(z, axis=1, keepdims=False)
)

# Batch size = 2, sequence length = 4, vector size = 3
a = tf.convert_to_tensor(
    [[[1., 1., 1.], [1., 1., 1.], [1., 0., 0.], [0., 0., 0.]], 
     [[1., 1., 1.], [0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]],
    dtype=tf.float32
)

mask = tf.not_equal(a, tf.zeros_like(a))        

non_zero = tf.ragged.boolean_mask(a, mask)  

# 平均池化可正常运行
averages = av_pool(non_zero)

# 标准差池化运行失败
standard_deviations = sd_pool(non_zero)

运行后触发报错：

ValueError: keepdims=True is not supported for RaggedTensors.

• 两个约简算子对RaggedTensor的适配程度不同：tf.math.reduce_mean 已经完成了RaggedTensor沿不规则轴（即代码中axis=1对应的序列维度，每个样本保留的非零元素数量不一致，属于不规则维度）的约简逻辑适配，无论是否设置keepdims都可以正常计算，因此平均池化能正常输出结果。
• 报错和你代码里显式传入的keepdims=False没有直接关系：tf.math.reduce_std的内部计算逻辑是「先求均值→再计算每个元素和均值的平方差→求平方差的均值→开方得到标准差」，其中内部调用约简方法时硬编码了keepdims=True的中间步骤，而当前TensorFlow版本中reduce_std本身没有对RaggedTensor做全流程适配，内部调用的约简步骤不支持RaggedTensor传入keepdims=True，因此哪怕外部显式传了keepdims=False，运行到内部逻辑时依然会触发报错。

不要直接调用tf.math.reduce_std处理RaggedTensor，手动用已经适配RaggedTensor的算子拼接标准差计算逻辑即可，示例代码如下：

def ragged_compatible_std(x, axis):
    # 计算均值，reduce_mean已适配RaggedTensor的keepdims参数
    elem_mean = tf.math.reduce_mean(x, axis=axis, keepdims=True)
    # 计算每个元素和均值的平方差
    square_diff = tf.math.square(x - elem_mean)
    # 求方差
    elem_var = tf.math.reduce_mean(square_diff, axis=axis)
    # 开方得到标准差
    return tf.math.sqrt(elem_var)

# 替换原有sd_pool即可正常运行
sd_pool_fixed = tf.keras.layers.Lambda(
    lambda z: ragged_compatible_std(z, axis=1)
)
standard_deviations = sd_pool_fixed(non_zero)

内容的提问来源于stack exchange，提问作者Lorcán