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
|
import numpy as np
from dlpack import from_numpy, to_numpy
def test_to_from_numpy_zero_copy():
"""Test the copy free conversion of numpy array via DLPack."""
np_ary = np.random.normal(size=[10, 10])
np_ary_big = np.random.normal(size=[12, 10])
dlpack_capsule = from_numpy(np_ary_big)
reconstructed_ary = to_numpy(dlpack_capsule)
del dlpack_capsule
np_ary_big[1:11] = np_ary
del np_ary_big
np.testing.assert_equal(actual=reconstructed_ary[1:11], desired=np_ary)
def test_to_from_numpy_memory():
"""Test that DLPack capsule keeps the source array alive"""
source_array = np.random.normal(size=[10, 10])
np_array_ref = source_array.copy()
dlpack_capsule = from_numpy(source_array)
del source_array
reconstructed_array = to_numpy(dlpack_capsule)
del dlpack_capsule
np.testing.assert_equal(actual=reconstructed_array, desired=np_array_ref)
if __name__ == "__main__":
"""
Run both tests a bunch of times to make
sure the conversions and memory management are stable.
"""
print("### Running `test_to_from_numpy_zero_copy`")
for i in range(10000):
test_to_from_numpy_zero_copy()
print("### Running `test_to_from_numpy_memory`")
for i in range(10000):
test_to_from_numpy_memory()
|
