DeviceArray: Improve support for copy, deepcopy, and pickle

[jakevdp]

Currently deep-copying or pickling of jax DeviceArrays is implemented by forwarding arr.__reduce__ to the ndarray value, meaning that the copied/unpickled result is a normal numpy array, and calling copy.copy is a no-op.

This change implements custom __copy__, __deepcopy__, and __reduce__ methods in order to properly copy and/or serialize DeviceArray objects in both traced and untraced contexts.

For __reduce__ (i.e. pickle), deserialization is done on the default device – I'm not sure whether it would make sense to try to persist the device; in any case it's not straightforward because jaxlib.xla_extension.Device is not serializable. Further, because deserialization may happen in a different runtime than the one where serialization took place, it's not clear how the non-default device would be identified.

Addresses #2632

Note that pickling of bfloat16 arrays is still broken due to #8505

[jakevdp]

pulling in for more testing - not sure whether there will be unintended side-effects.

[jakevdp]

Turns out there are a number of downstream users implicitly relying on pickling and/or deepcopy converting device arrays to numpy arrays. Perhaps we can start with a deprecation warning.

[jakevdp]

I'm working on fixing internal breakages; probably will not be able to merge this for a week or two. But requesting review new because I think this will be close to the final state of the PR. Thanks!

[jakevdp]

Ready for a review: PTAL

[hawkinsp]

A note, not necessarily for this change: I don't consider anything in jaxlib to be a public API. We should have public jax. names for anything we need to refer to.

[jakevdp]

Yeah, unfortunately array methods are only defined on the xla extension version of the DeviceArray class, so we can only cross-reference them that way: https://jax.readthedocs.io/en/latest/jax.numpy.html#jax.numpy.DeviceArray

Maybe the pjit array unification thing will let us clean this up?