用keras后端按行索引2D张量的最佳方法是什么？

Question

我的整数张量大小（n，）和2D张量大小=（n，k）。对于每个在2D张量中具有索引i的“行”，我想用在1D张量中的索引i中的数字对其进行索引。

这个想法是仅当我在GPU上运行时，才使用Keras后端功能。

有了列表提示，我会这样：

array = [1,2,0,0]
matrix = np.array([[1,2,3],[2,3,4],[2,3,1],[3,2,1]])
print([matrix[i, num] for i, num in enumerate(array)])