zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

divmod, partition, sort fixes (#2302)

Browse Source
This commit is contained in:
Awni Hannun
2025-06-16 18:49:32 -07:00
committed by GitHub
parent bc53f8293f
commit b8022c578a
8 changed files with 271 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/cuda/device/binary_ops.cuh
View File

@@ -22,7 +22,7 @@ struct FloorDivide {
if constexpr (cuda::std::is_integral_v<T>) {
return x / y;
} else {
return trunc(x / y);
return truncf(x / y);
}
}
};
@@ -132,7 +132,7 @@ struct LogAddExp {
cuda::std::numeric_limits<float>::quiet_NaN(),
cuda::std::numeric_limits<float>::quiet_NaN()};
}
constexpr float inf = cuda::std::numeric_limits<float>::infinity();
float inf = cuda::std::numeric_limits<float>::infinity();
auto maxval = x > y ? x : y;
auto minval = x < y ? x : y;
if (cuCrealf(minval) == -inf || cuCrealf(maxval) == inf)

2
mlx/backend/cuda/device/config.h
View File

@@ -5,7 +5,7 @@
#pragma once
// The maximum dimensions of shape/strides passed as kernel parameters.
#define MAX_NDIM 8
#define MAX_NDIM 10
// All existing NVIDIA hardware has a fixed 32 warp size. Though a built-in
// warpSize variable exists, using it would prevent compile-time optimizations.