Fix tape expand tests

pennylane/transforms/adjoint.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Code for the adjoint transform."""
-
+from collections.abc import Sequence
 from functools import wraps
 from pennylane.tape import QuantumTape, stop_recording
 
@@ -112,11 +112,11 @@ def circuit():
     @wraps(fn)
     def wrapper(*args, **kwargs):
         with stop_recording(), QuantumTape() as tape:
-            fn(*args, **kwargs)
+            res = fn(*args, **kwargs)
 
         if not tape.operations:
             # we called op.expand(): get the outputted tape
-            tape = fn(*args, **kwargs)
+            tape = res
 
         adjoint_ops = []
         for op in reversed(tape.operations):
@@ -130,6 +130,9 @@ def wrapper(*args, **kwargs):
                 if isinstance(new_ops, QuantumTape):
                     new_ops = new_ops.operations
 
+                if not isinstance(new_ops, Sequence):
+                    new_ops = [new_ops]
+
                 adjoint_ops.extend(new_ops)
 
         if len(adjoint_ops) == 1:

pennylane/transforms/tape_expand.py

@@ -225,20 +225,13 @@ def _custom_decomposition(obj, fn):
 
         original_decomp_method = obj.decomposition
 
-        # This is the method that will override the operations .decompose method
-        def new_decomp_method(self):
-            with NonQueuingTape():
-                if self.num_params == 0:
-                    return fn(self.wires)
-                return fn(*self.parameters, self.wires)
-
         try:
             # Explicitly set the new .decompose method
-            obj.decomposition = new_decomp_method
+            obj.compute_decomposition = staticmethod(fn)
             yield
 
         finally:
-            obj.decomposition = original_decomp_method
+            obj.compute_decomposition = original_decomp_method
 
     # Loop through the decomposition dictionary and create all the contexts
     try:

tests/transforms/test_control.py

@@ -108,7 +108,7 @@ def test_template(self):
         adjoint_of_controlled_op = qml.ctrl(op, control=control_wires)(par, wires=wires).adjoint()
 
         assert adjoint_of_controlled_op.control_wires == control_wires
-        res_ops = adjoint_of_controlled_op.subtape.operations[0].operations
+        res_ops = adjoint_of_controlled_op.subtape.operations
         expected = qml.adjoint(op)(par, wires=wires)
 
         for op1, op2 in zip(res_ops, expected):

tests/transforms/test_tape_expand.py

@@ -363,9 +363,9 @@ def custom_rot(phi, theta, omega, wires):
 
 
 # Decompose a template into another template
-def custom_basic_entangler_layers(weights, wires):
+def custom_basic_entangler_layers(weights, wires, **kwargs):
     return [
-        qml.AngleEmbedding(weights, wires=wires),
+        qml.AngleEmbedding(weights[0], wires=wires),
         qml.broadcast(qml.CNOT, pattern="ring", wires=wires),
     ]
 
@@ -495,7 +495,6 @@ def circuit(x):
         assert np.allclose(original_grad, decomp_grad)
 
         expected_ops = ["Hadamard", "RZ", "PauliX", "RY", "PauliX", "RZ", "Hadamard"]
-        print(decomp_qnode.qtape.operations)
         assert all(
             [op.name == name for op, name in zip(decomp_qnode.qtape.operations, expected_ops)]
         )
@@ -624,12 +623,12 @@ def circuit():
             return qml.expval(qml.PauliZ(0))
 
         # BasicEntanglerLayers custom decomposition involves AngleEmbedding. If
-        # expansion depth is 1, the AngleEmbedding will still be decomposed into
+        # expansion depth is 2, the AngleEmbedding will still be decomposed into
         # RX (since it's not a supported operation on the device), but the RX will
         # not be further decomposed even though the custom decomposition is specified.
         custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
         decomp_dev = qml.device(
-            "default.qubit", wires=2, custom_decomps=custom_decomps, decomp_depth=1
+            "default.qubit", wires=2, custom_decomps=custom_decomps, decomp_depth=2
         )
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()