Function _array_api_for_tests

sklearn/utils/_testing.py:1308–1409

Return (xp, device) for array API testing. Parameters ---------- array_namespace : str The importable name of the array namespace module. device_name : str or None, default=None The device name for array allocation. Can be None for default device. Returns --

(array_namespace, device_name=None, dtype_name=None)

def _array_api_for_tests(array_namespace, device_name=None, dtype_name=None):
    """Return (xp, device) for array API testing.

    Parameters
    ----------
    array_namespace : str
        The importable name of the array namespace module.
    device_name : str or None, default=None
        The device name for array allocation. Can be None for default device.

    Returns
    -------
    xp : module
        The module object for the requested array namespace.
    device : object, str or None
        The library specific device object that can be passed to
        xp.asarray(..., device=device). This might be a string and not
        a library specific device object.
    """
    try:
        array_mod = importlib.import_module(array_namespace)
    except (ModuleNotFoundError, ImportError):
        raise SkipTest(
            f"{array_namespace} is not installed: not checking array_api input"
        )

    if os.environ.get("SCIPY_ARRAY_API") is None:
        raise SkipTest("SCIPY_ARRAY_API is not set: not checking array_api input")

    from sklearn.externals.array_api_compat import get_namespace

    # First create an array using the chosen array module and then get the
    # corresponding (compatibility wrapped) array namespace based on it.
    # This is because `cupy` is not the same as the compatibility wrapped
    # namespace of a CuPy array.
    device = None
    xp = get_namespace(array_mod.asarray(1))
    if (
        array_namespace == "torch"
        and device_name == "cuda"
        and not xp.backends.cuda.is_built()
    ):
        raise SkipTest("PyTorch test requires cuda, which is not available")
    elif array_namespace == "dpnp":  # pragma: nocover
        dpctl = pytest.importorskip("dpctl")
        if device_name is None:
            available_devices = dpctl.get_devices()
            if not available_devices:
                raise SkipTest("Skipping dpnp test because no SYCL devices found")
            else:
                device = available_devices[0]
        elif not dpctl.get_devices(device_type=device_name):
            raise SkipTest(f"Skipping dpnp test because no {device_name} device found")

    elif array_namespace == "torch" and device_name == "mps":
        if os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") != "1":
            # For now we need PYTORCH_ENABLE_MPS_FALLBACK=1 for all estimators to work
            # when using the MPS device.