如何使用tf.gather或tf.gather_nd实现tf.batch_gather

Question

我有一个 A = batchx300x6x4 张量和一个 B = batchx300x1 索引张量。

如何在A和B中获得形状为batchx300x1x4的输出

没有tf.batch_gather（）吗？

例如：

假设A= [[[[1, 2, 3],[4,5,6]]]]和B = [[[0]]]是我想要的是能够使用相应的B采样A。这应该导致C = [[[[1, 2, 3]]]].