泡菜更改类型JAX

Question

我有一个包含jax numpy数组的亚麻结构数据类。

当我腌制此对象并再次加载它时，数组不再是jax numpy数组，并且转换为numpy数组，这是重现它的代码：

import flax
import jax.numpy as jnp
import pickle

@flax.struct.dataclass
class A:
    data: jnp.ndarray

a = A(data=jnp.zeros((2,2)))
print(a, type(a.data))



with open('file.pickle', 'wb') as handle:
    pickle.dump(a, handle)
    
with open('file.pickle', 'rb') as handle:
    loaded_a = pickle.load(handle)

print(loaded_a, type(loaded_a.data))

我不想要这种行为，我希望它能保持其原始类型，有可能吗？

Answer 1

更新：此错误已在。从JAX（v。0.3.14）的下一个版本

---

开始在jax; https://github.com/google/google/jax/jax/issues/2632

请参阅 图书馆开发人员认为是一种不幸的行为，但尚未确定修复程序。如果您有兴趣，您可能会在这个问题上权衡。