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 Packages Projects Releases Wiki Activity

Adding benchmarks and testing for max op nanpropagation

Browse Source
This commit is contained in:
Joona Havukainen
2025-07-06 14:27:40 -07:00
parent 0d30e9e8ec
commit af74818528
4 changed files with 43 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
benchmarks/cpp/single_ops.cpp
View File

@@ -192,6 +192,17 @@ void time_reductions() {
auto argmin_along_1 = [&a]() { return mx::argmin(a, 1, false); };
TIME(argmin_along_1);
auto indices = mlx::core::array({1});
auto updates = mlx::core::reshape(mlx::core::array({NAN}), {1, 1, 1});
std::vector<int> axes{0};
auto b = scatter(a, {indices}, updates, axes);
mx::eval(b);
auto max_along_0 = [&b]() { return mx::max(b, 0, false); };
TIME(max_along_0);
auto max_along_1 = [&b]() { return mx::max(b, 1, false); };
TIME(max_along_1);
}
void time_gather_scatter() {

6
benchmarks/python/single_ops.py
View File

@@ -50,6 +50,11 @@ def time_maximum():
mx.eval(a, b)
time_fn(mx.maximum, a, b)
def time_max():
a = mx.random.uniform(shape=(32, 1024, 1024))
a[1,1] = mx.nan
mx.eval(a)
time_fn(mx.max, a, 0)
def time_negative():
a = mx.random.uniform(shape=(10000, 1000))
@@ -108,6 +113,7 @@ if __name__ == "__main__":
time_add()
time_matmul()
time_max()
time_maximum()
time_exp()
time_negative()