如何在tf.function中用TensorFlow张量索引tf.keras.layers.Layer列表？

我拥有一个tf.keras.layers.Layer列表，希望在tf.function中通过索引张量对其进行索引，理想实现代码如下：

import tensorflow as tf
from typing import List

@tf.function
def compute_single_output(heads: List[tf.keras.layers.Layer], index_tensor: tf.Tensor, input: tf.Tensor):
    logit = heads[index_tensor](input)
    return logit

但运行时出现错误：

TypeError: list indices must be integers or slices, not Tensor

我尝试过tf.gather、TF Lookup Table，甚至将heads改为字典，均无法解决问题。唯一可行的方法是使用Python for循环遍历列表，但该方法每次迭代都会积累CPU内存（由于Autograph无法将heads转换为张量，无法使用tf.while_loop）：

@tf.function
def compute_single_output(heads: List[tf.keras.layers.Layer], index_tensor: tf.Tensor, input: tf.Tensor):
    logit = tf.constant(-1, dtype=tf.float32)
    m = tf.constant(0)
    for head in heads:
        if m == index_tensor:
            logit = head(input)
        else:
            pass
        m += 1
    return logit

由于需要长期运行该函数，无法承受内存无限积累的问题，求可行的解决办法？

方法一：使用tf.switch_case实现高效分支选择 #

tf.switch_case是图模式友好的分支控制API，支持基于张量值选择执行对应分支，不会产生Python循环带来的内存泄漏问题，且只会执行选中层的计算，效率更高。

import tensorflow as tf
from typing import List

@tf.function
def compute_single_output(heads: List[tf.keras.layers.Layer], index_tensor: tf.Tensor, input: tf.Tensor):
    # 为每个层构建对应的分支函数
    branches = {}
    for idx, head in enumerate(heads):
        def branch(input_tensor=input, layer=head):
            return layer(input_tensor)
        branches[idx] = branch
    
    # 根据索引张量选择对应分支执行
    logit = tf.switch_case(
        index_tensor,
        branch_fns=branches,
        default=lambda: tf.constant(-1, dtype=tf.float32)
    )
    return logit

方法二：堆叠所有层输出后用tf.gather选择（适合层数量较少场景） #

如果层的数量不多、计算成本较低，可以先计算所有层的输出并堆叠，再用tf.gather根据索引张量选择对应结果。这种方法实现更简洁，但会一次性计算所有层的输出。

import tensorflow as tf
from typing import List

@tf.function
def compute_single_output(heads: List[tf.keras.layers.Layer], index_tensor: tf.Tensor, input: tf.Tensor):
    # 计算所有层的输出并沿第0维度堆叠
    all_logits = tf.stack([head(input) for head in heads])
    # 根据索引张量获取对应输出
    logit = tf.gather(all_logits, index_tensor)
    return logit

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