L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix sort NaN handling for float16 and bfloat16 (#3269)

Browse Source
This commit is contained in:
Long Yixing
2026-03-19 15:19:41 -07:00
committed by GitHub
parent 5fa1a8d59f
commit 82809ebd12
4 changed files with 26 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cpu/sort.cpp
+7 -3
View File
@@ -15,10 +15,14 @@ namespace mlx::core {
namespace {
template <typename T>
inline constexpr bool is_floating_v = std::is_floating_point_v<T> ||
std::is_same_v<T, float16_t> || std::is_same_v<T, bfloat16_t>;
// NaN-aware comparator that places NaNs at the end
template <typename T>
bool nan_aware_less(T a, T b) {
if constexpr (std::is_floating_point_v<T> || std::is_same_v<T, complex64_t>) {
if constexpr (is_floating_v<T> || std::is_same_v<T, complex64_t>) {
if (std::isnan(a))
return false;
if (std::isnan(b))
@@ -198,7 +202,7 @@ void argsort(const array& in, array& out, int axis) {
auto v2 = data_ptr[b * in_stride];
// Handle NaNs (place them at the end)
if (std::is_floating_point<T>::value) {
if constexpr (is_floating_v<T>) {
if (std::isnan(v1))
return false;
if (std::isnan(v2))
@@ -299,7 +303,7 @@ void argpartition(const array& in, array& out, int axis, int kth) {
auto v2 = data_ptr[b * in_stride];
// Handle NaNs (place them at the end)
if (std::is_floating_point<T>::value) {
if constexpr (is_floating_v<T>) {
if (std::isnan(v1))
return false;
if (std::isnan(v2))
mlx/backend/cuda/kernel_utils.cuh
+2 -1
View File
@@ -89,7 +89,8 @@ using cuda_type_t = typename CTypeToCudaType<T>::type;
template <typename T>
inline constexpr bool is_floating_v =
cuda::std::is_same_v<T, float> || cuda::std::is_same_v<T, double> ||
cuda::std::is_same_v<T, float16_t> || cuda::std::is_same_v<T, bfloat16_t>;
cuda::std::is_same_v<T, float16_t> || cuda::std::is_same_v<T, bfloat16_t> ||
cuda::std::is_same_v<T, __half> || cuda::std::is_same_v<T, __nv_bfloat16>;
// Type traits for detecting complex numbers.
template <typename T>
mlx/backend/cuda/sort.cu
+2 -2
View File
@@ -52,7 +52,7 @@ struct InitValue {
};
template <typename T>
struct InitValue<T, cuda::std::enable_if_t<std::is_floating_point_v<T>>> {
struct InitValue<T, cuda::std::enable_if_t<is_floating_v<T>>> {
__device__ __forceinline__ static T value() {
return nan_value<T>();
}
@@ -72,7 +72,7 @@ struct LessThan {
}
__device__ __forceinline__ bool operator()(T a, T b) const {
if constexpr (std::is_floating_point_v<T>) {
if constexpr (is_floating_v<T>) {
bool an = cuda::std::isnan(a);
bool bn = cuda::std::isnan(b);
if (an | bn) {
python/tests/test_ops.py
+15 -3
View File
@@ -3311,11 +3311,23 @@ class TestOps(mlx_tests.MLXTestCase):
mx.broadcast_shapes()
def test_sort_nan(self):
x = mx.array([3.0, mx.nan, 2.0, 0.0])
expected = mx.array([0.0, 2.0, 3.0, mx.nan])
self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
for dtype in [mx.float32, mx.float16, mx.bfloat16]:
with self.subTest(dtype=dtype):
x = mx.array([3.0, mx.nan, 2.0, 0.0], dtype=dtype)
expected = mx.array([0.0, 2.0, 3.0, mx.nan], dtype=dtype)
self.assertTrue(mx.array_equal(mx.sort(x), expected, equal_nan=True))
x = mx.array([3.0, mx.nan, 2.0, 0.0]) + 1j * mx.array([1.0] * 4)
def test_argsort_nan(self):
for dtype in [mx.float32, mx.float16, mx.bfloat16]:
with self.subTest(dtype=dtype):
x = mx.array([3.0, mx.nan, 2.0, 0.0], dtype=dtype)
expected = mx.array([0.0, 2.0, 3.0, mx.nan], dtype=dtype)
indices = mx.argsort(x)
sorted_x = mx.take(x, indices)
self.assertTrue(mx.array_equal(sorted_x, expected, equal_nan=True))
def test_to_from_fp8(self):
vals = mx.array(
[448, 256, 192, 128, 96, 64, 48, 32, 24, 16, 12, 8, 6, 4, 3, 2, 0.015625]