Repository: google/jax

Issue:
Batching rule for 'sort_key_val' not implemented
I have implemented argsort like this:
```
@jax.numpy.lax_numpy._wraps(onp.argsort)
def argsort(a, axis=-1, kind='quicksort', order=None):
    assert(a.ndim == 1) # starting simple
    vals = np.arange(len(a), dtype=int)
    sorted_keys, sorted_vals = jax.lax.sort_key_val(keys=a, values=vals)
    return sorted_vals
```

This works:
`argsort(np.ones(3))`

And I can also write an argsort that works with 2d arrays, but instead I would like to use vmap so I don't need to rewrite my code with a batch dimension.

But using vmap gives me an error:
`jax.vmap(argsort)(np.ones((3, 3)))`
`NotImplementedError: Batching rule for 'sort_key_val' not implemented`

So I am trying to implement the batch rule for `jax.lax.sort_key_val`, but it is more complicated than eg https://github.com/google/jax/pull/82 because `jax.lax.sort_key_val` returns multiple values, and the batching code assumes only a single return value.

I've tried a few things but it's tricky. Here's an attempt:
```
def sort_key_val_batched(batched_args, batch_dims, dimension=-1):
    keys, values = batched_args
    sorted_keys, sorted_values = jax.lax.sort_key_val_p.bind(
        keys=keys,
        values=values,
        dimension=dimension
    )
    return sorted_keys, 0
jax.batching.primitive_batchers[jax.lax.sort_key_val_p] = sort_key_val_batched
```

But:
`jax.vmap(argsort)(np.ones((3, 3)))`
`RuntimeError: Invalid argument: Sort keys and values dimensions must match. Keys shape is: f32[10,3], Values shape (operand index 1) is: s32[3]:`
In other words the batching rule recieves the full 2d keys array but only one row at a time of the values array. Not sure how to handle this.

One potential workaround would be to use `np.sort(order=...` to sort rows in place where one column is the values to sort by, and the other column is indices:
```
@jax.numpy.lax_numpy._wraps(onp.argsort)
def jax_argsort(a, axis=-1, kind='quicksort', order=None):
    assert(a.ndim == 1) # starting simple
    inds = np.arange(len(a), dtype=int)
    inds = np.expand_dims(inds, 1)
    a = np.expand_dims(a, 1)
    a = a.view('i8,i8')
    tmp = np.concatenate((a, inds), axis=1)
    tmp_sorted = np.sort(a, order=['f0'])
    sort_inds = tmp_sorted[:, 1]
    return sort_inds
```
But jax does not support the order keyword to np.sort.

Any ideas for a solution or workaround?
Batching rule for 'sort_key_val' not implemented
I have implemented argsort like this:
```
@jax.numpy.lax_numpy._wraps(onp.argsort)
def argsort(a, axis=-1, kind='quicksort', order=None):
    assert(a.ndim == 1) # starting simple
    vals = np.arange(len(a), dtype=int)
    sorted_keys, sorted_vals = jax.lax.sort_key_val(keys=a, values=vals)
    return sorted_vals
```

This works:
`argsort(np.ones(3))`

And I can also write an argsort that works with 2d arrays, but instead I would like to use vmap so I don't need to rewrite my code with a batch dimension.

But using vmap gives me an error:
`jax.vmap(argsort)(np.ones((3, 3)))`
`NotImplementedError: Batching rule for 'sort_key_val' not implemented`

So I am trying to implement the batch rule for `jax.lax.sort_key_val`, but it is more complicated than eg https://github.com/google/jax/pull/82 because `jax.lax.sort_key_val` returns multiple values, and the batching code assumes only a single return value.

I've tried a few things but it's tricky. Here's an attempt:
```
def sort_key_val_batched(batched_args, batch_dims, dimension=-1):
    keys, values = batched_args
    sorted_keys, sorted_values = jax.lax.sort_key_val_p.bind(
        keys=keys,
        values=values,
        dimension=dimension
    )
    return sorted_keys, 0
jax.batching.primitive_batchers[jax.lax.sort_key_val_p] = sort_key_val_batched
```

But:
`jax.vmap(argsort)(np.ones((3, 3)))`
`RuntimeError: Invalid argument: Sort keys and values dimensions must match. Keys shape is: f32[10,3], Values shape (operand index 1) is: s32[3]:`
In other words the batching rule recieves the full 2d keys array but only one row at a time of the values array. Not sure how to handle this.

