Repository: PennyLaneAI/pennylane

Issue:
Adding `PSWAP` operation
### Context

We would like to define a new operation called `PSWAP`, as [defined in the PennyLane-Braket plugin](https://amazon-braket-pennylane-plugin-python.readthedocs.io/en/latest/code/api/braket.pennylane_plugin.PSWAP.html).

### **Steps to take**

 **Adding the operation**

- Adding the `PSWAP` class and its attributes to `pennylane/ops/qubit/parametric_ops.py`. 
- Complete its docstring in a similar fashion as for other operations. Make sure the class is accessible through `qml.PSWAP`.
- Add its decomposition, matrix, eigenvalues, adjoint and generator.
- Make sure the operation works with parameters from all interfaces (autograd, Torch, TensorFlow and JAX)

**Adding devices support**

- Add support for the operation in the `default.qubit` and `default.mixed` devices

**Testing**

- Add the matrix to `pennylane/tests/gate_data.py`
- Test the gate itself
- Test the gradients with multiple interfaces (autograd, torch, tensorflow and jax) with an **analytical** formula.

### **Notes**

You can refer to the implementation and testings of `IsingXX`, `IsingYY` and `IsingZZ`.

As an example, here is the [implementation](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/ops/qubit/parametric_ops.py#L2246) and [the testing of the decomposition](https://github.com/PennyLaneAI/pennylane/blob/master/tests/ops/qubit/test_parametric_ops.py#L307) for the `IsingYY` operation. The linked test file contains further test cases for `IsingYY`.



Discussion:

Resolving patch:
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -110,6 +110,7 @@ class DefaultQubit(QubitDevice):
         "CNOT",
         "SWAP",
         "ISWAP",
+        "PSWAP",
         "Adjoint(ISWAP)",
         "SISWAP",
         "Adjoint(SISWAP)",
diff --git a/pennylane/ops/qubit/__init__.py b/pennylane/ops/qubit/__init__.py
--- a/pennylane/ops/qubit/__init__.py
+++ b/pennylane/ops/qubit/__init__.py
@@ -59,6 +59,7 @@
     "SISWAP",
     "SQISW",
     "CSWAP",
+    "PSWAP",
     "ECR",
     "Toffoli",
     "RX",
diff --git a/pennylane/ops/qubit/parametric_ops.py b/pennylane/ops/qubit/parametric_ops.py
--- a/pennylane/ops/qubit/parametric_ops.py
+++ b/pennylane/ops/qubit/parametric_ops.py
@@ -3038,3 +3038,144 @@ def adjoint(self):
 
     def pow(self, z):
         return [IsingXY(self.data[0] * z, wires=self.wires)]
+
+
+class PSWAP(Operation):
+    r"""Phase SWAP gate
+
+    .. math:: PSWAP(\phi) = \begin{bmatrix}
+            1 & 0 & 0 & 0 \\
+            0 & 0 & e^{i \phi} & 0 \\
+            0 & e^{i \phi} & 0 & 0 \\
+            0 & 0 & 0 & 1
+        \end{bmatrix}.
+
+    **Details:**
+
+    * Number of wires: 2
+    * Number of parameters: 1
+    * Gradient recipe:
+
+    .. math::
+        \frac{d}{d \phi} PSWAP(\phi)
+        = \frac{1}{2} \left[ PSWAP(\phi + \pi / 2) - PSWAP(\phi - \pi / 2) \right]
+
+    Args:
+        phi (float): the phase angle
+        wires (int): the subsystem the gate acts on
+        do_queue (bool): Indicates whether the operator should be
+            immediately pushed into the Operator queue (optional)
+        id (str or None): String representing the operation (optional)
+    """
+    num_wires = 2
+    num_params = 1
+    """int: Number of trainable parameters that the operator depends on."""
+
+    grad_method = "A"
+    grad_recipe = ([[0.5, 1, np.pi / 2], [-0.5, 1, -np.pi / 2]],)
+
+    def __init__(self, phi, wires, do_queue=True, id=None):
+        super().__init__(phi, wires=wires, do_queue=do_queue, id=id)
+
+    @staticmethod
+    def compute_decomposition(phi, wires):
+        r"""Representation of the operator as a product of other operators (static method). :
+
+        .. math:: O = O_1 O_2 \dots O_n.
+
+
+        .. seealso:: :meth:`~.PSWAP.decomposition`.
+
+        Args:
+            phi (float): the phase angle
+            wires (Iterable, Wires): the subsystem the gate acts on
+
+        Returns:
+            list[Operator]: decomposition into lower level operations
+
+        **Example:**
+
+        >>> qml.PSWAP.compute_decomposition(1.23, wires=(0,1))
+        [SWAP(wires=[0, 1]), CNOT(wires=[0, 1]), PhaseShift(1.23, wires=[1]), CNOT(wires=[0, 1])]
+        """
+        return [
+            qml.SWAP(wires=wires),
+            qml.CNOT(wires=wires),
+            qml.PhaseShift(phi, wires=[wires[1]]),
+            qml.CNOT(wires=wires),
+        ]
+
+    @staticmethod
+    def compute_matrix(phi):  # pylint: disable=arguments-differ
+        r"""Representation of the operator as a canonical matrix in the computational basis (static method).
+
+        The canonical matrix is the textbook matrix representation that does not consider wires.
+        Implicitly, this assumes that the wires of the operator correspond to the global wire order.
+
+        .. seealso:: :meth:`~.PSWAP.matrix`
+
+
+        Args:
+           phi (tensor_like or float): phase angle
+
+        Returns:
+           tensor_like: canonical matrix
+
+        **Example**
+
+        >>> qml.PSWAP.compute_matrix(0.5)
+        array([[1.        +0.j, 0.        +0.j        , 0.        +0.j        , 0.        +0.j],
+              [0.        +0.j, 0.        +0.j        , 0.87758256+0.47942554j, 0.        +0.j],
+              [0.        +0.j, 0.87758256+0.47942554j, 0.        +0.j        , 0.        +0.j],
+              [0.        +0.j, 0.        +0.j        , 0.        +0.j        , 1.        +0.j]])
+        """
+        if qml.math.get_interface(phi) == "tensorflow":
+            phi = qml.math.cast_like(phi, 1j)
+
+        e = qml.math.exp(1j * phi)
+
+        return qml.math.stack(
+            [
+                stack_last([1, 0, 0, 0]),
+                stack_last([0, 0, e, 0]),
+                stack_last([0, e, 0, 0]),
+                stack_last([0, 0, 0, 1]),
+            ],
+            axis=-2,
+        )
+
+    @staticmethod
+    def compute_eigvals(phi):  # pylint: disable=arguments-differ
+        r"""Eigenvalues of the operator in the computational basis (static method).
+
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U`,
+        the operator can be reconstructed as
+
+        .. math:: O = U \Sigma U^{\dagger},
+
+        where :math:`\Sigma` is the diagonal matrix containing the eigenvalues.
+
+        Otherwise, no particular order for the eigenvalues is guaranteed.
+
+        .. seealso:: :meth:`~.PSWAP.eigvals`
+
+
+        Args:
+            phi (tensor_like or float): phase angle
+
+        Returns:
+            tensor_like: eigenvalues
+
+        **Example**
+
+        >>> qml.PSWAP.compute_eigvals(0.5)
+        array([ 1.        +0.j        ,  1.        +0.j,       -0.87758256-0.47942554j,  0.87758256+0.47942554j])
+        """
+        if qml.math.get_interface(phi) == "tensorflow":
+            phi = qml.math.cast_like(phi, 1j)
+
+        return qml.math.stack([1, 1, -qml.math.exp(1j * phi), qml.math.exp(1j * phi)])
+
+    def adjoint(self):
+        (phi,) = self.parameters
+        return PSWAP(-phi, wires=self.wires)


Test patch:
diff --git a/pennylane/devices/tests/test_gates.py b/pennylane/devices/tests/test_gates.py
--- a/pennylane/devices/tests/test_gates.py
+++ b/pennylane/devices/tests/test_gates.py
@@ -69,6 +69,7 @@
     "Adjoint(S)": qml.adjoint(qml.S(wires=[0])),
     "SWAP": qml.SWAP(wires=[0, 1]),
     "ISWAP": qml.ISWAP(wires=[0, 1]),
+    "PSWAP": qml.PSWAP(0, wires=[0, 1]),
     "ECR": qml.ECR(wires=[0, 1]),
     "Adjoint(ISWAP)": qml.adjoint(qml.ISWAP(wires=[0, 1])),
     "T": qml.T(wires=[0]),
@@ -229,6 +230,15 @@
     ]
 )
 
