diff --git a/CHANGELOG.md b/CHANGELOG.md index b4771ac5ae86..2c03ac9d8040 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/doc/reference/ndarray.rst b/doc/reference/ndarray.rst index 046648665f93..3be973566a75 100644 --- a/doc/reference/ndarray.rst +++ b/doc/reference/ndarray.rst @@ -399,7 +399,7 @@ For standard library functions: ndarray.__copy__ ndarray.__deepcopy__ - .. ndarray.__reduce__ + ndarray.__reduce__ ndarray.__setstate__ Basic customization: diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py index c7f9fe035943..ae0ad6c8e42f 100644 --- a/dpnp/dpnp_array.py +++ b/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): diff --git a/dpnp/tests/test_ndarray.py b/dpnp/tests/test_ndarray.py index 94331d34269f..5ebe61632c7c 100644 --- a/dpnp/tests/test_ndarray.py +++ b/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) diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py index e0f882f2175d..3faae4e660b1 100644 --- a/dpnp/tests/test_sycl_queue.py +++ b/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) diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py index 567078c3eff2..3d374583f6ff 100644 --- a/dpnp/tests/test_usm_type.py +++ b/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 diff --git a/dpnp/tests/third_party/cupy/core_tests/test_ndarray.py b/dpnp/tests/third_party/cupy/core_tests/test_ndarray.py index 65ff56e71f30..2c2b8f82189b 100644 --- a/dpnp/tests/third_party/cupy/core_tests/test_ndarray.py +++ b/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 diff --git a/dpnp/tests/third_party/cupy/io_tests/test_npz.py b/dpnp/tests/third_party/cupy/io_tests/test_npz.py index 7d42fd24d947..7fff2ad93f9a 100644 --- a/dpnp/tests/third_party/cupy/io_tests/test_npz.py +++ b/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)