softmax probability clipping at the logit level

softmax_clip_logits.py

import numpy

s = numpy.random.randn(10)

def clip_probability_logits(logit, epsilon=0.):
    s_max = numpy.max(logit)
    s_transf = logit - s_max

    normalization_const = numpy.exp(s_transf).sum()
    p_transf = numpy.exp(s_transf) / normalization_const

    p_transf = numpy.maximum(epsilon, p_transf)

    s_back = numpy.log(p_transf) + numpy.log(normalization_const) + s_max

    return s_back

print(numpy.abs(clip_probability_logits(s, epsilon=1e-3) - s_back).sum())