mha_ffnn_helper_functions.py

def build_multi_head_attention_layers(num_neurons, num_heads):
    multi_head_attention_layer = MultiHeadAttentionLayer(num_neurons, num_heads)   
    dropout = tf.keras.layers.Dropout(0.1)
    normalization = LayerNormalization(epsilon=1e-6)
    return multi_head_attention_layer, dropout, normalization

def build_feed_forward_layers(num_neurons, num_hidden_neurons):
    feed_forward_layer = tf.keras.Sequential()
    feed_forward_layer.add(Dense(num_hidden_neurons, activation='relu'))
    feed_forward_layer.add(Dense(num_neurons))
        
    dropout = Dropout(0.1)
    normalization = LayerNormalization(epsilon=1e-6)
    return feed_forward_layer, dropout, normalization