Repository: PennyLaneAI/pennylane
Pull request #7228 (2025-04-08 07:18:41)

Fix `PSWAP` eigvals axis ordering with broadcasting 
**Context:**
#7179 added broadcasting support to the matrix and eigvals of `PSWAP`. However, the broadcasting axis is currently the second (i.e. last) axis in broadcasted eigvals, although PL convention is to have the broadcasting axis first.

**Description of the Change:**
Transpose the eigenvalues. This does not change them when there is no broadcasting, and fixed the axis ordering whenever there is broadcasting.
Added tests for broadcasted eigenvalue computation, including ML interface tests.

**Benefits:**
Fixes eigenvalue broadcasting axis ordering.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
[sc-88504]

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.
dwierichs: Thanks @mudit2812, added that test! :+1: 

Patch:
diff --git a/doc/releases/changelog-0.41.0.md b/doc/releases/changelog-0.41.0.md
index 71a02cd48..ec717b33c 100644
--- a/doc/releases/changelog-0.41.0.md
+++ b/doc/releases/changelog-0.41.0.md
@@ -713,6 +713,7 @@ With `qml.decompositions.enable_graph()`, the following new features are availab
 
 * `PSWAP.matrix()` and `PSWAP.eigvals()` now support parameter broadcasting.
   [(#7179)](https://github.com/PennyLaneAI/pennylane/pull/7179)
+  [(#7228)](https://github.com/PennyLaneAI/pennylane/pull/7228)
 
 * `Device.eval_jaxpr` now accepts an `execution_config` keyword argument.
   [(#6991)](https://github.com/PennyLaneAI/pennylane/pull/6991)
diff --git a/pennylane/ops/qubit/parametric_ops_multi_qubit.py b/pennylane/ops/qubit/parametric_ops_multi_qubit.py
index 9462e87a3..f039b9092 100644
--- a/pennylane/ops/qubit/parametric_ops_multi_qubit.py
+++ b/pennylane/ops/qubit/parametric_ops_multi_qubit.py
@@ -1702,7 +1702,7 @@ class PSWAP(Operation):
 
         e = qml.math.exp(1j * phi)
         one = qml.math.ones_like(phi)
-        return qml.math.stack([one, one, -e, e])
+        return qml.math.transpose(qml.math.stack([one, one, -e, e]))
 
     def adjoint(self) -> "PSWAP":
         (phi,) = self.parameters


Test patch:
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
index 8e57e73c9..01959a507 100644
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -699,6 +699,9 @@ class TestDecompositions:
         assert np.allclose(decomposed_matrix, exp)
 
 
+pswap_angles = list(np.linspace(-np.pi, np.pi, 11)) + [np.linspace(-1, 1, 11)]
+
+
 class TestMatrix:
     def test_phase_shift(self, tol):
         """Test phase shift is correct"""
@@ -1001,44 +1004,89 @@ class TestMatrix:
             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_broadcasted(self, tol):
+        """Test that the broadcasted PSWAP operation is correct"""
+        batch_size = 3
+        z = np.zeros(batch_size)
+        mat = qml.PSWAP.compute_matrix(z)
+        assert qml.math.shape(mat) == (batch_size, 4, 4)
+        assert np.allclose(mat, np.diag([1, 1, 1, 1])[[0, 2, 1, 3]], atol=tol, rtol=0)
+        mat = qml.PSWAP(z, wires=[0, 1]).matrix()
+        assert qml.math.shape(mat) == (batch_size, 4, 4)
+        assert np.allclose(
+            mat,
+            np.diag([1, 1, 1, 1])[[0, 2, 1, 3]],
+            atol=tol,
+            rtol=0,
+        )
+
+        def get_expected(theta):
+            return np.array(
+                [np.diag([1, np.exp(1j * t), np.exp(1j * t), 1])[[0, 2, 1, 3]] for t in theta]
+            )
+
+        param = np.array([np.pi / 2, 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
+        )
+
+        param = np.array([2.152, np.pi / 2, 0.213])
+        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", pswap_angles)
     def test_pswap_eigvals(self, phi):
         """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)]
+        if len(qml.math.shape(phi)) > 0:
+            evs_expected = np.stack([[1, 1, -exp, exp] for exp in qml.math.exp(1j * phi)])
+        else:
+            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))
+    @pytest.mark.parametrize("phi", pswap_angles)
     def test_pswap_eigvals_tf(self, phi):
         """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)]
+        if len(qml.math.shape(phi)) > 0:
+            evs_expected = np.stack([[1, 1, -exp, exp] for exp in qml.math.exp(1j * phi)])
+        else:
+            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))
+    @pytest.mark.parametrize("phi", pswap_angles)
     def test_pswap_eigvals_torch(self, phi):
         """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)]
+        if len(qml.math.shape(phi)) > 0:
+            evs_expected = np.stack([[1, 1, -exp, exp] for exp in qml.math.exp(1j * phi)])
+        else:
+            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))
+    @pytest.mark.parametrize("phi", pswap_angles)
     def test_pswap_eigvals_jax(self, phi):
         """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)]
+        if len(qml.math.shape(phi)) > 0:
+            evs_expected = np.stack([[1, 1, -exp, exp] for exp in qml.math.exp(1j * phi)])
+        else:
+            evs_expected = [1, 1, -qml.math.exp(1j * phi), qml.math.exp(1j * phi)]
         assert qml.math.allclose(evs, evs_expected)
 
     @pytest.mark.tf