deepcopy of jax.numpy array generates regular numpy array

While updating some existing code that used copy.deepcopy I got what was a rather unexpected result namely that deepcopy returns a regular numpy array when I passed in a jax.numpy array. Is this expected behaviour? If so perhaps a note in the docs to indicated this would be good.

The following is a snippet that replicates the problem.

for jax.version.__version__ == '0.1.62''

import jax.numpy as np

import copy

array_in = np.array([1,2,3], dtype= np.int32) array_out = copy.deepcopy(array_in)

print(“type in:” + str(type(array_in))+ “type out:” +str(type(array_out)))

give type in:<class ‘jax.interpreters.xla.DeviceArray’>type out:<class ‘numpy.ndarray’>

[shoyer]

There are at least two surprising behaviors here that should be fixed:

• Pickling/unpickling a JAX DeviceArray objects should return another DeviceArray. To do so, we’ll need to define our own methods for pickling instead of using NumPy’s __reduce__ directly.
• Deep-copying JAX arrays should possibly copy device buffers, but certainly should not do so via serializing to the host with pickle. Here we also need to define __deepcopy__ instead of relying on Python’s default behavior. (There is some ambiguousness about whether we need to copy device buffers with deepcopy, given that JAX arrays themselves are immutable, but I think doing so would be within the spirit of deep copy.)

@tawe141 I’m not sure what’s going on with example. I believe deepcopy() on JAX tracers works just fine (preserving gradients) – the surprising implementation of copy only arises for arrays that are actually loaded into memory (DeviceArray).

[jakevdp]

Fixed in #10659