如何在@tf.function修饰的AutoGraph函数中创建动态形状张量

1. AutoGraph函数对动态形状张量的支持性 #

完全支持在启用AutoGraph的tf.function函数中声明动态形状张量，你遇到的报错不是动态形状本身不被支持，而是代码写法在静态图跟踪阶段触发了不符合静态图规则的隐式numpy调用。

2. 报错根因 #

当你给tf.zeros传递Python列表形式的形状参数[2*N, 2*N]时，部分旧版本TensorFlow会在静态图构建阶段调用np.prod计算形状总大小，而此时N是仅在运行时才会确定具体值的符号张量，无法被隐式转换为numpy数组，因此抛出NotImplementedError。

3. 正确的代码实现 #

你只需要将形状参数先构造为1D整数张量，再传给tf.zeros即可规避该问题，修改后的可运行代码如下：

import numpy as np
import tensorflow as tf

@tf.function
def make_matrix(a):
    N = tf.shape(a)[0]
    # 用tf.stack构造张量类型的形状参数
    shape = tf.stack([2*N, 2*N])
    A = tf.zeros(shape, dtype=a.dtype)
    return A

a = np.array([1,2,3,4], dtype=np.float32)
A = make_matrix(a)
print(A.numpy())

4. AutoGraph中动态形状张量分配通用规则 #

• 所有动态变化的尺寸参数，统一用Tensor类型传递，不要和Python原生的列表/元组混合作为形状入参
• 涉及形状拼接、构造的操作，用TensorFlow原生API（如tf.stack、tf.concat）完成，不要用Python列表推导、numpy操作处理动态尺寸
• 如果你使用TensorFlow 2.4及以上版本，也可以直接把形状参数写成元组(2*N, 2*N)，新版本已经做了兼容处理，但用tf.stack构造形状张量是跨版本兼容性最好的写法

内容的提问来源于stack exchange，提问作者capt