1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
from __future__ import annotations
import numpy as np
import pytest
from packaging.version import parse as parse_version
import awkward as ak
if parse_version(np.__version__) < parse_version("1.23.0"):
pytest.skip(
"NumPy 1.23 or greater is required for DLPack testing", allow_module_level=True
)
def test_from_dlpack_numpy():
np_array = np.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5)
array = ak.from_dlpack(np_array)
np_from_ak = ak.to_numpy(array)
assert np.shares_memory(np_array, np_from_ak)
class DLPackOf:
def __init__(self, array):
self._array = array
def __dlpack__(self, stream=None):
if stream is None:
return self._array.__dlpack__()
else:
return self._array.__dlpack__(stream)
def __dlpack_device__(self):
return self._array.__dlpack_device__()
def test_to_layout():
np_array = np.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5)
dlpack_array = DLPackOf(np_array)
layout = ak.to_layout(dlpack_array)
assert layout.is_numpy
np_from_ak = ak.to_numpy(layout)
assert np.shares_memory(np_array, np_from_ak)
def test_invalid_argument():
with pytest.raises(
TypeError, match=r"Expected an object that implements the DLPack protocol"
):
ak.from_dlpack([1, 2, 3])
|
