From ac99ecd94ed27db261cf2d09897145b29bcf55a0 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Sat, 20 Jul 2024 08:37:32 -0700
Subject: [PATCH] use statistics.median for timing (#5606)

---
 extra/mcts_search.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/extra/mcts_search.py b/extra/mcts_search.py
index 1b0fb9eea9..8181c22241 100644
--- a/extra/mcts_search.py
+++ b/extra/mcts_search.py
@@ -1,5 +1,5 @@
 from typing import List, Optional
-import math, functools, time, random
+import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put
 from tinygrad.codegen.kernel import Kernel
 from tinygrad.device import Buffer, Device
@@ -61,7 +61,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
       continue
 
     p, lib, _ = compile_ret
-    try: tm = min(_time_program(p, lib, var_vals, rawbufs, early_stop=best_tm*10/1e6))*1e6
+    try: tm = statistics.median(_time_program(p, lib, var_vals, rawbufs, cnt=5, early_stop=best_tm*10/1e6))*1e6
     except RuntimeError:
       remove_node(node)
       continue