Skip to content

Commit d594ff5

Browse files
committed
Cover everything in test_unique_values()
1 parent c0a47fd commit d594ff5

File tree

1 file changed

+27
-4
lines changed

1 file changed

+27
-4
lines changed

array_api_tests/test_set_functions.py

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
1-
from hypothesis import given
1+
import math
2+
3+
from hypothesis import assume, given
24

35
from . import _array_module as xp
6+
from . import array_helpers as ah
7+
from . import dtype_helpers as dh
48
from . import hypothesis_helpers as hh
9+
from . import pytest_helpers as ph
510
from . import xps
611

712

@@ -23,7 +28,25 @@ def test_unique_inverse(x):
2328
# TODO
2429

2530

26-
@given(xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes()))
31+
@given(xps.arrays(dtype=xps.scalar_dtypes(), shape=hh.shapes(min_side=1)))
2732
def test_unique_values(x):
28-
xp.unique_values(x)
29-
# TODO
33+
out = xp.unique_values(x)
34+
ph.assert_dtype("unique_values", x.dtype, out.dtype)
35+
scalar_type = dh.get_scalar_type(x.dtype)
36+
distinct = set(scalar_type(x[idx]) for idx in ah.ndindex(x.shape))
37+
vals_idx = {}
38+
nans = 0
39+
for idx in ah.ndindex(out.shape):
40+
val = scalar_type(out[idx])
41+
if math.isnan(val):
42+
nans += 1
43+
else:
44+
assert val in distinct, f"out[{idx}]={val}, but {val} not in input array"
45+
assert (
46+
val not in vals_idx.keys()
47+
), f"out[{idx}]={val}, but {val} is also in out[{vals_idx[val]}]"
48+
vals_idx[val] = idx
49+
if dh.is_float_dtype(out.dtype):
50+
assume(x.size <= 128) # may not be representable
51+
expected_nans = xp.sum(xp.astype(xp.isnan(x), xp.uint8))
52+
assert nans == expected_nans, f"{nans} NaNs in out, expected {expected_nans}"

0 commit comments

Comments
 (0)