+PSWAP = lambda phi: np.array(
+    [
+        [1, 0, 0, 0],
+        [0, 0, exp(1.0j * phi), 0],
+        [0, exp(1.0j * phi), 0, 0],
+        [0, 0, 0, 1],
+    ]
+)
+
 
 def adjoint_tuple(op, orig_mat):
     """Returns op constructor and matrix for provided base ops."""
@@ -278,6 +288,7 @@ def adjoint_tuple(op, orig_mat):
     (qml.IsingXY, IsingXY),
     (qml.IsingYY, IsingYY),
     (qml.IsingZZ, IsingZZ),
+    (qml.PSWAP, PSWAP),
 ]
 two_qubit_multi_param = [(qml.CRot, crot)]
 # list of all three-qubit gates
diff --git a/tests/gate_data.py b/tests/gate_data.py
--- a/tests/gate_data.py
+++ b/tests/gate_data.py
@@ -589,3 +589,20 @@ def OrbitalRotation(phi):
             [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
         ]
     )
+
+
+def PSWAP(phi):
+    r"""Phase SWAP gate
+    .. math:: \mathtt{PSWAP}(\phi) = \begin{bmatrix}
+            1 & 0 & 0 & 0 \\
+            0 & 0 & e^{i \phi} & 0 \\
+            0 & e^{i \phi} & 0 & 0 \\
+            0 & 0 & 0 & 1
+        \end{bmatrix}.
+    Args:
+        phi (float): rotation angle :math:`\phi`
+    Returns:
+        array[complex]: unitary 4x4 rotation matrix
+    """
+    e = math.exp(1j * phi)
+    return math.array([[1, 0, 0, 0], [0, 0, e, 0], [0, e, 0, 0], [0, 0, 0, 1]])
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -55,6 +55,7 @@
     qml.DoubleExcitation(0.123, wires=[0, 1, 2, 3]),
     qml.DoubleExcitationPlus(0.123, wires=[0, 1, 2, 3]),
     qml.DoubleExcitationMinus(0.123, wires=[0, 1, 2, 3]),
