tfd.TruncatedNormal Produces Samples Outside of Support

tfd.TruncatedNormal occasionally produce samples slightly outside of its support. I am not sure if this is an expected behavior and just a limitation of the rejection sampling implementation or if there is a bug in determining whether a proposed sample is in bounds.

I conducted a quick test on a singly truncated normal, tfd.TruncatedNormal(loc, scale, 0., np.inf), with randomly generated loc and scale parameters. The frequency at which samples are generated out of range appears to depend on the particular values of the location and scale parameters. Having the location parameter outside of [low, high] exacerbates the problem, but it isn't a prerequisite. Here is a plot of the samples with sample_value < TruncatedNormal.low from my experiments:

The values of the outlying samples are discretized in a suggestive way.

If this is a bug, triggering it is quite rare. Under the conditions tested, out of bounds samples were only generated in about 1 in 15 million cases. Here is the source code I used to generate the plots:

import numpy as np
from matplotlib import pyplot as plt
import tensorflow as tf
import pandas as pd
from tqdm import trange
from tensorflow_probability import distributions as tfd


low = 0.
high = np.inf
batch_size = 10000000
batches = 10000

loc_range = [-1., 1.]
scale_range = [0.01, 1.]

ct = 0
outliers = None
for i in trange(batches):
    loc   = tfd.Uniform(*loc_range).sample(batch_size)
    scale = tfd.Uniform(*scale_range).sample(batch_size)
    q = tfd.TruncatedNormal(loc, scale, low, high)
    z = q.sample()
    idx = (z < low).numpy()
    if len(idx) > 0:
        df  = pd.DataFrame({
            "z" : z[idx],
            'location' :  loc.numpy()[idx],
            'scale' :  scale.numpy()[idx],
        })
        outliers = df.append(outliers)

plt.figure()
plt.scatter(outliers.location, outliers.scale, c=outliers.z)
plt.colorbar(label='sample value')
plt.xlabel("loc")
plt.ylabel("scale")

plt.figure()
plt.hist(outliers.z, 100)
plt.xlabel("sample value")
plt.semilogy()
plt.show()

print(len(outliers))

Start by running the supplied Python reproduction for tfd.TruncatedNormal with the stated loc, scale, and support ranges. Inspect the TruncatedNormal sampling and rejection-sampling behavior, then add coverage showing that generated samples remain within [low, high] and confirm the rare out-of-bounds case is resolved.