从 tf.function 中的 tensorflow 2.x 中任何类型的对象列表中选择项目

如何解决从 tf.function 中的 tensorflow 2.x 中任何类型的对象列表中选择项目

我有一个 tensorflow 模型列表，我想使用索引选择其中一个。但是，如果调用发生在用@tf.function 修饰的函数中，我将无法使其工作。有人找到了方法吗？

Select an item from a list of object of any type when using tensorflow 2.x 是一个类似的问题，但它们不向函数传递参数。不幸的是，使用 tf.switch_case，列表元素需要是可调用的，并且会在执行时被调用

class an_object():
    def __init__(self,a):
        self.a = a
    def __call__(self,b):
        print("called ",self.a,b)

@tf.function
def a_function(a_list):
    idx = tf.constant(1)
    # return a_list[idx] # longer error / bug message. Cause: module 'gast' has no attribute 'Index'
    # return tf.switch_case(idx,a_list) # TypeError: __call__() missing 1 required positional argument: 'b'
    
a_list = [an_object(1),an_object(2),an_object(3)]       
x = a_function(a_list)