     else {return max(a, b);}
 }
 
+#if __CUDACC_VER_MAJOR__ >= 9
+// specialization for half for handling NaNs
+template <>
+__host__ __device__ __forceinline__ __half Max::operator()(const __half &a, const __half &b) const
+{
+    // NumPy behavior: NaN is always chosen!
+    if (isnan(a)) {return a;}
+    else if (isnan(b)) {return b;}
+    else {return CUB_MAX(a, b);}
+}
+#endif
+
 //
 // Min()
 //

/* ------------------------------------ "Patches" to CUB ------------------------------------
   These stubs are needed because CUB does not handle NaNs properly, while NumPy has certain
   behaviors with which we must comply.
   TODO(leofang): support half precision?
*/

//

    else {return min(a, b);}
}

#if __CUDACC_VER_MAJOR__ >= 9
// specialization for half for handling NaNs
template <>
__host__ __device__ __forceinline__ __half Min::operator()(const __half &a, const __half &b) const
{
    // NumPy behavior: NaN is always chosen!
    if (isnan(a)) {return a;}
    else if (isnan(b)) {return b;}
    else {return CUB_MIN(a, b);}
}
#endif

//
// ArgMax()
//
