Python torch.gather的tensorflow等效物

我有一个形状的张量

（164096,3）

。我有另一个形状指数的张量

（1632768,3）

。我试图收集沿

dim=1

的值。这最初是在pytorch中使用如下所示的方式完成的-

# a.shape (16L, 4096L, 3L)
# idx.shape (16L, 32768L, 3L)
b = a.gather(1, idx)
# b.shape (16L, 32768L, 3L)

b = tf.gather(a, idx, axis=1)
# b.shape (16, 16, 32768, 3, 3)

请注意，输出

b

的大小与

idx

的大小相同。然而，当我应用tensorflow的

聚集

函数时，我得到了完全不同的输出。发现输出维度不匹配，如下所示-

# a.shape (16L, 4096L, 3L)
# idx.shape (16L, 32768L, 3L)
b = a.gather(1, idx)
# b.shape (16L, 32768L, 3L)

b = tf.gather(a, idx, axis=1)
# b.shape (16, 16, 32768, 3, 3)

我还尝试使用

tf.gather\n

，但没有成功。见下文-

b = tf.gather_nd(a, idx)
# b.shape (16, 32768)

为什么我会得到不同形状的张量？我想得到和pytorch计算的形状相同的张量

---

换句话说，我想知道torch.gather的tensorflow等价物。

对于2D情况，有一种方法可以做到：

# a.shape (16L, 10L)
# idx.shape (16L,1)
idx = tf.stack([tf.range(tf.shape(idx)[0]),idx[:,0]],axis=-1)
b = tf.gather_nd(a,idx)

但是，对于ND情况，该方法可能非常复杂

这“应该”是使用tf.gather\n的一般解决方案（我只测试了沿最后一个轴的秩2和秩3张量）：

---

对于最后一个轴聚集，我们可以对一般ND情况使用2D整形技巧，然后使用上面的@Lishaoyan 2D代码

        # last-axis gathering only - use 2D-reshape-trick for Torch's style nD gathering
        def torch_gather(param, id_tensor):

            # 2d-gather torch equivalent from @LiShaoyuan above 
            def gather2d(target, id_tensor):
                idx = tf.stack([tf.range(tf.shape(id_tensor)[0]),id_tensor[:,0]],axis=-1)
                result = tf.gather_nd(target,idx)
                return tf.expand_dims(result,axis=-1)

            target = tf.reshape(param, (-1, param.shape[-1])) # reshape 2D
            target_shape = id_tensor.shape

            id_tensor = tf.reshape(id_tensor, (-1, 1)) # also 2D-index
            result = gather2d(target, id_tensor)
            return tf.reshape(result, target_shape)

---

有什么建议吗？你能试着结合上面给出的例子吗？