Python 使用索引上的基调的张量中的Tensorflow访问元素

如何使用张量索引访问tensorflow

Tensor

中的tenor元素，如下所示：

import tensorflow as tf
import numpy as np

# indexing in numpy [Working]
matrix = np.random.randint(0, 10, [100, 100])
indices = np.random.randint(0, 100, [1000, 100])
elements = matrix[indices[:, 0], indices[:, 1]]

# indexing in tensorflow [Not working]
tf_matrix = tf.constant(matrix, dtype=tf.int32)
tf_indices = tf.constant(indices, dtype=tf.int32)
tf_elements = tf_matrix[tf_indices[:, 0], tf_indices[:, 1]]  # Error

session = tf.Session()
session.run(tf_elements)

我发现以下错误：

tensorflow.python.framework.errors\u impl.InvalidArgumentError:形状 必须为等级1，但对于“跨步切片”为等级2（op: 输入形状为[100100]、[21000]、[21000]、[2]的“Straddslice”）

ValueError:形状必须为秩1，但对于“跨步切片”为秩2 （op:'Straddslice'），输入形状：[100100]，[21000]，[21000]， [2]

tf_elements = tf.gather_nd(tf_matrix, tf_indices[:, 0:2])