array-api-tests/array_api_tests/test_indexing_functions.py at 91dc500f076edb5b99b253a9778ddaa60faaa9d9 · asmeurer/array-api-tests · GitHub

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
import pytest
from hypothesis import given, note
from hypothesis import strategies as st

from . import _array_module as xp
from . import dtype_helpers as dh
from . import hypothesis_helpers as hh
from . import pytest_helpers as ph
from . import shape_helpers as sh


@pytest.mark.unvectorized
@pytest.mark.min_version("2022.12")
@given(
    x=hh.arrays(hh.all_dtypes, hh.shapes(min_dims=1, min_side=1)),
    data=st.data(),
)
def test_take(x, data):
    # TODO:
    # * negative axis
    # * negative indices
    # * different dtypes for indices

    # axis is optional but only if x.ndim == 1
    _axis_st = st.integers(0, max(x.ndim - 1, 0))
    if x.ndim == 1:
        kw = data.draw(hh.kwargs(axis=_axis_st))
    else:
        kw = {"axis": data.draw(_axis_st)}
    axis = kw.get("axis", 0)
    _indices = data.draw(
        st.lists(st.integers(0, x.shape[axis] - 1), min_size=1, unique=True),
        label="_indices",
    )
    indices = xp.asarray(_indices, dtype=dh.default_int)
    note(f"{indices=}")

    out = xp.take(x, indices, **kw)

    ph.assert_dtype("take", in_dtype=x.dtype, out_dtype=out.dtype)
    ph.assert_shape(
        "take",
        out_shape=out.shape,
        expected=x.shape[:axis] + (len(_indices),) + x.shape[axis + 1 :],
        kw=dict(
            x=x,
            indices=indices,
            axis=axis,
        ),
    )
    out_indices = sh.ndindex(out.shape)
    axis_indices = list(sh.axis_ndindex(x.shape, axis))
    for axis_idx in axis_indices:
        f_axis_idx = sh.fmt_idx("x", axis_idx)
        for i in _indices:
            f_take_idx = sh.fmt_idx(f_axis_idx, i)
            indexed_x = x[axis_idx][i, ...]
            for at_idx in sh.ndindex(indexed_x.shape):
                out_idx = next(out_indices)
                ph.assert_0d_equals(
                    "take",
                    x_repr=sh.fmt_idx(f_take_idx, at_idx),
                    x_val=indexed_x[at_idx],
                    out_repr=sh.fmt_idx("out", out_idx),
                    out_val=out[out_idx],
                )
    # sanity check
    with pytest.raises(StopIteration):
        next(out_indices)


@pytest.mark.unvectorized
@pytest.mark.min_version("2024.12")
@given(
    x=hh.arrays(hh.all_dtypes, hh.shapes(min_dims=1, min_side=1)),
    data=st.data(),
)
def test_take_along_axis(x, data):
    # TODO
    # 2. negative indices
    # 3. different dtypes for indices
    # 4. "broadcast-compatible" indices
    axis = data.draw(
        st.integers(-x.ndim, max(x.ndim - 1, 0)) | st.none(),
        label="axis"
    )
    if axis is None:
        axis_kw = {}
        n_axis = x.ndim - 1
    else:
        axis_kw = {"axis": axis}
        n_axis = axis + x.ndim if axis < 0 else axis

    new_len = data.draw(st.integers(0, 2*x.shape[n_axis]), label="new_len")
    idx_shape = x.shape[:n_axis] + (new_len,) + x.shape[n_axis+1:]
    indices = data.draw(
        hh.arrays(
            shape=idx_shape,
            dtype=dh.default_int,
            elements={"min_value": 0, "max_value": x.shape[n_axis]-1}
        ),
        label="indices"
    )
    note(f"{indices=}  {idx_shape=}")

    out = xp.take_along_axis(x, indices, **axis_kw)

    ph.assert_dtype("take_along_axis", in_dtype=x.dtype, out_dtype=out.dtype)
    ph.assert_shape(
        "take_along_axis",
        out_shape=out.shape,
        expected=x.shape[:n_axis] + (new_len,) + x.shape[n_axis+1:],
        kw=dict(
            x=x,
            indices=indices,
            axis=axis,
        ),
    )

    # value test: notation is from `np.take_along_axis` docstring
    Ni, Nk = x.shape[:n_axis], x.shape[n_axis+1:]
    for ii in sh.ndindex(Ni):
        for kk in sh.ndindex(Nk):
            a_1d = x[ii + (slice(None),) + kk]
            i_1d = indices[ii + (slice(None),) + kk]
            o_1d = out[ii + (slice(None),) + kk]
            for j in range(new_len):
                assert o_1d[j] == a_1d[i_1d[j]], f'{ii=}, {kk=}, {j=}'