Backpropagation across two parallel layers in Keras

Question

I want to create a network with two parallel layers (same input is given to two different layers and output of them is combined with some mathematical operations). Having said that, I am not sure the back-propagation will be done by Keras automatically. As a simple example of custom RNN cell,

class Example(keras.layers.Layer):

    def __init__(self, units, **kwargs):
        super(Example, self).__init__(**kwargs)
        self.units = units
        self.state_size = units
        self.la = keras.layers.Dense(self.units)
        self.lb = keras.layers.Dense(self.units)

    def call(self, inputs, states):
        prev_output = states[0]
        # parallel layers
        a = tf.sigmoid(self.la(inputs)) 
        b = tf.sigmoid(self.lb(inputs))
        # combined using mathematical operation
        output = (-1 * prev_output * a) + (prev_output * b)
        return output, [output]

Now, the loss gradient to `la` and `lb` layers are different (gradient of loss wrt `a`, should be `-output` but wrt `b` should be `output`), will this be taken care by Keras automatically or should we create custom gradient functions? 
Any insights and suggestions are much appreciated :)

Answer 1

Check the answer by Daniel Möller

Back-propagation will be taken care by Keras as long as all calculation is linked by tensor object, i.e don't cast tensor to another type like array, so don't worry about it.

For example of gradient tape, you can check the gradient at each layer by:

gradients = grad_tape.gradient(total_loss, model.trainable_variables)
gradient_of_last_layer = tf.reduce_max(gradients[-1])