test_set_functions.py

# TODO: disable if opted out, refactor things
importcmath
importmath
fromcollectionsimportCounter, defaultdict

importpytest
fromhypothesisimportassume, given

from . import_array_moduleasxp
from . importdtype_helpersasdh
from . importhypothesis_helpersashh
from . importpytest_helpersasph
from . importshape_helpersassh

pytestmark= [pytest.mark.data_dependent_shapes, pytest.mark.unvectorized]


@given(hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_side=1)))
deftest_unique_all(x):
out=xp.unique_all(x)

asserthasattr(out, "values")
asserthasattr(out, "indices")
asserthasattr(out, "inverse_indices")
asserthasattr(out, "counts")

ph.assert_dtype(
"unique_all", in_dtype=x.dtype, out_dtype=out.values.dtype, repr_name="out.values.dtype"
 )
ph.assert_default_index(
"unique_all", out.indices.dtype, repr_name="out.indices.dtype"
 )
ph.assert_default_index(
"unique_all", out.inverse_indices.dtype, repr_name="out.inverse_indices.dtype"
 )
ph.assert_default_index(
"unique_all", out.counts.dtype, repr_name="out.counts.dtype"
 )

assert (
out.indices.shape==out.values.shape
 ), f"{out.indices.shape=}, but should be {out.values.shape=}"
ph.assert_shape(
"unique_all",
out_shape=out.inverse_indices.shape,
expected=x.shape,
repr_name="out.inverse_indices.shape",
 )
assert (
out.counts.shape==out.values.shape
 ), f"{out.counts.shape=}, but should be {out.values.shape=}"

scalar_type=dh.get_scalar_type(out.values.dtype)
counts=defaultdict(int)
firsts= {}
fori, idxinenumerate(sh.ndindex(x.shape)):
val=scalar_type(x[idx])
ifcounts[val] ==0:
firsts[val] =i
counts[val] +=1

foridxinsh.ndindex(out.indices.shape):
val=scalar_type(out.values[idx])
ifcmath.isnan(val):
break
i=int(out.indices[idx])
expected=firsts[val]
asserti==expected, (
f"out.values[{idx}]={val} and out.indices[{idx}]={i}, "
f"but first occurence of {val} is at {expected}"
 )

foridxinsh.ndindex(out.inverse_indices.shape):
ridx=int(out.inverse_indices[idx])
val=out.values[ridx]
expected=x[idx]
msg= (
f"out.inverse_indices[{idx}]={ridx} results in out.values[{ridx}]={val}, "
f"but should result in x[{idx}]={expected}"
 )
ifdh.is_float_dtype(out.values.dtype) andxp.isnan(expected):
assertxp.isnan(val), msg
else:
assertval==expected, msg

vals_idx= {}
nans=0
foridxinsh.ndindex(out.values.shape):
val=scalar_type(out.values[idx])
count=int(out.counts[idx])
ifcmath.isnan(val):
nans+=1
assertcount==1, (
f"out.counts[{idx}]={count} for out.values[{idx}]={val}, "
"but count should be 1 as NaNs are distinct"
 )
else:
expected=counts[val]
assert (
expected>0
 ), f"out.values[{idx}]={val}, but {val} not in input array"
count=int(out.counts[idx])
assertcount==expected, (
f"out.counts[{idx}]={count} for out.values[{idx}]={val}, "
f"but should be {expected}"
 )
assert (
valnotinvals_idx.keys()
 ), f"out[{idx}]={val}, but {val} is also in out[{vals_idx[val]}]"
vals_idx[val] =idx

ifdh.is_float_dtype(out.values.dtype):
assume(math.prod(x.shape) <=128) # may not be representable
expected=sum(vfork, vincounts.items() ifcmath.isnan(k))
assertnans==expected, f"{nans} NaNs in out, but should be {expected}"


@given(hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_side=1)))
deftest_unique_counts(x):
out=xp.unique_counts(x)
asserthasattr(out, "values")
asserthasattr(out, "counts")
ph.assert_dtype(
"unique_counts", in_dtype=x.dtype, out_dtype=out.values.dtype, repr_name="out.values.dtype"
 )
ph.assert_default_index(
"unique_counts", out.counts.dtype, repr_name="out.counts.dtype"
 )
assert (
out.counts.shape==out.values.shape
 ), f"{out.counts.shape=}, but should be {out.values.shape=}"