+    qml.PSWAP(0.123, wires=[0, 1]),
 ]
 
 BROADCASTED_OPERATIONS = [
@@ -482,6 +483,36 @@ def test_U3_decomposition_broadcasted(self):
                 assert isinstance(op, c)
                 assert np.allclose(op.parameters, p)
 
+    def test_pswap_decomposition(self, tol):
+        """Tests that the decomposition of the PSWAP gate is correct"""
+        param = 0.1234
+        op = qml.PSWAP(param, wires=[0, 1])
+        res = op.decomposition()
+
+        assert len(res) == 4
+
+        assert res[0].wires == Wires([0, 1])
+        assert res[1].wires == Wires([0, 1])
+        assert res[2].wires == Wires([1])
+        assert res[3].wires == Wires([0, 1])
+
+        assert res[0].name == "SWAP"
+        assert res[1].name == "CNOT"
+        assert res[2].name == "PhaseShift"
+        assert res[3].name == "CNOT"
+
+        mats = []
+        for i in reversed(res):
+            if i.wires == Wires([1]):
+                # PhaseShift gate
+                mats.append(np.kron(np.eye(2), i.matrix()))
+            else:
+                mats.append(i.matrix())
+
+        decomposed_matrix = np.linalg.multi_dot(mats)
+
+        assert np.allclose(decomposed_matrix, op.matrix(), atol=tol, rtol=0)
+
     def test_isingxx_decomposition(self, tol):
         """Tests that the decomposition of the IsingXX gate is correct"""
         param = 0.1234
@@ -902,6 +933,73 @@ def get_expected(theta):
             qml.IsingXX(param, wires=[0, 1]).matrix(), get_expected(param), atol=tol, rtol=0
         )
 
