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

Code before:
    >>> rankdata([0, 2, 3, np.nan, -2, np.nan], nan_policy="omit")
    array([ 2.,  3.,  4., nan,  1., nan])

    """
    methods = ('average', 'min', 'max', 'dense', 'ordinal')
    if method not in methods:
        raise ValueError(f'unknown method "{method}"')

    xp = array_namespace(a)
    x = xp.asarray(a)

    if axis is None:
        x = xp_ravel(x)
        axis = -1

    if xp_size(x) == 0:
        dtype = xp.asarray(1.).dtype if method == 'average' else xp.asarray(1).dtype
        return xp.empty(x.shape, dtype=dtype)

    contains_nan = _contains_nan(x, nan_policy)

    x = xp_swapaxes(x, axis, -1, xp=xp)
    ranks = _rankdata(x, method, xp=xp)

    # JIT won't allow use of `contains_nan` for control flow here, so we have to choose
    # whether to always or never run this block wit JIT.
    # For now, *never* run it; otherwise, it would change dtype of `ranks`.
    # When gh-19889 is resolved, dtype will already be `float`, so *always* run it.
    # TODO then: broadcast `i_nan` to the shape of ranks before using `at.set`
    if not is_lazy_array(x) and contains_nan:
        default_float = xp_default_dtype(xp)
        i_nan = (xp.isnan(x) if nan_policy == 'omit'
                 else xp.any(xp.isnan(x), axis=-1))
        ranks = xp.asarray(ranks, dtype=default_float)  # copy=False when implemented
        ranks = xpx.at(ranks)[i_nan].set(xp.nan)

    ranks = xp_swapaxes(ranks, axis, -1, xp=xp)
    return ranks


def _order_ranks(ranks, j, *, xp):
    # Reorder ascending order `ranks` according to `j`
    xp = array_namespace(ranks) if xp is None else xp
    if is_numpy(xp) or is_cupy(xp):
        ordered_ranks = xp.empty(j.shape, dtype=ranks.dtype)
        xp.put_along_axis(ordered_ranks, j, ranks, axis=-1)
    else:
        # `put_along_axis` not in array API (data-apis/array-api#177)
        #  so argsort the argsort and take_along_axis...
        j_inv = xp.argsort(j, axis=-1, stable=True)
        ordered_ranks = xp.take_along_axis(ranks, j_inv, axis=-1)

Review comment (suggestion):
```suggestion
    # whether to always or never run this block with JIT.
```
Let's not hold this back for a typo, though.

Code after:
    >>> rankdata([0, 2, 3, np.nan, -2, np.nan], nan_policy="omit")
    array([ 2.,  3.,  4., nan,  1., nan])

    """
    methods = ('average', 'min', 'max', 'dense', 'ordinal')
    if method not in methods:
        raise ValueError(f'unknown method "{method}"')

    xp = array_namespace(a)
    x = xp.asarray(a)

    if axis is None:
        x = xp_ravel(x)
        axis = -1

    if xp_size(x) == 0:
        dtype = xp.asarray(1.).dtype if method == 'average' else xp.asarray(1).dtype
        return xp.empty(x.shape, dtype=dtype)

    contains_nan = _contains_nan(x, nan_policy)

    x = xp_swapaxes(x, axis, -1, xp=xp)
    ranks = _rankdata(x, method, xp=xp)

    # JIT won't allow use of `contains_nan` for control flow here, so we have to choose
    # whether to always or never run this block with JIT.
    # For now, *never* run it; otherwise, it would change dtype of `ranks`.
    # When gh-19889 is resolved, dtype will already be `float`, so *always* run it.
    # TODO then: broadcast `i_nan` to the shape of ranks before using `at.set`
    if not is_lazy_array(x) and contains_nan:
        default_float = xp_default_dtype(xp)
        i_nan = (xp.isnan(x) if nan_policy == 'omit'
                 else xp.any(xp.isnan(x), axis=-1))
        ranks = xp.asarray(ranks, dtype=default_float)  # copy=False when implemented
        ranks = xpx.at(ranks)[i_nan].set(xp.nan)

    ranks = xp_swapaxes(ranks, axis, -1, xp=xp)
    return ranks


def _order_ranks(ranks, j, *, xp):
    # Reorder ascending order `ranks` according to `j`
    xp = array_namespace(ranks) if xp is None else xp
    if is_numpy(xp) or is_cupy(xp):
        ordered_ranks = xp.empty(j.shape, dtype=ranks.dtype)
        xp.put_along_axis(ordered_ranks, j, ranks, axis=-1)
    else:
        # `put_along_axis` not in array API (data-apis/array-api#177)
        #  so argsort the argsort and take_along_axis...
        j_inv = xp.argsort(j, axis=-1, stable=True)
        ordered_ranks = xp.take_along_axis(ranks, j_inv, axis=-1)