def softmax_with_temperature(logits, temperature):
    """
    :param logits:      shape == (batch_size, seq_len, vocab_size), float32
    :param temperature: shape == (batch_size, 1), float32
    :return:
        transformed tokens probs, shape == (batch_size, seq_len, vocab_size), float32
    """

    def softmax_with_temp(args):
        logits, temperature = args
        repeat_num = K.shape(logits)[1]
        temperature_repeated = RepeatVector(repeat_num)(temperature)
        # shape == (batch_size, seq_len)
        scaled_logits = logits / temperature_repeated
        # shape == (batch_size, seq_len, vocab_size)

        # for numerical stability (e.g. for low temperatures):
        scaled_logits = scaled_logits - K.max(scaled_logits, axis=2, keepdims=True)
        # shape == (batch_size, seq_len, vocab_size)
        transformed_probs = K.softmax(scaled_logits)
        # shape == (batch_size, seq_len, vocab_size)
        return transformed_probs

    # wrap transformation in Lambda to turn the result to Keras layer
    transformed_probs = Lambda(
        function=softmax_with_temp,
        mask=lambda inputs, inputs_masks: inputs_masks[0],  # function to get mask of the first input
        name='softmax_with_temperature')([logits, temperature])
    # output shape == (batch_size, seq_len, vocab_size)

    return transformed_probs

    :param logits:      shape == (batch_size, seq_len, vocab_size), float32
    :param temperature: shape == (batch_size, 1), float32
    :return:
        transformed tokens probs, shape == (batch_size, seq_len, vocab_size), float32
    
cakechat/dialog_model/layers.py