如何使用TensorFlow对正态分布在指定范围内采样并获取对数概率

直接使用TensorFlow Probability内置的tfd.TruncatedNormal截断正态分布即可，原生支持区间约束，自带采样和对数概率计算接口，效率远高于手动实现的拒绝采样。

原有代码的问题 #

• 逻辑错误1：调用参数up=-1、down=1，区间上下界颠倒导致采样区间为空，永远无法得到符合要求的样本
• 逻辑错误2：tf.where条件写反，当前逻辑是符合区间要求的样本会被重新采样，不符合的样本会被保留，和预期完全相反
• 性能问题：使用Python原生while循环，无法被TensorFlow图优化，每次循环都需要在Python runtime和TF内核间交互，批量采样时只要有一个样本不满足条件就需要全量重采，算力浪费严重

正确实现代码 #

import tensorflow as tf
from tensorflow_probability import distributions as tfd

# mean和stddev为形状(N,)的张量，对应N组正态分布的参数
mean = ... 
stddev = ...
low = 1  # 区间下界
up = 5  # 需保证up >= low，否则区间无效

# 直接初始化截断正态分布，自动约束采样区间在[low, up]
truncated_dist = tfd.TruncatedNormal(
    loc=mean,
    scale=stddev,
    low=low,
    high=up,
    validate_args=True,
    allow_nan_stats=False
)

# 每组分布采1个样本，返回形状为(N,)的采样值和对应截断分布的对数概率
samples, log_probs = truncated_dist.experimental_sample_and_log_prob()

# 如果需要获取原正态分布（未截断）下的对数概率，使用原Normal分布计算即可
original_dist = tfd.Normal(loc=mean, scale=stddev)
original_log_probs = original_dist.log_prob(samples)

补充说明 #

该实现是原生TFP优化的实现，底层使用逆变换采样直接生成符合区间约束的样本，无循环开销，支持批量参数和图编译，性能远高于手动拒绝采样。若你确实需要使用拒绝采样实现（比如适配自定义分布），请使用tf.while_loop编写可图优化的循环逻辑，避免Python级循环的性能损耗。

内容的提问来源于stack exchange，提问作者f_s