Tensorflow Masked Softmax with mask after the exp operation.

masked_softmax_tf.py

import tensorflow as tf

def masked_softmax(logits, mask, axis):
    """softmax with mask after the exp operation"""
    e_logits = tf.exp(logits)
    masked_e = tf.multiply(e_logits, mask)
    sum_masked_e = tf.reduce_sum(masked_e, axis, keep_dims=True)
    # be careful when the sum is 0
    ones = tf.ones_like(sum_masked_e)
    sum_masked_e_safe = tf.where(tf.equal(sum_masked_e, 0), ones, sum_masked_e)
    return masked_e / sum_masked_e_safe