如何利用JAX在GPU上加速PyMC的sample_posterior_predictive采样？

我在Linux环境下使用pm.sampling_jax.sample_numpyro_nuts可借助GPU加速PyMC模型采样，效果良好，但执行pm.sample_posterior_predictive扩展idata时只能用CPU，成为性能瓶颈。相关代码如下：

# Sampling
if gpu_available:
    idata = pm.sampling_jax.sample_numpyro_nuts(draws=draws_def, 
                                                tune=tune_def, 
                                                target_accept=targ_acc_def, 
                                                chain_method='vectorized', 
                                                idata_kwargs={"log_likelihood": True})

else:
    idata = pm.sample(draws=draws_def, tune=tune_def, target_accept=targ_acc_def, idata_kwargs={"log_likelihood": True})


idata.extend(pm.sample_posterior_predictive(idata, var_names=["y_obs"]))

解决步骤 #

• 替换原生sample_posterior_predictive为PyMC JAX模块下的sample_posterior_predictive_jax，该方法会利用JAX的GPU加速能力
• 确保模型构建和采样全程使用JAX后端（若之前用CPU后端，需先执行pm.set_backend("jax")）

修改后的代码 #

# 确保使用JAX后端（若未设置）
pm.set_backend("jax")

# Sampling
if gpu_available:
    idata = pm.sampling_jax.sample_numpyro_nuts(draws=draws_def, 
                                                tune=tune_def, 
                                                target_accept=targ_acc_def, 
                                                chain_method='vectorized', 
                                                idata_kwargs={"log_likelihood": True})
    # 使用JAX版本的后验预测采样
    ppc_idata = pm.sampling_jax.sample_posterior_predictive_jax(idata, var_names=["y_obs"])
else:
    idata = pm.sample(draws=draws_def, tune=tune_def, target_accept=targ_acc_def, idata_kwargs={"log_likelihood": True})
    ppc_idata = pm.sample_posterior_predictive(idata, var_names=["y_obs"])

idata.extend(ppc_idata)

额外说明 #

• sample_posterior_predictive_jax会和前面的numpyro采样共享GPU上下文，无需额外配置
• 可通过jax.device_count()和jax.devices()验证JAX的GPU支持是否正常

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