问题：tf.function 发出警告并因 tf.stack 失败

Question

我在张量流中遇到tf.function 的问题。似乎无法转换包含指令tf.stack()的函数。

这里我写了一个简单的代码来突出这个问题：

import tensorflow as tf
c = tf.Variable([[1., 5.], [2., 4.]])
@tf.function
def toy_fct(x):
    y = tf.stack([x[0,:], x[1,:]], axis=0)
    return y
toy_fct(c)

我收到的消息如下：

WARNING:tensorflow:AutoGraph could not transform <function toy_fct at 0x000001DC9A13A670> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: module 'gast' has no attribute 'Index'
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
WARNING: AutoGraph could not transform <function toy_fct at 0x000001DC9A13A670> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: module 'gast' has no attribute 'Index'
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert

<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[1., 5.],
       [2., 4.]], dtype=float32)>

有人知道吗？

Answer 1

我已经在tf 2.0 到tf 2.4 中测试了您的代码，没有任何问题。但是当我使用gast==0.4.0 时它会出现。试试这个

pip install --upgrade pip
pip install gast==0.2.2

---

仅供参考，使用gast==0.4.0 和最新的tf 2.5.0-rc0 运行良好。