+    def test_pswap(self, tol):
+        """Test that the PSWAP operation is correct"""
+        assert np.allclose(
+            qml.PSWAP.compute_matrix(0), np.diag([1, 1, 1, 1])[[0, 2, 1, 3]], atol=tol, rtol=0
+        )
+        assert np.allclose(
+            qml.PSWAP(0, wires=[0, 1]).matrix(),
+            np.diag([1, 1, 1, 1])[[0, 2, 1, 3]],
+            atol=tol,
+            rtol=0,
+        )
+
+        def get_expected(theta):
+            return np.diag([1, np.exp(1j * theta), np.exp(1j * theta), 1])[[0, 2, 1, 3]]
+
+        param = np.pi / 2
+        assert np.allclose(qml.PSWAP.compute_matrix(param), get_expected(param), atol=tol, rtol=0)
+        assert np.allclose(
+            qml.PSWAP(param, wires=[0, 1]).matrix(), get_expected(param), atol=tol, rtol=0
+        )
+
+        param = np.pi
+        assert np.allclose(qml.PSWAP.compute_matrix(param), get_expected(param), atol=tol, rtol=0)
+        assert np.allclose(
+            qml.PSWAP(param, wires=[0, 1]).matrix(), get_expected(param), atol=tol, rtol=0
+        )
+
+    @pytest.mark.parametrize("phi", np.linspace(-np.pi, np.pi, 10))
+    def test_pswap_eigvals(self, phi, tol):
+        """Test eigenvalues computation for PSWAP"""
+        evs = qml.PSWAP.compute_eigvals(phi)
+        evs_expected = [1, 1, -qml.math.exp(1j * phi), qml.math.exp(1j * phi)]
+        assert qml.math.allclose(evs, evs_expected)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("phi", np.linspace(-np.pi, np.pi, 10))
+    def test_pswap_eigvals_tf(self, phi, tol):
+        """Test eigenvalues computation for PSWAP using Tensorflow interface"""
+        import tensorflow as tf
+
+        param_tf = tf.Variable(phi)
+        evs = qml.PSWAP.compute_eigvals(param_tf)
+        evs_expected = [1, 1, -qml.math.exp(1j * phi), qml.math.exp(1j * phi)]
+        assert qml.math.allclose(evs, evs_expected)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("phi", np.linspace(-np.pi, np.pi, 10))
+    def test_pswap_eigvals_torch(self, phi, tol):
+        """Test eigenvalues computation for PSWAP using Torch interface"""
+        import torch
+
+        param_torch = torch.tensor(phi)
+        evs = qml.PSWAP.compute_eigvals(param_torch)
+        evs_expected = [1, 1, -qml.math.exp(1j * phi), qml.math.exp(1j * phi)]
+        assert qml.math.allclose(evs, evs_expected)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("phi", np.linspace(-np.pi, np.pi, 10))
+    def test_pswap_eigvals_jax(self, phi, tol):
+        """Test eigenvalues computation for PSWAP using JAX interface"""
+        import jax
+
+        param_jax = jax.numpy.array(phi)
+        evs = qml.PSWAP.compute_eigvals(param_jax)
+        evs_expected = [1, 1, -qml.math.exp(1j * phi), qml.math.exp(1j * phi)]
+        assert qml.math.allclose(evs, evs_expected)
+
     def test_isingxy(self, tol):
         """Test that the IsingXY operation is correct"""
         assert np.allclose(qml.IsingXY.compute_matrix(0), np.identity(4), atol=tol, rtol=0)
