add register pass

Author: erick-xanadu

pennylane/compiler/python_compiler/transforms/__init__.py

@@ -1,7 +1,8 @@
-from .apply_transform_sequence import ApplyTransformSequence
+from .apply_transform_sequence import ApplyTransformSequence, register_pass
 from .transform_interpreter import TransformInterpreterPass
 
 __all__ = [
         "ApplyTransformSequence",
         "TransformInterpreterPass",
+        "register_pass",
 ]

pennylane/compiler/python_compiler/transforms/apply_transform_sequence.py

@@ -6,6 +6,11 @@
 
 from .transform_interpreter import TransformInterpreterPass
 
+available_passes = {}
+
+def register_pass(name, _callable):
+    available_passes[name] = _callable
+
 @dataclass(frozen=True)
 class ApplyTransformSequence(ModulePass):
     name = "apply-transform-sequence"
@@ -18,7 +23,7 @@ def apply(self, ctx: MLContext, module: builtin.ModuleOp) -> None:
                     if isinstance(op, builtin.ModuleOp):
                         nested_modules.append(op)
 
-        pipeline = PipelinePass((TransformInterpreterPass(passes={}),))
+        pipeline = PipelinePass((TransformInterpreterPass(passes=available_passes),))
         for op in nested_modules:
             pipeline.apply(ctx, op)