Fix sort when dtype is uint16 (#563)

Prior to this fix the code was using the mp_float_t data type for uint16 and producing incorrect sort results. Signed-off-by: Damien George <damien@micropython.org> Signed-off-by: Damien George <damien@micropython.org>
v923z · Nov 29, 2022 · 1a440d7 · 1a440d7
1 parent 25a825e
commit 1a440d7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/code/numpy/numerical.c b/code/numpy/numerical.c
@@ -650,7 +650,7 @@ static mp_obj_t numerical_sort_helper(mp_obj_t oin, mp_obj_t axis, uint8_t inpla
     if(ndarray->shape[ax]) {
         if((ndarray->dtype == NDARRAY_UINT8) || (ndarray->dtype == NDARRAY_INT8)) {
             HEAPSORT(ndarray, uint8_t, array, shape, strides, ax, increment, ndarray->shape[ax]);
-        } else if((ndarray->dtype == NDARRAY_INT16) || (ndarray->dtype == NDARRAY_INT16)) {
+        } else if((ndarray->dtype == NDARRAY_UINT16) || (ndarray->dtype == NDARRAY_INT16)) {
             HEAPSORT(ndarray, uint16_t, array, shape, strides, ax, increment, ndarray->shape[ax]);
         } else {
             HEAPSORT(ndarray, mp_float_t, array, shape, strides, ax, increment, ndarray->shape[ax]);

diff --git a/tests/2d/numpy/sort.py.exp b/tests/2d/numpy/sort.py.exp
@@ -11,7 +11,7 @@ array([1, 2, 3, 4], dtype=int8)
 
 array([], dtype=uint16)
 []
-array([0, 0, 0, 0], dtype=uint16)
+array([1, 2, 3, 4], dtype=uint16)
 [1, 3, 2, 0]
 
 array([], dtype=int16)