Repository: PennyLaneAI/pennylane

Issue:
[BUG] qml.transforms.transpile does not work with MultiRZ
### Expected behavior

qml.transforms.transpile does not add SWAP gates for MultiRZ operations between 2 qubits and just ignores it. This might be due to MultiRZ being able to act on more than 2 qubits at a time. However, it's currently the only way to add RZZ gates between 2 qubits.

This makes it difficult to see the transpiled circuit for QAOA setups as attempting to transpile a circuit constructed from `qml.qaoa` functions just produces the decomposition of the problem unitary (aka U_H) into RZZ gates that don't match the coupling map. 

Short example (expected behaviour using CX gates):
```python
def nwise(lst, k=2):
    # Function to obtain [(0,1), (1,2), (2,3), ... ], i.e. linear coupling map
    return list(zip(*[lst[i:] for i in range(k)])) 

dev = qml.device('default.qubit', wires=4)

@qml.qnode(dev)
@qml.transforms.transpile(coupling_map=nwise(np.arange(4)))
def circuit(param):
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[0, 2])
    qml.CNOT(wires=[0, 3])
    qml.PhaseShift(param, wires=0)
    return qml.probs(wires=[0, 1, 2, 3])
print(qml.draw(circuit)(0.6))

>>>

0: ─╭●───────╭●──────────╭●──Rϕ(0.60)─┤ ╭Probs
1: ─╰X─╭SWAP─╰X────╭SWAP─╰X───────────┤ ├Probs
2: ────╰SWAP─╭SWAP─╰SWAP──────────────┤ ├Probs
3: ──────────╰SWAP────────────────────┤ ╰Probs

```






### Actual behavior

But changing to MultiRZ (which is currently the only way to implement RZZ unless I manually decomp them into CX-RZ-CX) gives:

```python

def nwise(lst, k=2):
    # Function to obtain [(0,1), (1,2), (2,3), ... ], i.e. linear coupling map
    return list(zip(*[lst[i:] for i in range(k)])) 

dev = qml.device('default.qubit', wires=4)

@qml.qnode(dev)
@qml.transforms.transpile(coupling_map=nwise(np.arange(4)))
def circuit(param):
    qml.MultiRZ(param, wires=[0, 1])
    qml.MultiRZ(param, wires=[0, 2])
    qml.MultiRZ(param, wires=[0, 3])
    qml.PhaseShift(param, wires=0)
    return qml.probs(wires=[0, 1, 2, 3])
print(qml.draw(circuit)(0.6))

0: ─╭MultiRZ(0.60)─╭MultiRZ(0.60)─╭MultiRZ(0.60)──Rϕ(0.60)─┤ ╭Probs
1: ─╰MultiRZ(0.60)─│──────────────│────────────────────────┤ ├Probs
2: ────────────────╰MultiRZ(0.60)─│────────────────────────┤ ├Probs
3: ───────────────────────────────╰MultiRZ(0.60)───────────┤ ╰Probs
```

### Additional information

Full QAOA example below:

### Source code