scalar_type=dh.get_scalar_type(out.values.dtype)
counts=Counter(scalar_type(x[idx]) foridxinsh.ndindex(x.shape))
vals_idx= {}
nans=0
foridxinsh.ndindex(out.values.shape):
val=scalar_type(out.values[idx])
count=int(out.counts[idx])
ifcmath.isnan(val):
nans+=1
assertcount==1, (
f"out.counts[{idx}]={count} for out.values[{idx}]={val}, "
"but count should be 1 as NaNs are distinct"
 )
else:
expected=counts[val]
assert (
expected>0
 ), f"out.values[{idx}]={val}, but {val} not in input array"
count=int(out.counts[idx])
assertcount==expected, (
f"out.counts[{idx}]={count} for out.values[{idx}]={val}, "
f"but should be {expected}"
 )
assert (
valnotinvals_idx.keys()
 ), f"out[{idx}]={val}, but {val} is also in out[{vals_idx[val]}]"
vals_idx[val] =idx
ifdh.is_float_dtype(out.values.dtype):
assume(math.prod(x.shape) <=128) # may not be representable
expected=sum(vfork, vincounts.items() ifcmath.isnan(k))
assertnans==expected, f"{nans} NaNs in out, but should be {expected}"


@given(hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_side=1)))
deftest_unique_inverse(x):
out=xp.unique_inverse(x)
asserthasattr(out, "values")
asserthasattr(out, "inverse_indices")
ph.assert_dtype(
"unique_inverse", in_dtype=x.dtype, out_dtype=out.values.dtype, repr_name="out.values.dtype"
 )
ph.assert_default_index(
"unique_inverse",
out.inverse_indices.dtype,
repr_name="out.inverse_indices.dtype",
 )
ph.assert_shape(
"unique_inverse",
out_shape=out.inverse_indices.shape,
expected=x.shape,
repr_name="out.inverse_indices.shape",
 )
scalar_type=dh.get_scalar_type(out.values.dtype)
distinct=set(scalar_type(x[idx]) foridxinsh.ndindex(x.shape))
vals_idx= {}
nans=0
foridxinsh.ndindex(out.values.shape):
val=scalar_type(out.values[idx])
ifcmath.isnan(val):
nans+=1
else:
assert (
valindistinct
 ), f"out.values[{idx}]={val}, but {val} not in input array"
assert (
valnotinvals_idx.keys()
 ), f"out.values[{idx}]={val}, but {val} is also in out[{vals_idx[val]}]"
vals_idx[val] =idx
foridxinsh.ndindex(out.inverse_indices.shape):
ridx=int(out.inverse_indices[idx])
val=out.values[ridx]
expected=x[idx]
msg= (
f"out.inverse_indices[{idx}]={ridx} results in out.values[{ridx}]={val}, "
f"but should result in x[{idx}]={expected}"
 )
ifdh.is_float_dtype(out.values.dtype) andxp.isnan(expected):
assertxp.isnan(val), msg
else:
assertval==expected, msg
ifdh.is_float_dtype(out.values.dtype):
assume(math.prod(x.shape) <=128) # may not be representable
expected=xp.sum(xp.astype(xp.isnan(x), xp.uint8))
assertnans==expected, f"{nans} NaNs in out.values, but should be {expected}"


@given(hh.arrays(dtype=hh.all_dtypes, shape=hh.shapes(min_side=1)))
deftest_unique_values(x):
out=xp.unique_values(x)
ph.assert_dtype("unique_values", in_dtype=x.dtype, out_dtype=out.dtype)
scalar_type=dh.get_scalar_type(x.dtype)
distinct=set(scalar_type(x[idx]) foridxinsh.ndindex(x.shape))
vals_idx= {}
nans=0
foridxinsh.ndindex(out.shape):
val=scalar_type(out[idx])
ifcmath.isnan(val):
nans+=1
else:
assertvalindistinct, f"out[{idx}]={val}, but {val} not in input array"
assert (
valnotinvals_idx.keys()
 ), f"out[{idx}]={val}, but {val} is also in out[{vals_idx[val]}]"
vals_idx[val] =idx
ifdh.is_float_dtype(out.dtype):
assume(math.prod(x.shape) <=128) # may not be representable
expected=xp.sum(xp.astype(xp.isnan(x), xp.uint8))
assertnans==expected, f"{nans} NaNs in out, but should be {expected}"