PyMC 数学工具库 pymc.math 完全指南在贝叶斯模型中正确使用 PyTensor 数学表达式【免费下载链接】pymcBayesian Modeling and Probabilistic Programming in Python项目地址: https://gitcode.com/GitHub_Trending/py/pymc导读pymc.math是 PyMC 的数学函数子模块它把 PyTensorTheano 的后继框架的几乎所有张量数学运算重新导出并补充了 logit、logdiffexp、Kronecker 运算等贝叶斯建模专用函数。本文围绕 math.rst 这一 API 文档展开深入讲解「为什么模型内部必须使用 PyTensor 表达式而非 NumPy 代码」这一核心原则并基于 pymc/math.py 源码逐一拆解重导出函数与自定义函数的用法、实现原理和仓库内真实应用场景。读完本文你将能熟练地在自定义似然、先验和确定性变量中编写可被自动微分与符号优化的数学表达式。一、核心原则为什么模型中的数学必须用 PyTensor 表达式pymc.math文档开门见山地指出两个事实该子模块包含各种数学函数其中绝大多数直接从pytensor.tensor与pytensor.tensor.linalg重新导出完整签名与细节见 PyTensor 文档对 PyMC 随机变量做任何数学运算或者定义自定义似然likelihood与先验prior时必须使用这些 PyTensor 表达式而不是 NumPy 或纯 Python 代码。理解这一点是使用整个模块的前提。PyMC 的模型是一个符号计算图随机变量的分布参数、Deterministic变量、Potential以及自定义分布的logp都会被编译成 PyTensor 计算图随后经过梯度推导自动微分、重写优化rewrites和 JAX/Numba 后端编译。只有 PyTensor 张量表达式才能保留变量之间的依赖关系供 MCMC如 NUTS计算梯度支持形状推断与广播配合dims系统工作被 PyTensor 的优化器重写例如对三角矩阵打标签后启用快速求解在采样时按需惰性求值。如果在模型内混入 NumPy 调用例如np.logPyTensor 将无法对数值结果求导模型构建或采样会直接报错。这正是pymc.math存在的原因它让用户始终以pm.math.xxx的形式书写张量运算语义与 NumPy 一一对应学习成本极低。二、模块组成重导出层 PyMC 自定义层从源码结构看pymc/math.py 由三部分组成来源内容源码位置pytensor.tensor重导出约 100 个基础张量函数算术、比较、归约、形状、特殊函数pymc/math.pypytensor.tensor.linalg重导出线性代数函数cholesky、solve、det、eigh等pymc/math.pyPyMC 自定义实现kronecker、cartesian、logdiffexp、expand_packed_triangular、BatchedDiag、LogDet等pymc/math.py此外模块通过__all__见 pymc/math.py显式声明了全部公开符号覆盖约 130 个函数名。这一列表也是from pymc.math import *时的完整导出清单。2.1 重导出的基础张量函数文档说明这些函数直接来自pytensor.tensor按功能可分组如下算术运算abs、exp、expm1、log、log1p、log1pexp、log2、log10、sqr、sqrt、ceil、floor、round、clip、sgn、maximum、minimum、pow相关比较与逻辑eq、neq、lt、gt、le、ge、and_、or_、switch、where、all、any归约与统计sum、prod、mean、std、var、max、min、cumsum、cumprod、logsumexp、logaddexp形状与索引reshape、flatten、squeeze、expand_dims、broadcast_to、broadcast_arrays、transpose、moveaxis、swapaxes、stack、concatenate、repeat、tile、take、sort、argsort、argmax、argmin、unique、diff、diag、trace、tril、triu、linspace、arange创建与常量zeros、ones、zeros_like、ones_like、full、full_like、eye、constant、as_tensor张量运算dot、matmul。2.2 重导出的特殊函数与线性代数函数pytensor.tensor还贡献了大量统计建模常用的特殊函数均以pm.math.直接可用特殊函数erf、erfc、erfinv、erfcinv、gamma、gammaln、gammainc、gammaincc、digamma、polygamma、betainc、iv、kv、i0、i1概率变换sigmoid、softplus、softmax、log_softmax后两者来自pytensor.tensor.special见 pymc/math.py。线性代数部分来自pytensor.tensor.linalg见 pymc/math.py函数说明cholesky/cho_solveCholesky 分解与基于它的线性求解solve/solve_triangular通用线性求解 / 三角矩阵快速求解det/slogdet行列式 / 符号对数行列式eighHermitian/对称矩阵特征分解norm范数matrix_inverse即inv矩阵求逆block_diag分块对角矩阵三、PyMC 自定义数学函数实现原理与真实应用pymc.math的独到之处在于 PyMC 团队针对贝叶斯建模的数值稳定性与性能痛点实现了若干 NumPy/PyTensor 都没有直接提供的高阶工具。以下逐一解读其源码与仓库内的调用场景。3.1 数值稳定的对数域运算logdiffexp与log1mexp在对数概率域中log(exp(a) - exp(b))这类表达式极易发生下溢/上溢。logdiffexp的完整实现如下pymc/math.pydef logdiffexp(a, b): Return log(exp(a) - exp(b)). # Handle cases of -inf, -inf safely return pt.switch( pt.isneginf(b), a, a pt.log1mexp(b - a), )它利用pt.switch处理b -inf的边界情况此时直接返回a其余情况借助log1mexp实现数值稳定计算。配套的log1mexppymc/math.py实现了log(1 - exp(-x))其算法参考了 R 语言 Rmpfr 包中 Martin Mächler2012的稳定计算方法。兼容性提示旧版log1mexp支持negative_input参数当前版本已统一为“输入必须已取负”的语义显式传negative_inputFalse会抛出ValueError传negative_inputTrue则触发FutureWarning未来版本将完全移除该参数。写新代码时直接调用pm.math.log1mexp(x)并自行对输入取负即可。真实应用这两个函数在仓库中广泛用于分布的对数 CDF 计算。例如 Gumbel 分布的 logcdf 就用logdiffexp组合两个正态对数 CDF同时被 truncated.py、censoring.py、order.py 和 transforms.py 引用是截断、审查和顺序统计量对数概率计算的基础。3.2 概率链接函数logit/invlogit/probit/invprobit这四个函数覆盖了贝叶斯建模中最常用的链接函数pymc/math.py、pymc/math.pyinvlogit sigmoid # logit 的逆等价于 sigmoid def logit(p): return pt.log(p / (floatX(1) - p)) def probit(p): return -sqrt(2.0) * erfcinv(2.0 * p) def invprobit(x): return 0.5 * erfc(-x / sqrt(2.0))logit在实现中使用floatX(1)而非字面量1确保与 PyTensor 配置的浮点精度float32/float64一致invlogit直接是sigmoid的别名probit/invprobit基于erfcinv/erfc实现是高斯累积分布函数的逆变换。真实应用logit是 Beta 回归等模型中均值参数化 的核心invlogit则在 Bernoulli 等离散分布的对数概率 与 Logistic 回归确定性变量 中大量使用。3.3 Kronecker 运算家族kronecker/kron_diag/kron_dot/kron_solve_lower/kron_solve_upper这是pymc.math最具特色的部分专门服务于可分离协方差结构Kronecker 结构的高斯过程与多元分布。基础 Kronecker 积pymc/math.pydef kronecker(*Ks): rReturn the Kronecker product of arguments: K_1 ⊗ K_2 ⊗ ... ⊗ K_D return reduce(pt.linalg.kron, Ks) kron kronecker # 别名避免构造巨型 Kronecker 矩阵的惰性算子直接计算K_1 ⊗ K_2会得到尺寸平方爆炸的稠密矩阵。kron_matrix_oppymc/math.py实现了“不显式展开 Kronecker 积而是按分量依次作用”的算法其三个partial变体kron_dot partial(kron_matrix_op, oppt.dot) kron_solve_lower partial(kron_matrix_op, oppartial(solve_triangular, lowerTrue)) kron_solve_upper partial(kron_matrix_op, oppartial(solve_triangular, lowerFalse))分别对应 Kronecker 积的矩阵乘法、下三角/上三角求解。此外kron_diagpymc/math.py通过flat_outer即pt.outer(a, b).ravel()直接返回 Kronecker 积的对角线避免构造整个矩阵。真实应用这些函数是 pymc/gp/gp.py 中LatentKron与MarginalKronKronecker 结构高斯过程的核心引擎——先验用kron_dot(chols, v)实现旋转采样gp.py条件分布用kron_solve_lower/kron_solve_upper求解线性系统gp.py同时被 多变量分布的 MatrixNormal 实现 引用multivariate.py。测试用例 tests/gp/test_cov.py 验证了kronecker与直接计算的一致性。3.4 网格坐标生成cartesiandef cartesian(*arrays): Make the Cartesian product of arrays. N len(arrays) arrays_np [np.asarray(x) for x in arrays] arrays_2d [x[:, None] if np.asarray(x).ndim 1 else x for x in arrays_np] arrays_integer [np.arange(len(x)) for x in arrays_2d] product_integers np.stack(np.meshgrid(*arrays_integer, indexingij), -1).reshape(-1, N) return np.concatenate( [array[product_integers[:, i]] for i, array in enumerate(arrays_2d)], axis-1 )它生成多个一维/二维数组的完整笛卡尔积网格前面的数组循环更慢符合meshgrid的indexingij语义。真实应用Kronecker GP 要求输入位于完整网格上gp.py 的示例即Xnew1 np.linspace(-1, 2, 10)[:, None] Xnew2 np.linspace(0, 3, 10)[:, None] Xnew pm.math.cartesian(Xnew1, Xnew2) # 完整网格LatentKron 专用LatentKron.prior的参数说明中也明确指出总协方差函数是在cartesian(*Xs)完整网格上测量的gp.py并被 tests/gp/test_gp.py 与 tests/gp/test_cov.py 等测试覆盖。3.5 打包三角矩阵expand_packed_triangular三角矩阵可以只存储非零元素每行长度递增打包成一维数组。expand_packed_triangularpymc/math.py把这种打包格式还原为完整二维矩阵源码注释给出了编号规则[[0 - - -] [[0 1 2 3] [1 2 - -] [- 4 5 6] [3 4 5 -] [- - 7 8] [6 7 8 9]] [- - - 9]n三角矩阵行数必须是 Python 整数packed一维打包向量lowerTrue按下三角解包默认diagonal_onlyTrue时只提取对角线元素关键实现细节解包后通过out.tag.lower_triangular True或upper_triangular给张量打标签以启用 PyTensor 针对三角矩阵的优化重写如更快的三角求解。真实应用多变量正态的 Cholesky 因子经常以打包形式参数化。测试 tests/distributions/test_multivariate.py 与采样夹具 tests/sampler_fixtures.py 均展示了该用法L pm.math.expand_packed_triangular(D, packedL, lowerTrue)3.6 批量对角操作batched_diag与BatchedDiagbatched_diagpymc/math.py在二维与三维之间转换输入为二维矩阵若干对角向量形状(batch, dim)时通过自定义BatchedDiagOppymc/math.py输出三维对角块张量(batch, dim, dim)输入为三维张量时提取每个矩阵的对角线返回二维。BatchedDiag作为Op实现了make_node/perform/grad/infer_shapegrad通过对角索引提取实现infer_shape推导输出形状为(batch, dim, dim)。这样既保证前向高效也让梯度可自动微分。3.7 稳定的对数行列式LogDet与logdetlogdet LogDet() # 模块级单例LogDetpymc/math.py计算log(|det(M)|)通过 SVD 奇异值实现s np.linalg.svd(x, compute_uvFalse) log_det np.sum(np.log(np.abs(s)))相比直接算det(M)再取对数它避免了行列式的上溢/下溢且提供解析梯度gz * matrix_inverse(x).Tpymc/math.py。源码注释说明该实现源于 Theano 时代的历史遗留PR #3959现仍保留在 PyMC 中以服务特定需求。3.8 其他工具函数logbern(log_p, rngNone)pymc/math.py从对数概率采样一个伯努利样本log(u) log_plog_p为 NaN 时抛出FloatingPointErrorflatten_list(tensors)pymc/math.py将多个张量各自ravel()后拼接常用于把分块参数合并成单一向量flat_outer(a, b)pt.outer(a, b).ravel()kron_diag的底层原语。四、实战示例用 pm.math 编写自定义似然与确定性变量综合以上函数一个典型实践是在自定义模型中使用pm.math构造可微分的似然/先验表达式。仓库测试 tests/dims/test_model.py 展示了用softmax、concatenate、diff、maximum组合构造 B 样条基函数的设计模式delta_factors pm.math.softmax(z, axis-1) # 分组、节点上的权重 spline_slopes pm.math.concatenate(...) # 拼接各段斜率 beta pm.math.concatenate([beta0[:, None], pm.math.diff(spline_slopes)], axis-1) X pm.math.maximum(0, x[:, None] - knots[None, :]) # 截断幂基再如带约束参数的常见写法——用pm.math.exp保证正数、用pm.math.abs或Assert表达约束见 tests/model/test_core.pywith pm.Model(): log_sigma pm.Normal(log_sigma, 0, 1) sigma pm.Deterministic(sigma, pm.math.exp(log_sigma)) # 恒正 mu pm.Deterministic(mu, pm.math.abs(x)) # 无梯度断点安全这些写法之所以成立正是因为pm.math中的所有函数都返回 PyTensor 符号张量可被 NUTS 的梯度计算与 PyTensor 优化器正确处理。五、完整 API 速查表基于__all__pymc/math.py将pymc.math公开函数按用途整理如下类别函数基础算术abs, ceil, clip, exp, expm1, floor, log, log1p, log1pexp, log2, log10, maximum, minimum, round, sgn, sqr, sqrt比较/逻辑all, and_, any, eq, ge, gt, le, lt, neq, or_, switch, where归约/统计cumprod, cumsum, max, mean, min, prod, std, sum, var对数域logaddexp, logdiffexp, logsumexp, log1mexp, log_softmax概率/链接sigmoid, softplus, softmax, logit, invlogit, probit, invprobit, logbern特殊函数betainc, digamma, erf, erfc, erfcinv, erfinv, gamma, gammainc, gammaincc, gammaln, i0, i1, iv, kv, polygamma形状/索引arange, argsort, argmax, argmin, broadcast_arrays, broadcast_to, concatenate, diag, diff, expand_dims, flatten, linspace, moveaxis, repeat, reshape, sort, squeeze, stack, swapaxes, take, tile, trace, transpose, tril, triu, unique创建/常量as_tensor, constant, eye, full, full_like, ones, ones_like, zeros, zeros_like线性代数block_diag, cho_solve, cholesky, det, eigh, matmul, dot, norm, slogdet, solve, solve_triangular, matrix_inverseKronecker 族kronecker, kron, kron_diag, kron_dot, kron_solve_lower, kron_solve_upper, kron_matrix_opPyMC 专用cartesian, expand_packed_triangular, batched_diag, flat_outer, flatten_list, logdet六、进一步阅读API 文档源文件docs/source/api/math.rst完整实现源码pymc/math.py使用pm.math构建高斯过程的示例pymc/gp/gp.py分布内部对pymc.math的依赖pymc/distributions/continuous.py、pymc/distributions/multivariate.py、pymc/distributions/truncated.py相关测试tests/distributions/test_multivariate.py、tests/gp/test_cov.py、tests/dims/test_model.py总结pymc.math是 PyMC 与 PyTensor 之间的桥梁它把 PyTensor 的完整张量运算以pm.math命名空间暴露给用户并针对贝叶斯建模补充了数值稳定的对数域运算logdiffexp、log1mexp、链接函数logit/probit族、Kronecker 结构算子与打包三角矩阵工具。牢记文档的核心忠告——对随机变量的任何数学处理都必须走 PyTensor 表达式——你就掌握了在 PyMC 中编写自定义先验、似然与确定性变量的正确姿势从而充分利用符号图自动微分、重写优化与多后端编译的全部能力。【免费下载链接】pymcBayesian Modeling and Probabilistic Programming in Python项目地址: https://gitcode.com/GitHub_Trending/py/pymc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考