benchmark

e3nn · Jun 23, 2023 · 25b1ddf · 25b1ddf
1 parent 8611cab
commit 25b1ddf
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/examples/tensor_product_benchmark.py b/examples/tensor_product_benchmark.py
@@ -1,4 +1,5 @@
 import argparse
+import re
 import time
 
 import haiku as hk
@@ -160,7 +161,10 @@ def tp(x1, x2):
     print(f"{1e3 * perloop:.2f} ms")
 
     with open("xla.txt", "wt") as file:
-        file.write(jit_code(f, w, *inputs))
+        code = jit_code(f, w, *inputs)
+        code = re.sub(r"\d", "", code)
+
+        file.write(code)
 
 
 if __name__ == "__main__":