Implementing Additive (Bahdanau) Attention mechanism with keras functional API

I was looking for implementing attention in my seq2seq model. I know that keras has built-in attention layers (both bahdanau and luong), but I'm trying to fit a custom attention layer. I understand the math behind it but I'm not able to implement it in the desired way. I found out some custom made Additive attention layer but they only work with gradient tape and not with keras functional API.

Here's the custom attention layer -

class Attention(tf.keras.layers.Layer):
    def __init__(self, units, **kwargs):
        super(Attention, self).__init__(**kwargs)
        self.units = units

    def call(self, features, hidden):
        hidden_with_time_axis = tf.expand_dims(hidden, 1)
        score = tf.nn.tanh(self.W1(features) + self.W2(hidden_with_time_axis))
        attention_weights = tf.nn.softmax(self.V(score), axis=1)
        context_vector = attention_weights * features
        context_vector = tf.reduce_sum(context_vector, axis=1)
        return context_vector

    def build(self, input_shape):
        self.W1 = tf.keras.layers.Dense(self.units)
        self.W2 = tf.keras.layers.Dense(self.units)
        self.V = tf.keras.layers.Dense(1)

    def get_config(self):
        config = super().get_config()
        config.update({'units': self.units})
        return config

Here feature refers to the encoder output and hidden means the previous decoder hidden state. The keras attention layer takes the decoder and encoder outputs (batch,steps,units) as the inputs and outputs a 3D vector (most probably the context vector) which is concatenated with the decoder output and then fed to the fully connected layer. But in the case of the custom layer, it outputs a 2D vector which doesn't seem to work, unlike the former.