Support pickle and copy.deepcopy of dpnp.ndarray

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 * Added support for boolean input in `dpnp.positive` and the `+` operator of `dpnp.ndarray` [#3128](https://github.com/IntelPython/dpnp/pull/3128)
+* Added support for `pickle` and `copy.deepcopy` of `dpnp.ndarray`; an unpickled array is allocated on the default device with its original USM type and data type [#3142](https://github.com/IntelPython/dpnp/pull/3142)
 
 ### Changed
 

doc/reference/ndarray.rst

@@ -399,7 +399,7 @@ For standard library functions:
 
    ndarray.__copy__
    ndarray.__deepcopy__
-   .. ndarray.__reduce__
+   ndarray.__reduce__
    ndarray.__setstate__
 
 Basic customization:

dpnp/dpnp_array.py

@@ -95,6 +95,21 @@ def _get_unwrapped_index_key(key):
     return _unwrap_index_element(key)
 
 
+def _reconstruct_from_host(array_class, host_array, usm_type):
+    """
+    Rebuild an unpickled array from its host copy on the default device.
+
+    Pickles refer to this function, so keep its name and signature.
+
+    """
+
+    # explicit dtype raises instead of downcasting on devices without fp64
+    res = dpnp.asarray(host_array, dtype=host_array.dtype, usm_type=usm_type)
+    if array_class is not dpnp_array:
+        res = res.view(type=array_class)
+    return res
+
+
 # pylint: disable=too-many-public-methods
 class dpnp_array:
     """
@@ -259,7 +274,16 @@ def __copy__(self):
         """
         return self.copy(order="K")
 
-    # '__deepcopy__',
+    def __deepcopy__(self, memo, /):  # pylint: disable=unused-argument
+        """
+        Used if :func:`copy.deepcopy` is called on an array. Return a copy of
+        the array.
+
+        Equivalent to ``a.copy(order="K")``.
+
+        """
+        return self.copy(order="K")
+
     # '__dir__',
 
     def __divmod__(self, other, /):
@@ -549,7 +573,25 @@ def __rdivmod__(self, other, /):
         """Return ``divmod(value, self)``."""
         return dpnp.divmod(other, self)
 
-    # '__reduce__',
+    def __reduce__(self):
+        """
+        Used by :mod:`pickle` to serialize the array.
+
+        The data is copied to the host and pickled together with the array
+        type and the USM type. Unpickling allocates the array on the default
+        SYCL device with the original USM type, so it may end up on a
+        different device than the one it was pickled from.
+
+        Unpickling raises ``ValueError`` if the default device does not
+        support the data type, e.g. ``float64`` on a device without
+        double-precision support.
+
+        """
+        return (
+            _reconstruct_from_host,
+            (type(self), self.asnumpy(order="A"), self.usm_type),
+        )
+
     # '__reduce_ex__',
 
     def __repr__(self):

dpnp/tests/test_ndarray.py

@@ -1,3 +1,5 @@
+import copy
+import pickle
 import warnings
 
 import numpy
@@ -881,6 +883,90 @@ def test_error(self):
             ia.item()
 
 
+# pickle looks a class up by name, so it can't be local to a test
+class _PickleSubclass(dpnp.ndarray):
+    pass
+
+
+_layouts = pytest.mark.parametrize(
+    "view",
+    [lambda x: x, lambda x: x.T, lambda x: x[:, ::2]],
+    ids=["C", "F", "strided"],
+)
+
+
+class TestPickle:
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    @pytest.mark.parametrize(
+        "shape", [(), (0,), (5,), (2, 3)], ids=["()", "(0,)", "(5,)", "(2, 3)"]
+    )
+    def test_basic(self, dtype, shape):
+        a = generate_random_numpy_array(shape, dtype)
+        ia = dpnp.array(a)
+
+        expected = pickle.loads(pickle.dumps(a))
+        result = pickle.loads(pickle.dumps(ia))
+        assert_array_equal(result, expected)
+
+    @pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1))
+    def test_protocol(self, protocol):
+        a = numpy.arange(6, dtype="i4").reshape(2, 3)
+        ia = dpnp.array(a)
+
+        result = pickle.loads(pickle.dumps(ia, protocol=protocol))
+        assert_array_equal(result, a)
+
+    @_layouts
+    def test_layout(self, view):
+        a = numpy.arange(12, dtype="i4").reshape(3, 4)
+        ia = dpnp.array(a)
+
+        expected = pickle.loads(pickle.dumps(view(a)))
+        result = pickle.loads(pickle.dumps(view(ia)))
+        assert_array_equal(result, expected)
+        assert result.flags.c_contiguous == expected.flags.c_contiguous
+        assert result.flags.f_contiguous == expected.flags.f_contiguous
+
+    def test_subclass(self):
+        ia = dpnp.arange(3).view(type=_PickleSubclass)
+
+        result = pickle.loads(pickle.dumps(ia))
+        assert type(result) is _PickleSubclass
+        assert_array_equal(result, ia)
+
+    @pytest.mark.skipif(has_support_aspect64(), reason="device supports fp64")
+    def test_no_fp64(self):
+        # unpickling keeps the dtype instead of downcasting float64 data
+        func, args = dpnp.ones(3, dtype="f4").__reduce__()
+        args = (args[0], numpy.ones(3, dtype="f8"), args[2])
+        with pytest.raises(ValueError, match="double-precision"):
+            func(*args)
+
+
+class TestDeepCopy:
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    def test_basic(self, dtype):
+        a = generate_random_numpy_array((2, 3), dtype)
+        ia = dpnp.array(a)
+
+        result = copy.deepcopy(ia)
+        assert_array_equal(result, copy.deepcopy(a))
+
+        result[...] = 0
+        assert_array_equal(ia, a)
+
+    @_layouts
+    def test_layout(self, view):
+        a = numpy.arange(12, dtype="i4").reshape(3, 4)
+        ia = dpnp.array(a)
+
+        expected = copy.deepcopy(view(a))
+        result = copy.deepcopy(view(ia))
+        assert_array_equal(result, expected)
+        assert result.flags.c_contiguous == expected.flags.c_contiguous
+        assert result.flags.f_contiguous == expected.flags.f_contiguous
+
+
 class TestUsmNdarrayProtocol:
     def test_basic(self):
         a = dpnp.arange(256, dtype=dpnp.int64)

