Obtain ranks for each row in a 2D array in numpy

get_row_ranks.py

import numpy as np


def get_row_ranks(xs):
    """
    Args:
        xs: array of shape :math:`(N, E)`
        
    Returns:
         Ranks for each row of ``xs``. A rank of 0 is the highest value, and a rank of ``E - 1`` is the lowest value.
         
    Examples:
         >>> get_row_ranks(np.array([ ...
                 [1, 2, 3], ...
                 [1, 3, 2], ...
                 [2, 1, 3], ...
                 [2, 3, 1], ... 
                 [3, 1, 2], ...
                 [3, 2, 1], ...
             ]))
         array([[2, 1, 0],
                [2, 0, 1],
                [1, 2, 0],
                [1, 0, 2],
                [0, 2, 1],
                [0, 1, 2]])
    """
    assert xs.ndim <= 2
    n_elements = xs.shape[-1]
    sorted_idxs = xs.argsort(axis=-1)
    return (n_elements - 1) - sorted_idxs.argsort(axis=-1)