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

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]].

我目前有一个方法需要 reshape x 和 y，附加 y 的行索引，使用 gather_nd 并最终恢复到原来的形状。有没有更优雅的方法？有没有办法获取索引的张量(如 np.indices)，最好不需要了解满足上述条件的秩或形状？

Answer 1

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