Repository: PennyLaneAI/pennylane
Pull request #4966 (2023-12-21 14:51:36)

`Pow` operators are now pickleable.
[sc-49618]

We discovered that pow operators are currently not pickleable due to `__new__` black magic.  This PR makes it so `Pow` is updated to use the same behavior as `Adjoint`, which is picklable.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4966?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`89a7e77`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/89a7e775a7c8e444ee62fb411946bbdb15ecaa86?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`b407b3d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4966?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4966?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/ops/op\_math/pow.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4966?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Bvdy5weQ==) | 90.00% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4966?src=pr&el=tree&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    #4966      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         392      392              
  Lines       35556    35284     -272     
==========================================
- Hits        35442    35167     -275     
- Misses        114      117       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4966?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).


Patch:
diff --git a/doc/releases/changelog-0.34.0.md b/doc/releases/changelog-0.34.0.md
index bfb45ac32..cadafb0c7 100644
--- a/doc/releases/changelog-0.34.0.md
+++ b/doc/releases/changelog-0.34.0.md
@@ -568,6 +568,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `Pow` operators are now picklable.
+  [(#4966)](https://github.com/PennyLaneAI/pennylane/pull/4966)
+
 * Finite differences and SPSA can now be used with tensorflow-autograph on setups that were seeing a bus error.
   [(#4961)](https://github.com/PennyLaneAI/pennylane/pull/4961)
 
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 1e0baaec0..4cdf73bfa 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -86,7 +86,7 @@ def stopping_condition(op: qml.operation.Operator) -> bool:
         return False
     if op.name == "Snapshot":
         return True
-    if op.__class__.__name__ == "Pow" and qml.operation.is_trainable(op):
+    if op.__class__.__name__[:3] == "Pow" and qml.operation.is_trainable(op):
         return False
 
     return op.has_matrix
diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
index 3f2c5498c..35fc26fa9 100644
--- a/pennylane/devices/default_qubit_legacy.py
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -232,7 +232,7 @@ class DefaultQubitLegacy(QubitDevice):
             if getattr(obj, "has_matrix", False):
                 # pow operations dont work with backprop or adjoint without decomposition
                 # use class name string so we don't need to use isinstance check
-                return not (obj.__class__.__name__ == "Pow" and qml.operation.is_trainable(obj))
+                return not (obj.__class__.__name__[:3] == "Pow" and qml.operation.is_trainable(obj))
             return obj.name in self.observables.union(self.operations)
 
         return qml.BooleanFn(accepts_obj)
diff --git a/pennylane/ops/functions/assert_valid.py b/pennylane/ops/functions/assert_valid.py
index 5e6e288f5..514667169 100644
--- a/pennylane/ops/functions/assert_valid.py
+++ b/pennylane/ops/functions/assert_valid.py
@@ -102,7 +102,7 @@ def _check_matrix_matches_decomp(op):
     """Check that if both the matrix and decomposition are defined, they match."""
     if op.has_matrix and op.has_decomposition:
         mat = op.matrix()
-        decomp_mat = qml.matrix(op.decomposition, wire_order=op.wires)()
+        decomp_mat = qml.matrix(qml.tape.QuantumScript(op.decomposition()), wire_order=op.wires)
         failure_comment = (
             f"matrix and matrix from decomposition must match. Got \n{mat}\n\n {decomp_mat}"
         )
diff --git a/pennylane/ops/op_math/pow.py b/pennylane/ops/op_math/pow.py
index 089bb93fe..0dfd96e26 100644
--- a/pennylane/ops/op_math/pow.py
+++ b/pennylane/ops/op_math/pow.py
@@ -31,7 +31,7 @@ from pennylane.operation import (
 from pennylane.ops.identity import Identity
 from pennylane.queuing import QueuingManager, apply
 
-from .symbolicop import ScalarSymbolicOp, SymbolicOp
+from .symbolicop import ScalarSymbolicOp
 
 _superscript = str.maketrans("0123456789.+-", "⁰¹²³⁴⁵⁶⁷⁸⁹⋅⁺⁻")
 
@@ -106,29 +106,6 @@ def pow(base, z=1, lazy=True, id=None):
     return pow_op
 
 
-# pylint: disable=no-member
-class PowOperation(Operation):
-    """Operation-specific methods and properties for the ``Pow`` class.
-
-    Dynamically mixed in based on the provided base operator.  If the base operator is an
-    Operation, this class will be mixed in.
-
-    When we no longer rely on certain functionality through `Operation`, we can get rid of this
-    class.
-    """
-
-    # until we add gradient support
-    grad_method = None
-
-    @property
-    def name(self):
-        return self._name
-
-    @property
-    def control_wires(self):
-        return self.base.control_wires
-
-
 class Pow(ScalarSymbolicOp):
     """Symbolic operator denoting an operator raised to a power.
 
@@ -160,10 +137,6 @@ class Pow(ScalarSymbolicOp):
     def _unflatten(cls, data, _):
         return pow(data[0], z=data[1])
 
-    _operation_type = None  # type if base inherits from operation and not observable
-    _operation_observable_type = None  # type if base inherits from both operation and observable
-    _observable_type = None  # type if base inherits from observable and not oepration
-
     # pylint: disable=unused-argument
     def __new__(cls, base=None, z=1, id=None):
         """Mixes in parents based on inheritance structure of base.
@@ -187,22 +160,13 @@ class Pow(ScalarSymbolicOp):
 
         if isinstance(base, Operation):
             if isinstance(base, Observable):
-                if cls._operation_observable_type is None:
-                    base_classes = (PowOperation, Pow, SymbolicOp, Observable, Operation)
-                    cls._operation_observable_type = type("Pow", base_classes, dict(cls.__dict__))
-                return object.__new__(cls._operation_observable_type)
+                return object.__new__(PowOpObs)
 
             # not an observable
-            if cls._operation_type is None:
-                base_classes = (PowOperation, Pow, SymbolicOp, Operation)
-                cls._operation_type = type("Pow", base_classes, dict(cls.__dict__))
-            return object.__new__(cls._operation_type)
+            return object.__new__(PowOperation)
 
         if isinstance(base, Observable):
-            if cls._observable_type is None:
-                base_classes = (Pow, SymbolicOp, Observable)
-                cls._observable_type = type("Pow", base_classes, dict(cls.__dict__))
-            return object.__new__(cls._observable_type)
+            return object.__new__(PowObs)
 
         return object.__new__(Pow)
 
@@ -370,3 +334,46 @@ class Pow(ScalarSymbolicOp):
             return op.simplify()
         except PowUndefinedError:
             return Pow(base=base, z=self.z)
+
+
+# pylint: disable=no-member
+class PowOperation(Pow, Operation):
+    """Operation-specific methods and properties for the ``Pow`` class.
+
+    Dynamically mixed in based on the provided base operator.  If the base operator is an
+    Operation, this class will be mixed in.
+
+    When we no longer rely on certain functionality through `Operation`, we can get rid of this
+    class.
+    """
+
+    def __new__(cls, *_, **__):
+        return object.__new__(cls)
+
+    # until we add gradient support
+    grad_method = None
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def control_wires(self):
+        return self.base.control_wires
+
+
+class PowObs(Pow, Observable):
+    """A child class of ``Pow`` that also inherits from ``Observable``."""
+
+    def __new__(cls, *_, **__):
+        return object.__new__(cls)
+
+
+# pylint: disable=too-many-ancestors
+class PowOpObs(PowOperation, Observable):
+    """A child class of ``Pow`` that inherits from both
+    ``Observable`` and ``Operation``.
+    """
+
+    def __new__(cls, *_, **__):
+        return object.__new__(cls)


Test patch:
diff --git a/tests/ops/functions/conftest.py b/tests/ops/functions/conftest.py
index c0e35cb2e..276013fc1 100644
--- a/tests/ops/functions/conftest.py
+++ b/tests/ops/functions/conftest.py
@@ -22,6 +22,7 @@ import pytest
 import pennylane as qml
 from pennylane.operation import Operator, Operation, Observable, Tensor, Channel
 from pennylane.ops.op_math.adjoint import Adjoint, AdjointObs, AdjointOperation, AdjointOpObs
+from pennylane.ops.op_math.pow import PowObs, PowOperation, PowOpObs
 
 # if you would like to validate one of these operators, add an instance to the parametrization
 # of `test_explicit_list_of_ops` in `test_assert_valid.py`, and kindly move it below the
@@ -81,6 +82,9 @@ _ABSTRACT_OR_META_TYPES = {
     AdjointOpObs,
     AdjointOperation,
     AdjointObs,
+    PowOpObs,
+    PowOperation,
+    PowObs,
 }
 
 
diff --git a/tests/ops/op_math/test_pow_op.py b/tests/ops/op_math/test_pow_op.py
index 415b283b8..dbaafab80 100644
--- a/tests/ops/op_math/test_pow_op.py
+++ b/tests/ops/op_math/test_pow_op.py
@@ -39,12 +39,17 @@ def pow_using_dunder_method(base, z, id=None):
     return base**z
 
 
-@pytest.mark.xfail  # TODO: remove xfail as part of story 49618
 def test_basic_validity():
     """Run basic operator validity checks."""
     op = qml.pow(qml.RX(1.2, wires=0), 3)
     qml.ops.functions.assert_valid(op)
 
+    op = qml.pow(qml.PauliX(0), 2.5)
+    qml.ops.functions.assert_valid(op)
+
+    op = qml.pow(qml.Hermitian(np.eye(2), 0), 2)
+    qml.ops.functions.assert_valid(op)
+
 
 class TestConstructor:
     def test_lazy_mode(self):