Function test_average

sklearn/utils/tests/test_array_api.py:240–259

(
    array_namespace, device_name, dtype_name, weights, axis, normalize, expected
)

    ],
)
def test_average(
    array_namespace, device_name, dtype_name, weights, axis, normalize, expected
):
    xp, device = _array_api_for_tests(array_namespace, device_name, dtype_name)
    array_in = numpy.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype_name)
    array_in = xp.asarray(array_in, device=device)
    if weights is not None:
        weights = numpy.asarray(weights, dtype=dtype_name)
        weights = xp.asarray(weights, device=device)

    with config_context(array_api_dispatch=True):
        result = _average(array_in, axis=axis, weights=weights, normalize=normalize)

        if np_version < parse_version("2.0.0") or np_version >= parse_version("2.1.0"):
            # NumPy 2.0 has a problem with the device attribute of scalar arrays:
            # https://github.com/numpy/numpy/issues/26850
            assert array_api_device(array_in) == array_api_device(result)

    result = move_to(result, xp=numpy, device="cpu")
    assert_allclose(result, expected, atol=_atol_for_type(dtype_name))


@pytest.mark.parametrize(