One potential workaround would be to use `np.sort(order=...` to sort rows in place where one column is the values to sort by, and the other column is indices:
```
@jax.numpy.lax_numpy._wraps(onp.argsort)
def jax_argsort(a, axis=-1, kind='quicksort', order=None):
    assert(a.ndim == 1) # starting simple
    inds = np.arange(len(a), dtype=int)
    inds = np.expand_dims(inds, 1)
    a = np.expand_dims(a, 1)
    a = a.view('i8,i8')
    tmp = np.concatenate((a, inds), axis=1)
    tmp_sorted = np.sort(a, order=['f0'])
    sort_inds = tmp_sorted[:, 1]
    return sort_inds
```
But jax does not support the order keyword to np.sort.

Any ideas for a solution or workaround?


Discussion:
Thanks for digging into this! You're on the right track, but you're right that `sort_key_val` is especially tricky.

I don't have any specific suggestions yet, but I'll try to poke at it this week when I can.
Thanks for digging into this! You're on the right track, but you're right that `sort_key_val` is especially tricky.

I don't have any specific suggestions yet, but I'll try to poke at it this week when I can.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -240,7 +240,7 @@ def vmap(fun, in_axes=0, out_axes=0):
   @wraps(fun, docstr=docstr)
   def batched_fun(*args, **kwargs):
     if not isinstance(fun, lu.WrappedFun):
-      f = lu.wrap_init(fun)
+      f = lu.wrap_init(fun, kwargs)
     in_axes_ = (in_axes,) * len(args) if type(in_axes) is int else in_axes
     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -391,7 +391,7 @@ def sort(operand, dimension=-1):
   return sort_p.bind(operand, dimension=dimension)
 
 def sort_key_val(keys, values, dimension=-1):
-  # TODO new sort_key_val is variadic
+  # TODO(mattjj): new sort_key_val is variadic
   result = sort_key_val_p.bind(keys, values, dimension=dimension)
   sorted_keys, sorted_values = result
   return sorted_keys, sorted_values
@@ -2275,19 +2275,48 @@ def sort_key_val_jvp(primals, tangents, dimension):
 def sort_key_val_transpose_rule(t, keys, values, dimension):
   t_keys, t_values = t
   assert t_keys is ad_util.zero
-  broadcasted_iota = broadcast_in_dim(
-      onp.arange(keys.shape[dimension]), keys.shape, [dimension % keys.ndim])
-  _, perm = sort_key_val(keys, broadcasted_iota)
+  iota = broadcasted_iota(onp.int32, keys.shape, dimension % keys.ndim)
+  _, perm = sort_key_val(keys, iota)
   keys_result = ad_util.zero if keys is None else None
   values_result = sort_key_val(perm, t_values)[1] if values is None else None
   return [keys_result, values_result]
 
+def sort_key_val_batch_rule(batched_args, batch_dims, dimension):
+  keys, values = batched_args
+  keys_bdim, values_bdim = batch_dims
+  assert keys_bdim is not None or values_bdim is not None
+  if keys_bdim == values_bdim:
+    new_dimension = dimension + (keys_bdim <= dimension)
+    out = sort_key_val(keys, values, new_dimension)
+    return core.pack(out), keys_bdim
+  elif keys_bdim is not None and values_bdim is not None:
+    keys_trans = batching.moveaxis(keys.shape[keys_bdim], values_bdim,
+                                   keys_bdim, keys)
+    new_dimension = dimension + (values_bdim <= dimension)
+    out = sort_key_val(keys_trans, values, new_dimension)
+    return core.pack(out), values_bdim
+  elif keys_bdim is None:
+    broadcast_dimensions = onp.delete(onp.arange(values.ndim), values_bdim)
+    new_keys = broadcast_in_dim(keys, values.shape, broadcast_dimensions)
+    new_dimension = dimension + (values_bdim <= dimension)
+    out = sort_key_val(new_keys, values, new_dimension)
+    return core.pack(out), values_bdim
+  elif values_bdim is None:
+    broadcast_dimensions = onp.delete(onp.arange(keys.ndim), keys_bdim)
+    new_values = broadcast_in_dim(values, keys.shape, broadcast_dimensions)
+    new_dimension = dimension + (keys_bdim <= dimension)
+    out = sort_key_val(keys, new_values, new_dimension)
+    return core.pack(out), keys_bdim
+  else:
+    raise Exception  # unreachable
+
 sort_key_val_p = Primitive('sort_key_val')
 sort_key_val_p.def_impl(sort_key_val_impl)
 sort_key_val_p.def_abstract_eval(sort_key_val_abstract_eval)
 xla.translations[sort_key_val_p] = partial(standard_translate, 'sort_key_val')
 ad.primitive_jvps[sort_key_val_p] = sort_key_val_jvp
 ad.primitive_transposes[sort_key_val_p] = sort_key_val_transpose_rule
