在维度内建立索引的最简单方法

Question

我有两个张量x和y在前一个shape维度上具有相等的k。第二张量包含索引，以从第一维到最后一维检索值。对于rank为3的情况，则输出z应为z[i_1, i_2,...,i_k, j] = x[i_1, i_2,...,i_k, y[i_1, i_2, ...,i_k, j]]。

我目前有一种方法，需要重塑x和y，使用y附加gather_nd的行索引，最后返回到原始形状。有没有更优雅的方法？有没有一种方法可以获取索引的张量（如np.indices），最好是不需要满足秩或形状的知识，只要它们满足上述条件即可。

Answer 1

找到了！ tf.batch_gather和tf.batch_scatter。