def sample_truncated_gaussian(mu=0, sigma=1, lb=-np.Inf, ub=np.Inf):
    """
    Sample a truncated normal with the specified params. This
    is not the most stable way but it works as long as the
    truncation region is not too far from the mean.
    """
    # Broadcast arrays to be of the same shape
    mu, sigma, lb, ub = np.broadcast_arrays(mu, sigma, lb, ub)
    shp = mu.shape
    if np.allclose(sigma, 0.0):
        return mu

    cdflb = normal_cdf(lb, mu, sigma)
    cdfub = normal_cdf(ub, mu, sigma)

    # Sample uniformly from the CDF
    cdfsamples = cdflb + np.random.rand(*shp) * (cdfub-cdflb)

    # Clip the CDF samples so that we can invert them
    cdfsamples = np.clip(cdfsamples, 1e-15, 1-1e-15)
    zs = -np.sqrt(2) * special.erfcinv(2 * cdfsamples)

    # Transform the standard normal samples
    xs = sigma * zs + mu
    xs = np.clip(xs, lb, ub)

    return xs
Sample a truncated normal with the specified params. This
    is not the most stable way but it works as long as the
    truncation region is not too far from the mean.