tf.switch_case使用Keras Input作分支索引报类型错误的问题咨询

tf.switch_case是TensorFlow底层实现的控制流算子，入参校验环节只认可原生的Tensor类型；而Keras的Input层生成的KerasTensor是Keras框架内部用来做静态图结构推导的符号占位符，不是可直接被底层算子识别的真实张量对象。直接在模型搭建阶段把KerasTensor传给tf.switch_case，就会触发类型校验不通过的报错。

1. 用自定义Keras层封装分支逻辑 #

这是稳定性最好的实现方式：把tf.switch_case的调用逻辑放到自定义层的call方法里。Keras在层的call方法执行上下文里，会自动完成KerasTensor到底层运算张量的转换，这时候传入tf.switch_case的入参类型就符合要求了。
示例代码：

import tensorflow as tf
from tensorflow.keras.layers import Input, Layer

class SwitchCaseLayer(Layer):
    def __init__(self, branch_fns, default_fn, **kwargs):
        super().__init__(**kwargs)
        self.branch_fns = branch_fns
        self.default_fn = default_fn

    def call(self, inputs):
        # 压缩多余维度，匹配switch_case需要的标量索引输入
        branch_index = tf.squeeze(inputs)
        return tf.switch_case(
            branch_index=branch_index,
            branch_fns=self.branch_fns,
            default=self.default_fn
        )

# 定义分支计算逻辑
def f1(): return tf.constant(17)
def f2(): return tf.constant(31)
def f3(): return tf.constant(-1)

# 构建可正常运行的模型
t_input = Input(shape=(1,), name="t_input")
switch_layer = SwitchCaseLayer(branch_fns={0: f1, 1: f2}, default_fn=f3)
r = switch_layer(t_input)

2. 用Lambda层快速实现 #

如果分支逻辑较简单，不需要追踪额外权重，可以直接把tf.switch_case的逻辑封装在Lambda层中，利用Lambda层的输入适配机制完成类型转换。
示例代码：

t_input = Input(shape=(1,), name="t_input")
def switch_logic(x):
    idx = tf.squeeze(x)
    return tf.switch_case(
        branch_index=idx,
        branch_fns={0: f1, 1: f2},
        default=f3
    )
r = tf.keras.layers.Lambda(switch_logic)(t_input)

• 分支索引需要提前处理为0维整数张量：由于Input层的输入默认带批次维度，需要用tf.squeeze去掉多余的维度，再做整数类型转换，否则会出现索引维度不匹配的运行错误。
• 所有分支的输出形状、数据类型必须完全统一，否则Keras在推导模型结构时会触发形状不匹配的报错。
• 如果分支内需要接入完整的神经网络子模块，要把子模块作为自定义层的类属性进行实例化，保证模型的权重能被正常追踪、序列化加载。
• tf.switch_case的原生逻辑是对整个批次统一选择分支，如果需要实现按单样本粒度独立选择分支的逻辑，需要搭配tf.vectorized_map或者tf.map_fn做逐样本的分支判断，否则会出现索引维度和批次不匹配的错误。

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