Repository: scipy/scipy
File: scipy/stats/_quantile.py (Python)

Code before:
    # JIX, so we can come back to this later.
    if not is_lazy_array(y) and contains_nans:
        nans = xp.isnan(y)

        # Note that if length along `axis` were 0 to begin with,
        # it is now length 1 and filled with NaNs.
        if nan_policy == 'propagate':
            nan_out = xp.any(nans, axis=-1)
        else:  # 'omit'
            n_int = n - xp.count_nonzero(nans, axis=-1, keepdims=True)
            n = xp.astype(n_int, dtype)
            # NaNs are produced only if slice is empty after removing NaNs
            nan_out = xp.any(n == 0, axis=-1)
            n = xpx.at(n, nan_out).set(y.shape[-1])  # avoids pytorch/pytorch#146211

        if xp.any(nan_out):
            y = xp.asarray(y, copy=True)  # ensure writable
            y = xpx.at(y, nan_out).set(xp.nan)
        elif xp.any(nans) and method == 'harrell-davis':
            y = xp.asarray(y, copy=True)  # ensure writable
            y = xpx.at(y, nans).set(0)  # any non-nan will prevent NaN from propagating

    n = xp.asarray(n, dtype=dtype, device=xp_device(y))

    p_mask = (p > 1) | (p < 0) | xp.isnan(p)
    if is_lazy_array(y) or xp.any(p_mask):
        p = xp.asarray(p, copy=True)
        p = xpx.at(p, p_mask).set(0.5)  # these get NaN-ed out at the end

    return (y, p, method, axis, nan_policy, keepdims,
            n, axis_none, ndim, p_mask, weights, xp)


@xp_capabilities(skip_backends=[("dask.array", "No take_along_axis yet.")])
def quantile(x, p, *, method='linear', axis=0, nan_policy='propagate', keepdims=None,
             weights=None):
    """
    Compute the p-th quantile of the data along the specified axis.

    Parameters
    ----------
    x : array_like of real numbers
        Data array.
    p : array_like of float
        Probability or sequence of probabilities of the quantiles to compute.
        Values must be between 0 and 1 (inclusive).
        While `numpy.quantile` can only compute quantiles according to the Cartesian
        product of the first two arguments, this function enables calculation of
        quantiles at different probabilities for each axis slice by following
        broadcasting rules like those of `scipy.stats` reducing functions.
        See `axis`, `keepdims`, and the examples.

Review comment (suggestion):
We could just eliminate the `xp.any`. And probably should just use `where` instead of copying then `xpx.at.set`. I can do that in a separate PR, when I address the code comment above.

Code after:
    # JIX, so we can come back to this later.
    if not is_lazy_array(y) and contains_nans:
        nans = xp.isnan(y)

        # Note that if length along `axis` were 0 to begin with,
        # it is now length 1 and filled with NaNs.
        if nan_policy == 'propagate':
            nan_out = xp.any(nans, axis=-1)
        else:  # 'omit'
            n_int = n - xp.count_nonzero(nans, axis=-1, keepdims=True)
            n = xp.astype(n_int, dtype)
            # NaNs are produced only if slice is empty after removing NaNs
            nan_out = xp.any(n == 0, axis=-1)
            n = xpx.at(n, nan_out).set(y.shape[-1])  # avoids pytorch/pytorch#146211

        if xp.any(nan_out):
            y = xp.asarray(y, copy=True)  # ensure writable
            y = xpx.at(y, nan_out).set(xp.nan)
        elif xp.any(nans) and method == 'harrell-davis':
            y = xp.asarray(y, copy=True)  # ensure writable
            y = xpx.at(y, nans).set(0)  # any non-nan will prevent NaN from propagating

    n = xp.asarray(n, dtype=dtype, device=xp_device(y))

    p_mask = (p > 1) | (p < 0) | xp.isnan(p)
    p = xp.where(p_mask, 0.5, p)

    return (y, p, method, axis, nan_policy, keepdims,
            n, axis_none, ndim, p_mask, weights, xp)


@xp_capabilities(skip_backends=[("dask.array", "No take_along_axis yet.")])
def quantile(x, p, *, method='linear', axis=0, nan_policy='propagate', keepdims=None,
             weights=None):
    """
    Compute the p-th quantile of the data along the specified axis.

    Parameters
    ----------
    x : array_like of real numbers
        Data array.
    p : array_like of float
        Probability or sequence of probabilities of the quantiles to compute.
        Values must be between 0 and 1 (inclusive).
        While `numpy.quantile` can only compute quantiles according to the Cartesian
        product of the first two arguments, this function enables calculation of
        quantiles at different probabilities for each axis slice by following
        broadcasting rules like those of `scipy.stats` reducing functions.
        See `axis`, `keepdims`, and the examples.
    method : str, default: 'linear'
        The method to use for estimating the quantile.