From 71500bab827d3b1a4780f8ec17a6e85e87206353 Mon Sep 17 00:00:00 2001
From: Erick Ochoa Lopez <erick.ochoalopez@xanadu.ai>
Date: Thu, 9 Jan 2025 16:25:37 -0500
Subject: [PATCH] Documentation

---
 frontend/catalyst/passes/pass_api.py          | 32 ++++++++++++++++++-
 frontend/test/lit/test_mlir_plugin.py         |  2 ++
 .../test/pytest/test_mlir_plugin_interface.py |  6 +++-
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/frontend/catalyst/passes/pass_api.py b/frontend/catalyst/passes/pass_api.py
index 4d797088c5..09784ce7be 100644
--- a/frontend/catalyst/passes/pass_api.py
+++ b/frontend/catalyst/passes/pass_api.py
@@ -239,6 +239,20 @@ class Pass:
     :class:`Pass` will be used when generating `ApplyRegisteredPassOp`s.
     The attribute `pass_name` corresponds to the field `name`.
     The attribute `options` is generated by the `get_options` method.
+
+    People working on MLIR plugins may use this or :class:`PassPlugin` to
+    schedule their compilation pass. E.g.,
+
+    .. code-block:: python
+
+        def an_optimization(qnode):
+            @functools.wraps(qnode)
+            def wrapper(*args, **kwargs):
+                pass_pipeline = kwargs.pop("pass_pipeline", [])
+                pass_pipeline.append(Pass("my_library.my_optimization", *args, **kwargs))
+                kwargs["pass_pipeline"] = pass_pipeline
+                return qnode(*args, **kwargs)
+        return wrapper
     """
 
     def __init__(self, name: str, *options: list[str], **valued_options: dict[str, str]):
@@ -277,7 +291,9 @@ def __repr__(self):
         return (
             self.name
             + " ".join(f"--{str(option)}" for option in self.options)
-            + " ".join([f"--{str(option)}={str(value)}" for option, value in self.valued_options.items()])
+            + " ".join(
+                [f"--{str(option)}={str(value)}" for option, value in self.valued_options.items()]
+            )
         )
 
 
@@ -288,6 +304,20 @@ class PassPlugin(Pass):
     E.g.,
 
       --pass-plugin=path/to/plugin --dialect-plugin=path/to/plugin
+
+    People working on MLIR plugins may use this or :class:`Pass` to
+    schedule their compilation pass. E.g.,
+
+    .. code-block:: python
+
+        def an_optimization(qnode):
+            @functools.wraps(qnode)
+            def wrapper(*args, **kwargs):
+                pass_pipeline = kwargs.pop("pass_pipeline", [])
+                pass_pipeline.append(PassPlugin(path_to_plugin, "my_optimization", *args, **kwargs))
+                kwargs["pass_pipeline"] = pass_pipeline
+                return qnode(*args, **kwargs)
+        return wrapper
     """
 
     def __init__(
diff --git a/frontend/test/lit/test_mlir_plugin.py b/frontend/test/lit/test_mlir_plugin.py
index 26b8daa2f8..ccc2e04b6d 100644
--- a/frontend/test/lit/test_mlir_plugin.py
+++ b/frontend/test/lit/test_mlir_plugin.py
@@ -64,6 +64,7 @@ def module():
 from pathlib import Path
 
 import pennylane as qml
+
 import catalyst
 from catalyst.compiler import CompileOptions, Compiler
 from catalyst.utils.filesystem import WorkspaceManager
@@ -112,4 +113,5 @@ def example():
 
     print(example.mlir)
 
+
 test_pass_options()
diff --git a/frontend/test/pytest/test_mlir_plugin_interface.py b/frontend/test/pytest/test_mlir_plugin_interface.py
index 20bee0f63c..d476fc3032 100644
--- a/frontend/test/pytest/test_mlir_plugin_interface.py
+++ b/frontend/test/pytest/test_mlir_plugin_interface.py
@@ -67,6 +67,7 @@ def example():
 
         assert example.mlir
 
+
 def test_get_options():
     """
       ApplyRegisteredPassOp expects options to be a single StringAttr
@@ -82,5 +83,8 @@ def test_get_options():
     However, experimentally we found that single-options also work without values.
     """
     assert catalyst.passes.Pass("example-pass", "single-option").get_options() == "single-option"
-    assert catalyst.passes.Pass("example-pass", "an-option", "bn-option").get_options() == "an-option bn-option"
+    assert (
+        catalyst.passes.Pass("example-pass", "an-option", "bn-option").get_options()
+        == "an-option bn-option"
+    )
     assert catalyst.passes.Pass("example-pass", option=True).get_options() == "option=True"