运行MERF时出现JAX相关AttributeError问题求助

核心错误定位 #

关键报错为：

AttributeError: 'jaxlib.xla_extension.DeviceList' object has no attribute 'split'

该错误出现在grid_utils.py第175行，代码尝试对变量y调用split方法，但y实际是DeviceList类型而非预期的JAX数组，导致方法调用失败。

可能原因 #

• JAX版本不兼容：MERF代码基于特定版本的JAX编写，新版本JAX的API或内部类型处理逻辑变化，导致变量类型异常。
• 设备上下文异常：尽管日志显示回退到CPU运行，但设备分配过程中可能出现错误，导致变量被错误绑定为DeviceList而非JAX数组。
• GPU环境缺失间接影响：日志中提示CUDA-enabled jaxlib未安装，虽然当前用CPU运行，但JAX的设备管理逻辑可能因此出现异常。

解决方案 #

1. 安装MERF官方指定的JAX版本 #

多数NeRF类项目会在README或requirements.txt中指定兼容的jax/jaxlib版本，找到对应版本后执行安装：

# 示例：安装0.4.13版本，需替换为MERF指定版本
pip install jax==0.4.13 jaxlib==0.4.13

如果需要GPU支持，根据本地CUDA版本安装对应cuda版jaxlib：

# 示例：CUDA 11.8 + cuDNN 8.6
pip install jaxlib==0.4.13+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

2. 调试变量来源 #

在grid_utils.py第175行前添加调试代码，确认y的类型和来源：

print(f"Type of y: {type(y)}")
print(f"Value of y: {y}")

运行后根据输出追溯上游代码，找到y被错误赋值为DeviceList的位置，修正设备分配逻辑。

3. 启用完整回溯信息 #

设置环境变量禁用JAX的简化回溯，获取更详细的错误链：

export JAX_TRACEBACK_FILTERING=off

重新运行代码，查看完整的调用栈，定位变量类型异常的根源。

4. 修复GPU依赖（可选） #

如果计划使用GPU，安装匹配版本的CUDA和cuDNN，再重新安装cuda版jaxlib，确保JAX能正确识别GPU设备，避免设备管理逻辑异常。

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