dpnp/tests/test_sycl_queue.py

@@ -1,4 +1,5 @@
 import copy
+import pickle
 import tempfile
 
 import dpctl
@@ -210,13 +211,24 @@ def test_copy_method(device_x, device_y):
     assert_sycl_queue_equal(y.sycl_queue, q)
 
 
+@pytest.mark.parametrize(
+    "func", [copy.copy, copy.deepcopy], ids=["copy", "deepcopy"]
+)
 @pytest.mark.parametrize("device", valid_dev, ids=dev_ids)
-def test_copy_operation(device):
+def test_copy_operation(func, device):
     x = dpnp.array([[1, 2, 3], [4, 5, 6]], device=device)
-    y = copy.copy(x)
+    y = func(x)
     assert_sycl_queue_equal(y.sycl_queue, x.sycl_queue)
 
 
+@pytest.mark.parametrize("device", valid_dev, ids=dev_ids)
+def test_pickle(device):
+    x = dpnp.array([[1, 2, 3], [4, 5, 6]], device=device)
+    y = pickle.loads(pickle.dumps(x))
+    # unpickled arrays go to the default device, not the original one
+    assert_sycl_queue_equal(y.sycl_queue, dpnp.empty(0).sycl_queue)
+
+
 @pytest.mark.parametrize("device", valid_dev, ids=dev_ids)
 def test_extract(device):
     x = dpnp.arange(3, device=device)

dpnp/tests/test_usm_type.py

@@ -1,4 +1,5 @@
 import copy
+import pickle
 import tempfile
 from math import prod
 
@@ -270,10 +271,15 @@ def test_copy_method(usm_type_x, usm_type_y):
     assert y.usm_type == usm_type_y
 
 
+@pytest.mark.parametrize(
+    "func",
+    [copy.copy, copy.deepcopy, lambda x: pickle.loads(pickle.dumps(x))],
+    ids=["copy", "deepcopy", "pickle"],
+)
 @pytest.mark.parametrize("usm_type", list_of_usm_types)
