使用张量流收集 3D 张量

Question

我有一个名为 V 的 3D 张量，形状为 (P,Q,R)，还有一个名为 W 的 2D 张量，形状为 (P,S)。 我想创建形状为 (P,Q,R) 的张量 Z，其中 Z[i,j,k] = W[i,V[i,j,k]]。换句话说，W 的每一行对应于 V 的一行/切片。我尝试过使用 Gather 和 Gather_nd，但输出是 2D 或 4D。我尝试过一些变体：

import tensorflow as tf
V=tf.random.uniform((2,3,4),minval=0,maxval=2,dtype=tf.int32)
W=tf.random.uniform((2,20), minval=0, maxval=4, dtype=tf.int32)
Z=tf.gather(params=W, indices=V,axis=1,batch_dims=0)
print(Z.shape)

我有张量流2.15。有人可以帮忙吗？

Answer 1

我意识到我应该将批次维度的数量设置为 1。现在可以了：

import tensorflow as tf
V=tf.random.uniform((2,3,4),minval=0,maxval=2,dtype=tf.int32)
print(V[0,:])
W=tf.random.uniform((2,5), minval=0, maxval=4, dtype=tf.int32)
print('W=',W)
Z=tf.gather(params=W, indices=V, axis=1, batch_dims=1)
print(Z.shape)

形状为(2,3,4)，如所希望的。最后很简单，但弄清楚很痛苦！