```python


import numpy as np
import networkx as nx

import pennylane as qml
from pennylane import numpy as plnp
from pennylane import qaoa


def rand_mc(nv, deg, seed):
    # Function to generate max-cut matrices of random n-variable d-regular graphs
    graph = nx.generators.random_graphs.random_regular_graph(deg, nv, seed=seed)
    matrix = nx.to_numpy_array(graph, nodelist=np.arange(nv))
    
    for k in range(nv):
        matrix[k][k] = -np.sum(matrix[k])
    return matrix, seed


def get_hvals(matrix): 
    # For max cut problem, this is already written in a form to be minimized.
    # So the ground state of the resulting Hamiltonian will provide the solution to the max cut problem.
    
    # Gets the Ising coefficients, NOT the 2^n by 2^n Ising Hamiltonian.
    # Reparameterizes the QUBO problem into {+1, -1}
    # Reparameterization done using x=(Z+1)/2
    # Minimize z.J.z + z.h + offset
    nvars = len(matrix)
    jmat = np.zeros(shape=(nvars,nvars))
    hvec = np.zeros(nvars)
    
    for i in range(nvars):
        for j in range(nvars):
            if i == j:
                hvec[i] = np.sum(matrix[i])/2
            else:
                jmat[i][j] = matrix[i][j]/4
                jmat[j][i] = matrix[i][j]/4              
    
    # Gives the correct offset value to the CF
    offset = (np.sum(matrix)/4 + np.trace(matrix)/4)
    return jmat, hvec, offset


def get_qml_ising(matrix):
    # Function to obtain qml.Hamiltonian from Ising coefficients
    nq = len(matrix)
    jj, hh, oo = get_hvals(matrix)
    h_coeff = []
    h_obs = []

    for i in range(nq):
        for j in range(i,nq):
            if i == j:
                h_coeff.append(hh[i])
                h_obs.append(qml.PauliZ(i))
            else:
                h_coeff.append(2*jj[i,j])
                h_obs.append( qml.PauliZ(i) @ qml.PauliZ(j))

    h_coeff.append(oo)
    h_coeff = np.array(h_coeff)
    h_obs.append(qml.Identity(0))
    hamiltonian = qml.Hamiltonian(h_coeff, h_obs)
    return hamiltonian

def get_qaoa_mixer(matrix):
    # Function to obtain QAOA mixer Hamiltonian
    nq = len(matrix)
    h_coeff = []
    h_obs = []
    
    for i in range(nq):
        h_coeff.append(1)
        h_obs.append(qml.PauliX(i))
    hamiltonian = qml.Hamiltonian(h_coeff, h_obs)
    return hamiltonian


def qaoa_circ(params, wires = None, problem_hamiltonian = None, mixing_hamiltonian = None):
    # QAOA Circuit
    P = len(params)//2
    nq = len(problem_hamiltonian.wires)

    def qaoa_layer(gamma, beta):
        qaoa.cost_layer(gamma, problem_hamiltonian)
        qaoa.mixer_layer(beta, mixing_hamiltonian)

    # g1, b1, g2, b2, etc....
    gamma_list = params[::2][:P] # for U_H
    beta_list = params[1::2][:P] # for U_X

    for nq_ in range(nq): # apply Hadamards to get all in x-basis
        qml.Hadamard(wires=nq_)
    qml.layer(qaoa_layer, P, gamma_list, beta_list) # p instances of unitary operators

    # Returns probs because measurement of H not supported according to documents
    # https://docs.pennylane.ai/en/stable/code/api/pennylane.transforms.transpile.html
    return qml.probs(wires=np.arange(nq))                               


def nwise(lst, k=2):
    # Function to obtain [(0,1), (1,2), (2,3), ... ], i.e. linear coupling map
    return list(zip(*[lst[i:] for i in range(k)])) 

#+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

nqq = 4                                                                 # Number of qubits
tmat = rand_mc(nqq, 3, seed=1029)[0]                                    # n-variable 3-regular graph
n_measurements = None                                                   # Use statevector
coupling_map = nwise(np.arange(nqq))                                    # Get coupling map

hamiltonian = get_qml_ising(tmat)                                       # Get Pennylane Hamiltonian
mixing_hamiltonian = get_qaoa_mixer(tmat)


dev = qml.device('default.qubit', wires=nqq)                            # Initialize device
transpiled_qaoa = qml.transforms.transpile(coupling_map=coupling_map)(qaoa_circ)
transpiled_qaoa_qnode = qml.QNode(transpiled_qaoa, dev)
print(qml.draw(transpiled_qaoa_qnode)([1,1], wires = None, problem_hamiltonian = hamiltonian, mixing_hamiltonian = mixing_hamiltonian))

# Output is supposed to be transpiled according to linear coupling map which includes a bunch of SWAP gates.
>>>
0: ──H──RZ(0.00)─╭RZZ(1.00)─╭RZZ(1.00)─╭RZZ(1.00)──RX(2.00)────────────────────────────────
1: ──H───────────╰RZZ(1.00)─│──────────│───────────RZ(0.00)─╭RZZ(1.00)─╭RZZ(1.00)──RX(2.00)
2: ──H──────────────────────╰RZZ(1.00)─│────────────────────╰RZZ(1.00)─│───────────RZ(0.00)
3: ──H─────────────────────────────────╰RZZ(1.00)──────────────────────╰RZZ(1.00)──────────

─────────────────────────────────┤ ╭Probs
─────────────────────────────────┤ ├Probs
──╭RZZ(1.00)──RX(2.00)───────────┤ ├Probs
──╰RZZ(1.00)──RZ(0.00)──RX(2.00)─┤ ╰Probs



```


### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.26.0

