Repository: cupy/cupy

Issue:
Lexsort output is incorrect
The output of cupy.lexsort is incorrect. I can't see any pattern as to what is wrong with the output.

### Code to reproduce
```
import cupy
import numpy

# numpy output is correct:
a = (numpy.random.random((10,2)) * 10).astype(int)  # random ints between 0 and 9
print(a[numpy.lexsort(a.T[::-1])])                  # sorted by first column, not last

# cupy output is unsorted:
b = cupy.array(a)
print(b[cupy.lexsort(b.T[::-1])])
```
The same thing happens for floats, so that's not the reason.

### Conditions
(output of `python -c 'import cupy; cupy.show_config()'`):

CuPy Version          : 7.1.1
CUDA Root             : /opt/cuda
CUDA Build Version    : 10020
CUDA Driver Version   : 10020
CUDA Runtime Version  : 10020
cuBLAS Version        : 10202
cuFFT Version         : 10102
cuRAND Version        : 10102
cuSOLVER Version      : (10, 3, 0)
cuSPARSE Version      : 10301
NVRTC Version         : (10, 2)
cuDNN Build Version   : 7605
cuDNN Version         : 7605
NCCL Build Version    : 2507
NCCL Runtime Version  : 2507


Discussion:
Thanks for the report, I can reproduce the result.
Not sure if this is relevant, but lexsort would fail these two tests too:
```diff
diff --git a/tests/cupy_tests/sorting_tests/test_sort.py b/tests/cupy_tests/sorting_tests/test_sort.py
index b00856866..5b0c67d58 100644
--- a/tests/cupy_tests/sorting_tests/test_sort.py
+++ b/tests/cupy_tests/sorting_tests/test_sort.py
@@ -202,6 +202,20 @@ class TestLexsort(unittest.TestCase):
         with self.assertRaises(TypeError):
             return cupy.lexsort(a)
 
+    @testing.for_dtypes('fdFD')
+    @testing.numpy_cupy_allclose()
+    def test_nan2(self, xp, dtype):
+        a = testing.shaped_random((2, 10), xp, dtype)
+        a[0, 2] = a[1, 6] = xp.nan
+        return xp.lexsort(a)
+
+    @testing.for_dtypes('fdFD')
+    @testing.numpy_cupy_allclose()
+    def test_nan3(self, xp, dtype):
+        a = testing.shaped_random((2, 10), xp, dtype)
+        a[1, 2] = a[1, 6] = xp.nan
+        return xp.lexsort(a)
+
 
 @testing.parameterize(*testing.product({
     'external': [False, True],
```
lexsort with `a[0, 2] = a[0, 6] = xp.nan` is correct, though.

Resolving patch:
diff --git a/cupy/_sorting/sort.py b/cupy/_sorting/sort.py
--- a/cupy/_sorting/sort.py
+++ b/cupy/_sorting/sort.py
@@ -73,6 +73,10 @@ def lexsort(keys):
         raise NotImplementedError('Keys with the rank of three or more is not '
                                   'supported in lexsort')
 
+    # thrust.lexsort() assumes a C-contiguous array
+    if not keys.flags.c_contiguous:
+        keys = keys.copy('C')
+
     idx_array = cupy.ndarray(keys._shape[1:], dtype=numpy.intp)
     k = keys._shape[0]
     n = keys._shape[1]
