如何在TensorFlow Probability中对独立二项随机变量求和并生成分布对象

多个独立、参数不同的二项分布变量之和没有标准闭式分布，可根据总试验次数规模选择以下两种方案实现，最终都返回原生TensorFlow Probability（TFP）的Distribution类型对象。

前置依赖导入 #

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions

方案1：精确PMF实现（总试验次数≤1000时优先选用） #

原理是通过概率生成函数（PGF）的卷积计算和的精确概率质量函数，最终封装为Categorical分布。如果所有变量都是试验次数为1的伯努利分布，直接调用TFP内置的tfd.PoissonBinomial(probs=p_list)即可。

通用二项分布（试验次数可大于1）的实现代码：

def get_exact_sum_dist(n_list: tf.Tensor, p_list: tf.Tensor) -> tfd.Distribution:
    """
    计算独立非等同二项变量和的精确分布
    Args:
        n_list:  shape为[N]的张量，存储每个二项分布的试验次数
        p_list:  shape为[N]的张量，存储每个二项分布的成功概率
    Returns:
        TFP Distribution类型的和分布对象
    """
    # 初始化PGF系数：0个变量时和为0的概率为1
    pgf_coeff = tf.constant([1.], dtype=p_list.dtype)
    
    for n, p in zip(tf.unstack(n_list), tf.unstack(p_list)):
        # 生成单个二项分布的PGF系数 (1-p + p*z)^n
        k_vals = tf.range(n + 1, dtype=tf.int32)
        single_coeff = tf.cast(tf.math.comb(n, k_vals), p_list.dtype) * \
                       tf.math.pow(1 - p, tf.cast(n - k_vals, p_list.dtype)) * \
                       tf.math.pow(p, tf.cast(k_vals, p_list.dtype))
        # 卷积更新总PGF系数
        pgf_coeff = tf.nn.conv1d(
            pgf_coeff[tf.newaxis, :, tf.newaxis],
            single_coeff[::-1, tf.newaxis, tf.newaxis],
            stride=1,
            padding="VALID"
        )[0, :, 0]
    
    return tfd.Categorical(probs=pgf_coeff)

# 调用示例
n_list = tf.constant([2, 3, 5], dtype=tf.int32)
p_list = tf.constant([0.2, 0.5, 0.7], dtype=tf.float32)
exact_sum_dist = get_exact_sum_dist(n_list, p_list)

方案特点：

• 输出的分布所有方法（sample()、log_prob()、mean()、variance()等）结果完全精确
• 总试验次数过大会导致卷积计算复杂度和内存占用过高，不适合总规模超过1000的场景

方案2：近似变换分布实现（总试验次数较大时选用） #

通过TFP的变换分布接口直接对独立二项联合分布做求和操作，不需要预计算全量PMF。
实现代码：

def get_approx_sum_dist(n_list: tf.Tensor, p_list: tf.Tensor) -> tfd.Distribution:
    # 构造独立二项联合分布
    joint_binom = tfd.Independent(
        tfd.Binomial(total_count=n_list, probs=p_list),
        reinterpreted_batch_ndims=1
    )
    # 沿最后一维求和得到和分布
    return tfd.TransformedDistribution(
        distribution=joint_binom,
        bijector=tfp.bijectors.ReduceSum(axes=-1)
    )

# 调用示例
approx_sum_dist = get_approx_sum_dist(n_list, p_list)

方案特点：

• 实现简单，支持任意总试验次数，采样速度快
• 大样本下如果需要快速计算log_prob，可改用正态近似：构造tfd.Normal(loc=tf.reduce_sum(n_list*p_list), scale=tf.sqrt(tf.reduce_sum(n_list*p_list*(1-p_list))))，计算速度最快

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