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

Code before:
        g = (1 - xp.astype((g == 0) & (j % 2 == 1), jg.dtype))
    if method in {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation'}:
        g = xp.asarray(g)
        g = xpx.at(g, jg < 0).set(0)

    g = xpx.at(g)[j < 0].set(0)
    j = xp.clip(j, 0., n - 1)
    jp1 = xp.clip(jp1, 0., n - 1)

    return ((1 - g) * xp.take_along_axis(y, xp.astype(j, xp.int64), axis=-1)
            + g * xp.take_along_axis(y, xp.astype(jp1, xp.int64), axis=-1))


def _quantile_hd(y, p, n, xp):
    # RE axis handling: We need to perform a reducing operation over rows of `y` for
    # each element in the corresponding row of `p` (a la Cartesian product). Strategy:
    # move rows of `p` to an axis at the front that is orthogonal to all the rest,
    # perform the reducing operating over the last axis, then move the front axis back
    # to the end.
    p = xp.moveaxis(p, -1, 0)[..., xp.newaxis]
    a = p * (n + 1)
    b = (1 - p) * (n + 1)
    i = xp.arange(y.shape[-1] + 1, dtype=y.dtype, device=xp_device(y))
    w = betainc(a, b, i / n)
    w = w[..., 1:] - w[..., :-1]
    w = xpx.at(w, xp.isnan(w)).set(0)
    res = xp.vecdot(w, y, axis=-1)
    return xp.moveaxis(res, 0, -1)


def _quantile_bc(y, p, n, method, xp):
    # Methods retained for backward compatibility. NumPy documentation is not
    # quite right about what these methods do: if `p * (n - 1)` is integral,
    # that is used as the index. See numpy/numpy#28910.
    ij = p * (n - 1)
    if method == '_midpoint':
        return (xp.take_along_axis(y, xp.astype(xp.floor(ij), xp.int64), axis=-1)
                + xp.take_along_axis(y, xp.astype(xp.ceil(ij), xp.int64), axis=-1)) / 2
    elif method == '_lower':
        k = xp.floor(ij)
    elif method == '_higher':
        k = xp.ceil(ij)
    elif method == '_nearest':
        k = xp.round(ij)
    return xp.take_along_axis(y, xp.astype(k, xp.int64), axis=-1)


@xp_capabilities(skip_backends=[("dask.array", "No take_along_axis yet.")])
def _xp_searchsorted(x, y, *, side='left', xp=None):
    # Vectorize np.searchsorted. Assumes search is along last axis, which is always
    # preserved in the output. Does not support zero-length `x`. For side='right',

Review comment (suggestion):
For `nan_policy='omit'` and `marray`, `n` is an array the is the number of valid entries in a slice; that's why we can't just take the last element. (`marray` tests are skipped right now because I'm running into a mutability issue. I think it's an `marray` issue and the implementation will probably "just work" when that gets resolved because we're already accounting for `n` depending on the slice.)

Code after:
        res = xp.reshape(res, shape)
        axis = -1

    res = xp.moveaxis(res, -1, axis)

    if not keepdims:
        res = xp.squeeze(res, axis=axis)

    return res[()] if res.ndim == 0 else res


def _quantile_hf(y, p, n, method, weights, xp):
    ms = dict(inverted_cdf=0, averaged_inverted_cdf=0, closest_observation=-0.5,
              interpolated_inverted_cdf=0, hazen=0.5, weibull=p, linear=1 - p,
              median_unbiased=p/3 + 1/3, normal_unbiased=p/4 + 3/8)
    m = ms[method]

    if weights is None:
        jg = p * n + m
        jp1 = jg // 1
        j = jp1 - 1
    else:
        cumulative_weights = xp.cumulative_sum(weights, axis=-1)
        n_int = xp.asarray(n, dtype=xp.int64)
        n_int = xp.broadcast_to(n_int, cumulative_weights.shape[:-1] + (1,))
        total_weight = xp.take_along_axis(cumulative_weights, n_int-1, axis=-1)
        jg = p * total_weight + m
        jp1 = _xp_searchsorted(cumulative_weights, jg, side='right')
        j = _xp_searchsorted(cumulative_weights, jg-1, side='right')
        j, jp1 = xp.astype(j, y.dtype), xp.astype(jp1, y.dtype)

    g = jg % 1
    if method == 'inverted_cdf':
        g = xp.astype((g > 0), jg.dtype)
    elif method == 'averaged_inverted_cdf':
        g = (1 + xp.astype((g > 0), jg.dtype)) / 2
    elif method == 'closest_observation':
        g = (1 - xp.astype((g == 0) & (j % 2 == 1), jg.dtype))
    if method in {'inverted_cdf', 'averaged_inverted_cdf', 'closest_observation'}:
        g = xp.asarray(g)
        g = xpx.at(g, jg < 0).set(0)

    g = xpx.at(g)[j < 0].set(0)
    j = xp.clip(j, 0., n - 1)
    jp1 = xp.clip(jp1, 0., n - 1)

    return ((1 - g) * xp.take_along_axis(y, xp.astype(j, xp.int64), axis=-1)
            + g * xp.take_along_axis(y, xp.astype(jp1, xp.int64), axis=-1))


def _quantile_hd(y, p, n, xp):