+batching.primitive_batchers[sort_key_val_p] = sort_key_val_batch_rule
 
 
 def while_loop_abstract_eval(init_val, opaque_params):
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -19,6 +19,7 @@
 import collections
 import itertools
 import string
+import warnings
 
 import numpy as onp
 import opt_einsum
@@ -235,7 +236,6 @@ def _one_to_one_binop(numpy_fn, lax_fn, promote_like=False):
 fabs = _one_to_one_unop(onp.fabs, lax.abs, True)
 bitwise_not = _one_to_one_unop(onp.bitwise_not, lax.bitwise_not)
 negative = _one_to_one_unop(onp.negative, lax.neg)
-sort = _one_to_one_unop(onp.sort, lax.sort)
 sign = _one_to_one_unop(onp.sign, lax.sign)
 
 floor = _one_to_one_unop(onp.floor, lax.floor, True)
@@ -1523,6 +1523,36 @@ def _argminmax(op, a, axis):
   mask_idxs = where(lax._eq_meet(a, op(a, axis, keepdims=True)), idxs, maxval)
   return min(mask_idxs, axis)
 
+
+@_wraps(onp.sort)
+def sort(a, axis=-1, kind='quicksort', order=None):
+  if kind != 'quicksort':
+    warnings.warn("'kind' argument to sort is ignored.")
+  if order is not None:
+    raise ValueError("'order' argument to sort is not supported.")
+
+  if axis is None:
+    return lax.sort(a.ravel(), 0)
+  else:
+    return lax.sort(a, axis % ndim(a))
+
+
+@_wraps(onp.argsort)
+def argsort(a, axis=-1, kind='quicksort', order=None):
+  if kind != 'quicksort':
+    warnings.warn("'kind' argument to argsort is ignored.")
+  if order is not None:
+    raise ValueError("'order' argument to argsort is not supported.")
+
+  if axis is None:
+    return argsort(a.ravel(), 0)
+  else:
+    axis = axis % ndim(a)
+    iota = lax.broadcasted_iota(onp.int64, shape(a), axis)
+    _, perm = lax.sort_key_val(a, iota, dimension=axis)
+    return perm
+
+
 ### Indexing
 
 
@@ -1722,7 +1752,8 @@ def _static_idx(idx, size):
 def _not_implemented(fun):
   @_wraps(fun)
   def wrapped(*args, **kwargs):
-    raise Exception("Numpy function {} not yet implemented".format(fun))
+    msg = "Numpy function {} not yet implemented"
+    raise NotImplementedError(msg.format(fun))
   return wrapped
 
 # Build a set of all unimplemented NumPy functions.


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -335,6 +335,36 @@ def testRandom(self):
     self.assertAllClose(ans, expected, check_dtypes=False)
     assert len(onp.unique(ans)) == 10 * 3 * 2
 
