Python TensorFlow 2 中自定义层的 call() 中的 tf.while_loop()

我想编写一个自定义层，将密集层和一些指定函数应用于该计算的输出.我想指定应用于列表中各个输出的函数，这样我就可以轻松地更改它们.

我试图在tf.while_loop中应用函数，但我不知道如何访问和写入dense_output_nodes的各个元素.

dense_output_nodes[i] = ...不能像它告诉我的那样起作用

TypeError:"Tensor"对象不支持项分配

所以我之前try 过tf.unstack，这是下面的代码，但是现在当用hidden_1 = ArithmeticLayer(unit_types=['id', 'sin', 'cos'])(inputs)创建图层时，我得到了一个错误:

类型错误:列表索引必须是整数或切片，而不是张量

因为显然TensorFlow将i从tf.constant转换为tf.Tensor.

到现在为止，我真的很难找到解决这个问题的方法.有什么办法可以让它发挥作用吗？

class ArithmeticLayer(layers.Layer):
    # u = number of units
    
    def __init__(self, name=None, regularizer=None, unit_types=['id', 'sin', 'cos']):
        self.regularizer=regularizer
        super().__init__(name=name)
        self.u = len(unit_types)
        self.u_types = unit_types

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.u),
                                 initializer='random_normal',
                                 regularizer=self.regularizer,
                                 trainable=True)
        self.b = self.add_weight(shape=(self.u,),
                                 initializer='random_normal',
                                 regularizer=self.regularizer,
                                 trainable=True)


    def call(self, inputs):
        # get the output nodes of the dense layer as a list
        dense_output_nodes = tf.matmul(inputs, self.w) + self.b
        dense_output_list = tf.unstack(dense_output_nodes, axis=1)
        
        # apply the function units
        i = tf.constant(0)
        def c(i):
            return tf.less(i, self.u)
        def b(i):
            dense_output_list[i] = tf.cond(self.u_types[i] == 'sin',
                                            lambda: tf.math.sin(dense_output_list[i]),
                                            lambda: dense_output_list[i]
                                           )
            dense_output_list[i] = tf.cond(self.u_types[i] == 'cos',
                                            lambda: tf.math.cos(dense_output_list[i]),
                                            lambda: dense_output_list[i]
                                           )
            return (tf.add(i, 1), )
        [i] = tf.while_loop(c, b, [i])
        
        final_output_nodes = tf.stack(dense_output_list, axis=1)
        return final_output_nodes

谢谢你的建议！