-#include <cupy/complex.cuh>
 #include <cupy/type_dispatcher.cuh>
 #include <thrust/device_ptr.h>
 #include <thrust/device_vector.h>

    }

    // pick 1 from 3 and compare the numerical values (nan+nan*I compares to itself as false)
    return (((lhsIm && rhsIm) && (lhs.x < rhs.x)) || ((lhsRe && rhsRe) && (lhs.y < rhs.y)));
}

/*
 * Unfortunately we need explicit (instead of templated) definitions here, because the template specializations would
 * go through some wild routes in Thrust that passing by reference to device functions is not working...
 */

__host__ __device__ __forceinline__ bool operator<(const cuComplex& lhs, const cuComplex& rhs) {
    return _cmp_less(lhs, rhs);
}

__host__ __device__ __forceinline__ bool operator<(const cuDoubleComplex& lhs, const cuDoubleComplex& rhs) {
    return _cmp_less(lhs, rhs);
}

/*

 * Ref: https://docs.scipy.org/doc/numpy/reference/generated/numpy.sort.html
 */

template <typename T>
__host__ __device__ __forceinline__ bool _less(const T& lhs, const T& rhs) {
    return false;  // just a placeholder to be overwritten by specializations
}

template <typename T>
__host__ __device__ __forceinline__ bool _tuple_less(const tuple<size_t, T>& lhs,
                                                     const tuple<size_t, T>& rhs) {
    const size_t& lhs_k = lhs.template get<0>();
    const size_t& rhs_k = rhs.template get<0>();
    const T& lhs_v = lhs.template get<1>();
    const T& rhs_v = rhs.template get<1>();

    // tuple's comparison rule: compare the 1st member, then 2nd, then 3rd, ...,
    // which should be respected
    if (lhs_k < rhs_k) {
        return true;
    } else if (lhs_k == rhs_k) {
        // same key, compare values
        // note that we can't rely on native operator< due to NaN, so use _less() above
        return _less<T>(lhs_v, rhs_v);
    } else {
        return false;
    }
}

/*
 * ********** complex numbers **********
 * We need to specialize thrust::less because obviously we can't overload operator< for complex numbers...
 */

template <typename T>
__host__ __device__ __forceinline__ bool _cmp_less(const T& lhs, const T& rhs) {
    bool lhsRe = isnan(lhs.real());
    bool lhsIm = isnan(lhs.imag());
    bool rhsRe = isnan(rhs.real());
    bool rhsIm = isnan(rhs.imag());

    // neither side has nan
    if (!lhsRe && !lhsIm && !rhsRe && !rhsIm) {
        return lhs < rhs;
    }

    // one side has nan, and the other does not
