Repository: scikit-learn/scikit-learn
File: sklearn/utils/_array_api.py (Python)

Code before:

    Returns
    -------
    supported_dtypes : tuple
        Tuple of real floating data types supported by the provided array namespace,
        ordered from the highest precision to lowest.

    See Also
    --------
    max_precision_float_dtype : Maximum float dtype for a namespace/device pair.

    Notes
    -----
    `float16` is not officially part of the Array API spec at the
    time of writing but scikit-learn estimators and functions can choose
    to accept it when xp.float16 is defined.

    Additionally, some devices available within a namespace may not support
    all floating-point types that the namespace provides.

    https://data-apis.org/array-api/latest/API_specification/data_types.html
    """
    # TODO: Update to use `__array_namespace__info__()` from array-api v2023.12
    #       when/if that becomes more widespread.
    if xp.__name__ in {"array_api_compat.torch", "torch"} and device == "mps":
        # N.B. Yanked from pull/27232
        dtypes = (xp.float32,)
    else:
        dtypes = (xp.float64, xp.float32)

    if hasattr(xp, "float16"):
        return (*dtypes, xp.float16)

    return dtypes


def ensure_common_namespace_device(reference, *arrays):
    """Ensure that all arrays use the same namespace and device as reference.

    If necessary the arrays are moved to the same namespace and device as
    the reference array.

    Parameters
    ----------
    reference : array
        Reference array.

    *arrays : array
        Arrays to check.

    Returns

Review comment (security):
I am not sure if we still want to make `supported_float_dtypes` device sensitive: the new `_max_precision_float_dtype` in `main` does not delegate to `supported_float_dtypes` so we do not really have a case for this change anymore.

Code after:
    https://data-apis.org/array-api/latest/API_specification/data_types.html
    """
    dtypes_dict = xp.__array_namespace_info__().dtypes(
        kind="real floating", device=device
    )
    valid_float_dtypes = []
    for dtype_key in ("float64", "float32"):
        if dtype_key in dtypes_dict:
            valid_float_dtypes.append(dtypes_dict[dtype_key])

    if hasattr(xp, "float16"):
        valid_float_dtypes.append(xp.float16)

    return tuple(valid_float_dtypes)


def ensure_common_namespace_device(reference, *arrays):
    """Ensure that all arrays use the same namespace and device as reference.

    If necessary the arrays are moved to the same namespace and device as
    the reference array.

    Parameters
    ----------
    reference : array
        Reference array.

    *arrays : array
        Arrays to check.

    Returns
    -------
    arrays : list
        Arrays with the same namespace and device as reference.
    """
    xp, is_array_api = get_namespace(reference)

    if is_array_api:
        device_ = device(reference)
        # Move arrays to the same namespace and device as the reference array.
        return [
            xp.asarray(a, device=device_) if a is not None else None for a in arrays
        ]
    else:
        return arrays


def _remove_non_arrays(*arrays, remove_none=True, remove_types=(str,)):
    """Filter arrays to exclude None and/or specific types.

    Sparse arrays are always filtered out.