tensorflow如何在段求和后获得与输入张量相同大小的输出

Question

我正在使用TensorFlow的tf.unsorted_segment_sum方法，它可以工作。 例如：

tf.unsorted_segment_sum(tf.constant([0.2, 0.1, 0.5, 0.7, 0.8]),
                        tf.constant([0, 0, 1, 2, 2]), 3)

给出正确的结果：

array([ 0.3,  0.5 , 1.5 ], dtype=float32)

我想得到：

array([0.3, 0.3, 0.5, 1.5, 1.5], dtype=float32)

Answer 1

我已经解决了。

data = tf.constant([0.2, 0.1, 0.5, 0.7, 0.8])
gr_idx = tf.constant([0, 0, 1, 2, 2])
y, idx, count = tf.unique_with_count(gr_idx)
group_sum = tf.segment_sum(data, gr_idx)
group_sup = tf.gather(group_sum, idx)

answer：

array([0.3, 0.3, 0.5, 1.5, 1.5], dtype=float32)