Tensorflow 2+. Использование альтернативных функций активации в сетях с использованием градиентов

Всех приветствую!

Для решения дифференциальных уравнения я пытаюсь применить методику PINN, которая использует стандартный перцептрон, задаваемый функцией

def create_standart_model(layers, activation='tanh', 
                          kernel_initializer='glorot_uniform'):
    model = keras.models.Sequential()
    model.add(keras.layers.InputLayer(input_shape=(layers[0])))
    for i in range(1, len(layers) - 1):
        model.add(keras.layers.Dense(layers[i], activation=activation, 
                                     kernel_initializer=kernel_initializer))
    model.add(keras.layers.Dense(layers[-1], activation='linear'))
    return model

В процессе обучения сети используются следующие функции (применительно к одномерному уравнению Пуассона):

def __loss(self, u, u_pred):
    " Функция ошибки "
    f_pred = self.f_model()
    return tf.reduce_mean(tf.square(u - u_pred)) + \
           tf.reduce_mean(tf.square(f_pred))

def __grad(self, X, u):
    " Расчет градиентов "
    with tf.GradientTape() as tape:
        loss_value = self.__loss(u, self.model(X))
    return loss_value, tape.gradient(loss_value, self.model.trainable_variables)

def f_model(self):
    " Расчет градиентов и невязки уравнения "
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(self.x_f)                                    
        u = self.model(self.x_f)                                
        u_x = tape.gradient(u, self.x_f)                        
    u_xx = tape.gradient(u_x, self.x_f)                         
    del tape                                                    
    return u_xx - self.x_f_init   

Далее производится обучение:

def fit(self, epochs=5000):
    X_u = tf.convert_to_tensor(X_u_global, dtype=self.dtype)
    u = tf.convert_to_tensor(u_global, dtype=self.dtype)
    for epoch in range(epochs):                                 
        loss_value, grads = self.__grad(X_u, u)
        self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

Если кратко, то численно рассчитывается второй градиент по x, а ошибка представляет из себя сумму невязки решения уравнения и выполнения ГУ.

Так вот, если при задании модели указать activation='tanh', то никаких абсолютно проблем не встречается и все великолепным образом рассчитывается.

Если же я захочу другую функцию активации, например softsign, то, по идее, пользуясь Keras, я просто должен задать activation='softsign', но, к сожалению, выскакивает следующая ошибка:

LookupError: gradient registry has no entry for: SoftsignGrad

Если кто-то знает, каким образом можно легко решить эту проблему, то я буду искренне благодарен. Я слышал про возможность задавания активаций через @tf.custom_gradient, но никаких значительных примеров в сети не нашел.