How to sample from a normal distribution random values inside a range using Tensorflow?

Question

I have two variables mean and stddev which are tensors of shape (1,) and they represent many normal distributions with mean lets say mean[i] and stardard deviation stddev[i].

From these distributions I want to sample one value within a range in [low,up,] for everyones and then I want to get the log probabilities of the sampled values.

From the docs I found that the experimental_sample_and_log_prob method is almost for me because it does not sample elements within a range of values (low, up) that I would like to have.

So I coded few lines but it doesn't work very well naturally ad it is so computationally expensive.

import tensorflow as tf
from tensorflow_probability import distributions as tfd


def sample_and_log_prob(dist, up, down):
    samples = dist.sample()
    accepted = False
    print("Is {} accepted? {}".format(samples, accepted))
    while not accepted:
        # sample < up
        cond1 = tf.less_equal(samples, up)
        # sample > down
        cond2 = tf.greater_equal(samples, down)
        # if down < sample < up
        accepted = tf.logical_and(cond1, cond2) 
        samples = tf.where(
            tf.logical_not(accepted),
            samples,
            dist.sample())
        print("Is {} accepted? {}".format(samples, accepted))
    
    return samples, dist.log_prob(samples)


distribution = tfd.Normal(
    loc=mean ,
    scale=stddev,
    validate_args=True,
    allow_nan_stats=False)

samples, log_probs = sample_and_log_prob(distribution, up=-1, down=1)

Any tips to solve it?

Answer 1

It sounds like you want a TruncatedNormal distribution.