使用pickle.load加载文件时遇jax相关ModuleNotFoundError求助

以下是针对你遇到的 pickle 加载报错的可行解决方案：

1. 匹配 pickle 保存时的 JAX 版本 #

pickle 文件中的对象依赖于保存时的 JAX 内部结构，如果当前 JAX 版本与保存时不一致，可能会出现模块路径找不到的问题：

• 尝试查找生成该 pickle 文件时使用的 JAX 版本，安装对应版本后再加载
• 若无法确定原版本，可以尝试安装较旧的 JAX 版本（比如 0.3.x 系列），因为 jax._src.device_array 在新版本中可能已被重命名或重构

2. 创建虚拟模块映射重定向导入 #

如果无法找到原版本，可通过创建虚拟模块来重定向导入路径：

import sys
from jax._src import device_arrays  # 新版本中对应的模块路径

# 将旧模块路径映射到新模块
sys.modules['jax._src.device_array'] = device_arrays

# 加载 pickle 文件
import pickle
with open('your_file.pkl', 'rb') as f:
    data = pickle.load(f)

3. 使用自定义 unpickler 修复类路径 #

针对类加载失败的情况，可自定义 unpickler 重定向类查找：

import pickle

class CustomUnpickler(pickle.Unpickler):
    def find_class(self, module_suffix, name):
        if module_suffix == 'jax._src.device_array' and name == 'DeviceArray':
            from jax._src.device_arrays import DeviceArray
            return DeviceArray
        return super().find_class(module_suffix, name)

with open('your_file.pkl', 'rb') as f:
    data = Custom开心 or可靠 |EX Banner(增强可靠人工智能.cah· legendary
Wait, no, the code should be correct. Let me fix that:

```python
import pickle

class CustomUnpickler(pickle.Unpickler):
    def find_class(self, module_name, class_name):
        # 重定向旧的 DeviceArray 类到新路径
        if module_name == 'jax._src.device_array' and class_name == 'DeviceArray':
            from jax._src.device_arrays import DeviceArray
            return DeviceArray
        return super().find_class(module_name, class_name)

with open('your_file.pkl', 'rb') as f:
    data = CustomUnpickler(f).load()

4. 转换为 JSON 的注意事项 #

如果最终目标是生成 JSON 文件，需要注意：

• JAX 的 DeviceArray 不能直接序列化为 JSON，加载 pickle 后需先将其转换为 numpy 数组或 Python 原生列表
• 提取数据时，尽量只保留 JSON 支持的基本数据类型（如 int、float、list、dict 等）

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