import tensorflow as tf



def CE_loss(word_probs, answers, loss_weights):
    '''
    word_probs: [batch_size, max_dec_steps, vocab]
    answers: [batch_size, max_dec_steps]
    loss_weigts: [batch_size, max_dec_steps]
    '''
    #word_probs = tf.nn.softmax(word_probs, dim=-1)
    input_shape = tf.shape(word_probs)
    vsize = input_shape[2]

    epsilon = 1.0e-6
    word_probs = _clip_and_normalize(word_probs, epsilon)

    one_hot_spare_rep = tf.one_hot(answers, vsize)

    xent = -tf.reduce_sum(one_hot_spare_rep * tf.log(word_probs), axis=-1) # [batch_size, max_dec_steps]
