.

Author: hylin2002

src/benchmark_utils.py

@@ -31,9 +31,15 @@ def simple_timeit(f, *args, matrix_dim=None, tries=10, task=None, trace_dir=None
     assert task is not None
 
     if trace_dir:
-        outcomes_ms = timeit_from_trace(f, *args, matrix_dim=matrix_dim, tries=tries, task=task, trace_dir=trace_dir)
-        if outcomes_ms is not None:
-            return outcomes_ms
+        try:
+            outcomes_ms = timeit_from_trace(
+                f, *args, matrix_dim=matrix_dim, tries=tries, task=task, trace_dir=trace_dir
+            )
+            if outcomes_ms is not None:
+                return outcomes_ms
+            print("Warning: timeit_from_trace returned empty results. Falling back to manual timing.")
+        except Exception as e:
+            print(f"Warning: Failed to get metrics from trace due to: {e}. Falling back to manual timing.")
 
     outcomes_ms = []
     jax.block_until_ready(f(*args))  # warm it up!