-def test_copy_operation(usm_type):
+def test_copy_operation(func, usm_type):
     x = dpnp.array([[1, 2, 3], [4, 5, 6]], usm_type=usm_type)
-    y = copy.copy(x)
+    y = func(x)
     assert x.usm_type == y.usm_type == usm_type
 
 

dpnp/tests/third_party/cupy/core_tests/test_ndarray.py

@@ -232,7 +232,6 @@ def test_excessive_ndim(self):
         }
     )
 )
-@pytest.mark.skip("deepcopy() is not supported")
 class TestNdarrayDeepCopy(unittest.TestCase):
 
     def _check_deepcopy(self, arr, arr2):
@@ -244,15 +243,15 @@ def _check_deepcopy(self, arr, arr2):
         testing.assert_array_equal(arr, arr2)
 
     def test_deepcopy(self):
-        arr = _core.ndarray(self.shape)
+        arr = cupy.ndarray(self.shape)
         arr2 = copy.deepcopy(arr)
         self._check_deepcopy(arr, arr2)
 
     @testing.multi_gpu(2)
     def test_deepcopy_multi_device(self):
-        arr = _core.ndarray(self.shape)
-        with cuda.Device(1):
-            arr2 = copy.deepcopy(arr)
+        arr = cupy.ndarray(self.shape)
+        # with cuda.Device(1):
+        arr2 = copy.deepcopy(arr)
         self._check_deepcopy(arr, arr2)
         assert arr2.device == arr.device
 

dpnp/tests/third_party/cupy/io_tests/test_npz.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import io
 import pickle
 import unittest
@@ -7,11 +9,10 @@
 import dpnp as cupy
 from dpnp.tests.third_party.cupy import testing
 
-pytest.skip("text functions are not supported yet", allow_module_level=True)
-
 
 class TestNpz(unittest.TestCase):
 
+    @pytest.mark.skip("save() and load() are not supported yet")
     @testing.for_all_dtypes()
     def test_save_load(self, dtype):
         a = testing.shaped_arange((2, 3, 4), dtype=dtype)
@@ -26,6 +27,7 @@ def test_save_load(self, dtype):
 
         testing.assert_array_equal(a, b)
 
+    @pytest.mark.skip("save() is not supported yet")
     def test_save_pickle(self):
         data = object()
 
@@ -38,6 +40,7 @@ def test_save_pickle(self):
         cupy.save(sio, data, allow_pickle=True)
         sio.close()
 
+    @pytest.mark.skip("dump() and load() are not supported yet")
     def test_load_pickle(self):
         a = testing.shaped_arange((2, 3, 4), dtype=cupy.float32)
 
@@ -75,19 +78,25 @@ def check_savez(self, savez, dtype):
         testing.assert_array_equal(a1, b1)
         testing.assert_array_equal(a2, b2)
 
+    @pytest.mark.skip("savez() and load() are not supported yet")
     def test_savez(self):
         self.check_savez(cupy.savez)
 
+    @pytest.mark.skip("savez_compressed() and load() are not supported yet")
     def test_savez_compressed(self):
         self.check_savez(cupy.savez_compressed)
 
     @testing.for_all_dtypes()
     def test_pickle(self, dtype):
-        a = testing.shaped_arange((2, 3, 4), dtype=dtype)
-        s = pickle.dumps(a)
-        b = pickle.loads(s)
-        testing.assert_array_equal(a, b)
-
+        for order in ["C", "F"]:
+            a = testing.shaped_arange((2, 3, 4), dtype=dtype, order=order)
+            s = pickle.dumps(a)
+            b = pickle.loads(s)
+            testing.assert_array_equal(a, b)
+            assert a.flags.c_contiguous == b.flags.c_contiguous
+            assert a.flags.f_contiguous == b.flags.f_contiguous
+
+    @pytest.mark.skip("dump() and load() are not supported yet")
     @testing.for_all_dtypes()
     def test_dump(self, dtype):
         a = testing.shaped_arange((2, 3, 4), dtype=dtype)