Repository: cupy/cupy

Issue:
Make cupy.sort support arrays with rank two or more.
# Background
Arrays sorted with `cupy.sort` operation have some properties such as dtype, rank, sorting axis and C/F-contiguousness. Currently, `cupy.sort` supports sorting arrays only with the rank of one because of its implementation reason, see #55.

# Problem
This issue addresses a problem that makes `cupy.sort` support sorting arrays with the rank of two or more, with the last axis and C-contiguousness.

# Approach

**Rank two**

For an array with the rank of two, 

```
[[4, 3]
 [2, 1]]
```

treating the array as flattened one, `[4, 3, 2 ,1]`, and providing the following comparator in pseudo code to underlying Thrust library:

```
if floor(i / 2) < floor(j / 2) then return true;
else if floor(i / 2) > floor(j / 2) then return false;
else return data[i] < data[j];
```

where `i` and `j` are array indices, and `data[i]` represents `i` th element of array `data`,

we get the C-contiguous array sorted with the last axis.

```
[[3, 4]
 [1, 2]]
```

**Rank N**

Generalized to the rank of N with shape `(d_0, d_1, ..., d_n-1)`, the following comparator works:

```
if floor(i / d_n-1) < floor(j / d_n-1) then return true;
else if floor(i / d_n-1) > floor(j / d_n-1) then return false;
else return data[i] < data[j];
```



Discussion:
For sorting along an axis in middle of axes, we need modulo. 

I wrote a simple python code of the compare function and a simple test.

```python
def compare(data, i, j, axis):
    size_k = np.prod(data.shape[axis:])
    size_k1 = np.prod(data.shape[axis + 1:])

    if i // size_k < j // size_k:
        return True
    elif i // size_k > j // size_k:
        return False
    else:
        if i % size_k1 == j % size_k1:
            return data.reshape(-1)[i] < data.reshape(-1)[j]
        else:
            return i < j


# This is a quick test to verify that the compare function produces
# correct expected outputs using the values of sorted arrays.
import numpy as np
x = np.arange(250)
np.random.shuffle(x)
x = x.reshape(2, 5, 5, 5)

for k in range(4):
    sorted_x = np.sort(x, axis=k)
    for i in range(sorted_x.size):
        for j in range(sorted_x.size):
            if i < j:
                assert compare(sorted_x, i, j, axis=k), '{} {}'.format(i, j)
            if i > j:
                assert not compare(sorted_x, i, j, axis=k), '{} {}'.format(i, j)
```
The last branch of the comparer might be also with modulo as following?

```Python
else:
    return i % size_k1 < j % size_k1
```
Really?
I simplified the compare. I think the logic is clearer.

```python
def compare(data, i, j, axis):
    size_k = np.prod(data.shape[axis:])
    size_k1 = np.prod(data.shape[axis + 1:])

    if (i // size_k == j // size_k) and (i % size_k1 == j % size_k1):
        return data.reshape(-1)[i] < data.reshape(-1)[j]
    else:
        return i < j
```

Basically, when there is no need to compare by values in a list, the ordering follows the order of indices.

I look at this problem by reducing general rank N case to rank 3, where sort axis is in the middle.
By looking at it that way, the relation between linear indexing and three dim indexing is as follows.

```python
x.reshape(-1)[i] == x[i // size_k, i // size_k1, i % size_k1]
```
It becomes clear that the aforementioned condition `(i // size_k == j // size_k) and (i % size_k1 == j % size_k1)` is True only when the first and the last index are same for `i` and `j`.
Why I asked it is that sorting compares are required to satisfy not only relations between item pairs but relations across all items without inconsistency.

For example, when the comparer is applied to the following array:

    [[4, 3],
     [2, 1]]

it leads the inconsistent returns below:

- Compare 4 and 2 to get [2, 4] and compare 3 and 2 to get [3, 2], which lead to [3, 2, 4]
- Compare 4 and 3 to get [4, 3]

I actually tried the comparer with Python's built-in sort operation,

```Python
import numpy as np

def make_comparer(data, axis):
    def comparer(i, j):
        size_k = np.prod(data.shape[axis:])
        size_k1 = np.prod(data.shape[axis+1:])

        if i // size_k < j // size_k:
            return -1
        elif i // size_k > j // size_k:
            return 1
        else:
            if i % size_k1 == j % size_k1:
                return cmp(data.reshape(-1)[i], data.reshape(-1)[j])
            else:
                return cmp(i, j)
    return comparer

def my_sorted(data, cmp):
    idx = range(data.size)
    sorted_idx = sorted(idx, cmp=cmp)
    _data = data.reshape(-1)
    return np.array([_data[i] for i in sorted_idx]).reshape(data.shape)

x = np.array([4, 3, 2, 1])
x = x.reshape(2, 2)

x1 = my_sorted(x, cmp=make_comparer(x, 0))
x2 = np.sort(x, axis=0)

assert np.array_equal(x1, x2)
```

but the assertion fails.

Fixing the comparer's last line gives consistent returns, but it also causes unexpected result:

    [[2, 4],
     [1, 3]]

where the values with the same modulo are placed along the last axis, it's not the case.




Oh. You are right. The function I posted does not work because the order is not transitive.
@okuta suggested the second answer of the following SOF is a quite cleaver way to sort multiple arrays at once.

https://stackoverflow.com/questions/8298728/concurrently-sorting-many-arrays-with-cuda-thrust
For sorting along the last axis, it would be good to use the Back-to-back approach, because  division and modulo for each comparison and random memory access will suffer its efficiency.

For sorting along other axes but the last one, I have no idea but using rolling axes, so I will start implementation with the approach.

Resolving patch:
diff --git a/cupy/sorting/sort.py b/cupy/sorting/sort.py
--- a/cupy/sorting/sort.py
+++ b/cupy/sorting/sort.py
@@ -15,9 +15,9 @@ def sort(a):
         cupy.ndarray: Array of the same type and shape as ``a``.
 
     .. note::
-       For its implementation reason, ``cupy.sort`` currently supports only
-       arrays with their rank of one and does not support ``axis``, ``kind``
-       and ``order`` parameters that ``numpy.sort`` does support.
+       For its implementation reason, ``cupy.sort`` currently does not support
+       ``axis``, ``kind`` and ``order`` parameters that ``numpy.sort`` does
+       support.
 
     .. seealso:: :func:`numpy.sort`
 


Test patch:
diff --git a/tests/cupy_tests/sorting_tests/test_sort.py b/tests/cupy_tests/sorting_tests/test_sort.py
--- a/tests/cupy_tests/sorting_tests/test_sort.py
+++ b/tests/cupy_tests/sorting_tests/test_sort.py
@@ -23,15 +23,16 @@ def test_external_sort_zero_dim(self, xp):
         a = testing.shaped_random((), xp)
         return xp.sort(a)
 
-    def test_sort_two_or_more_dim(self):
-        a = testing.shaped_random((2, 3), cupy)
-        with self.assertRaises(ValueError):
-            a.sort()
+    @testing.numpy_cupy_array_equal()
+    def test_sort_two_or_more_dim(self, xp):
+        a = testing.shaped_random((2, 3, 3), xp)
+        a.sort()
+        return a
 
-    def test_external_sort_two_or_more_dim(self):
-        a = testing.shaped_random((2, 3), cupy)
-        with self.assertRaises(ValueError):
-            return cupy.sort(a)
+    @testing.numpy_cupy_array_equal()
+    def test_external_sort_two_or_more_dim(self, xp):
+        a = testing.shaped_random((2, 3, 3), xp)
+        return xp.sort(a)
 
     # Test dtypes