[export] Conditionally load decomp table (#8209)

Co-authored-by: Siyuan Liu <lsiyuan@google.coim>
pytorch · 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)