如何使用tensorflow 1.14为keras层定义自定义渐变

我试图在tf.keras（tfver1.14）中实现梯度反转层。我使用tf.custom_gradient定义了一个自定义渐变，但从未调用grad函数，结果是错误的

我正在使用下面的测试层

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras import backend as k
import tensorflow as tf
tf.enable_eager_execution()
import numpy as np

@tf.custom_gradient
def custom_op(x):
    result = x
    def custom_grad(dy):
        grad = tf.negative(dy)
        return grad
    return result, custom_grad

class GradientReversal(tf.keras.layers.Layer):
  def __init__(self):
    super(GradientReversal, self).__init__()

  def call(self, inputs):
    return custom_op(inputs)


a = Input(shape=(2,))
b = GradientReversal()(a)
c = Dense(1)(b)
model = Model(inputs=a, outputs=c)
model.compile(loss='binary_crossentropy', optimizer='sgd')

@tf.function
def test():  
  with tf.GradientTape() as tape:
    out = model(np.ones((1, 2)))
  gradients = tape.gradient(out, model.trainable_variables)
  return gradients

print(test())

---

没关系，我应该在密度之后使用渐变层来获得我想要的结果