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

Reductions update (#1351)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-11-04 22:25:16 -08:00
committed by GitHub
parent 76f275b4df
commit 248431eb3c
8 changed files with 462 additions and 206 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
benchmarks/python/comparative/bench_mlx.py
View File

@@ -144,6 +144,13 @@ def reduction(op, axis, x):
mx.eval(ys)
def sum_and_add(axis, x, y):
z = x.sum(axis=axis, keepdims=True)
for i in range(50):
z = (z + y).sum(axis=axis, keepdims=True)
mx.eval(z)
def softmax(axis, x):
ys = []
for i in range(100):
@@ -505,5 +512,8 @@ if __name__ == "__main__":
elif args.benchmark == "selu":
print(bench(selu, x))
elif args.benchmark == "sum_and_add":
print(bench(sum_and_add, axis, *xs))
else:
raise ValueError("Unknown benchmark")