Add expander_name to stablehlo printer (#80)

pearu · Dec 21, 2024 · 7a84fef · 7a84fef
1 parent 2615811
commit 7a84fef
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/functional_algorithms/targets/stablehlo.py b/functional_algorithms/targets/stablehlo.py
@@ -143,8 +143,9 @@ def tostring(self, expr, tab=""):
             sargs = ", ".join(sargs)
 
             lines = []
+            expander_name = expr.props.get("expander_name", "")
             chlo_name = expr.props.get("name", f"CHLO_{name.ref.title()}")
-            lines.append(f"{tab}def : Pat<({chlo_name} {sargs}),")
+            lines.append(f"{tab}def {expander_name}: Pat<({chlo_name} {sargs}),")
 
             for line in self.tostring(body, tab=tab + "  ").splitlines():
                 sline = line.lstrip()

diff --git a/results/README.md b/results/README.md
@@ -84,7 +84,7 @@ Finally,
 | square | complex128 | 995705 | 6296 | - | - | - | - |
 | sqrt | float32 | 1000001 | - | - | - | - | using native sqrt |
 | sqrt | complex64 | 639573 | 362328 | 100 | - | - | using native sqrt |
-| sqrt<sub>2</sub> | complex64 | 639573 | 362328 | 100 | - | - | using native sqrt |
+| sqrt<sub>2</sub> | complex64 | 644997 | 356784 | 212 | 8 | - | - |
 | sqrt | complex128 | 637905 | 364008 | 88 | - | - | - |
 | angle | complex64 | 940289 | 61332 | 376 | 4 | - | - |
 | angle | complex128 | 989787 | 12214 | - | - | - | - |

diff --git a/results/estimate_accuracy.py b/results/estimate_accuracy.py
@@ -83,8 +83,8 @@ def get_inputs():
         # ("sqrt", np.float32, {}),     # real_sqrt is not implemented
         ("sqrt", np.float32, dict(use_native_sqrt=True)),
         # ("sqrt", np.float64, {}),     # real_sqrt is not implemented
-        ("sqrt", np.complex64, {}),
         ("sqrt", np.complex64, dict(use_native_sqrt=True)),
+        ("sqrt", np.complex64, {}),
         ("sqrt", np.complex128, {}),
         ("angle", np.complex64, {}),
         ("angle", np.complex128, {}),