TensorFlow中tf.cond与tf.case无视条件执行全部分支解决方案问询

1. 原生Python操作在编译期全量执行 #

你第一个示例中看到所有分支的print都被触发，是因为tf.function首次调用时会遍历所有分支的Python代码完成计算图构建，这个阶段所有原生Python操作（比如print、普通变量赋值）都会执行，和运行时的条件判断无关。实际运行推理时只会执行符合条件的TensorFlow运算，返回值符合预期就是证明。如果要验证运行时的分支执行情况，把print替换为tf.print即可，运行时只会触发匹配条件分支的tf.print。

2. 第二个示例报错的核心原因 #

• 代码存在笔误：default=d的写法错误，d需要传入参数t才能返回无参的执行函数，和前面a(t)、b(t)的逻辑对齐，应该写为default=d(t)。
• 提前触发了分支运算：你在tf.case的参数中直接写a(t)，相当于编译阶段就执行了a函数，不管条件是否成立都会触发内部einsum的形状校验，而tf.shape(t)是动态张量，编译阶段无法确定具体值，不同分支的einsum输出形状不匹配就会抛出错误。
• 普通if tf.equal(x,2):报错的原因：tf.equal返回的是张量对象，不是Python布尔值，无法直接作为Python if的判断条件。

场景1：张量形状静态可确定（输入张量形状在定义时固定，你的示例属于该场景） #

直接使用张量的静态形状属性t.shape，它返回Python层面的TensorShape对象，不是张量，可以直接用普通Python if判断，不会触发布尔值报错，编译时也只会构建匹配条件的分支计算图：

x = tf.ones((3,2,1))
y = tf.ones((1,2,3))
z = tf.ones((4,3,5))

@tf.function
def f(t):
    if t.shape[0] == 3:
        return tf.einsum("ijk,lmi", t, y)
    elif t.shape[1] == 3:
        return tf.einsum("ijk,ljm", t, z)
    else:
        return tf.einsum("ijk,klm", t, z)

print(f(x))

该写法和你去掉tf.function的运行结果完全一致，不会触发其他分支的运算。

场景2：张量形状动态可变（比如输入张量的维度大小在编译阶段无法确定，比如batch维度动态） #

调整tf.case的传参逻辑，传递无参可调用对象，避免编译阶段提前执行分支运算，同时保证所有分支的输出形状、数据类型完全一致：

@tf.function
def f(t):
    return tf.case([
        (tf.equal(tf.shape(t)[0], 3), lambda: tf.einsum("ijk,lmi", t, y)),
        (tf.equal(tf.shape(t)[1], 3), lambda: tf.einsum("ijk,ljm", t, z)),
    ], default=lambda: tf.einsum("ijk,klm", t, z), exclusive=True)

使用lambda包装分支运算后，只有运行时条件匹配时才会执行对应分支的einsum运算，不会提前触发全部分支的形状校验。

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