Platform info:           Windows-10-10.0.19041-SP0
Python version:          3.7.13
Numpy version:           1.21.5
Scipy version:           1.7.3
Installed devices:
- default.gaussian (PennyLane-0.26.0)
- default.mixed (PennyLane-0.26.0)
- default.qubit (PennyLane-0.26.0)
- default.qubit.autograd (PennyLane-0.26.0)
- default.qubit.jax (PennyLane-0.26.0)
- default.qubit.tf (PennyLane-0.26.0)
- default.qubit.torch (PennyLane-0.26.0)
- default.qutrit (PennyLane-0.26.0)
- lightning.qubit (PennyLane-Lightning-0.26.0)
- cirq.mixedsimulator (PennyLane-Cirq-0.22.0)
- cirq.pasqal (PennyLane-Cirq-0.22.0)
- cirq.qsim (PennyLane-Cirq-0.22.0)
- cirq.qsimh (PennyLane-Cirq-0.22.0)
- cirq.simulator (PennyLane-Cirq-0.22.0)
- qiskit.aer (PennyLane-qiskit-0.21.0)
- qiskit.basicaer (PennyLane-qiskit-0.21.0)
- qiskit.ibmq (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.21.0)
- forest.numpy_wavefunction (PennyLane-Forest-0.21.0.dev0)
- forest.qvm (PennyLane-Forest-0.21.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.21.0.dev0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Discussion:
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://codecov.io/gh/PennyLaneAI/pennylane/pull/3104?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3104](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3104?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (73a2b51) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c158bcdde048dd140a0e70a9d44e2123792572cb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c158bcd) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3104   +/-   ##
=======================================
  Coverage   99.67%   99.67%           
=======================================
  Files         273      273           
  Lines       23546    23546           
=======================================
  Hits        23470    23470           
  Misses         76       76           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3104?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/transpile.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3104/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvdHJhbnNwaWxlLnB5) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7425dbcea..89991a3b0 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -99,6 +99,9 @@ keyword argument when using `GellMann`, which determines which of the 8 Gell-Man
 
 <h3>Bug fixes</h3>
 
