class CustomConstSoftmaxLayer(Layer):

    def __init__(self, **kwargs):
        super(CustomConstSoftmaxLayer, self).__init__(**kwargs)

    def build(self, input_shape):
        self.const = self.add_weight(name='const',
                                     shape=tuple([1] * len(input_shape[0])),
                                     initializer='zero',
                                     trainable=True)
        super(CustomConstSoftmaxLayer, self).build(input_shape)

    def call(self, args):
        [x, logits] = args
        tmp = K.tile(self.const, K.shape(x))
        return K.sum(K.softmax(logits) * K.concatenate([x, tmp], axis=-1), axis=-1, keepdims=True)

    def compute_output_shape(self, input_shape):
        return input_shape[0]