forked from data-apis/array-api-tests
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathpytest_helpers.py
More file actions
270 lines (231 loc) · 7.89 KB
/
Copy pathpytest_helpers.py
File metadata and controls
270 lines (231 loc) · 7.89 KB
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
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
import math
from inspect import getfullargspec
from typing import Any, Dict, Optional, Sequence, Tuple, Union
from . import _array_module as xp
from . import array_helpers as ah
from . import dtype_helpers as dh
from . import function_stubs
from . import shape_helpers as sh
from .typing import Array, DataType, Scalar, ScalarType, Shape
__all__ = [
"raises",
"doesnt_raise",
"nargs",
"fmt_kw",
"is_pos_zero",
"is_neg_zero",
"assert_dtype",
"assert_kw_dtype",
"assert_default_float",
"assert_default_int",
"assert_default_index",
"assert_shape",
"assert_result_shape",
"assert_keepdimable_shape",
"assert_0d_equals",
"assert_fill",
"assert_array",
]
def raises(exceptions, function, message=""):
"""
Like pytest.raises() except it allows custom error messages
"""
try:
function()
except exceptions:
return
except Exception as e:
if message:
raise AssertionError(
f"Unexpected exception {e!r} (expected {exceptions}): {message}"
)
raise AssertionError(f"Unexpected exception {e!r} (expected {exceptions})")
raise AssertionError(message)
def doesnt_raise(function, message=""):
"""
The inverse of raises().
Use doesnt_raise(function) to test that function() doesn't raise any
exceptions. Returns the result of calling function.
"""
if not callable(function):
raise ValueError("doesnt_raise should take a lambda")
try:
return function()
except Exception as e:
if message:
raise AssertionError(f"Unexpected exception {e!r}: {message}")
raise AssertionError(f"Unexpected exception {e!r}")
def nargs(func_name):
return len(getfullargspec(getattr(function_stubs, func_name)).args)
def fmt_kw(kw: Dict[str, Any]) -> str:
return ", ".join(f"{k}={v}" for k, v in kw.items())
def is_pos_zero(n: float) -> bool:
return n == 0 and math.copysign(1, n) == 1
def is_neg_zero(n: float) -> bool:
return n == 0 and math.copysign(1, n) == -1
def assert_dtype(
func_name: str,
in_dtype: Union[DataType, Sequence[DataType]],
out_dtype: DataType,
expected: Optional[DataType] = None,
*,
repr_name: str = "out.dtype",
):
in_dtypes = in_dtype if isinstance(in_dtype, Sequence) else [in_dtype]
f_in_dtypes = dh.fmt_types(tuple(in_dtypes))
f_out_dtype = dh.dtype_to_name[out_dtype]
if expected is None:
expected = dh.result_type(*in_dtypes)
f_expected = dh.dtype_to_name[expected]
msg = (
f"{repr_name}={f_out_dtype}, but should be {f_expected} "
f"[{func_name}({f_in_dtypes})]"
)
assert out_dtype == expected, msg
def assert_kw_dtype(func_name: str, kw_dtype: DataType, out_dtype: DataType):
f_kw_dtype = dh.dtype_to_name[kw_dtype]
f_out_dtype = dh.dtype_to_name[out_dtype]
msg = (
f"out.dtype={f_out_dtype}, but should be {f_kw_dtype} "
f"[{func_name}(dtype={f_kw_dtype})]"
)
assert out_dtype == kw_dtype, msg
def assert_default_float(func_name: str, dtype: DataType):
f_dtype = dh.dtype_to_name[dtype]
f_default = dh.dtype_to_name[dh.default_float]
msg = (
f"out.dtype={f_dtype}, should be default "
f"floating-point dtype {f_default} [{func_name}()]"
)
assert dtype == dh.default_float, msg
def assert_default_int(func_name: str, dtype: DataType):
f_dtype = dh.dtype_to_name[dtype]
f_default = dh.dtype_to_name[dh.default_int]
msg = (
f"out.dtype={f_dtype}, should be default "
f"integer dtype {f_default} [{func_name}()]"
)
assert dtype == dh.default_int, msg
def assert_default_index(func_name: str, dtype: DataType, repr_name="out.dtype"):
f_dtype = dh.dtype_to_name[dtype]
msg = (
f"{repr_name}={f_dtype}, should be the default index dtype, "
f"which is either int32 or int64 [{func_name}()]"
)
assert dtype in (xp.int32, xp.int64), msg
def assert_shape(
func_name: str,
out_shape: Union[int, Shape],
expected: Union[int, Shape],
/,
repr_name="out.shape",
**kw,
):
if isinstance(out_shape, int):
out_shape = (out_shape,)
if isinstance(expected, int):
expected = (expected,)
msg = (
f"{repr_name}={out_shape}, but should be {expected} [{func_name}({fmt_kw(kw)})]"
)
assert out_shape == expected, msg
def assert_result_shape(
func_name: str,
in_shapes: Sequence[Shape],
out_shape: Shape,
/,
expected: Optional[Shape] = None,
*,
repr_name="out.shape",
**kw,
):
if expected is None:
expected = sh.broadcast_shapes(*in_shapes)
f_in_shapes = " . ".join(str(s) for s in in_shapes)
f_sig = f" {f_in_shapes} "
if kw:
f_sig += f", {fmt_kw(kw)}"
msg = f"{repr_name}={out_shape}, but should be {expected} [{func_name}({f_sig})]"
assert out_shape == expected, msg
def assert_keepdimable_shape(
func_name: str,
out_shape: Shape,
in_shape: Shape,
axes: Tuple[int, ...],
keepdims: bool,
/,
**kw,
):
if keepdims:
shape = tuple(1 if axis in axes else side for axis, side in enumerate(in_shape))
else:
shape = tuple(side for axis, side in enumerate(in_shape) if axis not in axes)
assert_shape(func_name, out_shape, shape, **kw)
def assert_0d_equals(
func_name: str, x_repr: str, x_val: Array, out_repr: str, out_val: Array, **kw
):
msg = (
f"{out_repr}={out_val}, should be {x_repr}={x_val} "
f"[{func_name}({fmt_kw(kw)})]"
)
if dh.is_float_dtype(out_val.dtype) and xp.isnan(out_val):
assert xp.isnan(x_val), msg
else:
assert x_val == out_val, msg
def assert_scalar_equals(
func_name: str,
type_: ScalarType,
idx: Shape,
out: Scalar,
expected: Scalar,
/,
repr_name: str = "out",
**kw,
):
repr_name = repr_name if idx == () else f"{repr_name}[{idx}]"
f_func = f"{func_name}({fmt_kw(kw)})"
if type_ is bool or type_ is int:
msg = f"{repr_name}={out}, but should be {expected} [{f_func}]"
assert out == expected, msg
elif math.isnan(expected):
msg = f"{repr_name}={out}, but should be {expected} [{f_func}]"
assert math.isnan(out), msg
else:
msg = f"{repr_name}={out}, but should be roughly {expected} [{f_func}]"
assert math.isclose(out, expected, rel_tol=0.25, abs_tol=1), msg
def assert_fill(
func_name: str, fill_value: Scalar, dtype: DataType, out: Array, /, **kw
):
msg = f"out not filled with {fill_value} [{func_name}({fmt_kw(kw)})]\n{out=}"
if math.isnan(fill_value):
assert ah.all(ah.isnan(out)), msg
else:
assert ah.all(ah.equal(out, ah.asarray(fill_value, dtype=dtype))), msg
def assert_array(func_name: str, out: Array, expected: Array, /, **kw):
assert_dtype(func_name, out.dtype, expected.dtype)
assert_shape(func_name, out.shape, expected.shape, **kw)
f_func = f"[{func_name}({fmt_kw(kw)})]"
if dh.is_float_dtype(out.dtype):
for idx in sh.ndindex(out.shape):
at_out = out[idx]
at_expected = expected[idx]
msg = (
f"{sh.fmt_idx('out', idx)}={at_out}, should be {at_expected} "
f"{f_func}"
)
if xp.isnan(at_expected):
assert xp.isnan(at_out), msg
elif at_expected == 0.0 or at_expected == -0.0:
scalar_at_expected = float(at_expected)
scalar_at_out = float(at_out)
if is_pos_zero(scalar_at_expected):
assert is_pos_zero(scalar_at_out), msg
else:
assert is_neg_zero(scalar_at_expected) # sanity check
assert is_neg_zero(scalar_at_out), msg
else:
assert at_out == at_expected, msg
else:
assert xp.all(out == expected), (
f"out not as expected {f_func}\n" f"{out=}\n{expected=}"
)