如何在Tensorflow中进行分组操作？

Question

我有一个tensor数据点,points其形状是(1, #data, #dimension)另一个dist,其形状是(#cluster, 1, #dimension),最后一个是cluster,具有形状(#data, ),表示points所属的簇的索引.现在我想计算每个集群的平均值.所以预期的产量应该有形(#cluster, 1, #dimension).我做了类似的东西,tf.gather(points, cluster)并得到了一个输出形状(#data, #data, #dimension).现在我不知道应该怎么处理它.请你帮助我好吗？

编辑:

例如,pointsis [[2, 3], [1, 4], [8, 10]]和clusteris [1, 0, 2],预期输出应该是[[[0,0], [2, 3], [0, 0]], [[1, 4], [0, 0], [0, 0]], [[0, 0], [0, 0], [8, 10]]]这样,便于计算平均值.