expectation of x * gelu(x) where x ~ N(mu, sigma)

x_gelu_expectation.py

def x_gelu_expectation(mu, sigma):
    """Compute E[x * gelu(x)] for x ~ N(mu, sigma^2) analytically."""
    evCDF = norm.cdf(mu / np.sqrt(1 + sigma**2))
    evPDF = norm.pdf(mu / np.sqrt(1 + sigma**2)) / np.sqrt(1 + sigma**2)
    evZPDF = -mu*sigma/np.sqrt(1 + sigma**2)**3 * norm.pdf(mu / np.sqrt(1 + sigma**2))

    # linearity
    evXPDF = mu * evPDF + sigma * evZPDF

    # identity (first time)
    evXCDF = mu * evCDF + sigma**2 * evPDF

    # by defn
    evGELU = evXCDF
    evDGELU = evCDF + evXPDF

    # identity (second time)
    return mu * evGELU + sigma**2 * evDGELU