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

Code before:
    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 preserved
    # in the output unless x is 1d and y is 0d.
    xp = array_namespace(x, y) if xp is None else xp
    y_0d = xp.asarray(y).ndim == 0
    x, y = _broadcast_arrays((x, y), axis=-1, xp=xp)
    x_1d = x.ndim <= 1

    if x_1d or is_torch(xp):
        y = xp.reshape(y, ()) if (y_0d and x_1d) else y
        return xp.searchsorted(x, y, side=side)

    a = xp.full(y.shape, 0, device=xp_device(x))

    if x.shape[-1] == 0:
        return a

    n = xp.count_nonzero(~xp.isnan(x), axis=-1, keepdims=True)
    b = xp.broadcast_to(n, y.shape)

    compare = xp.less_equal if side == 'left' else xp.less

    # while xp.any(b - a > 1):
    # refactored to for loop with ~log2(n) iterations for JAX JIT
    for i in range(int(math.log2(x.shape[-1])) + 1):
        c = (a + b) // 2
        x0 = xp.take_along_axis(x, c, axis=-1)

Review comment (suggestion):
```suggestion
    # Vectorized xp.searchsorted. Assumes search is along last axis, which is preserved
```

Code after:
    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):
    # Vectorized xp.searchsorted. Search is performed along last axis. The shape of the
    # output is that of `y`, broadcasting the batch dimensions with those of `x` if
    # necessary.
    xp = array_namespace(x, y) if xp is None else xp
    xp_default_int = xp.asarray(1).dtype
    y_0d = xp.asarray(y).ndim == 0
    x, y = _broadcast_arrays((x, y), axis=-1, xp=xp)
    x_1d = x.ndim <= 1

    if x_1d or is_torch(xp):
        y = xp.reshape(y, ()) if (y_0d and x_1d) else y
        out = xp.searchsorted(x, y, side=side)
        out = xp.astype(out, xp_default_int, copy=False)
        return out

    a = xp.full(y.shape, 0, device=xp_device(x))

    if x.shape[-1] == 0:
        return a

    n = xp.count_nonzero(~xp.isnan(x), axis=-1, keepdims=True)
    b = xp.broadcast_to(n, y.shape)

    compare = xp.less_equal if side == 'left' else xp.less

    # while xp.any(b - a > 1):