Behaviour of pickling a JAX array

[ayaka14732]

I am wondering what the behaviour of pickling a JAX array is.

On 10 Nov 2021, @jakevdp commented in #8494 (reply in thread):

Regarding general use of pickle with JAX arrays: JAX delegates pickle to numpy. So any devicearray that you pickle will be saved (and re-loaded) as a standard numpy array, e.g.

import pickle
import jax.numpy as jnp
type(pickle.loads(pickle.dumps(jnp.arange(10))))
# numpy.ndarray

But it seems that the behaviour has been changed. In JAX v0.4.11, pickle returns a JAX array.

>>> import pickle
>>> import jax.numpy as jnp
>>> type(pickle.loads(pickle.dumps(jnp.arange(10))))
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
<class 'jaxlib.xla_extension.ArrayImpl'>

What I am trying to achieve is to save JAX arrays stored in nested namedtuples to disk. For example:

from jax import Array
import jax.numpy as jnp
from typing import NamedTuple

class Apple(NamedTuple):
    a: Array
    b: Array

class Banana(NamedTuple):
    a: Apple
    b: Array

a0 = jnp.asarray(np.random.rand(1, 5))
a1 = jnp.asarray(np.random.rand(2, 7))
a2 = jnp.asarray(np.random.rand(4, 8))

params = Banana(a=Apple(a=a0, b=a1), b=a2)  # save this to disk

My question:

1. Is pickle the intended tool to achieve this?
2. What is the behaviour of pickling a JAX array?

[jakevdp]

The intended behavior of pickle with JAX arrays is described here: https://jax.readthedocs.io/en/latest/jax.numpy.html#copying-and-serialization

In short, as of #10659 it is expected that pickling and unpickling a JAX array will return a JAX array.