我想在keras图层中实现自然渐变。这必须在已经到位的自定义渐变中发生。我希望能够在调用优化器时选择应该计算哪个实现（常规或自然梯度）。

我面临的问题是，当我nat_grad=True在训练时（而不是在图形生成时）将布尔值传递给Op时，AutoGraph并不满意。

目前，正在发生的事情的伪代码如下：

@tf.custom_gradient
def MyOp(inputs, w, nat_grad=False):
    output = w*inputs
    def grad(dy):
        if nat_grad:
            return dy, 1.0
        else:
            return -dy, -1.0
    return output, grad


class MyKerasLayer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.nat_grad = False
        
    def build(self, input_shape):
        self.w = self.add_weight("w", dtype=tf.float32, trainable=True, initializer=tf.random_normal_initializer)
        super().build(input_shape)

    def call(self, inputs):
        return MyOp(inputs, self.w, self.nat_grad)


class MyModel(tf.keras.Sequential):
    def __init__(self, num_layers):
        super().__init__([tf.keras.Input(shape=[1], batch_size=None, dtype=tf.float32)]+[MyKerasLayer() for _ in range(num_layers)])


def optimize(model, X, Y, nat_grad:bool):
    for layer in model.layers:
        layer.nat_grad = nat_grad
    model.fit(x=X, y=Y)

    
model = MyModel(5)
model.compile(optimizer='SGD', loss=lambda x,y:x-y, metrics=[])
X = np.array([1.0, 2.0, 3.0])
Y = np.array([1.0, 2.0, 3.0])
optimize(model, X, Y, nat_grad=True)
>>> OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

什么是正确的方法？