Update on "Don't decompose functional ops in predispatch functionalization"

tugsbayasgalan · tugsbayasgalan · commit 852881b96872 · 2023-12-27T18:58:59.000-08:00
[ghstack-poisoned]
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
@@ -375,7 +375,6 @@ def _export_non_strict(
     *,
     transform=lambda x: x,  # TODO(zhxchen17) Revisit if this is needed later.
     pre_dispatch=False,
-    decomp_table=None,
 ):
     # This _reparametrize_module makes sure inputs and module.params/buffers have the same fake_mode,
     # otherwise aot_export_module will error out because it sees a mix of fake_modes.
@@ -386,7 +385,6 @@ def _export_non_strict(
             (*fake_args, *fake_kwargs.values()),
             trace_joint=False,
             pre_dispatch=pre_dispatch,
-            decompositions=decomp_table,
         )
 
     # NOTE: aot_export adds symint metadata for placeholders with int values;
@@ -710,7 +708,6 @@ def _aot_export_strict(gm_torch_level: torch.fx.GraphModule, args, **kwargs):
         fake_params_buffers,
         transform=_process_user_inputs,
         pre_dispatch=pre_dispatch,
-        decomp_table=decomp_table,
     )
 
     gm = ep_non_strict.gm
