Repository: PennyLaneAI/pennylane
Pull request #5669 (2024-05-07 18:21:50)

Allow specifying error type to be raised by decompose
**Context:**
As the `decompose` transform is used in various contexts, it would be nice to specify the error type that should be raised if its not possible to get a suitable decomposition for one of the operators  (i.e. devices raise `DeviceError`, transforms and gradients raise `DecompositionUndefinedError`, catalyst raises `CompileError`)

**Description of the Change:**
An error kwarg is added to `decompose` and its helper function `_operator_decomposition_gen` so that the error type can be specified. It continues to default to `DeviceError`

**Benefits:**
We don't have to duplicate many lines of code to get an identical function with a different error type

**Possible Drawbacks:**
More kwargs

**Related GitHub Issues:**
N/A

**Related Shortcut Stories:**
[[sc-61450]](https://app.shortcut.com/xanaduai/story/61450/catalyst-can-decompose-nested-operations-hybrid-op-part-2)

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5669?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.68%. Comparing base [(`cfcef86`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/cfcef8612d72869f43c9c95e28128ae0fc36d653?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`8cbf93f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5669?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5669      +/-   ##
==========================================
- Coverage   99.68%   99.68%   -0.01%     
==========================================
  Files         412      412              
  Lines       38676    38385     -291     
==========================================
- Hits        38556    38264     -292     
- Misses        120      121       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5669?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

lillian542: [sc-61450]
trbromley: Thanks @lillian542!

> it would be nice to specify the error type that should be raised

Are you able to speak more towards that motivation? I guess this is needed for Catalyst integration, but I'm just curious to understand more.
lillian542: > Also, do we want a basic test that asserts the error kwarg being used correctly?
> And a changelog entry? 🤔

@dwierichs - why yes. Yes we do want those things 🙈 
lillian542: > Are you able to speak more towards that motivation? I guess this is needed for Catalyst integration, but I'm just curious to understand more.

@trbromley - this function was originally used only in preprocessing for standard devices (which we want to raise a `DeviceError` if it gets an operation it can't use or decompose). 

As of this week, it's also used in transforms, where no device is present, and it would be better to raise the original `DecompositionUndefinedError`. For catalyst, the issue that we hit in `decompose` is whether or not it can be compiled via `qjit` for the given device, and everything in that part of the pipe line raises a `CompileError`. 

So the motivation is to keep error types consistent with what process is failing, even though at the end of the day they all come down to a `DecompositionUndefinedError` arising in different contexts.

An alternative would be to always use `DecompositionUndefinedError`. That would be accurate, but inconvenient in `try-except` blocks where we would like to, for example, try to run preprocessing and have special handling if a `DeviceError` is raised.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f206fb023..0a09e4048 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -63,6 +63,10 @@
 * Sets up the framework for the development of an `assert_equal` function for testing operator comparison.
   [(#5634)](https://github.com/PennyLaneAI/pennylane/pull/5634)
 
+* The `decompose` transform has an `error` kwarg to specify the type of error that should be raised, 
+  allowing error types to be more consistent with the context the `decompose` function is used in.
+  [(#5669)](https://github.com/PennyLaneAI/pennylane/pull/5669)
+
 <h3>Breaking changes 💔</h3>
 
 * `qml.is_commuting` no longer accepts the `wire_map` argument, which does not bring any functionality.
diff --git a/pennylane/devices/preprocess.py b/pennylane/devices/preprocess.py
index 98d6442f4..882ed8445 100644
--- a/pennylane/devices/preprocess.py
+++ b/pennylane/devices/preprocess.py
@@ -45,6 +45,7 @@ def _operator_decomposition_gen(
     max_expansion: Optional[int] = None,
     current_depth=0,
     name: str = "device",
+    error: Exception = DeviceError,
 ) -> Generator[qml.operation.Operator, None, None]:
     """A generator that yields the next operation that is accepted."""
     max_depth_reached = False
@@ -57,7 +58,7 @@ def _operator_decomposition_gen(
             decomp = decomposer(op)
             current_depth += 1
         except qml.operation.DecompositionUndefinedError as e:
-            raise DeviceError(
+            raise error(
                 f"Operator {op} not supported with {name} and does not provide a decomposition."
             ) from e
 
@@ -69,6 +70,7 @@ def _operator_decomposition_gen(
                 max_expansion=max_expansion,
                 current_depth=current_depth,
                 name=name,
+                error=error,
             )
 
 
@@ -253,6 +255,7 @@ def decompose(
     ] = None,
     max_expansion: Union[int, None] = None,
     name: str = "device",
+    error: Exception = DeviceError,
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Decompose operations until the stopping condition is met.
 
@@ -260,14 +263,20 @@ def decompose(
         tape (QuantumTape or QNode or Callable): a quantum circuit.
         stopping_condition (Callable): a function from an operator to a boolean. If ``False``, the operator
             should be decomposed. If an operator cannot be decomposed and is not accepted by ``stopping_condition``,
-            a ``DecompositionUndefinedError`` will be raised.
+            an ``Exception`` will be raised (of a type specified by the ``error`` kwarg).
+
+    Keyword Args:
         stopping_condition_shots (Callable): a function from an operator to a boolean. If ``False``, the operator
             should be decomposed. If an operator cannot be decomposed and is not accepted by ``stopping_condition``,
-            a ``DecompositionUndefinedError`` will be raised. This replaces stopping_condition if and only if the tape has shots.
-        skip_initial_state_prep=True (bool): If ``True``, the first operator will not be decomposed if it inherits from :class:`~.StatePrepBase`.
+            an ``Exception`` will be raised (of a type specified by the ``error`` kwarg). This replaces stopping_condition if and only if the tape has shots.
+        skip_initial_state_prep (bool): If ``True``, the first operator will not be decomposed if it inherits
+            from :class:`~.StatePrepBase`. Defaults to ``True``.
         decomposer (Callable): an optional callable that takes an operator and implements the relevant decomposition.
             If None, defaults to using a callable returning ``op.decomposition()`` for any :class:`~.Operator` .
-        max_expansion (int): The maximum depth of the expansion.
+        max_expansion (int): The maximum depth of the expansion. Defaults to None.
+        name (str): The name of the transform, process or device using decompose. Used in the error message. Defaults to "device".
+        error (Error): An error type to raise if it is not possible to obtain a decomposition that fulfills
+            the ``stopping_condition``. Defaults to ``DeviceError``.
 
 
     Returns:
@@ -276,9 +285,9 @@ def decompose(
         The decomposed circuit. The output type is explained in :func:`qml.transform <pennylane.transform>`.
 
     Raises:
-        DecompositionUndefinedError: if an operator is not accepted and does not define a decomposition
-
-        DeviceError: If the decomposition enters and infinite loop and raises a ``RecursionError``.
+        Exception: Type defaults to ``DeviceError`` but can be modified via keyword argument. Raised if
+            an operator is not accepted and does not define a decomposition, or if the decomposition
+            enters an infinite loop and raises a ``RecursionError``.
 
     **Example:**
 
@@ -343,10 +352,11 @@ def decompose(
                 decomposer=decomposer,
                 max_expansion=max_expansion,
                 name=name,
+                error=error,
             )
         ]
     except RecursionError as e:
-        raise DeviceError(
+        raise error(
             "Reached recursion limit trying to decompose operations. "
             "Operator decomposition may have entered an infinite loop."
         ) from e
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index e478bc653..05a353b44 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -203,6 +203,7 @@ def _expand_transform_finite_diff(
         stopping_condition=_finite_diff_stopping_condition,
         skip_initial_state_prep=False,
         name="finite_diff",
+        error=qml.operation.DecompositionUndefinedError,
     )
     if new_tape is tape:
         return [tape], postprocessing
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index ef2709e7a..9036378f7 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -780,6 +780,7 @@ def _expand_transform_param_shift(
         stopping_condition=_param_shift_stopping_condition,
         skip_initial_state_prep=False,
         name="param_shift",
+        error=qml.operation.DecompositionUndefinedError,
     )
     if new_tape is tape:
         return [tape], postprocessing
diff --git a/pennylane/transforms/compile.py b/pennylane/transforms/compile.py
index a2421a984..50bb71049 100644
--- a/pennylane/transforms/compile.py
+++ b/pennylane/transforms/compile.py
@@ -190,7 +190,11 @@ def compile(
             return obj.name in basis_set and (not getattr(obj, "only_visual", False))
 
         [expanded_tape], _ = qml.devices.preprocess.decompose(
-            tape, stopping_condition=stop_at, max_expansion=expand_depth, name="compile"
+            tape,
+            stopping_condition=stop_at,
+            max_expansion=expand_depth,
+            name="compile",
+            error=qml.operation.DecompositionUndefinedError,
         )
 
         # Apply the full set of compilation transforms num_passes times
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index e3a7362ae..2c8178c19 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -123,7 +123,10 @@ def merge_rotations(
         return not isinstance(obj, Adjoint)
 
     [expanded_tape], _ = qml.devices.preprocess.decompose(
-        tape, stopping_condition=stop_at, name="merge_rotations"
+        tape,
+        stopping_condition=stop_at,
+        name="merge_rotations",
+        error=qml.operation.DecompositionUndefinedError,
     )
     list_copy = expanded_tape.operations
     new_operations = []
diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py
index f2f6e4e08..8e9fa68e6 100644
--- a/pennylane/transforms/transpile.py
+++ b/pennylane/transforms/transpile.py
@@ -168,7 +168,10 @@ def transpile(
             return (obj.name in all_ops) and (not getattr(obj, "only_visual", False))
 
         [expanded_tape], _ = qml.devices.preprocess.decompose(
-            tape, stopping_condition=stop_at, name="transpile"
+            tape,
+            stopping_condition=stop_at,
+            name="transpile",
+            error=qml.operation.DecompositionUndefinedError,
         )
         # make copy of ops
         list_op_copy = expanded_tape.operations.copy()


Test patch:
diff --git a/tests/devices/test_preprocess.py b/tests/devices/test_preprocess.py
index fb5897041..38c3d4837 100644
--- a/tests/devices/test_preprocess.py
+++ b/tests/devices/test_preprocess.py
@@ -56,6 +56,15 @@ class NoMatNoDecompOp(Operation):
         return False
 
 
+class InfiniteOp(qml.operation.Operation):
+    """An op with an infinite decomposition."""
+
+    num_wires = 1
+
+    def decomposition(self):
+        return [InfiniteOp(*self.parameters, self.wires)]
+
+
 class TestPrivateHelpers:
     """Test the private helpers for preprocessing."""
 
@@ -205,18 +214,28 @@ class TestDecomposeValidation:
     def test_infinite_decomposition_loop(self):
         """Test that a device error is raised if decomposition enters an infinite loop."""
 
-        class InfiniteOp(qml.operation.Operation):
-            """An op with an infinite decomposition."""
-
-            num_wires = 1
-
-            def decomposition(self):
-                return [InfiniteOp(*self.parameters, self.wires)]
-
         qs = qml.tape.QuantumScript([InfiniteOp(1.23, 0)])
         with pytest.raises(DeviceError, match=r"Reached recursion limit trying to decompose"):
             decompose(qs, lambda obj: obj.has_matrix)
 
+    @pytest.mark.parametrize(
+        "error_type", [RuntimeError, qml.operation.DecompositionUndefinedError]
+    )
+    def test_error_type_can_be_set(self, error_type):
+        """Test that passing a class of Error the ``decompose`` transform allows raising another type
+        of error instead of the default ``DeviceError``."""
+
+        decomp_error_tape = QuantumScript(
+            ops=[NoMatNoDecompOp(0)], measurements=[qml.expval(qml.Hadamard(0))]
+        )
+        recursion_error_tape = qml.tape.QuantumScript([InfiniteOp(1.23, 0)])
+
+        with pytest.raises(error_type, match="not supported with abc"):
+            decompose(decomp_error_tape, lambda op: op.has_matrix, name="abc", error=error_type)
+
+        with pytest.raises(error_type, match=r"Reached recursion limit trying to decompose"):
+            decompose(recursion_error_tape, lambda obj: obj.has_matrix, error=error_type)
+
 
 class TestValidateObservables:
     """Tests for the validate observables transform."""