Merge branch 'master' into test_pow_with_scalars

Author: ev-br

array_api_tests/test_indexing_functions.py

@@ -60,3 +60,63 @@ def test_take(x, data):
     # 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=}'

array_api_tests/test_operators_and_elementwise_functions.py

@@ -1597,6 +1597,26 @@ def test_not_equal(ctx, data):
     )
 
 
+@pytest.mark.min_version("2024.12")
+@given(
+    shapes=hh.two_mutually_broadcastable_shapes,
+    dtype=hh.real_floating_dtypes,
+    data=st.data() 
+)
+def test_nextafter(shapes, dtype, data):
+    x1 = data.draw(hh.arrays(dtype=dtype, shape=shapes[0]), label="x1")
+    x2 = data.draw(hh.arrays(dtype=dtype, shape=shapes[0]), label="x2")
+
+    out = xp.nextafter(x1, x2)
+    _assert_correctness_binary(
+        "nextafter",
+        math.nextafter,
+        in_dtypes=[x1.dtype, x2.dtype],
+        in_shapes=[x1.shape, x2.shape],
+        in_arrs=[x1, x2],
+        out=out
+    )
+
 @pytest.mark.parametrize("ctx", make_unary_params("positive", dh.numeric_dtypes))
 @given(data=st.data())
 def test_positive(ctx, data):
@@ -1815,6 +1835,7 @@ def _filter_zero(x):
         ("divide", operator.truediv, {"filter_": lambda s: s != 0}, None),
         ("hypot", math.hypot, {}, None),
         ("logaddexp", logaddexp_refimpl, {}, None),
+        ("nextafter", math.nextafter, {}, None),
         ("maximum", max, {'strict_check': True}, None),
         ("minimum", min, {'strict_check': True}, None),
         ("multiply", operator.mul, {}, None),
@@ -1900,3 +1921,34 @@ def test_binary_with_scalars_bitwise_shifts(func_data, x1x2):
     refimpl_ = lambda l, r: mock_int_dtype(refimpl(l, r), xp.int32 )
     _check_binary_with_scalars((func_name, refimpl_, kwargs, expected), x1x2)
 
+
+@pytest.mark.unvectorized
+@given(
+    x1x2=hh.array_and_py_scalar([xp.int32]),
+    data=st.data()
+)
+def test_where_with_scalars(x1x2, data):
+    x1, x2 = x1x2
+
+    if dh.is_scalar(x1):
+        dtype, shape = x2.dtype, x2.shape
+        x1_arr, x2_arr = xp.broadcast_to(xp.asarray(x1), shape), x2
+    else:
+        dtype, shape = x1.dtype, x1.shape
+        x1_arr, x2_arr = x1, xp.broadcast_to(xp.asarray(x2), shape)
+
+    condition = data.draw(hh.arrays(shape=shape, dtype=xp.bool))
+
+    out = xp.where(condition, x1, x2)
+
+    assert out.dtype == dtype, f"where: got {out.dtype = } for {dtype=}, {x1=} and {x2=}"
+    assert out.shape == shape, f"where: got {out.shape = } for {shape=}, {x1=} and {x2=}"
+
+    # value test
+    for idx in sh.ndindex(shape):
+        if condition[idx]:
+            assert out[idx] == x1_arr[idx]
+        else:
+            assert out[idx] == x2_arr[idx]
+
+

array_api_tests/test_utility_functions.py

@@ -63,3 +63,79 @@ def test_any(x, data):
         expected = any(elements)
         ph.assert_scalar_equals("any", type_=scalar_type, idx=out_idx,
                                 out=result, expected=expected, kw=kw)
+
+
+@pytest.mark.unvectorized
+@pytest.mark.min_version("2024.12")
+@given(
+    x=hh.arrays(hh.numeric_dtypes, hh.shapes(min_dims=1, min_side=1)),
+    data=st.data(),
+)
+def test_diff(x, data):
+    axis = data.draw(
+        st.integers(-x.ndim, max(x.ndim - 1, 0)) | st.none(),
+        label="axis"
+    )
+    if axis is None:
+        axis_kw = {"axis": -1}
+        n_axis = x.ndim - 1
+    else:
+        axis_kw = {"axis": axis}
+        n_axis = axis + x.ndim if axis < 0 else axis
+
+    n = data.draw(st.integers(1, min(x.shape[n_axis], 3)))
+
+    out = xp.diff(x, **axis_kw, n=n)
+
+    expected_shape = list(x.shape)
+    expected_shape[n_axis] -= n
+
+    assert out.shape == tuple(expected_shape)
+
+    # value test
+    if n == 1:
+        for idx in sh.ndindex(out.shape):
+            l = list(idx)
+            l[n_axis] += 1
+            assert out[idx] == x[tuple(l)] - x[idx], f"diff failed with {idx = }"
+
+
+@pytest.mark.min_version("2024.12")
+@pytest.mark.unvectorized
+@given(
+    x=hh.arrays(hh.numeric_dtypes, hh.shapes(min_dims=1, min_side=1)),
+    data=st.data(),
+)
+def test_diff_append_prepend(x, data):
+    axis = data.draw(
+        st.integers(-x.ndim, max(x.ndim - 1, 0)) | st.none(),
+        label="axis"
+    )
+    if axis is None:
+        axis_kw = {"axis": -1}
+        n_axis = x.ndim - 1
+    else:
+        axis_kw = {"axis": axis}
+        n_axis = axis + x.ndim if axis < 0 else axis
+
+    n = data.draw(st.integers(1, min(x.shape[n_axis], 3)))
+
+    append_shape = list(x.shape)
+    append_axis_len = data.draw(st.integers(1, 2*append_shape[n_axis]), label="append_axis")
+    append_shape[n_axis] = append_axis_len
+    append = data.draw(hh.arrays(dtype=x.dtype, shape=tuple(append_shape)), label="append")
+
+    prepend_shape = list(x.shape)
+    prepend_axis_len = data.draw(st.integers(1, 2*prepend_shape[n_axis]), label="prepend_axis")
+    prepend_shape[n_axis] = prepend_axis_len
+    prepend = data.draw(hh.arrays(dtype=x.dtype, shape=tuple(prepend_shape)), label="prepend")
+
+    out = xp.diff(x, **axis_kw, n=n, append=append, prepend=prepend)
+
+    in_1 = xp.concat((prepend, x, append), **axis_kw)
+    out_1 = xp.diff(in_1, **axis_kw, n=n)
+
+    assert out.shape == out_1.shape
+    for idx in sh.ndindex(out.shape):
+        assert out[idx] == out_1[idx], f"{idx = }"
+