update compile benchmark

ml-explore · Feb 14, 2024 · 523b2d1 · 523b2d1
1 parent 469c5b2
commit 523b2d1
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/benchmarks/python/compile_bench.py b/benchmarks/python/compile_bench.py
@@ -17,7 +17,7 @@ def gelu(x):
 
     def gen_fun(fun):
         def bench_fun(x):
-            for _ in range(30):
+            for _ in range(10):
                 x = fun(x)
             return x
 
@@ -32,10 +32,9 @@ def randint():
     def gen_fun(fun):
         def bench_fun(x, y):
             x = x[: randint()]
-            y = y[: randint()]
-            for _ in range(30):
-                y = fun(y)
+            for _ in range(10):
                 x = fun(x)
+                y = fun(y)
             return x, y
 
         return bench_fun
@@ -65,11 +64,11 @@ def layernorm(x):
         x = x.astype(mx.float16)
         return weight * x + bias
 
-    x = mx.random.uniform(shape=(1000, 4096))
+    x = mx.random.uniform(shape=(1000, 4096)).astype(mx.float16)
 
     def gen_fun(fun):
         def bench_fun(x):
-            for _ in range(30):
+            for _ in range(10):
                 x = fun(x)
             return x
 
@@ -84,7 +83,7 @@ def randint():
     def gen_fun(fun):
         def bench_fun(x):
             x = x[: randint()]
-            for _ in range(30):
+            for _ in range(10):
                 x = fun(x)
             return x