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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
|
from __future__ import annotations
from typing import TYPE_CHECKING, Literal
import numpy as np
import pytest
import zarr
from zarr.abc.buffer import ArrayLike, BufferPrototype, NDArrayLike
from zarr.buffer import cpu, gpu
from zarr.codecs.blosc import BloscCodec
from zarr.codecs.crc32c_ import Crc32cCodec
from zarr.codecs.gzip import GzipCodec
from zarr.codecs.transpose import TransposeCodec
from zarr.codecs.zstd import ZstdCodec
from zarr.errors import ZarrUserWarning
from zarr.storage import MemoryStore, StorePath
from zarr.testing.buffer import (
NDBufferUsingTestNDArrayLike,
StoreExpectingTestBuffer,
TestBuffer,
TestNDArrayLike,
)
from zarr.testing.utils import gpu_mark, gpu_test, skip_if_no_gpu
if TYPE_CHECKING:
import types
try:
import cupy as cp
except ImportError:
cp = None
import zarr.api.asynchronous
if TYPE_CHECKING:
import types
def test_nd_array_like(xp: types.ModuleType) -> None:
ary = xp.arange(10)
assert isinstance(ary, ArrayLike)
assert isinstance(ary, NDArrayLike)
@pytest.mark.asyncio
async def test_async_array_prototype() -> None:
"""Test the use of a custom buffer prototype"""
expect = np.zeros((9, 9), dtype="uint16", order="F")
a = await zarr.api.asynchronous.create_array(
StorePath(StoreExpectingTestBuffer()) / "test_async_array_prototype",
shape=expect.shape,
chunks=(5, 5),
dtype=expect.dtype,
fill_value=0,
)
expect[1:4, 3:6] = np.ones((3, 3))
my_prototype = BufferPrototype(buffer=TestBuffer, nd_buffer=NDBufferUsingTestNDArrayLike)
await a.setitem(
selection=(slice(1, 4), slice(3, 6)),
value=np.ones((3, 3)),
prototype=my_prototype,
)
got = await a.getitem(selection=(slice(0, 9), slice(0, 9)), prototype=my_prototype)
# ignoring a mypy error here that TestNDArrayLike doesn't meet the NDArrayLike protocol
# The test passes, so it clearly does.
assert isinstance(got, TestNDArrayLike)
assert np.array_equal(expect, got) # type: ignore[unreachable]
@gpu_test
@pytest.mark.asyncio
async def test_async_array_gpu_prototype() -> None:
"""Test the use of the GPU buffer prototype"""
expect = cp.zeros((9, 9), dtype="uint16", order="F")
a = await zarr.api.asynchronous.create_array(
StorePath(MemoryStore()) / "test_async_array_gpu_prototype",
shape=expect.shape,
chunks=(5, 5),
dtype=expect.dtype,
fill_value=0,
)
expect[1:4, 3:6] = cp.ones((3, 3))
await a.setitem(
selection=(slice(1, 4), slice(3, 6)),
value=cp.ones((3, 3)),
prototype=gpu.buffer_prototype,
)
got = await a.getitem(selection=(slice(0, 9), slice(0, 9)), prototype=gpu.buffer_prototype)
assert isinstance(got, cp.ndarray)
assert cp.array_equal(expect, got)
@pytest.mark.asyncio
async def test_codecs_use_of_prototype() -> None:
expect = np.zeros((10, 10), dtype="uint16", order="F")
a = await zarr.api.asynchronous.create_array(
StorePath(StoreExpectingTestBuffer()) / "test_codecs_use_of_prototype",
shape=expect.shape,
chunks=(5, 5),
dtype=expect.dtype,
fill_value=0,
compressors=[BloscCodec(), Crc32cCodec(), GzipCodec(), ZstdCodec()],
filters=[TransposeCodec(order=(1, 0))],
)
expect[:] = np.arange(100).reshape(10, 10)
my_prototype = BufferPrototype(buffer=TestBuffer, nd_buffer=NDBufferUsingTestNDArrayLike)
await a.setitem(
selection=(slice(0, 10), slice(0, 10)),
value=expect[:],
prototype=my_prototype,
)
got = await a.getitem(selection=(slice(0, 10), slice(0, 10)), prototype=my_prototype)
# ignoring a mypy error here that TestNDArrayLike doesn't meet the NDArrayLike protocol
# The test passes, so it clearly does.
assert isinstance(got, TestNDArrayLike)
assert np.array_equal(expect, got) # type: ignore[unreachable]
@gpu_test
@pytest.mark.asyncio
async def test_codecs_use_of_gpu_prototype() -> None:
expect = cp.zeros((10, 10), dtype="uint16", order="F")
a = await zarr.api.asynchronous.create_array(
StorePath(MemoryStore()) / "test_codecs_use_of_gpu_prototype",
shape=expect.shape,
chunks=(5, 5),
dtype=expect.dtype,
fill_value=0,
compressors=[BloscCodec(), Crc32cCodec(), GzipCodec(), ZstdCodec()],
filters=[TransposeCodec(order=(1, 0))],
)
expect[:] = cp.arange(100).reshape(10, 10)
msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
with pytest.warns(ZarrUserWarning, match=msg):
await a.setitem(
selection=(slice(0, 10), slice(0, 10)),
value=expect[:],
prototype=gpu.buffer_prototype,
)
with pytest.warns(ZarrUserWarning, match=msg):
got = await a.getitem(
selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype
)
assert isinstance(got, cp.ndarray)
assert cp.array_equal(expect, got)
@gpu_test
@pytest.mark.asyncio
async def test_sharding_use_of_gpu_prototype() -> None:
with zarr.config.enable_gpu():
expect = cp.zeros((10, 10), dtype="uint16", order="F")
a = await zarr.api.asynchronous.create_array(
StorePath(MemoryStore()) / "test_codecs_use_of_gpu_prototype",
shape=expect.shape,
chunks=(5, 5),
shards=(10, 10),
dtype=expect.dtype,
fill_value=0,
)
expect[:] = cp.arange(100).reshape(10, 10)
msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path"
with pytest.warns(ZarrUserWarning, match=msg):
await a.setitem(
selection=(slice(0, 10), slice(0, 10)),
value=expect[:],
prototype=gpu.buffer_prototype,
)
with pytest.warns(ZarrUserWarning, match=msg):
got = await a.getitem(
selection=(slice(0, 10), slice(0, 10)), prototype=gpu.buffer_prototype
)
assert isinstance(got, cp.ndarray)
assert cp.array_equal(expect, got)
def test_numpy_buffer_prototype() -> None:
buffer = cpu.buffer_prototype.buffer.create_zero_length()
ndbuffer = cpu.buffer_prototype.nd_buffer.create(shape=(1, 2), dtype=np.dtype("int64"))
assert isinstance(buffer.as_array_like(), np.ndarray)
assert isinstance(ndbuffer.as_ndarray_like(), np.ndarray)
with pytest.raises(ValueError, match="Buffer does not contain a single scalar value"):
ndbuffer.as_scalar()
@gpu_test
def test_gpu_buffer_prototype() -> None:
buffer = gpu.buffer_prototype.buffer.create_zero_length()
ndbuffer = gpu.buffer_prototype.nd_buffer.create(shape=(1, 2), dtype=cp.dtype("int64"))
assert isinstance(buffer.as_array_like(), cp.ndarray)
assert isinstance(ndbuffer.as_ndarray_like(), cp.ndarray)
with pytest.raises(ValueError, match="Buffer does not contain a single scalar value"):
ndbuffer.as_scalar()
# TODO: the same test for other buffer classes
def test_cpu_buffer_as_scalar() -> None:
buf = cpu.buffer_prototype.nd_buffer.create(shape=(), dtype="int64")
assert buf.as_scalar() == buf.as_ndarray_like()[()] # type: ignore[index]
@pytest.mark.parametrize(
"prototype",
[
cpu.buffer_prototype,
pytest.param(
gpu.buffer_prototype,
marks=[gpu_mark, skip_if_no_gpu],
),
BufferPrototype(
buffer=cpu.Buffer,
nd_buffer=NDBufferUsingTestNDArrayLike,
),
],
)
@pytest.mark.parametrize(
"shape",
[
(1, 2),
(1, 2, 3),
],
)
@pytest.mark.parametrize("dtype", ["int32", "float64"])
@pytest.mark.parametrize("order", ["C", "F"])
def test_empty(
prototype: BufferPrototype, shape: tuple[int, ...], dtype: str, order: Literal["C", "F"]
) -> None:
buf = prototype.nd_buffer.empty(shape=shape, dtype=dtype, order=order)
result = buf.as_ndarray_like()
assert result.shape == shape
assert result.dtype == dtype
if order == "C":
assert result.flags.c_contiguous # type: ignore[attr-defined]
else:
assert result.flags.f_contiguous # type: ignore[attr-defined]
|
