From 31ef8ddb8c4467f5b8698ef1eb9bb8bab7056855 Mon Sep 17 00:00:00 2001
From: Wei Wei <wwei6@fb.com>
Date: Wed, 7 Sep 2022 17:21:27 +0000
Subject: [PATCH] add option to remove passes (#84425)

Summary:
Add a remove_pass method in pass_manager to provide user option to remove any pass.

Reviewed By: wushirong

Differential Revision: D39080077

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84425
Approved by: https://github.com/yinghai
---
 torch/fx/passes/pass_manager.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/torch/fx/passes/pass_manager.py b/torch/fx/passes/pass_manager.py
index 52f1170290db5..5a34c5bca3621 100644
--- a/torch/fx/passes/pass_manager.py
+++ b/torch/fx/passes/pass_manager.py
@@ -212,6 +212,16 @@ def add_constraint(self, constraint):
         self.constraints.append(constraint)
         self._validated = False
 
+    def remove_pass(self, _passes: List[Callable]):
+        if _passes is None:
+            return
+        passes_left = []
+        for ps in self.passes:
+            if ps.__name__ not in _passes:
+                passes_left.append(ps)
+        self.passes = passes_left
+        self._validated = False
+
     def validate(self):
         """
         Validates that current pass schedule defined by `self.passes` is valid