Implement NumPy search routines

[shoyer]

xref https://docs.scipy.org/doc/numpy/reference/routines.sort.html

Extracted from #70:

• np.argmax
• np.argmin
• np.argwhere
• np.compress
• np.nonzero
• np.flatnonzero
• np.where
• np.searchsorted
• np.extract

(see #2077 for nanargmax/nanargmin)

Notes:

• argwhere/flatnonzero/extractcould should just be minimal wrappers around nonzero
• searchsorted is harder: it need would need low-level routines/wrappers for efficient binary search, which to my knowledge do not currently existing in XLA

[djdongjin]

@shoyer I'd like to work on the last searchsorted. Where and how should I implement the efficient binary search you mentioned? Should I implement it in tensorflow.xla first and wrap it in Jax when it is merged into tensorflow?

And a more general question, if we want to add some functions using ops that don't exist in XLA (like here), how should we do? Can we implement the ops directly in python, or we need to implement it in XLA first?

[terhorst]

If your data don't have ties, then the following should work:

def hacky_searchsorted(a, v):
    i = np.argsort(np.concatenate([a, v]))
    ranks = i.argsort()
    A = len(a)
    V = len(v)
    return (np.arange(A + V) < A)[i].cumsum()[ranks[A:]]

[hawkinsp]

#2938 implemented searchsorted.