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

Question

我有两个张量 x 和 y 在第一个 k 维度中具有相等的 shape。第二个张量包含从第一个维度到最后一个维度检索值的索引。对于 3 的 rank，则输出 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。