如何将旧版Jax的DeviceArray转换为新版Jax兼容的Array？

方法1：在旧环境转成Numpy数组再导出 #

你已经有jax=0.3.25的环境了，加载.pkl后先把所有DeviceArray递归转成标准Numpy数组，再存成.npy，这样新版加载时会自动转为jax.Array：

import numpy as np
import pickle

# 加载旧pkl文件
with open('filename.pkl', 'rb') as f:
    file_content = pickle.load(f)

# 递归处理嵌套结构里的DeviceArray
def convert_to_numpy(obj):
    if hasattr(obj, '__array__'):  # 匹配DeviceArray类型
        return np.asarray(obj)
    elif isinstance(obj, dict):
        return {k: convert_to_numpy(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return type(obj)(convert_to_numpy(item) for item in obj)
    return obj

converted_content = convert_to_numpy(file_content)
np.save('filename_converted.npy', converted_content)

新版环境加载代码：

import jax.numpy as jnp
file_content = jnp.load('filename_converted.npy', allow_pickle=True)

方法2：新版环境加兼容补丁 #

不想重新导出的话，在v0.4.31环境里加载.npy前，先给pickle补个DeviceArray的映射类：

import sys
import jax
import jax.numpy as jnp

# 用新版jax.Array模拟旧版DeviceArray
class MockDeviceArray(jax.Array):
    pass

# 注册到模块路径，让pickle能识别该类
sys.modules['jax._src.device_array'] = type('jax._src.device_array', (), {'DeviceArray': MockDeviceArray})()

# 正常加载文件
file_content = jnp.load('filename.npy', allow_pickle=True)

方法3：旧环境导出兼容格式的Pickle #

如果后续还要用pickle格式，在旧环境里把DeviceArray转成主机数组（脱离设备）再存：

import jax
import pickle

with open('filename.pkl', 'rb') as f:
    file_content = pickle.load(f)

# 递归转成主机可序列化的数组
def convert_to_host(obj):
    if isinstance(obj, jax._src.device_array.DeviceArray):
        return jax.device_get(obj)  # 转成主机上的numpy数组
    elif isinstance(obj, dict):
        return {k: convert_to_host(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return type(obj)(convert_to_host(item) for item in obj)
    return obj

host_content = convert_to_host(file_content)
with open('filename_compatible.pkl', 'wb') as f:
    pickle.dump(host_content, f)

新版环境直接加载这个新pkl文件即可。

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