Windows纯CPU环境下BayesNF包安装使用问题及解决方案咨询

问题背景 #

已在Windows 10系统的Python 3.10新虚拟环境中，按项目指引完成bayesnf包的安装，但设备仅配备CPU（无GPU）。运行import bayesnf时触发JAX相关的循环导入错误，错误栈如下：

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[7], line 1
----> 1 import bayesnf

File ~\pyenv\lib\site-packages\bayesnf\__init__.py:21
     17 # A new PyPI release will be pushed every time __version__ is increased.
     18 # When changing this, also update the CHANGELOG.md
     19 __version__ = '0.1.3'
---> 21 from .spatiotemporal import BayesianNeuralFieldMAP
     22 from .spatiotemporal import BayesianNeuralFieldMLE
     23 from .spatiotemporal import BayesianNeuralFieldVI

File ~\pyenv\lib\site-packages\bayesnf\spatiotemporal.py:19
     15 """API for Bayesian Neural Field estimators."""
     17 from collections.abc import Sequence
---> 19 import jax
     20 import jax.numpy as jnp
     21 import numpy as np

File ~\pyenv\lib\site-packages\jax\__init__.py:25
     22 from jax.version import __version_info__ as __version_info__
     24 # Set Cloud TPU env vars if necessary before transitively loading C++ backend
---> 25 from jax._src.cloud_tpu_init import cloud_tpu_init as _cloud_tpu_init
     26 try:
     27   _cloud_tpu_init()

File ~\pyenv\lib\site-packages\jax\_src\cloud_tpu_init.py:17
     15 import os
     16 from jax import version
---> 17 from jax._src import config
     18 from jax._src import hardware_utils
     20 running_in_cloud_tpu_vm: bool = False

File ~\pyenv\lib\site-packages\jax\_src\config.py:27
     24 import threading
     25 from typing import Any, Generic, NamedTuple, NoReturn, Protocol, TypeVar, cast
---> 27 from jax._src import lib
     28 from jax._src.lib import jax_jit
     29 from jax._src.lib import transfer_guard_lib

File ~\pyenv\lib\site-packages\jax\_src\lib\__init__.py:76
     71   return _jaxlib_version
     74 version_str = jaxlib.version.__version__
     75 version = check_jaxlib_version(
---> 76   jax_version=jax.version.__version__,
     77   jaxlib_version=jaxlib.version.__version__,
     78   minimum_jaxlib_version=jax.version._minimum_jaxlib_version)
     80 # Before importing any C compiled modules from jaxlib, first import the CPU
     81 # feature guard module to verify that jaxlib was compiled in a way that only
     82 # uses instructions that are present on this machine.
     83 import jaxlib.cpu_feature_guard as cpu_feature_guard

AttributeError: partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)

项目GitHub页面相关说明：

该系统基于JAX机器学习平台构建。
典型安装时间为1分钟。本软件已在Python 3.10的标准Debian GNU/Linux环境中测试。scripts/中的大规模实验使用TPU v3-8加速器运行。要在本地处理中大规模数据，至少需要GPU。

解决方案 #

1. 修复JAX循环导入错误 #

当前错误是JAX版本兼容或安装不完整导致的循环导入，需重新安装适配Windows CPU的JAX版本：

pip uninstall -y jax jaxlib
pip install jax[cpu] jaxlib

此命令会安装仅支持CPU的JAX版本，避免GPU相关组件引发的兼容性问题。

2. Windows纯CPU环境可用性说明 #

• 小规模测试场景下，纯CPU环境可以运行bayesnf，但处理速度会远慢于GPU/TPU环境。
• 但JAX对Windows的原生支持不如Linux完善，后续可能遇到其他兼容性问题。

3. 备选优化方案 #

• 若CPU性能不足，可使用WSL2在Windows上搭建Linux虚拟环境，项目在Linux环境的兼容性更稳定。
• 或使用云平台的免费CPU/GPU实例运行大规模任务。

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