pymc模型中，当使用自定义pytensor op定义对数似然并尝试结合blackjax采样器时，可能遭遇jax转换兼容性错误。本文将深入探讨如何实现自定义对数似然，分析blackjax集成时的挑战，并提供一种通过数学表达式重构来显著提升核心计算函数性能的通用优化策略，即使无法利用jax加速，也能有效缩短采样时间。
1. PyMC中自定义对数似然函数的实现
在贝叶斯建模中，有时标准分布无法满足特定需求，需要引入自定义的对数似然函数。PyMC（基于PyTensor）提供了一种机制，允许用户通过定义自定义的pytensor.Op来集成任意Python函数及其梯度。
1.1 定义自定义PyTensor Op
要将一个复杂的Python函数（例如，涉及外部库或数值求解器）集成到PyMC的计算图中，需要创建两个pytensor.Op类：一个用于计算函数值（对数似然），另一个用于计算其梯度。
LogLikeWithGrad (对数似然函数)
这个Op负责计算给定参数的对数似然值。它需要实现perform方法来执行实际的对数似然计算，并重载grad方法来指定如何计算梯度。
import pytensor.tensor as pt
import numpy as np
from scipy.optimize import approx_fprime # 用于数值梯度
class LogLikeWithGrad(pt.Op):
itypes = [pt.dvector] # 输入是一个参数向量
otypes = [pt.dscalar] # 输出是一个标量（对数似然值）
def __init__(self, loglike_function):
self.likelihood = loglike_function
self.loglike_grad_op = LogLikeGrad(loglike_function) # 初始化梯度Op
def perform(self, node, inputs, outputs):
(theta,) = inputs
logl = self.likelihood(theta)
outputs[0][0] = np.array(logl)
def grad(self, inputs, grad_outputs):
(theta,) = inputs
# 调用自定义的梯度Op来计算梯度
grads = self.loglike_grad_op(theta)
return [grad_outputs[0] * grads]LogLikeGrad (对数似然梯度函数)
这个Op专门用于计算对数似然函数相对于其输入的梯度。在缺乏解析梯度的情况下，可以使用数值近似方法，例如scipy.optimize.approx_fprime。
class LogLikeGrad(pt.Op):
itypes = [pt.dvector] # 输入是一个参数向量
otypes = [pt.dvector] # 输出是一个梯度向量
def __init__(self, loglike_function):
self.likelihood = loglike_function
def perform(self, node, inputs, outputs):
(theta,) = inputs
# 使用数值方法近似梯度
grads = approx_fprime(theta, self.likelihood, epsilon=1e-8)
outputs[0][0] = grads1.2 在PyMC模型中集成自定义似然
一旦定义了自定义的LogLikeWithGrad Op，就可以将其作为pm.Potential添加到PyMC模型中。pm.Potential允许用户在模型中引入任意的对数概率贡献。
import pymc as pm
# 假设 applyMCMC 是你的核心对数似然计算函数
# 并且 param_names 和 lower/upper_boundaries 已经定义
# logl = LogLikeWithGrad(applyMCMC)
with pm.Model() as model:
# 定义模型参数
for i, name in enumerate(param_names):
pm.Uniform(name, lower=lower_boundaries[0][i], upper=upper_boundaries[0][i])
# 将所有参数组合成一个PyTensor向量
theta = pt.as_tensor_variable([model[param] for param in param_names])
# 将自定义对数似然作为潜力项添加到模型中
pm.Potential("likelihood", logl(theta))
# 执行采样
# trace = pm.sample(draws=niter, step=pm.NUTS(), tune=500, cores=64, init="jitter+adapt_diag", progressbar=True)2. Blackjax采样器与PyTensor Op的JAX兼容性挑战
PyMC 5.x 版本支持多种NUTS采样器后端，包括其默认的PyTensor后端以及基于JAX的Blackjax采样器，后者在GPU等加速设备上表现出色。然而，当模型中包含自定义的pytensor.Op时，尝试使用Blackjax采样器可能会遇到兼容性问题。
2.1 NotImplementedError 的根源
当尝试通过 pm.sample(nuts_sampler="blackjax") 使用Blackjax时，如果自定义的LogLikeWithGrad Op没有对应的JAX转换实现，PyTensor的JAX后端会抛出 NotImplementedError，错误信息通常为 No JAX conversion for the given Op: LogLikeWithGrad。
这是因为Blackjax采样器依赖于JAX的即时编译（JIT）能力，而JAX只能编译其能够理解的操作。自定义的pytensor.Op本质上是一个Python对象，PyTensor需要一个明确的规则来告诉JAX如何将其转换为JAX可执行的操作。对于标准PyTensor操作，这些转换已经内置，但对于用户自定义的Op，则需要手动提供。
2.2 解决方案方向
解决此问题通常需要以下两种方法之一：
- 重写Op以完全使用JAX操作：如果自定义对数似然函数的核心逻辑可以用JAX原生的操作（如jax.numpy）表达，那么可以直接在JAX中构建这个似然函数，并将其集成到PyMC模型中。这通常涉及将所有PyTensor代码替换为JAX代码。
- 为自定义Op提供JAX转换规则：这是一种更高级的方法，涉及为自定义pytensor.Op编写一个JAX转换函数，告诉PyTensor的JAX后端如何将该Op转换为JAX的计算图。这通常通过注册JAX调度函数来实现，但实现起来较为复杂，且文档相对较少。
在许多情况下，特别是当自定义似然函数依赖于复杂外部库（如物理模拟器）时，直接将其完全转换为JAX操作可能非常困难或不可能。此时，即使无法利用Blackjax的JAX加速，我们仍然可以通过优化核心计算逻辑来提升采样性能。
3. 提升PyMC模型计算性能的通用策略：数学表达式优化
即使无法直接利用JAX的GPU加速，通过对核心数学计算函数进行细致的优化，也能显著提升PyMC模型的采样速度。这种优化策略侧重于减少冗余计算、避免重复的函数调用以及利用局部变量缓存中间结果。
3.1 识别并消除冗余计算
在复杂的数学表达式中，往往存在重复计算相同子表达式的情况。通过将这些子表达式的结果存储在局部变量中，可以避免多次计算，从而提高效率。
以原始代码中的dH和du函数为例，它们包含大量重复的幂运算和乘法：
- (1 + z) 的不同幂次 ((1 + z)**2, (1 + z)**3, (1 + z)**4, (1 + z)**5)
- math.pi * Rho_m
- Omega_k * Phi
- Phi * Phi (即 Phi**2)
3.2 优化示例：dH 和 du 函数
以下是针对 dH 和 du 函数的优化版本，通过引入局部变量来缓存重复计算的中间结果：
import math
import timeit # 用于性能测试
# 假设 Rho_m, Phi, u, omega_BD, Omega_k, z 为示例输入
# (为了测试方便，这里使用任意值，实际应是模型参数)
Rho_m = -1.0
Phi = 0.1
u = 3.0
omega_BD = 4.0
Omega_k = -5.0
z = 6.0
# 原始 du 函数 (为对比而保留，实际代码中应替换为优化版本)
def du_original(Rho_m, Phi, u, omega_BD, Omega_k, z):
return (
24 * math.pi * Rho_m * Phi**3
+ (1 + z)
* u
* Phi**2
* (
8 * math.pi * (-3 + omega_BD) * Rho_m
- 3 * (1 + z) ** 2 * (3 + 2 * omega_BD) * Omega_k * Phi
)
- 3
* (1 + z) ** 2
* u**2
* Phi
* (
-4 * math.pi * omega_BD * Rho_m
+ (1 + z) ** 4 * (3 + 2 * omega_BD) * Omega_k * Phi
)
- omega_BD
* u**3
* (
4 * math.pi * (1 + z) ** 3 * (1 + omega_BD) * Rho_m
+ (1 + z) ** 5 * (3 + 2 * omega_BD) * Omega_k * Phi
)
) / (
(1 + z) ** 2
* (3 + 2 * omega_BD)
* Phi**2
* (8 * math.pi * Rho_m + 3 * (1 + z) ** 2 * Omega_k * Phi)
)
# 优化后的 du 函数
def du_optimized(Rho_m, Phi, u, omega_BD, Omega_k, z):
# 缓存幂次和重复乘法
Phi_pow2 = Phi * Phi
Phi_pow3 = Phi_pow2 * Phi
one_plus_z = 1 + z
one_plus_z_pow2 = one_plus_z * one_plus_z
one_plus_z_pow3 = one_plus_z_pow2 * one_plus_z
one_plus_z_pow4 = one_plus_z_pow3 * one_plus_z
one_plus_z_pow5 = one_plus_z_pow4 * one_plus_z
# 缓存其他重复子表达式
one_plus_z_pow2_times_3 = 3 * one_plus_z_pow2
pi_times_Rho_m = math.pi * Rho_m
Omega_k_times_Phi = Omega_k * Phi
u_pow2 = u * u
u_pow3 = u_pow2 * u
omg1 = (3 + 2 * omega_BD) # (3 + 2 * omega_BD)
omg = omg1 * Omega_k_times_Phi # (3 + 2 * omega_BD) * Omega_k * Phi
omg2 = omega_BD * pi_times_Rho_m # omega_BD * math.pi * Rho_m
return (
24 * pi_times_Rho_m * Phi_pow3
+ one_plus_z * u * Phi_pow2 * (8 * (-3 + omega_BD) * pi_times_Rho_m - one_plus_z_pow2_times_3 * omg)
- one_plus_z_pow2_times_3 * u_pow2 * Phi * (-4 * omg2 + one_plus_z_pow4 * omg)
- omega_BD * u_pow3 * (4 * one_plus_z_pow3 * (pi_times_Rho_m + omg2) + one_plus_z_pow5 * omg)
) / (
one_plus_z_pow2 * omg1 * Phi_pow2 * (8 * pi_times_Rho_m + one_plus_z_pow2_times_3 * Omega_k_times_Phi)
)
# 原始 dH 函数 (为对比而保留)
def dH_original(Rho_m, Phi, u, omega_BD, Omega_k, z):
val = (-16 * math.pi * Rho_m - 6 * (1 + z) ** 2 * Omega_k * Phi) / (
6 * (1 + z) * u + ((1 + z) ** 2 * omega_BD * u**2) / Phi - 6 * Phi
)
if val >= 0:
return -(
(
(1 + z)
* (16 * math.pi * Rho_m + 6 * (1 + z) ** 2 * Omega_k * Phi)
* (
(1 + z) * omega_BD * u**3
- 2
* omega_BD
* u
* ((1 + z) * du_original(Rho_m, Phi, u, omega_BD, Omega_k, z) + u)
* Phi
- 6 * du_original(Rho_m, Phi, u, omega_BD, Omega_k, z) * Phi**2
)
+ (
6
* Phi
* (
-8 * math.pi * Rho_m
+ (1 + z) ** 2 * Omega_k * ((1 + z) * u + 2 * Phi)
)
* (6 * Phi**2 - (1 + z) * u * ((1 + z) * omega_BD * u + 6 * Phi))
)
/ (1 + z)
)
/ (
2
* math.sqrt(val)
* (
(1 + z) ** 2 * omega_BD * u**2
+ 6 * (1 + z) * u * Phi
- 6 * Phi**2
)
** 2
)
)
else:
return None
# 优化后的 dH 函数
def dH_optimized(Rho_m, Phi, u, omega_BD, Omega_k, z):
# 缓存常用变量和幂次
Phi_pow2 = Phi * Phi
Phi_pow2_times_6 = Phi_pow2 * 6
Phi_times_6 = Phi * 6
one_plus_z = 1 + z
one_plus_z_pow2 = one_plus_z * one_plus_z
one_plus_z_times_u = one_plus_z * u
pi_times_Rho_m = math.pi * Rho_m
Omega_k_times_Phi = Omega_k * Phi
u_pow2 = u * u
u_pow3 = u_pow2 * u
# 重新计算 duu (如果 duu 在 dH 内部被调用多次，直接内联其计算可进一步优化)
# 此处为简洁起见，仍调用 du_optimized，但注意实际场景可内联
# 或者，如原答案所示，直接将 du_optimized 的计算逻辑复制到此处
# duu 的内联计算部分 (来自 du_optimized)
Phi_pow3_du = Phi_pow2 * Phi
one_plus_z_pow3_du = one_plus_z_pow2 * one_plus_z
one_plus_z_pow4_du = one_plus_z_pow3_du * one_plus_z
one_plus_z_pow5_du = one_plus_z_pow4_du * one_plus_z
one_plus_z_pow2_times_3_du = 3 * one_plus_z_pow2
omg1_du = (3 + 2 * omega_BD)
omg_du = omg1_du * Omega_k_times_Phi
omg2_du = omega_BD * pi_times_Rho_m
duu = (
24 * pi_times_Rho_m * Phi_pow3_du
+ one_plus_z * u * Phi_pow2 * (8 * (-3 + omega_BD) * pi_times_Rho_m - one_plus_z_pow2_times_3_du * omg_du)
- one_plus_z_pow2_times_3_du * u_pow2 * Phi * (-4 * omg2_du + one_plus_z_pow4_du * omg_du)
- omega_BD * u_pow3 * (4 * one_plus_z_pow3_du * (pi_times_Rho_m + omg2_du) + one_plus_z_pow5_du * omg_du)
) / (
one_plus_z_pow2 * omg1_du * Phi_pow2 * (8 * pi_times_Rho_m + one_plus_z_pow2_times_3_du * Omega_k_times_Phi)
)
# duu 内联计算结束
val1 = (-16 * pi_times_Rho_m - 6 * one_plus_z_pow2 * Omega_k_times_Phi)
val = val1 / (6 * one_plus_z_times_u + (one_plus_z_pow2 * omega_BD * u_pow2) / Phi - Phi_times_6)
if val >= 0:
Phi_times_2 = Phi + Phi
val2 = (one_plus_z_pow2 * omega_BD * u_pow2 + 6 * (one_plus_z_times_u * Phi - Phi_pow2))
# 优化分子中的复杂项
term1_numerator = one_plus_z_pow2 * val1 * (omega_BD * u * (one_plus_z * u_pow2 - Phi_times_2 * (one_plus_z * duu + u)) - duu * Phi_pow2_times_6)
term2_numerator = Phi_times_6 * (-8 * pi_times_Rho_m + one_plus_z_pow2 * Omega_k * (one_plus_z_times_u + Phi_times_2)) * (Phi_pow2_times_6 - one_plus_z_times_u * (one_plus_z_times_u * omega_BD + Phi_times_6))
return (term1_numerator - term2_numerator) / (2 * one_plus_z * math.sqrt(val) * val2 * val2)
else:
return None
# 性能测试
t_original = timeit.timeit('dH_original(Rho_m, Phi, u, omega_ 
