Mac M1环境下PyMC sample()函数编译错误排查求助

问题说明 #

在Mac M1设备上运行贝叶斯线性回归代码，调用PyMC的sample()生成后验分布时，抛出CompileError，核心报错为括号嵌套层级超过256上限。已尝试重装Jupyter及相关依赖（PyMC 5.6.1、NumPy 1.23.5、PyTensor 2.12.3），问题未解决。

复现代码 #

import pymc as pm
import numpy as np
import arviz as az

# 生成模拟数据
np.random.seed(42)
X = np.random.randn(100, 5)
beta_true = np.array([1, -2, 3, -4, 5])
y = X @ beta_true + np.random.randn(100) * 0.5

# 构建贝叶斯线性回归模型
with pm.Model() as linear_model:
    # 定义先验分布
    beta = pm.Normal('beta', mu=0, sigma=10, shape=X.shape[1])
    sigma = pm.HalfNormal('sigma', sigma=1)
    # 计算均值
    mu = pm.math.dot(X, beta)
    # 定义似然
    y_obs = pm.Normal('y_obs', mu=mu, sigma=sigma, observed=y)
    # 采样生成后验
    trace = pm.sample(2000, tune=1000, cores=2)

报错堆栈片段 #

CompileError: Compilation failed (return status=1):
/var/folders/xx/xxxxxx/T/pytensor_compilation_xxxxxx/model_xxxxxx.c:102345:1: error: bracket nesting level exceeded maximum of 256
}
^

可行解决方案 #

1. 调整PyTensor编译参数，提高嵌套上限 #

M1默认的clang编译器对括号嵌套层级有严格限制，直接修改PyTensor的编译配置即可解决：
在代码最开头添加以下配置：

import pytensor

# 设置clang编译器参数，将嵌套层级上限提高到1024
pytensor.config.cxx = "clang"
pytensor.config.cxxflags = "-fbracket-depth=1024"

添加后再运行模型代码即可。

2. 降级PyTensor到兼容版本 #

PyTensor 2.12.x在M1平台的编译优化存在bug，降级到2.11.0版本可规避该问题：
执行终端命令：

pip install pytensor==2.11.0

安装完成后重启Jupyter内核，重新运行代码。

3. 切换到PyMC的JAX采样后端 #

JAX在M1上的兼容性优于默认编译后端，PyMC 5支持直接切换：
首先安装M1兼容的JAX版本：

pip install "jax[cpu]" -f https://storage.googleapis.com/jax-releases/jax_macos_releases.html

然后修改采样代码的backend参数：

with linear_model:
    trace = pm.sample(2000, tune=1000, cores=2, target_accept=0.95, backend="jax")

验证修复 #

采样完成后，用ArviZ检查后验分布是否正常生成：

az.summary(trace)
az.plot_trace(trace)

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