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

Code before:
            "Complex floating point values are not supported by average."
        )

    output_dtype = _find_matching_floating_dtype(a, weights, xp=xp)
    a = xp.astype(a, output_dtype)

    if weights is None:
        return (xp.mean if normalize else xp.sum)(a, axis=axis)

    weights = xp.astype(weights, output_dtype)

    sum_ = xp.sum(xp.multiply(a, weights), axis=axis)

    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, _, 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 has no "median", as per the spec,
    # as `_convert_to_numpy` does not generically 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)



Review comment (refactor):
```suggestion
    # 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.
```

Code after:
            "Complex floating point values are not supported by average."
        )

    output_dtype = _find_matching_floating_dtype(a, weights, xp=xp)
    a = xp.astype(a, output_dtype)

    if weights is None:
        return (xp.mean if normalize else xp.sum)(a, axis=axis)

    weights = xp.astype(weights, output_dtype)

    sum_ = xp.sum(xp.multiply(a, weights), axis=axis)

    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)