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

Code before:

    if not normalize:
        return sum_

    scale = xp.sum(weights, axis=axis)
    if xp.any(scale == 0.0):
        raise ZeroDivisionError("Weights sum to zero, can't be normalized")

    return sum_ / scale


def _median(x, axis=None, keepdims=False, xp=None):
    # `median` is not included in the Array API spec, but is implemented in most
    # array libraries, and all that we support (as of May 2025).
    xp, _ = get_namespace(x, xp=xp)

    if _is_xp_namespace(xp, "array_api_strict"):
        x_np = numpy.asarray(xp.asarray(x, device=xp.Device("CPU_DEVICE")))
        return numpy.median(x_np, axis=axis, keepdims=keepdims)

    # `torch.median` takes the lower of the two medians when `x` has even number
    # of elements, thus we use `torch.quantile(q=0.5)`, which gives mean of the two
    if array_api_compat.is_torch_namespace(xp):
        return xp.quantile(x, q=0.5, dim=axis, keepdim=keepdims)

    return xp.median(x, axis=axis, keepdims=keepdims)


def _xlogy(x, y, xp=None):
    # TODO: Remove this once https://github.com/scipy/scipy/issues/21736 is fixed
    xp, _, device_ = get_namespace_and_device(x, y, xp=xp)

    with numpy.errstate(divide="ignore", invalid="ignore"):
        temp = x * xp.log(y)
    return xp.where(x == 0.0, xp.asarray(0.0, dtype=temp.dtype, device=device_), temp)


def _nanmin(X, axis=None, xp=None):
    # TODO: refactor once nan-aware reductions are standardized:
    # https://github.com/data-apis/array-api/issues/621
    xp, _, device_ = get_namespace_and_device(X, xp=xp)
    if _is_numpy_namespace(xp):
        return xp.asarray(numpy.nanmin(X, axis=axis))

    else:
        mask = xp.isnan(X)
        X = xp.min(
            xp.where(mask, xp.asarray(+xp.inf, dtype=X.dtype, device=device_), X),
            axis=axis,
        )
        # Replace Infs from all NaN slices with NaN again

Review comment (suggestion):
I think the code was rearranged and this will now break in case the array api does not support this. It is better to keep the default case to the conversion to numpy and we should still only call this if the array api actually contains this function. Please see `_bincount` or `_tolist` as examples of how we have done this. In case you still want to go with the check for `array-api-strict` then we should still use `_convert_to_numpy` to convert the original array to numpy, as the code above seems like a duplication of what is already being done there.

Code after:

    if not normalize:
        return sum_

    scale = xp.sum(weights, axis=axis)
    if xp.any(scale == 0.0):
        raise ZeroDivisionError("Weights sum to zero, can't be normalized")

    return sum_ / scale


def _median(x, axis=None, keepdims=False, xp=None):
    # XXX: `median` is not included in the array API spec, but is implemented
    # in most array libraries, and all that we support (as of May 2025).
    # TODO: consider simplifying this code to use scipy instead once the oldest
    # supported SciPy version provides `scipy.stats.quantile` with native array API
    # support (likely scipy 1.6 at the time of writing). Proper benchmarking of
    # either option with popular array namespaces is required to evaluate the
    # impact of this choice.
    xp, _, device = get_namespace_and_device(x, xp=xp)

    # `torch.median` takes the lower of the two medians when `x` has even number
    # of elements, thus we use `torch.quantile(q=0.5)`, which gives mean of the two
    if array_api_compat.is_torch_namespace(xp):
        return xp.quantile(x, q=0.5, dim=axis, keepdim=keepdims)

    if hasattr(xp, "median"):
        return xp.median(x, axis=axis, keepdims=keepdims)

    # Intended mostly for array-api-strict (which as no "median", as per the spec)
    # as `_convert_to_numpy` does not necessarily work for all array types.
    x_np = _convert_to_numpy(x, xp=xp)
    return xp.asarray(numpy.median(x_np, axis=axis, keepdims=keepdims), device=device)


def _xlogy(x, y, xp=None):
    # TODO: Remove this once https://github.com/scipy/scipy/issues/21736 is fixed
    xp, _, device_ = get_namespace_and_device(x, y, xp=xp)

    with numpy.errstate(divide="ignore", invalid="ignore"):
        temp = x * xp.log(y)
    return xp.where(x == 0.0, xp.asarray(0.0, dtype=temp.dtype, device=device_), temp)


def _nanmin(X, axis=None, xp=None):
    # TODO: refactor once nan-aware reductions are standardized:
    # https://github.com/data-apis/array-api/issues/621
    xp, _, device_ = get_namespace_and_device(X, xp=xp)
    if _is_numpy_namespace(xp):
        return xp.asarray(numpy.nanmin(X, axis=axis))