检查张量中是否包含元素（带有'@tf.function'的Python Tensorflow2）

Question

我要检查一个元素是否包含在张量中，我会遇到问题。

例如 1

def foo(a):
    if 5 in tf.constant([5, 7, 9]):
        tf.print(a)

foo(2)

# you'll get '2', and no erros

例如 2

@tf.function
def foo(a):
    if 5 in tf.constant([5, 7, 9]):
        tf.print(a)

foo(2)

# you'll get erros like "TypeError: argument of type 'Tensor' is not iterable"

显然，添加@tf.function后情况有所不同。 如果您能帮我解决这个问题，我将不胜感激：:)

Answer 1

当您使用@tf.function 装饰 function 时，它将以图形模式运行。 在图形模式下，您不能迭代 tf.Tensor （这就是您在 if 语句中所做的）。