[export] Conditionally load decomp table (pytorch#8209)

Co-authored-by: Siyuan Liu <[email protected]>
dvhg · Oct 3, 2024 · c22dfe9 · c22dfe9
1 parent 4813702
commit c22dfe9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torch_xla/stablehlo.py b/torch_xla/stablehlo.py
@@ -320,7 +320,10 @@ def _extract_input_args(exported_model, options):
 
 
 def _run_decompositions(exported_model):
-  decomp_table = core_aten_decompositions()
+  if hasattr(torch.export, 'core_op_decompositions'):
+    decomp_table = torch.export.core_op_decompositions()
+  else:
+    decomp_table = core_aten_decompositions()
   decomp_table.update(_extra_decompositions)
   decomp_table[torch.ops.aten._safe_softmax.default] = torch.softmax
   exported_model = exported_model.run_decompositions(decomp_table)