Function move_to

sklearn/utils/_array_api.py:519–628

Move all arrays to `xp` and `device`. Each array will be moved to the reference namespace and device if it is not already using it. Otherwise the array is left unchanged. `arrays` may contain `None` entries, these are left unchanged. Sparse arrays are accepted (as pass through) if

(*arrays, xp, device)

def move_to(*arrays, xp, device):
    """Move all arrays to `xp` and `device`.

    Each array will be moved to the reference namespace and device if
    it is not already using it. Otherwise the array is left unchanged.

    `arrays` may contain `None` entries, these are left unchanged.

    Sparse arrays are accepted (as pass through) if the reference namespace is
    NumPy, in which case they are returned unchanged. Otherwise a `TypeError`
    is raised.

    Parameters
    ----------
    *arrays : iterable of arrays
        Arrays to (potentially) move.

    xp : namespace
        Array API namespace to move arrays to.

    device : device
        Array API device to move arrays to.

    Returns
    -------
    arrays : tuple or array
        Tuple of arrays with the same namespace and device as reference. Single array
        returned if only one `arrays` input.
    """
    if isinstance(device, str) and device == "xpu":  # pragma: nocover
        # XXX: Workaround for PyTorch XPU bug for `from_dlpack` calls with
        # device strings that do not include any device number suffix.
        # https://github.com/pytorch/pytorch/issues/181140
        device += ":0"

    sparse_mask = [sp.issparse(array) for array in arrays]
    none_mask = [array is None for array in arrays]
    if any(sparse_mask) and not _is_numpy_namespace(xp):
        raise TypeError(
            "Sparse arrays are only accepted (and passed through) when the target "
            "namespace is Numpy"
        )

    arrays_ = arrays
    # Down cast float64 `arrays` when highest precision of `xp`/`device` is float32
    if _max_precision_float_dtype(xp, device) == xp.float32:
        arrays_ = []
        for array in arrays:
            xp_array, _ = get_namespace(array)
            if getattr(array, "dtype", None) == xp_array.float64:
                arrays_.append(xp_array.astype(array, xp_array.float32))
            else:
                arrays_.append(array)

    converted_arrays = []
    for array, is_sparse, is_none in zip(arrays_, sparse_mask, none_mask):
        if is_none:
            converted_arrays.append(None)