Tensorflow 如何根据2D索引收集日志？

假设序列

[[2,1,4]，[3,4,2]

是由预先训练的LSTM生成的。它的维度是

（2*3）

意味着每个样本中

批次大小=2

和

3个时间步

例如，总共有

5个功能

，因此登录可能是：

[[[0.2, 0.2, 0.1, 0.1, 0.2], 
  [0.3, 0.2, 0.1, 0.1, 0.1],
  [0.1, 0.2, 0.1, 0.1, 0.3]],
 [[0.2, 0.2, 0.1, 0.1, 0.2],
  [0.2, 0.2, 0.1, 0.1, 0.2],
  [0.2, 0.2, 0.1, 0.1, 0.2]]]

我想使用序列作为索引，从每个样本和每个时间步的logits中获得相应的概率。关于上面的例子，我想要得到的最终结果是

[[0.1, 0.2, 0.3],[0.1, 0.2, 0.1]]

我知道我可能需要tf.stack（），但我对如何处理维度感到困惑。谢谢你的帮助

我想我找到了办法

tf.losses.sparse_softmax_cross_entropy(labels = None, logits = None) 

labels

是索引，

logits

是模型的输出