+* Fixed the `qml.transforms.transpile` transform to work correctly for all two-qubit operations.
+  [(#3104)](https://github.com/PennyLaneAI/pennylane/pull/3104)
+
 * Fixed a bug with the control values of a controlled version of a `ControlledQubitUnitary`.
   [(#3119)](https://github.com/PennyLaneAI/pennylane/pull/3119)
 
@@ -114,4 +117,5 @@ Soran Jahangiri,
 Christina Lee,
 Mudit Pandey,
 Jay Soni,
+Antal Száva,
 David Wierichs,
diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py
index b85ed1b73..95ddbfc95 100644
--- a/pennylane/transforms/transpile.py
+++ b/pennylane/transforms/transpile.py
@@ -96,6 +96,11 @@ def transpile(tape: QuantumTape, coupling_map: Union[List, nx.Graph]):
             "Measuring expectation values of tensor products or Hamiltonians is not yet supported"
         )
 
+    if any(len(op.wires) > 2 for op in tape.operations):
+        raise NotImplementedError(
+            "The transpile transform only supports gates acting on 1 or 2 qubits."
+        )
+
     gates = []
 
     # we wrap all manipulations inside stop_recording() so that we don't queue anything due to unrolling of templates
@@ -115,7 +120,7 @@ def transpile(tape: QuantumTape, coupling_map: Union[List, nx.Graph]):
             op = list_op_copy[0]
 
             # gates which act only on one wire
-            if op.num_wires != 2:
+            if len(op.wires) == 1:
                 gates.append(op)
                 list_op_copy.pop(0)
                 continue


Test patch:
diff --git a/tests/transforms/test_transpile.py b/tests/transforms/test_transpile.py
index 5e59ecc96..d7b3f0ace 100644
--- a/tests/transforms/test_transpile.py
+++ b/tests/transforms/test_transpile.py
@@ -55,7 +55,7 @@ class TestTranspile:
         transpiled_qfunc = transpile(coupling_map=[(0, 1)])(original_qfunc)
         transpiled_qnode = qml.QNode(transpiled_qfunc, dev)
         err_msg = (
-            "Not all wires present in coupling map! wires: \[0, 2, 1\], coupling map: \[0, 1\]"
+            r"Not all wires present in coupling map! wires: \[0, 2, 1\], coupling map: \[0, 1\]"
         )
         with pytest.raises(ValueError, match=err_msg):
             transpiled_qnode(0.1, 0.2, 0.3)
@@ -150,3 +150,123 @@ class TestTranspile:
         transpiled_qnode = qml.QNode(transpiled_circ, dev, interface="autograd")
         params = np.array([0.5, 0.1, 0.2], requires_grad=True)
         qml.gradients.param_shift(transpiled_qnode)(params)
+
+    def test_more_than_2_qubits_raises_anywires(self):
+        """test that transpile raises an error for an operation with AnyWires that acts on more than 2 qubits"""
+        dev = qml.device("default.qubit", wires=[0, 1, 2])
+
+        def circuit(param):
+            qml.MultiRZ(param, wires=[0, 1, 2])
+            return qml.probs(wires=[0, 1])
+
+        param = 0.3
+
+        transpiled_qfunc = transpile(coupling_map=[(0, 1), (1, 2)])(circuit)
+        transpiled_qnode = qml.QNode(transpiled_qfunc, dev)
+        with pytest.raises(
+            NotImplementedError,
+            match="transpile transform only supports gates acting on 1 or 2 qubits",
+        ):
+            transpiled_expectation = transpiled_qnode(param)
+
+    def test_more_than_2_qubits_raises_3_qubit_gate(self):
+        """test that transpile raises an error for an operation that acts on more than 2 qubits"""
+        dev = qml.device("default.qubit", wires=[0, 1, 2])
+
+        def circuit():
+            qml.Toffoli(wires=[0, 1, 2])
+            return qml.probs(wires=[0, 1])
+
+        transpiled_qfunc = transpile(coupling_map=[(0, 1), (1, 2)])(circuit)
+        transpiled_qnode = qml.QNode(transpiled_qfunc, dev)
+        with pytest.raises(
+            NotImplementedError,
+            match="transpile transform only supports gates acting on 1 or 2 qubits",
+        ):
+            transpiled_expectation = transpiled_qnode()
+
+    def test_transpile_ops_anywires(self):
+        """test that transpile does not alter output for expectation value of an observable if the qfunc contains
+        operations that act on any number of wires"""
+        dev = qml.device("default.qubit", wires=[0, 1, 2])
+
+        def circuit(param):
+            qml.MultiRZ(param, wires=[0, 1])
+            qml.PhaseShift(param, wires=2)
+            qml.MultiRZ(param, wires=[0, 2])
+            return qml.probs(wires=[0, 1])
+
+        param = 0.3
+
+        # build circuit without transpile
+        original_qfunc = circuit
+        original_qnode = qml.QNode(original_qfunc, dev)
+        original_expectation = original_qnode(param)
+
+        # build circuit with transpile
+        transpiled_qfunc = transpile(coupling_map=[(0, 1), (1, 2)])(original_qfunc)
+        transpiled_qnode = qml.QNode(transpiled_qfunc, dev)
+        transpiled_expectation = transpiled_qnode(param)
+
+        original_ops = list(transpiled_qnode.qtape)
+        transpiled_ops = list(transpiled_qnode.qtape)
+        assert qml.equal(transpiled_ops[0], original_ops[0])
+        assert qml.equal(transpiled_ops[1], original_ops[1])
+
+        # SWAP to ensure connectivity
+        assert isinstance(transpiled_ops[2], qml.SWAP)
+        assert transpiled_ops[2].wires == qml.wires.Wires([1, 2])
+
+        assert isinstance(transpiled_ops[3], qml.MultiRZ)
+        assert transpiled_ops[3].data == [param]
+        assert transpiled_ops[3].wires == qml.wires.Wires([0, 1])
+
+        assert isinstance(transpiled_ops[4], qml.measurements.MeasurementProcess)
+        assert transpiled_ops[4].wires == qml.wires.Wires([0, 2])
+
+        assert qml.math.allclose(
+            original_expectation, transpiled_expectation, atol=np.finfo(float).eps
+        )
+
+    def test_transpile_ops_anywires_1_qubit(self):
+        """test that transpile does not alter output for expectation value of an observable if the qfunc contains
+        1-qubit operations with AnyWires defined for the operation"""
+        dev = qml.device("default.qubit", wires=[0, 1, 2])
+
+        def circuit(param):
+            qml.MultiRZ(param, wires=[0])
+            qml.PhaseShift(param, wires=2)
+            qml.MultiRZ(param, wires=[0, 2])
+            return qml.probs(wires=[0, 1])
+
+        param = 0.3
+
+        # build circuit without transpile
+        original_qfunc = circuit
+        original_qnode = qml.QNode(original_qfunc, dev)
+        original_expectation = original_qnode(param)
+
+        # build circuit with transpile
+        transpiled_qfunc = transpile(coupling_map=[(0, 1), (1, 2)])(original_qfunc)
+        transpiled_qnode = qml.QNode(transpiled_qfunc, dev)
+        transpiled_expectation = transpiled_qnode(param)
+
+        original_ops = list(transpiled_qnode.qtape)
+        transpiled_ops = list(transpiled_qnode.qtape)
+        assert qml.equal(transpiled_ops[0], original_ops[0])
+        assert qml.equal(transpiled_ops[1], original_ops[1])
+
+        # SWAP to ensure connectivity
+        assert isinstance(transpiled_ops[2], qml.SWAP)
+        assert transpiled_ops[2].wires == qml.wires.Wires([1, 2])
+
+        assert isinstance(transpiled_ops[3], qml.MultiRZ)
+        assert transpiled_ops[3].data == [param]
+        assert transpiled_ops[3].wires == qml.wires.Wires([0, 1])
+
+        assert isinstance(transpiled_ops[4], qml.measurements.MeasurementProcess)
+        assert transpiled_ops[4].wires == qml.wires.Wires([0, 2])
+
+        assert qml.math.allclose(
+            original_expectation, transpiled_expectation, atol=np.finfo(float).eps
+        )