+  def testSortKeyVal(self):
+    k = onp.arange(12)[::-1].reshape(3, 4)
+    v = onp.random.RandomState(0).permutation(12).reshape(3, 4)
+
+    sk, sv = vmap(partial(lax.sort_key_val, dimension=0), (0, 0))(k, v)
+    self.assertAllClose(sk, k[:, ::-1], check_dtypes=True)
+    self.assertAllClose(sv, v[:, ::-1], check_dtypes=True)
+
+    sk, sv = vmap(partial(lax.sort_key_val, dimension=0), (1, 1), 1)(k, v)
+    self.assertAllClose(sk, k[::-1, :], check_dtypes=True)
+    self.assertAllClose(sv, v[::-1, :], check_dtypes=True)
+
+    sk, sv = vmap(partial(lax.sort_key_val, dimension=0), (0, 1))(k, v.T)
+    self.assertAllClose(sk, k[:, ::-1], check_dtypes=True)
+    self.assertAllClose(sv, v[:, ::-1], check_dtypes=True)
+
+    sk, sv = vmap(partial(lax.sort_key_val, dimension=0), (1, 0))(k.T, v)
+    self.assertAllClose(sk, k[:, ::-1], check_dtypes=True)
+    self.assertAllClose(sv, v[:, ::-1], check_dtypes=True)
+
+    sk, sv = vmap(partial(lax.sort_key_val, dimension=0), (None, 0))(k[0], v)
+    self.assertAllClose(sk, onp.broadcast_to(k[0, ::-1], (3, 4)),
+                        check_dtypes=True)
+    self.assertAllClose(sv, v[:, ::-1], check_dtypes=True)
+
+    sk, sv = vmap(partial(lax.sort_key_val, dimension=0), (1, None))(k.T, v[0])
+    self.assertAllClose(sk, k[:, ::-1], check_dtypes=True)
+    self.assertAllClose(sv, onp.broadcast_to(v[0, ::-1], (3, 4)),
+                        check_dtypes=True)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -932,7 +932,6 @@ def testRot90(self, shape, dtype, k, axes, rng):
   # TODO(mattjj): test infix operator overrides
 
   def testRavel(self):
-    # TODO(mattjj): support this method-based syntax?
     rng = onp.random.RandomState(0)
     args_maker = lambda: [rng.randn(3, 4).astype("float32")]
     self._CompileAndCheck(lambda x: x.ravel(), args_maker, check_dtypes=True)
@@ -959,6 +958,54 @@ def testArangeOnFloats(self):
     ans = lnp.arange(0.0, 1.0, 0.1)
     self.assertAllClose(expected, ans, check_dtypes=True)
 
+  def testSortManually(self):
+    # manual tests for sort are nice because we don't have to worry about ties.
+    # lax.sort is tested combinatorially.
+    ans = lnp.sort(onp.array([16, 15, 23, 42, 8, 4]))
+    expected = onp.array([4, 8, 15, 16, 23, 42])
+    self.assertAllClose(expected, ans, check_dtypes=True)
+
+    a = onp.array([[1, 4], [3, 1]])
+    ans = lnp.sort(a, axis=None)
+    expected = onp.array([[1, 1, 3, 4]])
+    self.assertAllClose(expected, ans, check_dtypes=True)
+
+    a = onp.array([[1, 4], [3, 1]])
+    ans = lnp.sort(a)  # last axis
+    expected = onp.array([[1, 4], [1, 3]])
+    self.assertAllClose(expected, ans, check_dtypes=True)
+
+    a = onp.array([[1, 4], [3, 1]])
+    ans = lnp.sort(a, axis=0)
+    expected = onp.array([[1, 1], [3, 4]])
+    self.assertAllClose(expected, ans, check_dtypes=True)
+
+  def testArgsortManually(self):
+    x = onp.array([16, 15, 23, 42, 8, 4])
+    ans = lnp.argsort(x)
+    expected = onp.argsort(x)
+    self.assertAllClose(expected, ans, check_dtypes=False)
+
+    x = onp.array([[16, 15, 23], [42, 8, 4]])
+    ans = lnp.argsort(x, axis=0)
+    expected = onp.argsort(x, axis=0)
+    self.assertAllClose(expected, ans, check_dtypes=False)
+
+    x = onp.array([[16, 15, 23], [42, 8, 4]])
+    ans = lnp.argsort(x, axis=1)
+    expected = onp.argsort(x, axis=1)
+    self.assertAllClose(expected, ans, check_dtypes=False)
+
+    x = onp.array([[16, 15, 23], [42, 8, 4]])
+    ans = lnp.argsort(x, axis=None)
+    expected = onp.argsort(x, axis=None)
+    self.assertAllClose(expected, ans, check_dtypes=False)
+
+    x = onp.array([[16, 15, 23], [42, 8, 4]])
+    ans = lnp.argsort(x)
+    expected = onp.argsort(x)
+    self.assertAllClose(expected, ans, check_dtypes=False)
+
 
 if __name__ == "__main__":
   absltest.main()