@@ -120,7 +124,6 @@ def msort(a):
 
     """
 
-    # TODO(takagi): Support float16 and bool.
     return sort(a, axis=0)
 
 


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
@@ -38,21 +38,21 @@ def test_external_sort_two_or_more_dim(self, xp):
     # Test dtypes
 
     @testing.for_all_dtypes()
-    @testing.numpy_cupy_allclose()
+    @testing.numpy_cupy_array_equal()
     def test_sort_dtype(self, xp, dtype):
         a = testing.shaped_random((10,), xp, dtype)
         a.sort()
         return a
 
     @testing.for_all_dtypes()
-    @testing.numpy_cupy_allclose()
+    @testing.numpy_cupy_array_equal()
     def test_external_sort_dtype(self, xp, dtype):
         a = testing.shaped_random((10,), xp, dtype)
         return xp.sort(a)
 
     # Test contiguous arrays
 
-    @testing.numpy_cupy_allclose()
+    @testing.numpy_cupy_array_equal()
     def test_sort_contiguous(self, xp):
         a = testing.shaped_random((10,), xp)  # C contiguous view
         a.sort()
@@ -63,12 +63,12 @@ def test_sort_non_contiguous(self):
         with self.assertRaises(NotImplementedError):
             a.sort()
 
-    @testing.numpy_cupy_allclose()
+    @testing.numpy_cupy_array_equal()
     def test_external_sort_contiguous(self, xp):
         a = testing.shaped_random((10,), xp)  # C contiguous view
         return xp.sort(a)
 
-    @testing.numpy_cupy_allclose()
+    @testing.numpy_cupy_array_equal()
     def test_external_sort_non_contiguous(self, xp):
         a = testing.shaped_random((10,), xp)[::2]  # Non contiguous view
         return xp.sort(a)
@@ -76,11 +76,23 @@ def test_external_sort_non_contiguous(self, xp):
     # Test axis
 
     @testing.numpy_cupy_array_equal()
-    def test_sort_axis(self, xp):
-        a = testing.shaped_random((2, 3, 3), xp)
+    def test_sort_axis1(self, xp):
+        a = testing.shaped_random((2, 3, 4), xp)
         a.sort(axis=0)
         return a
 
+    @testing.numpy_cupy_array_equal()
+    def test_sort_axis2(self, xp):
+        a = testing.shaped_random((2, 3, 4), xp)
+        a.sort(axis=1)
+        return a
+
+    @testing.numpy_cupy_array_equal()
+    def test_sort_axis3(self, xp):
+        a = testing.shaped_random((2, 3, 4), xp)
+        a.sort(axis=2)
+        return a
+
     @testing.numpy_cupy_array_equal()
     def test_external_sort_axis(self, xp):
         a = testing.shaped_random((2, 3, 3), xp)
@@ -149,12 +161,36 @@ def test_external_sort_invalid_negative_axis2(self):
     # Test NaN ordering
 
     @testing.for_dtypes('efdFD')
-    @testing.numpy_cupy_allclose()
-    def test_nan(self, xp, dtype):
+    @testing.numpy_cupy_array_equal()
+    def test_nan1(self, xp, dtype):
         a = testing.shaped_random((10,), xp, dtype)
         a[2] = a[6] = xp.nan
-        a.sort()
-        return a
+        out = xp.sort(a)
+        return out
+
+    @testing.for_dtypes('efdFD')
+    @testing.numpy_cupy_array_equal()
+    def test_nan2(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        a[0, 2, 1] = a[1, 0, 3] = xp.nan
+        out = xp.sort(a, axis=0)
+        return out
+
+    @testing.for_dtypes('efdFD')
+    @testing.numpy_cupy_array_equal()
+    def test_nan3(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        a[0, 2, 1] = a[1, 0, 3] = xp.nan
+        out = xp.sort(a, axis=1)
+        return out
+
+    @testing.for_dtypes('efdFD')
+    @testing.numpy_cupy_array_equal()
+    def test_nan4(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        a[0, 2, 1] = a[1, 0, 3] = xp.nan
+        out = xp.sort(a, axis=2)
+        return out
 
 
 @testing.gpu
@@ -189,23 +225,51 @@ def test_lexsort_three_or_more_dim(self):
     # Test dtypes
 
     @testing.for_all_dtypes()
-    @testing.numpy_cupy_allclose()
+    @testing.numpy_cupy_array_equal()
     def test_lexsort_dtype(self, xp, dtype):
         a = testing.shaped_random((2, 10), xp, dtype)
         return xp.lexsort(a)
 
     # Test NaN ordering
-    # TODO(leofang): test two more cases after #3232 is fixed
-    # 1. a[1, 2] = a[1, 6] = xp.nan
-    # 2. a[0, 2] = a[1, 6] = xp.nan
 
     @testing.for_dtypes('efdFD')
-    @testing.numpy_cupy_allclose()
-    def test_nan(self, xp, dtype):
+    @testing.numpy_cupy_array_equal()
+    def test_nan1(self, xp, dtype):
         a = testing.shaped_random((2, 10), xp, dtype)
         a[0, 2] = a[0, 6] = xp.nan
         return xp.lexsort(a)
 
+    @testing.for_dtypes('efdFD')
+    @testing.numpy_cupy_array_equal()
+    def test_nan2(self, xp, dtype):
+        a = testing.shaped_random((2, 10), xp, dtype)
+        a[1, 2] = a[0, 6] = xp.nan
+        return xp.lexsort(a)
+
+    @testing.for_dtypes('efdFD')
+    @testing.numpy_cupy_array_equal()
+    def test_nan3(self, xp, dtype):
+        a = testing.shaped_random((2, 10), xp, dtype)
+        a[1, 2] = a[1, 6] = xp.nan
+        return xp.lexsort(a)
+
+    # Test non C-contiguous input
+
+    @testing.numpy_cupy_array_equal()
+    def test_view(self, xp):
+        # from #3232
+        a = testing.shaped_random((4, 8), xp, dtype=xp.float64)
+        a = a.T[::-1]
+        return xp.lexsort(a)
+
+    @testing.numpy_cupy_array_equal()
+    def test_F_order(self, xp):
+        a = testing.shaped_random((4, 8), xp, dtype=xp.float64)
+        a = xp.asfortranarray(a)
+        assert a.flags.f_contiguous
+        assert not a.flags.c_contiguous
+        return xp.lexsort(a)
+
 
 @testing.parameterize(*testing.product({
     'external': [False, True],
@@ -312,11 +376,18 @@ def test_argsort_original_array_not_modified_multi_dim(self):
     # Test NaN ordering
 
     @testing.for_dtypes('efdFD')
-    @testing.numpy_cupy_allclose()
-    def test_nan(self, xp, dtype):
+    @testing.numpy_cupy_array_equal()
+    def test_nan1(self, xp, dtype):
         a = testing.shaped_random((10,), xp, dtype)
         a[2] = a[6] = xp.nan
-        return xp.argsort(a)
+        return self.argsort(a)
+
+    @testing.for_dtypes('efdFD')
+    @testing.numpy_cupy_array_equal()
+    def test_nan2(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        a[0, 2, 1] = a[1, 1, 3] = xp.nan
+        return self.argsort(a)
 
 
 @testing.gpu