тензор потока Как получить выход такого же размера, как входной тензор после суммы сегмента

Question

Я использую метод tf.unsorted_segment_sum в TensorFlow, и он работает. Например:

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)

Ответ:

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