TensorFlow中按指定张量分桶高效累加浮点数值的方法咨询

你要的功能可以直接用TensorFlow原生的高度优化算子实现，完全满足并行计算、无额外复杂度开销的要求：

1. 稠密输出方案 #

直接使用tf.math.bincount的weights参数即可，该算子底层为并行实现，时间复杂度为O(N)（N为输入元素个数），无多余内存开销，是你描述场景的最优解：

import tensorflow as tf
import numpy as np

buckets = 8
indices = tf.convert_to_tensor([0, 0, 0, 1, 1, 3, 3, 5], tf.int32)
values = tf.convert_to_tensor([.5, .3, .2, .1, 1., 1., 1., .1], tf.float32)

# 核心实现，一行代码完成
result = tf.math.bincount(indices, weights=values, minlength=buckets, maxlength=buckets, dtype=tf.float32)

# 验证结果
expected_result_dense = [1., 1.1, 0., 2., 0., .1, 0., 0.]
np.testing.assert_array_almost_equal(expected_result_dense, result.numpy())

你也可以用tf.math.unsorted_segment_sum实现等价效果，适合分桶ID不连续从0开始的场景：

result = tf.math.unsorted_segment_sum(values, indices, num_segments=buckets)

2. 稀疏输出方案 #

如果需要直接输出非零分桶的索引和值，避免生成大尺寸的稠密张量，可以按如下方式实现：

# 方法1：基于稠密结果提取（适合分桶数不大的场景）
non_zero_mask = result != 0
expected_indices_sparse = tf.where(non_zero_mask)[:, 0]
expected_values_sparse = tf.boolean_mask(result, non_zero_mask)

# 方法2：直接生成稀疏结果（适合分桶数极大、绝大多数分桶为0的场景，无中间稠密张量）
unique_indices, idx_in_unique = tf.unique(indices)
sorted_order = tf.argsort(unique_indices)
unique_indices_sorted = tf.gather(unique_indices, sorted_order)
sum_values = tf.math.unsorted_segment_sum(values, idx_in_unique, num_segments=tf.shape(unique_indices)[0])
sum_values_sorted = tf.gather(sum_values, sorted_order)

两种方案得到的结果都和你给出的稀疏预期值完全匹配。

内容的提问来源于stack exchange，提问作者Bartłomiej Popielarz