code: def argmaxMulti(a, groupKeys, assumeSorted=False):
    if not assumeSorted:
      sorter = np.argsort(groupKeys, kind="mergesort")
      a = a[sorter]
      groupKeys = groupKeys[sorter]
    _, indices, lengths = np.unique(groupKeys, return_index=True,
                                    return_counts=True)
    maxValues = np.maximum.reduceat(a, indices)
    allMaxIndices = np.flatnonzero(np.repeat(maxValues, lengths) == a)
    indices = allMaxIndices[np.searchsorted(allMaxIndices, indices)]
    if assumeSorted:
      return indices
    else:
      return sorter[indices]