Use default method in stablehlo txt/bytecode getter of StableHLOGraph…

…Module (pytorch#5745) Co-authored-by: Siyuan Liu <[email protected]>
chunnienc · Dec 14, 2023 · 5c481db · 5c481db
1 parent 9a51ab5
commit 5c481db
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/torch_xla/stablehlo.py b/torch_xla/stablehlo.py
@@ -73,10 +73,14 @@ def evaluate(self, method_name, args):
       res = pytree.tree_unflatten(res, out_spec)
     return res
 
-  def get_stablehlo_bytecode(self, method_name):
+  def get_stablehlo_bytecode(self, method_name=None):
+    if method_name is None:
+      method_name = self._default_method
     return self._name_to_stablehlo[method_name].bytecode
 
-  def get_stablehlo_text(self, method_name):
+  def get_stablehlo_text(self, method_name=None):
+    if method_name is None:
+      method_name = self._default_method
     return self._name_to_stablehlo[method_name].text
 
   def save(self, directory_path):