Блог

Я пытаюсь перебрать тензор в режиме ожидания, но не могу.

Естественно, вы сделали бы что-то вроде:

 probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([1, 2, 3]))

@tf.function
def iterate_tensor(probs, indexs):
    return [output[label] for output, label in zip(probs, indexs)]
iterate_tensor(probs, indexs)
 

Но это приводит к ошибке OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed:

Еще одна вещь, которую я попробовал, была:

 probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([1, 2, 3]))

@tf.function
def iterate_tensor(probs, indexs):
    return tf.map_fn(lambda i: i[0][i[1]], (probs, indexs), dtype=(tf.int64, tf.int64))

iterate_tensor(probs, indexs)
 

Выдает ошибку ValueError: The two structures don't have the same nested structure.

Кажется, это работает:

 probs = tf.convert_to_tensor(np.array([[1,2,3], [4,5,6], [7,8,9]]))
indexs = tf.convert_to_tensor(np.array([1,1,1]))

@tf.function
def iterate_tensor(probs, indexs):
    return tf.linalg.diag_part(tf.gather(probs, indexs, axis=1))
iterate_tensor(probs, indexs)
 

Выход: <tf.Tensor: shape=(3,), dtype=int64, numpy=array([2, 5, 8])>