TensorFlow中基于张量与Mask生成目标新张量及流程优化问题

修正原有代码的逻辑问题，补全压缩张量生成逻辑，适配通用场景的实现如下：

import tensorflow as tf

def stc_compression(tensor, sparsification_rate, rate_is_ratio=False):
    """
    STC压缩实现
    :param tensor: 输入任意维度张量
    :param sparsification_rate: 稀疏率，若rate_is_ratio为True则为保留元素比例，否则为保留元素个数k
    :param rate_is_ratio: 标记sparsification_rate是否为比例值
    :return: 压缩后的张量
    """
    # 提前计算绝对值避免重复计算
    abs_tensor = tf.abs(tensor)
    # 若传入的是比例，先计算实际要保留的元素个数k
    if rate_is_ratio:
        total_elements = tf.reduce_prod(tf.shape(tensor))
        k = tf.cast(total_elements * sparsification_rate, tf.int32)
    else:
        k = sparsification_rate
    
    # 打平张量计算top_k，适配任意维度输入
    flat_abs = tf.reshape(abs_tensor, [-1])
    top_k_vals, top_k_indices = tf.math.top_k(flat_abs, k=k, sorted=False)
    # 用indices生成准确的mask，避免阈值相等时保留元素数超过k的问题
    flat_mask = tf.scatter_nd(
        indices=tf.expand_dims(top_k_indices, axis=-1),
        updates=tf.ones_like(top_k_indices, dtype=tf.float32),
        shape=tf.shape(flat_abs)
    )
    mask = tf.reshape(flat_mask, tf.shape(tensor))
    
    # 计算k个top元素的平均绝对值
    average = tf.reduce_sum(top_k_vals) / tf.cast(k, tf.float32)
    
    # 生成压缩张量：符号*平均绝对值*mask，非top位置自动为0
    compressed_tensor = tf.sign(tensor) * average * mask
    return compressed_tensor

核心实现逻辑：用tf.sign(tensor)直接获取原张量每个位置的符号，乘以平均绝对值后再乘mask，刚好满足「k_top位置正负对应±average、其余位置为0」的需求。

• 原有逻辑bug：原有代码计算average时用tf.reduce_mean(tf.abs(tensor_masked))是对整个张量所有元素求平均，包含大量0值，计算出的平均值远小于真实k个top元素的平均值，修正为直接对top_k的取值求和后除以k。
• 高维适配优化：原有tf.math.top_k默认仅处理张量最后一维，遇到2维及以上的张量（如batch数据、特征图）时生成的mask不符合预期，改为先打平张量计算top_k再恢复形状，适配任意维度输入。
• 计算效率优化：提前将tf.abs(tensor)存为中间变量，避免重复计算3次绝对值，减少运算开销。
• 鲁棒性优化：原有通过阈值比较生成mask的逻辑，在多个元素等于top_k最小阈值时，会保留超过k个元素，改为通过top_k返回的索引生成mask，严格保证保留的元素数符合要求。
• 功能扩展性优化：增加比例参数支持，可直接传入保留比例无需提前计算k值。

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