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

Code before:

    contains_nans = _contains_nan(x, nan_policy, xp_omit_okay=True, xp=xp)

    if keepdims not in {None, True, False}:
        message = "If specified, `keepdims` must be True or False."
        raise ValueError(message)

    # If data has length zero along `axis`, the result will be an array of NaNs just
    # as if the data had length 1 along axis and were filled with NaNs. This is treated
    # naturally below whether `nan_policy` is `'propagate'` or `'omit'`.
    if x.shape[axis] == 0:
        shape = list(x.shape)
        shape[axis] = 1
        x = xp.full(shape, xp.nan, dtype=dtype, device=xp_device(x))

    if weights is None:
        y = xp.sort(x, axis=axis, stable=False)
        y, p = _broadcast_arrays((y, p), axis=axis)
        n_zero_weight = None
    else:
        i_zero_weight = (weights == 0)
        n_zero_weight = xp.count_nonzero(i_zero_weight, axis=axis, keepdims=True)
        x = xpx.at(x)[i_zero_weight].set(xp.inf)
        i_y = xp.argsort(x, axis=axis, stable=False)
        y = xp.take_along_axis(x, i_y, axis=axis)
        y, p, weights, i_y, n_zero_weight = _broadcast_arrays(
            (y, p, weights, i_y, n_zero_weight), axis=axis)
        weights = xp.take_along_axis(weights, i_y, axis=axis)
        # When NaNs have zero weight, they shouldn't propagate
        # data-apis/array-api-extra#506 raises an error here, so use `where`
        # y = xpx.at(y)[xp.isnan(y) & (weights == 0)].set(1.)
        if nan_policy == 'propagate':
            y = xp.where(~xp.isfinite(y) & (weights == 0), 1., y)

    if (keepdims is False) and (p.shape[axis] != 1):
        message = "`keepdims` may be False only if the length of `p` along `axis` is 1."
        raise ValueError(message)
    keepdims = (p.shape[axis] != 1) if keepdims is None else keepdims

    y = xp.moveaxis(y, axis, -1)
    p = xp.moveaxis(p, axis, -1)
    weights = weights if weights is None else xp.moveaxis(weights, axis, -1)
    n_zero_weight = (n_zero_weight if n_zero_weight is None
                     else xp.moveaxis(n_zero_weight, axis, -1))

    n = _length_nonmasked(y, -1, xp=xp, keepdims=True)
    n = n if n_zero_weight is None else n - n_zero_weight

    if contains_nans:
        nans = xp.isnan(y)


Review comment (suggestion):
Hmm we may need to broadcast first depending on what `xpx.at` does when the arrays are of different (but broadcastable) shapes (next line)

Code after:
    if weights is not None and method in no_weights:
        message = f"`method='{method}'` does not support `weights`."
        raise ValueError(message)

    contains_nans = _contains_nan(x, nan_policy, xp_omit_okay=True, xp=xp)

    if keepdims not in {None, True, False}:
        message = "If specified, `keepdims` must be True or False."
        raise ValueError(message)

    # If data has length zero along `axis`, the result will be an array of NaNs just
    # as if the data had length 1 along axis and were filled with NaNs. This is treated
    # naturally below whether `nan_policy` is `'propagate'` or `'omit'`.
    if x.shape[axis] == 0:
        shape = list(x.shape)
        shape[axis] = 1
        x = xp.full(shape, xp.nan, dtype=dtype, device=xp_device(x))

    if weights is None:
        y = xp.sort(x, axis=axis, stable=False)
        y, p = _broadcast_arrays((y, p), axis=axis)
        n_zero_weight = None
    else:
        x, weights = xp.broadcast_arrays(x, weights)
        i_zero_weight = (weights == 0)
        n_zero_weight = xp.count_nonzero(i_zero_weight, axis=axis, keepdims=True)
        x = xpx.at(x)[i_zero_weight].set(xp.inf, copy=True)
        i_y = xp.argsort(x, axis=axis, stable=False)
        y = xp.take_along_axis(x, i_y, axis=axis)
        weights = xp.take_along_axis(weights, i_y, axis=axis)
        y, p, weights, i_y, n_zero_weight = _broadcast_arrays(
            (y, p, weights, i_y, n_zero_weight), axis=axis)

    if (keepdims is False) and (p.shape[axis] != 1):
        message = "`keepdims` may be False only if the length of `p` along `axis` is 1."
        raise ValueError(message)
    keepdims = (p.shape[axis] != 1) if keepdims is None else keepdims

    y = xp.moveaxis(y, axis, -1)
    p = xp.moveaxis(p, axis, -1)
    weights = weights if weights is None else xp.moveaxis(weights, axis, -1)
    n_zero_weight = (n_zero_weight if n_zero_weight is None
                     else xp.moveaxis(n_zero_weight, axis, -1))

    n = _length_nonmasked(y, -1, xp=xp, keepdims=True)
    n = n if n_zero_weight is None else n - n_zero_weight

    if contains_nans:
        nans = xp.isnan(y)

        # Note that if length along `axis` were 0 to begin with,