WeightedAverage merging layer for Keras

WeightedAverage_layer.py

from keras import backend as K
from keras.layers import Average
from keras.activations import softmax

class WeightedAverage(Average):
  
  def build(self, input_shape):
    self.kernel = self.add_weight(name='kernel', 
                                  shape=(1, len(input_shape)),
                                  initializer='ones',
                                  trainable=True)
    super(WeightedAverage, self).build(input_shape)
  
  def _merge_function(self, inputs):
    w = softmax(self.kernel)
    output = K.zeros_like(inputs[0])
    for i in range(len(inputs)):
      output += w[0,i] * inputs[i]
    return output