ENH: core: Update the sorting module for half/float16

2 files changed, 21 insertions, 7 deletions

diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index c2a607d62..b854e879f 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -662,7 +662,7 @@ def configuration(parent_package='',top_path=None):
                                   generate_numpyconfig_h,
                                   generate_numpy_api,
                                   ],
-                         )
+                         libraries=['npymath'])
 
     # npymath needs the config.h and numpyconfig.h files to be generated, but
     # build_clib cannot handle generate_config_h and generate_numpyconfig_h
diff --git a/numpy/core/src/_sortmodule.c.src b/numpy/core/src/_sortmodule.c.src
index 709a26a1c..527d0c402 100644
--- a/numpy/core/src/_sortmodule.c.src
+++ b/numpy/core/src/_sortmodule.c.src
@@ -30,6 +30,7 @@
 #include "Python.h"
 #include "numpy/noprefix.h"
 #include "numpy/npy_math.h"
+#include "numpy/halffloat.h"
 
 #include "npy_config.h"
 
@@ -48,12 +49,12 @@
 /**begin repeat
  *
  * #TYPE = BOOL, BYTE, UBYTE, SHORT, USHORT, INT, UINT, LONG, ULONG,
- *         LONGLONG, ULONGLONG, FLOAT, DOUBLE, LONGDOUBLE, CFLOAT,
+ *         LONGLONG, ULONGLONG, HALF, FLOAT, DOUBLE, LONGDOUBLE, CFLOAT,
  *         CDOUBLE,CLONGDOUBLE, INTP#
  * #type = npy_bool, npy_byte, npy_ubyte, npy_short, npy_ushort, npy_int,
  *         npy_uint, npy_long, npy_ulong, npy_longlong, npy_ulonglong,
- *         npy_float, npy_double, npy_longdouble, npy_cfloat, npy_cdouble,
- *         npy_clongdouble, npy_intp#
+ *         npy_half, npy_float, npy_double, npy_longdouble, npy_cfloat,
+ *         npy_cdouble, npy_clongdouble, npy_intp#
  */
 #define @TYPE@_SWAP(a,b) {@type@ tmp = (b); (b)=(a); (a) = tmp;}
 
@@ -92,6 +93,19 @@ NPY_INLINE static int
 }
 /**end repeat**/
 
+NPY_INLINE static int
+HALF_LT(npy_half a, npy_half b)
+{
+    int ret;
+
+    if (npy_half_isnan(b)) {
+        ret = !npy_half_isnan(a);
+    } else {
+        ret = !npy_half_isnan(a) && npy_half_lt_nonan(a, b);
+    }
+
+    return ret;
+}
 
 /*
  * For inline functions SUN recommends not using a return in the then part
@@ -217,10 +231,10 @@ UNICODE_LT(npy_ucs4 *s1, npy_ucs4 *s2, size_t len)
 /**begin repeat
  *
  * #TYPE = BOOL, BYTE, UBYTE, SHORT, USHORT, INT, UINT, LONG, ULONG,
- *         LONGLONG, ULONGLONG, FLOAT, DOUBLE, LONGDOUBLE,
+ *         LONGLONG, ULONGLONG, HALF, FLOAT, DOUBLE, LONGDOUBLE,
  *         CFLOAT, CDOUBLE, CLONGDOUBLE#
  * #type = Bool, byte, ubyte, short, ushort, int, uint, long, ulong,
- *         longlong, ulonglong, float, double, longdouble,
+ *         longlong, ulonglong, ushort, float, double, longdouble,
  *         cfloat, cdouble, clongdouble#
  */
 
@@ -965,7 +979,7 @@ add_sortfuncs(void)
     /**begin repeat
      *
      * #TYPE = BOOL, BYTE, UBYTE, SHORT, USHORT, INT, UINT, LONG, ULONG,
-     *         LONGLONG, ULONGLONG, FLOAT, DOUBLE, LONGDOUBLE,
+     *         LONGLONG, ULONGLONG, HALF, FLOAT, DOUBLE, LONGDOUBLE,
      *         CFLOAT, CDOUBLE, CLONGDOUBLE, STRING, UNICODE#
      */
     descr = PyArray_DescrFromType(PyArray_@TYPE@);