tensorflow - pytorchgather()的tensorflow等效实现

Question

如果我有一个大小为 (2,1,64,512) 的张量S和一个大小为 ( 2,5,64,512 ) 的索引张量D。在 pytorch 中，我可以根据 D by: 沿轴 = 1 从 S 中提取值o = torch.gather(S, D, axis=1)，这会产生一个大小为 (2,5,64,512) 的新张量。

我想知道是否有一种方法可以使用 tensorflow 产生相同的结果，即使张量的大小不同。

谢谢，

标签： tensorflowpytorch