@@ -1776,6 +1874,152 @@ class TestGrad:
         for device, method in device_methods:
             configuration.append([device, method, npp.array(phi, requires_grad=True)])
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("dev_name,diff_method,phi", configuration)
+    def test_pswap_autograd_grad(self, tol, dev_name, diff_method, phi):
+        """Test the gradient with Autograd for the gate PSWAP."""
+
+        if diff_method in {"adjoint"}:
+            # PSWAP does not have a generator defined
+            pytest.skip("PSWAP does not support adjoint")
+
+        dev = qml.device(dev_name, wires=2)
+
+        psi_0 = 0.1
+        psi_1 = 0.2
+        psi_2 = 0.3
+        psi_3 = 0.4
+
+        init_state = npp.array([psi_0, psi_1, psi_2, psi_3], requires_grad=False)
+        norm = np.linalg.norm(init_state)
+        init_state /= norm
+
+        @qml.qnode(dev, diff_method=diff_method, interface="autograd")
+        def circuit(phi):
+            qml.QubitStateVector(init_state, wires=[0, 1])
+            qml.PSWAP(phi, wires=[0, 1])
+            return qml.expval(qml.PauliY(0))
+
+        expected = 2 * np.cos(phi) * (psi_0 * psi_1 - psi_3 * psi_2) / norm**2
+
+        res = qml.grad(circuit)(phi)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("dev_name,diff_method,phi", configuration)
+    def test_pswap_torch_grad(self, tol, dev_name, diff_method, phi):
+        """Test the gradient with Torch for the gate PSWAP."""
+
+        if diff_method in {"adjoint"}:
+            # PSWAP does not have a generator defined
+            pytest.skip("PSWAP does not support adjoint")
+
+        import torch
+
+        dev = qml.device(dev_name, wires=2)
+
+        psi_0 = torch.tensor(0.1)
+        psi_1 = torch.tensor(0.2)
+        psi_2 = torch.tensor(0.3)
+        psi_3 = torch.tensor(0.4)
+
+        init_state = torch.tensor([psi_0, psi_1, psi_2, psi_3], requires_grad=False)
+        norm = torch.norm(init_state)
+        init_state /= norm
+
+        @qml.qnode(dev, diff_method=diff_method, interface="torch")
+        def circuit(phi):
+            qml.QubitStateVector(init_state, wires=[0, 1])
+            qml.PSWAP(phi, wires=[0, 1])
+            return qml.expval(qml.PauliY(0))
+
+        phi = torch.tensor(phi, requires_grad=True)
+
+        expected = 2 * torch.cos(phi) * (psi_0 * psi_1 - psi_3 * psi_2) / norm**2
+
+        result = circuit(phi)
+        result.backward()
+        res = phi.grad
+        assert np.allclose(res, expected.detach(), atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name,diff_method,phi", configuration)
+    def test_pswap_jax_grad(self, tol, dev_name, diff_method, phi):
+        """Test the gradient with JAX for the gate PSWAP."""
+
+        if diff_method in {"adjoint"}:
+            # PSWAP does not have a generator defined
+            pytest.skip("PSWAP does not support adjoint")
+
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device(dev_name, wires=2)
+
+        psi_0 = 0.1
+        psi_1 = 0.2
+        psi_2 = 0.3
+        psi_3 = 0.4
+
+        init_state = jnp.array([psi_0, psi_1, psi_2, psi_3])
+        norm = jnp.linalg.norm(init_state)
+        init_state = init_state / norm
+
+        @qml.qnode(dev, diff_method=diff_method, interface="jax")
+        def circuit(phi):
+            qml.QubitStateVector(init_state, wires=[0, 1])
+            qml.PSWAP(phi, wires=[0, 1])
+            return qml.expval(qml.PauliY(0))
+
+        phi = jnp.array(phi)
+
+        expected = 2 * np.cos(phi) * (psi_0 * psi_1 - psi_3 * psi_2) / norm**2
+
+        res = jax.grad(circuit, argnums=0)(phi)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("dev_name,diff_method,phi", configuration)
+    def test_pswap_tf_grad(self, tol, dev_name, diff_method, phi):
+        """Test the gradient with Tensorflow for the gate PSWAP."""
+
+        if diff_method in {"adjoint"}:
+            # PSWAP does not have a generator defined
+            pytest.skip("PSWAP does not support adjoint")
+
+        import tensorflow as tf
+
+        dev = qml.device(dev_name, wires=2)
+
+        psi_0 = tf.Variable(0.1, dtype=tf.complex128)
+        psi_1 = tf.Variable(0.2, dtype=tf.complex128)
+        psi_2 = tf.Variable(0.3, dtype=tf.complex128)
+        psi_3 = tf.Variable(0.4, dtype=tf.complex128)
+
+        init_state = tf.Variable([psi_0, psi_1, psi_2, psi_3], dtype=tf.complex128)
+        norm = tf.norm(init_state)
+        init_state = init_state / norm
+
+        @qml.qnode(dev, interface="tf", diff_method=diff_method)
+        def circuit(phi):
+            qml.QubitStateVector(init_state, wires=[0, 1])
+            qml.PSWAP(phi, wires=[0, 1])
+            return qml.expval(qml.PauliY(0))
+
+        phi = tf.Variable(phi, dtype=tf.complex128)
+
+        expected = 2 * tf.cos(phi) * (psi_0 * psi_1 - psi_3 * psi_2) / norm**2
+
+        with tf.GradientTape() as tape:
+            result = circuit(phi)
+
+        res = tape.gradient(result, phi)
+        if diff_method == "backprop":
+            # Check #2872 https://github.com/PennyLaneAI/pennylane/issues/2872
+            assert np.allclose(np.real(res), expected, atol=tol, rtol=0)
+        else:
+            assert np.allclose(res, expected, atol=tol, rtol=0)
+
     @pytest.mark.autograd
     @pytest.mark.parametrize("dev_name,diff_method,phi", configuration)
     def test_isingxx_autograd_grad(self, tol, dev_name, diff_method, phi):
@@ -3243,6 +3487,7 @@ def test_pow_matrix(self, op, n):
     (qml.IsingXY(1.234, wires=(0, 1)), Wires([])),
     (qml.IsingYY(np.array([-5.1, 0.219]), wires=(0, 1)), Wires([])),
     (qml.IsingZZ(1.234, wires=(0, 1)), Wires([])),
+    (qml.PSWAP(1.234, wires=(0, 1)), Wires([])),
     ### Controlled Ops
     (qml.ControlledPhaseShift(1.234, wires=(0, 1)), Wires(0)),
     (qml.CPhase(1.234, wires=(0, 1)), Wires(0)),