Skip to content

Commit

Permalink
[Export] Allow ExportedProgram to take empty decomp table (#126142)
Browse files Browse the repository at this point in the history
**As title.**
Still, `ep.run_decompositions()` will use `core_aten_decompositions()` by default. Cases like `ep.run_decompositions(get_decompositions([]))` will use empty table, and go with [`aot_autograd_decompositions`](https://github.com/pytorch/pytorch/blob/04877dc430a6e93765471b28f422bf3e81d02c9e/torch/_functorch/aot_autograd.py#L456-459) only.

**Motivation**
We didn't have a clean way to pass in an empty decomp table. Since we've made `pre_dispatch` export as default and `ep.run_decompositions` remains with `aot_export_module(..., pre_dispatch=False)`, allowing empty table would help make blank control easier.

**Testing**
CI
Also looked through all the references in fbcode. The only concern I have is whether we should update [this example](https://github.com/pytorch/pytorch/blob/04877dc430a6e93765471b28f422bf3e81d02c9e/torch/onnx/_internal/exporter.py#L817) or not.
Pull Request resolved: #126142
Approved by: https://github.com/angelayi
  • Loading branch information
StellarrZ authored and pytorchmergebot committed May 16, 2024
1 parent b3f1882 commit 51e9bb8
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion torch/export/exported_program.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,7 +547,8 @@ def _get_placeholders(gm):
placeholders.append(node)
return placeholders

decomp_table = decomp_table or core_aten_decompositions()
if decomp_table is None:
decomp_table = core_aten_decompositions()

old_placeholders = _get_placeholders(self.graph_module)
fake_args = [node.meta["val"] for node in old_placeholders]
Expand Down

0 comments on commit 51e9bb8

Please sign in to comment.