Python @tf.function错误地处理输入张量的无（批）维度

我想使用

@tf.function

decorator使用

tensorflow

的自动绘图功能将Python代码转换为tensorflow静态图，以便以后可以将其作为

.pb

静态图的一部分进行计算

我有

keras.layers.Input（…）

tensor作为输入。通常，keras的输入层具有未知（无）批次维度（第0个）

在我的修饰函数中，我在这个未知维度上使用python的for循环，期望tensorflow能够正确地推断每个反馈的真实输入的第0个批次维度大小，并正确地执行相应的已知循环迭代次数

但根据结果，tensorflow只是假设这个未知维度总是等于0，并且进行静态的0次循环迭代。因此，当未提供

llen

时，在下一个代码中打印

result[]

。但我希望打印的结果是

result[17]

，就像我显式地为函数输入提供等于2的llen一样

是否有任何方法可以迭代未知（无）维度并执行不同数量的循环迭代和张量变换

接下来是解释行为的代码：

import tensorflow as tf, numpy as np

def f(t, llen = None):
    if llen is None:
        llen = int(tf.shape(t)[0])
    l = [0] * llen
    for i in range(len(l)):
        l = l[:i] + [t[i, 0, 0]] + l[i + 1:]
    return tf.convert_to_tensor(l)
    
@tf.function
def ft(t, llen = None):
    return f(t, llen)

input_ = np.array([[[1,2,3],[4,5,6]],[[7,8,9],[10,11,12]]])

for eager in [True, False]:
    print('is eager:', eager)
    if eager:
        tf.compat.v1.enable_eager_execution()
        print('result', f(input_).numpy())
    else:
        tf.compat.v1.disable_eager_execution()
        for llen in [None, 2]:
            print('input llen', llen)
            i = tf.keras.layers.Input([2, 3])
            t = ft(i, llen)
            with tf.compat.v1.Session() as sess:
                print('result', sess.run(t, {i: input_}))