import os from functools import partial import numpy import pytest import scipy import scipy.sparse as sp from numpy.testing import assert_allclose from sklearn._config import config_context from sklearn._loss import HalfMultinomialLoss from sklearn.base import BaseEstimator from sklearn.utils._array_api import ( _add_to_diagonal, _asarray_with_order, _atol_for_type, _average, _convert_to_numpy, _count_nonzero, _estimator_with_converted_arrays, _fill_diagonal, _get_namespace_device_dtype_ids, _half_multinomial_loss, _is_numpy_namespace, _isin, _logsumexp, _max_precision_float_dtype, _median, _nanmax, _nanmean, _nanmin, _ravel, _validate_diagonal_args, device, get_namespace, get_namespace_and_device, indexing_dtype, move_to, np_compat, supported_float_dtypes, yield_namespace_device_dtype_combinations, ) from sklearn.utils._testing import ( SkipTest, _array_api_for_tests, _convert_container, assert_array_equal, skip_if_array_api_compat_not_configured, ) from sklearn.utils.fixes import _IS_32BIT, CSR_CONTAINERS, np_version, parse_version @pytest.mark.parametrize("X", [numpy.asarray([1, 2, 3]), [1, 2, 3]]) def test_get_namespace_ndarray_default(X): """Check that get_namespace returns NumPy wrapper""" xp_out, is_array_api_compliant = get_namespace(X) assert xp_out is np_compat assert not is_array_api_compliant def test_get_namespace_ndarray_creation_device(): """Check expected behavior with device and creation functions.""" X = numpy.asarray([1, 2, 3]) xp_out, _ = get_namespace(X) full_array = xp_out.full(10, fill_value=2.0, device="cpu") assert_allclose(full_array, [2.0] * 10) with pytest.raises(ValueError, match="Unsupported device"): xp_out.zeros(10, device="cuda") @skip_if_array_api_compat_not_configured def test_get_namespace_ndarray_with_dispatch(): """Test get_namespace on NumPy ndarrays.""" X_np = numpy.asarray([[1, 2, 3]]) with config_context(array_api_dispatch=True): xp_out, is_array_api_compliant = get_namespace(X_np) assert is_array_api_compliant # In the future, NumPy should become API compliant library and we should have # assert xp_out is numpy assert xp_out is np_compat @skip_if_array_api_compat_not_configured @pytest.mark.parametrize( "constructor_name", ["pyarrow", "dataframe", "polars", "series"] ) def test_get_namespace_df_with_dispatch(constructor_name): """Test get_namespace on dataframes and series.""" df = _convert_container([[1, 4, 2], [3, 3, 6]], constructor_name) with config_context(array_api_dispatch=True): xp_out, is_array_api_compliant = get_namespace(df) assert not is_array_api_compliant # When operating on dataframes or series the Numpy namespace is # the right thing to use. assert xp_out is np_compat @skip_if_array_api_compat_not_configured def test_get_namespace_sparse_with_dispatch(): """Test get_namespace on sparse arrays.""" with config_context(array_api_dispatch=True): xp_out, is_array_api_compliant = get_namespace(sp.csr_array([[1, 2, 3]])) assert not is_array_api_compliant # When operating on sparse arrays the Numpy namespace is # the right thing to use. assert xp_out is np_compat @skip_if_array_api_compat_not_configured def test_get_namespace_array_api(monkeypatch): """Test get_namespace for ArrayAPI arrays.""" xp = pytest.importorskip("array_api_strict") X_np = numpy.asarray([[1, 2, 3]]) X_xp = xp.asarray(X_np) with config_context(array_api_dispatch=True): xp_out, is_array_api_compliant = get_namespace(X_xp) assert is_array_api_compliant with pytest.raises(TypeError): xp_out, is_array_api_compliant = get_namespace(X_xp, X_np) def mock_getenv(key): if key == "SCIPY_ARRAY_API": return "0" monkeypatch.setattr("os.environ.get", mock_getenv) assert os.environ.get("SCIPY_ARRAY_API") != "1" with pytest.raises( RuntimeError, match="scipy's own support is not enabled.", ): get_namespace(X_xp) @pytest.mark.parametrize( "array_input, reference", [ pytest.param(("cupy", None), ("torch", "cuda"), id="cupy to torch cuda"), pytest.param(("torch", "mps"), ("numpy", None), id="torch mps to numpy"), pytest.param(("numpy", None), ("torch", "cuda"), id="numpy to torch cuda"), pytest.param(("numpy", None), ("torch", "mps"), id="numpy to torch mps"), pytest.param( ("array_api_strict", None), ("torch", "mps"), id="array_api_strict to torch mps", ), ], ) def test_move_to_array_api_conversions(array_input, reference): """Check conversion between various namespace and devices.""" if array_input[0] == "array_api_strict": array_api_strict = pytest.importorskip( "array_api_strict", reason="array-api-strict not available" ) xp = _array_api_for_tests(reference[0], reference[1]) xp_array = _array_api_for_tests(array_input[0], array_input[1]) with config_context(array_api_dispatch=True): device_ = device(xp.asarray([1], device=reference[1])) if array_input[0] == "array_api_strict": array_device = array_api_strict.Device("CPU_DEVICE") else: array_device = array_input[1] array = xp_array.asarray([1, 2, 3], device=array_device) array_out = move_to(array, xp=xp, device=device_) assert get_namespace(array_out)[0] == xp assert device(array_out) == device_ def test_move_to_sparse(): """Check sparse inputs are handled correctly.""" xp_numpy = _array_api_for_tests("numpy", None) xp_torch = _array_api_for_tests("torch", "cpu") sparse1 = sp.csr_array([0, 1, 2, 3]) sparse2 = sp.csr_array([0, 1, 0, 1]) numpy_array = numpy.array([1, 2, 3]) with config_context(array_api_dispatch=True): device_cpu = xp_torch.asarray([1]).device # sparse and None to NumPy result1, result2 = move_to(sparse1, None, xp=xp_numpy, device=None) assert result1 is sparse1 assert result2 is None # sparse to non-NumPy msg = r"Sparse arrays are only accepted \(and passed through\)" with pytest.raises(TypeError, match=msg): move_to(sparse1, numpy_array, xp=xp_torch, device=device_cpu) with pytest.raises(TypeError, match=msg): move_to(sparse1, None, xp=xp_torch, device=device_cpu) @pytest.mark.parametrize("array_api", ["numpy", "array_api_strict"]) def test_asarray_with_order(array_api): """Test _asarray_with_order passes along order for NumPy arrays.""" xp = pytest.importorskip(array_api) X = xp.asarray([1.2, 3.4, 5.1]) X_new = _asarray_with_order(X, order="F", xp=xp) X_new_np = numpy.asarray(X_new) assert X_new_np.flags["F_CONTIGUOUS"] @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) @pytest.mark.parametrize( "weights, axis, normalize, expected", [ # normalize = True (None, None, True, 3.5), (None, 0, True, [2.5, 3.5, 4.5]), (None, 1, True, [2, 5]), ([True, False], 0, True, [1, 2, 3]), # boolean weights ([True, True, False], 1, True, [1.5, 4.5]), # boolean weights ([0.4, 0.1], 0, True, [1.6, 2.6, 3.6]), ([0.4, 0.2, 0.2], 1, True, [1.75, 4.75]), ([1, 2], 0, True, [3, 4, 5]), ([1, 1, 2], 1, True, [2.25, 5.25]), ([[1, 2, 3], [1, 2, 3]], 0, True, [2.5, 3.5, 4.5]), ([[1, 2, 1], [2, 2, 2]], 1, True, [2, 5]), # normalize = False (None, None, False, 21), (None, 0, False, [5, 7, 9]), (None, 1, False, [6, 15]), ([True, False], 0, False, [1, 2, 3]), # boolean weights ([True, True, False], 1, False, [3, 9]), # boolean weights ([0.4, 0.1], 0, False, [0.8, 1.3, 1.8]), ([0.4, 0.2, 0.2], 1, False, [1.4, 3.8]), ([1, 2], 0, False, [9, 12, 15]), ([1, 1, 2], 1, False, [9, 21]), ([[1, 2, 3], [1, 2, 3]], 0, False, [5, 14, 27]), ([[1, 2, 1], [2, 2, 2]], 1, False, [8, 30]), ], ) def test_average( array_namespace, device_, dtype_name, weights, axis, normalize, expected ): xp = _array_api_for_tests(array_namespace, device_) 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 device(array_in) == device(result) result = _convert_to_numpy(result, xp) assert_allclose(result, expected, atol=_atol_for_type(dtype_name)) @pytest.mark.parametrize( "array_namespace, device, dtype_name", yield_namespace_device_dtype_combinations(include_numpy_namespaces=False), ids=_get_namespace_device_dtype_ids, ) def test_average_raises_with_wrong_dtype(array_namespace, device, dtype_name): xp = _array_api_for_tests(array_namespace, device) array_in = numpy.asarray([2, 0], dtype=dtype_name) + 1j * numpy.asarray( [4, 3], dtype=dtype_name ) complex_type_name = array_in.dtype.name if not hasattr(xp, complex_type_name): # This is the case for cupy as of March 2024 for instance. pytest.skip(f"{array_namespace} does not support {complex_type_name}") array_in = xp.asarray(array_in, device=device) err_msg = "Complex floating point values are not supported by average." with ( config_context(array_api_dispatch=True), pytest.raises(NotImplementedError, match=err_msg), ): _average(array_in) @pytest.mark.parametrize( "array_namespace, device, dtype_name", yield_namespace_device_dtype_combinations(include_numpy_namespaces=True), ids=_get_namespace_device_dtype_ids, ) @pytest.mark.parametrize( "axis, weights, error, error_msg", ( ( None, [1, 2], TypeError, "Axis must be specified", ), ( 0, [[1, 2]], # NumPy 2 raises ValueError, NumPy 1 raises TypeError (ValueError, TypeError), "weights", # the message is different for NumPy 1 and 2... ), ( 0, [1, 2, 3, 4], ValueError, "weights", ), (0, [-1, 1], ZeroDivisionError, "Weights sum to zero, can't be normalized"), ), ) def test_average_raises_with_invalid_parameters( array_namespace, device, dtype_name, axis, weights, error, error_msg ): xp = _array_api_for_tests(array_namespace, device) array_in = numpy.asarray([[1, 2, 3], [4, 5, 6]], dtype=dtype_name) array_in = xp.asarray(array_in, device=device) weights = numpy.asarray(weights, dtype=dtype_name) weights = xp.asarray(weights, device=device) with config_context(array_api_dispatch=True), pytest.raises(error, match=error_msg): _average(array_in, axis=axis, weights=weights) def test_device_none_if_no_input(): assert device() is None assert device(None, "name") is None @skip_if_array_api_compat_not_configured def test_device_inspection(): class Device: def __init__(self, name): self.name = name def __eq__(self, device): return self.name == device.name def __hash__(self): raise TypeError("Device object is not hashable") def __str__(self): return self.name class Array: def __init__(self, device_name): self.device = Device(device_name) # Sanity check: ensure our Device mock class is non hashable, to # accurately account for non-hashable device objects in some array # libraries, because of which the `device` inspection function shouldn't # make use of hash lookup tables (in particular, not use `set`) with pytest.raises(TypeError): hash(Array("device").device) # If array API dispatch is disabled the device should be ignored. Erroring # early for different devices would prevent the np.asarray conversion to # happen. For example, `r2_score(np.ones(5), torch.ones(5))` should work # fine with array API disabled. assert device(Array("cpu"), Array("mygpu")) is None # Test that ValueError is raised if on different devices and array API dispatch is # enabled. err_msg = "Input arrays use different devices: cpu, mygpu" with config_context(array_api_dispatch=True): with pytest.raises(ValueError, match=err_msg): device(Array("cpu"), Array("mygpu")) # Test expected value is returned otherwise array1 = Array("device") array2 = Array("device") assert array1.device == device(array1) assert array1.device == device(array1, array2) assert array1.device == device(array1, array1, array2) # TODO: add cupy to the list of libraries once the following upstream issue # has been fixed: # https://github.com/cupy/cupy/issues/8180 @skip_if_array_api_compat_not_configured @pytest.mark.parametrize("library", ["numpy", "array_api_strict", "torch"]) @pytest.mark.parametrize( "X,reduction,expected", [ ([1, 2, numpy.nan], _nanmin, 1), ([1, -2, -numpy.nan], _nanmin, -2), ([numpy.inf, numpy.inf], _nanmin, numpy.inf), ( [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]], partial(_nanmin, axis=0), [1.0, 2.0, 3.0], ), ( [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]], partial(_nanmin, axis=1), [1.0, numpy.nan, 4.0], ), ([1, 2, numpy.nan], _nanmax, 2), ([1, 2, numpy.nan], _nanmax, 2), ([-numpy.inf, -numpy.inf], _nanmax, -numpy.inf), ( [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]], partial(_nanmax, axis=0), [4.0, 5.0, 6.0], ), ( [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]], partial(_nanmax, axis=1), [3.0, numpy.nan, 6.0], ), ([1, 2, numpy.nan], _nanmean, 1.5), ([1, -2, -numpy.nan], _nanmean, -0.5), ([-numpy.inf, -numpy.inf], _nanmean, -numpy.inf), ( [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]], partial(_nanmean, axis=0), [2.5, 3.5, 4.5], ), ( [[1, 2, 3], [numpy.nan, numpy.nan, numpy.nan], [4, 5, 6.0]], partial(_nanmean, axis=1), [2.0, numpy.nan, 5.0], ), ], ) def test_nan_reductions(library, X, reduction, expected): """Check NaN reductions like _nanmin and _nanmax""" xp = pytest.importorskip(library) with config_context(array_api_dispatch=True): result = reduction(xp.asarray(X)) result = _convert_to_numpy(result, xp) assert_allclose(result, expected) @pytest.mark.parametrize( "namespace, _device, _dtype", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) def test_ravel(namespace, _device, _dtype): xp = _array_api_for_tests(namespace, _device) array = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] array_xp = xp.asarray(array, device=_device) with config_context(array_api_dispatch=True): result = _ravel(array_xp) result = _convert_to_numpy(result, xp) expected = numpy.ravel(array, order="C") assert_allclose(expected, result) if _is_numpy_namespace(xp): assert numpy.asarray(result).flags["C_CONTIGUOUS"] @skip_if_array_api_compat_not_configured @pytest.mark.parametrize("library", ["cupy", "torch"]) def test_convert_to_numpy_gpu(library): # pragma: nocover """Check convert_to_numpy for GPU backed libraries.""" xp = pytest.importorskip(library) if library == "torch": if not xp.backends.cuda.is_built(): pytest.skip("test requires cuda") X_gpu = xp.asarray([1.0, 2.0, 3.0], device="cuda") else: X_gpu = xp.asarray([1.0, 2.0, 3.0]) X_cpu = _convert_to_numpy(X_gpu, xp=xp) expected_output = numpy.asarray([1.0, 2.0, 3.0]) assert_allclose(X_cpu, expected_output) def test_convert_to_numpy_cpu(): """Check convert_to_numpy for PyTorch CPU arrays.""" torch = pytest.importorskip("torch") X_torch = torch.asarray([1.0, 2.0, 3.0], device="cpu") X_cpu = _convert_to_numpy(X_torch, xp=torch) expected_output = numpy.asarray([1.0, 2.0, 3.0]) assert_allclose(X_cpu, expected_output) class SimpleEstimator(BaseEstimator): def fit(self, X, y=None): self.X_ = X self.n_features_ = X.shape[0] return self @skip_if_array_api_compat_not_configured @pytest.mark.parametrize( "array_namespace, converter", [ ("torch", lambda array: array.cpu().numpy()), ("array_api_strict", lambda array: numpy.asarray(array)), ("cupy", lambda array: array.get()), ], ) def test_convert_estimator_to_ndarray(array_namespace, converter): """Convert estimator attributes to ndarray.""" xp = pytest.importorskip(array_namespace) X = xp.asarray([[1.3, 4.5]]) est = SimpleEstimator().fit(X) new_est = _estimator_with_converted_arrays(est, converter) assert isinstance(new_est.X_, numpy.ndarray) @skip_if_array_api_compat_not_configured def test_convert_estimator_to_array_api(): """Convert estimator attributes to ArrayAPI arrays.""" xp = pytest.importorskip("array_api_strict") X_np = numpy.asarray([[1.3, 4.5]]) est = SimpleEstimator().fit(X_np) new_est = _estimator_with_converted_arrays(est, lambda array: xp.asarray(array)) assert hasattr(new_est.X_, "__array_namespace__") @pytest.mark.parametrize( "namespace, _device, _dtype", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) def test_indexing_dtype(namespace, _device, _dtype): xp = _array_api_for_tests(namespace, _device) if _IS_32BIT: assert indexing_dtype(xp) == xp.int32 else: assert indexing_dtype(xp) == xp.int64 @pytest.mark.parametrize( "namespace, _device, _dtype", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) def test_max_precision_float_dtype(namespace, _device, _dtype): xp = _array_api_for_tests(namespace, _device) expected_dtype = xp.float32 if _device == "mps" else xp.float64 assert _max_precision_float_dtype(xp, _device) == expected_dtype @pytest.mark.parametrize( "array_namespace, device, _", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) @pytest.mark.parametrize("invert", [True, False]) @pytest.mark.parametrize("assume_unique", [True, False]) @pytest.mark.parametrize("element_size", [6, 10, 14]) @pytest.mark.parametrize("int_dtype", ["int16", "int32", "int64", "uint8"]) def test_isin( array_namespace, device, _, invert, assume_unique, element_size, int_dtype ): xp = _array_api_for_tests(array_namespace, device) r = element_size // 2 element = 2 * numpy.arange(element_size).reshape((r, 2)).astype(int_dtype) test_elements = numpy.array(numpy.arange(14), dtype=int_dtype) element_xp = xp.asarray(element, device=device) test_elements_xp = xp.asarray(test_elements, device=device) expected = numpy.isin( element=element, test_elements=test_elements, assume_unique=assume_unique, invert=invert, ) with config_context(array_api_dispatch=True): result = _isin( element=element_xp, test_elements=test_elements_xp, xp=xp, assume_unique=assume_unique, invert=invert, ) assert_array_equal(_convert_to_numpy(result, xp=xp), expected) @pytest.mark.skipif( os.environ.get("SCIPY_ARRAY_API") != "1", reason="SCIPY_ARRAY_API not set to 1." ) def test_get_namespace_and_device(): # Use torch as a library with custom Device objects: torch = pytest.importorskip("torch") from sklearn.externals.array_api_compat import torch as torch_compat some_torch_tensor = torch.arange(3, device="cpu") some_numpy_array = numpy.arange(3) # When dispatch is disabled, get_namespace_and_device should return the # default NumPy wrapper namespace and "cpu" device. Our code will handle such # inputs via the usual __array__ interface without attempting to dispatch # via the array API. namespace, is_array_api, device = get_namespace_and_device(some_torch_tensor) assert namespace is get_namespace(some_numpy_array)[0] assert not is_array_api assert device is None # Otherwise, expose the torch namespace and device via array API compat # wrapper. with config_context(array_api_dispatch=True): namespace, is_array_api, device = get_namespace_and_device(some_torch_tensor) assert namespace is torch_compat assert is_array_api assert device == some_torch_tensor.device @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) @pytest.mark.parametrize("csr_container", CSR_CONTAINERS) @pytest.mark.parametrize("axis", [0, 1, None, -1, -2]) @pytest.mark.parametrize("sample_weight_type", [None, "int", "float"]) def test_count_nonzero( array_namespace, device_, dtype_name, csr_container, axis, sample_weight_type ): from sklearn.utils.sparsefuncs import count_nonzero as sparse_count_nonzero xp = _array_api_for_tests(array_namespace, device_) array = numpy.array([[0, 3, 0], [2, -1, 0], [0, 0, 0], [9, 8, 7], [4, 0, 5]]) if sample_weight_type == "int": sample_weight = numpy.asarray([1, 2, 2, 3, 1]) elif sample_weight_type == "float": sample_weight = numpy.asarray([0.5, 1.5, 0.8, 3.2, 2.4], dtype=dtype_name) else: sample_weight = None expected = sparse_count_nonzero( csr_container(array), axis=axis, sample_weight=sample_weight ) array_xp = xp.asarray(array, device=device_) with config_context(array_api_dispatch=True): result = _count_nonzero( array_xp, axis=axis, sample_weight=sample_weight, xp=xp, device=device_ ) assert_allclose(_convert_to_numpy(result, xp=xp), expected) 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 device(array_xp) == device(result) @pytest.mark.parametrize( "array, value, match", [ (numpy.array([1, 2, 3]), 1, "`array` should be 2D"), (numpy.array([[1, 2], [3, 4]]), numpy.array([1, 2, 3]), "`value` needs to be"), (numpy.array([[1, 2], [3, 4]]), [1, 2, 3], "`value` needs to be"), ( numpy.array([[1, 2], [3, 4]]), numpy.array([[1, 2], [3, 4]]), "`value` needs to be a", ), ], ) def test_validate_diagonal_args(array, value, match): """Check `_validate_diagonal_args` raises the correct errors.""" xp = _array_api_for_tests("numpy", None) with pytest.raises(ValueError, match=match): _validate_diagonal_args(array, value, xp) @pytest.mark.parametrize("function", ["fill", "add"]) @pytest.mark.parametrize("c_contiguity", [True, False]) def test_fill_and_add_to_diagonal(c_contiguity, function): """Check `_fill/add_to_diagonal` behaviour correct with numpy arrays.""" xp = _array_api_for_tests("numpy", None) if c_contiguity: array = numpy.zeros((3, 4)) else: array = numpy.zeros((3, 4)).T assert array.flags["C_CONTIGUOUS"] == c_contiguity if function == "fill": func = _fill_diagonal else: func = _add_to_diagonal func(array, 1, xp) assert_allclose(array.diagonal(), numpy.ones((3,))) func(array, [0, 1, 2], xp) if function == "fill": expected_diag = numpy.arange(3) else: expected_diag = numpy.ones((3,)) + numpy.arange(3) assert_allclose(array.diagonal(), expected_diag) fill_array = numpy.array([11, 12, 13]) func(array, fill_array, xp) if function == "fill": expected_diag = fill_array else: expected_diag = fill_array + numpy.arange(3) + numpy.ones((3,)) assert_allclose(array.diagonal(), expected_diag) @pytest.mark.parametrize("array", ["standard", "transposed", "non-contiguous"]) @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) def test_fill_diagonal(array, array_namespace, device_, dtype_name): """Check array API `_fill_diagonal` consistent with `numpy._fill_diagonal`.""" xp = _array_api_for_tests(array_namespace, device_) array_np = numpy.zeros((4, 5), dtype=dtype_name) if array == "transposed": array_xp = xp.asarray(array_np.copy(), device=device_).T array_np = array_np.T elif array == "non-contiguous": array_xp = xp.asarray(array_np.copy(), device=device_)[::2, ::2] array_np = array_np[::2, ::2] else: array_xp = xp.asarray(array_np.copy(), device=device_) numpy.fill_diagonal(array_np, val=1) with config_context(array_api_dispatch=True): _fill_diagonal(array_xp, value=1, xp=xp) assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) def test_add_to_diagonal(array_namespace, device_, dtype_name): """Check `_add_to_diagonal` consistent between array API xp and numpy namespace.""" xp = _array_api_for_tests(array_namespace, device_) np_xp = _array_api_for_tests("numpy", None) array_np = numpy.zeros((3, 4), dtype=dtype_name) array_xp = xp.asarray(array_np.copy(), device=device_) add_val = [1, 2, 3] _fill_diagonal(array_np, value=add_val, xp=np_xp) with config_context(array_api_dispatch=True): _fill_diagonal(array_xp, value=add_val, xp=xp) assert_array_equal(_convert_to_numpy(array_xp, xp=xp), array_np) @pytest.mark.parametrize("csr_container", CSR_CONTAINERS) @pytest.mark.parametrize("dispatch", [True, False]) def test_sparse_device(csr_container, dispatch): np_arr = numpy.array([1]) # For numpy < 2, the device attribute is not available on numpy arrays expected_numpy_array_device = getattr(np_arr, "device", None) if dispatch else None a, b = csr_container(numpy.array([[1]])), csr_container(numpy.array([[2]])) if dispatch and os.environ.get("SCIPY_ARRAY_API") is None: raise SkipTest("SCIPY_ARRAY_API is not set: not checking array_api input") with config_context(array_api_dispatch=dispatch): assert device(a, b) is None assert device(a, np_arr) == expected_numpy_array_device assert get_namespace_and_device(a, b)[2] is None assert get_namespace_and_device(a, np_arr)[2] == expected_numpy_array_device @pytest.mark.parametrize( "namespace, device, dtype_name", yield_namespace_device_dtype_combinations(), ids=_get_namespace_device_dtype_ids, ) @pytest.mark.parametrize("axis", [None, 0, 1]) def test_median(namespace, device, dtype_name, axis): # Note: depending on the value of `axis`, this test will compare median # computations on arrays of even (4) or odd (5) numbers of elements, hence # will test for median computation with and without interpolation to check # that array API namespaces yield consistent results even when the median is # not mathematically uniquely defined. xp = _array_api_for_tests(namespace, device) rng = numpy.random.RandomState(0) X_np = rng.uniform(low=0.0, high=1.0, size=(5, 4)).astype(dtype_name) result_np = numpy.median(X_np, axis=axis) X_xp = xp.asarray(X_np, device=device) with config_context(array_api_dispatch=True): result_xp = _median(X_xp, axis=axis) if xp.__name__ != "array_api_strict": # We convert array-api-strict arrays to numpy arrays as `median` is not # part of the Array API spec assert get_namespace(result_xp)[0] == xp assert result_xp.device == X_xp.device assert_allclose(result_np, _convert_to_numpy(result_xp, xp=xp)) @pytest.mark.parametrize( "array_namespace, device_, dtype_name", yield_namespace_device_dtype_combinations() ) @pytest.mark.parametrize("axis", [0, 1, None]) def test_logsumexp_like_scipy_logsumexp(array_namespace, device_, dtype_name, axis): xp = _array_api_for_tests(array_namespace, device_) array_np = numpy.asarray( [ [0, 3, 1000], [2, -1, 1000], [-10, 0, 0], [-50, 8, -numpy.inf], [4, 0, 5], ], dtype=dtype_name, ) array_xp = xp.asarray(array_np, device=device_) res_np = scipy.special.logsumexp(array_np, axis=axis) rtol = 1e-6 if "float32" in str(dtype_name) else 1e-12 # if torch on CPU or array api strict on default device # check that _logsumexp works when array API dispatch is disabled if (array_namespace == "torch" and device_ == "cpu") or ( array_namespace == "array_api_strict" and "CPU" in str(device_) ): assert_allclose(_logsumexp(array_xp, axis=axis), res_np, rtol=rtol) with config_context(array_api_dispatch=True): res_xp = _logsumexp(array_xp, axis=axis) res_xp = _convert_to_numpy(res_xp, xp) assert_allclose(res_np, res_xp, rtol=rtol) # Test with NaNs and +np.inf array_np_2 = numpy.asarray( [ [0, numpy.nan, 1000], [2, -1, 1000], [numpy.inf, 0, 0], [-50, 8, -numpy.inf], [4, 0, 5], ], dtype=dtype_name, ) array_xp_2 = xp.asarray(array_np_2, device=device_) res_np_2 = scipy.special.logsumexp(array_np_2, axis=axis) with config_context(array_api_dispatch=True): res_xp_2 = _logsumexp(array_xp_2, axis=axis) res_xp_2 = _convert_to_numpy(res_xp_2, xp) assert_allclose(res_np_2, res_xp_2, rtol=rtol) @pytest.mark.parametrize( ("namespace", "device_", "expected_types"), [ ("numpy", None, ("float64", "float32", "float16")), ("array_api_strict", None, ("float64", "float32")), ("torch", "cpu", ("float64", "float32", "float16")), ("torch", "cuda", ("float64", "float32", "float16")), ("torch", "mps", ("float32", "float16")), ], ) def test_supported_float_types(namespace, device_, expected_types): xp = _array_api_for_tests(namespace, device_) float_types = supported_float_dtypes(xp, device=device_) expected = tuple(getattr(xp, dtype_name) for dtype_name in expected_types) assert float_types == expected @pytest.mark.parametrize("use_sample_weight", [False, True]) @pytest.mark.parametrize( "namespace, device_, dtype_name", yield_namespace_device_dtype_combinations() ) def test_half_multinomial_loss(use_sample_weight, namespace, device_, dtype_name): """Check that the array API version of :func:`_half_multinomial_loss` works correctly and matches the results produced by :class:`HalfMultinomialLoss` of the private `_loss` module. """ n_samples = 5 n_classes = 3 rng = numpy.random.RandomState(42) y = rng.randint(0, n_classes, n_samples).astype(dtype_name) pred = rng.rand(n_samples, n_classes).astype(dtype_name) xp = _array_api_for_tests(namespace, device_) y_xp = xp.asarray(y, device=device_) pred_xp = xp.asarray(pred, device=device_) if use_sample_weight: sample_weight = numpy.ones_like(y) sample_weight[1::2] = 2 sample_weight_xp = xp.asarray(sample_weight, device=device_) else: sample_weight, sample_weight_xp = None, None np_loss = HalfMultinomialLoss(n_classes=n_classes)( y_true=y, raw_prediction=pred, sample_weight=sample_weight ) with config_context(array_api_dispatch=True): xp_loss = _half_multinomial_loss( y=y_xp, pred=pred_xp, sample_weight=sample_weight_xp, xp=xp ) assert numpy.isclose(np_loss, xp_loss)