Repository: PennyLaneAI/pennylane
Pull request #7179 (2025-04-01 18:20:06)

Make `PSWAP` support parameter broadcasting
**Context:**
`PSWAP` does not support parameter broadcasting in its `compute_matrix` and `compute_eigvals` methods.

**Description of the Change:**
Add support for parameter broadcasting.

**Benefits:**
More support for broadcasting.

**Possible Drawbacks:**
Tiny code/computation growth.

**Related GitHub Issues:**
#7180 

[sc-87781]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7179?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.62%. Comparing base [(`470c585`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/470c585c859b65e2834215aa4ad32e2014f8b2c5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`2a847e2`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/2a847e20492162b9628bec9789edbcb89d1bf427?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #7179   +/-   ##
=======================================
  Coverage   99.62%   99.62%           
=======================================
  Files         509      509           
  Lines       49619    49625    +6     
=======================================
+ Hits        49434    49440    +6     
  Misses        185      185           
```

</details>

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

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
JerryChen97: Mind deleting the `xfail` in `tests/ops/qubit/test_parametric_ops.py::TestSimplify::test_simplify_rotations` in this PR? With both `qml.matrix` bug-fix merged and this PR implementing `PSWAP.matrix`, that `xfail`'d case should be released.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index da54d7335..d2210e529 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -195,6 +195,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `PSWAP.matrix()` and `PSWAP.eigvals()` now support parameter broadcasting.
+  [(#7179)](https://github.com/PennyLaneAI/pennylane/pull/7179)
+
 * `PrepSelPrep` now has a concise representation when drawn with `qml.draw` or `qml.draw_mpl`.
   [(#7164)](https://github.com/PennyLaneAI/pennylane/pull/7164)
 
diff --git a/pennylane/ops/qubit/parametric_ops_multi_qubit.py b/pennylane/ops/qubit/parametric_ops_multi_qubit.py
index fc996f19a..9462e87a3 100644
--- a/pennylane/ops/qubit/parametric_ops_multi_qubit.py
+++ b/pennylane/ops/qubit/parametric_ops_multi_qubit.py
@@ -1586,6 +1586,9 @@ class PSWAP(Operation):
     num_params = 1
     """int: Number of trainable parameters that the operator depends on."""
 
+    ndim_params = (0,)
+    """tuple[int]: Number of dimensions per trainable parameter that the operator depends on."""
+
     resource_keys = set()
 
     grad_method = "A"
@@ -1654,13 +1657,15 @@ class PSWAP(Operation):
             phi = qml.math.cast_like(phi, 1j)
 
         e = qml.math.exp(1j * phi)
+        zero = qml.math.zeros_like(phi)
+        one = qml.math.ones_like(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]),
+                stack_last([one, zero, zero, zero]),
+                stack_last([zero, zero, e, zero]),
+                stack_last([zero, e, zero, zero]),
+                stack_last([zero, zero, zero, one]),
             ],
             axis=-2,
         )
@@ -1695,7 +1700,9 @@ class PSWAP(Operation):
         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)])
+        e = qml.math.exp(1j * phi)
+        one = qml.math.ones_like(phi)
+        return 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 2103cb8da..8e57e73c9 100644
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -103,6 +103,7 @@ BROADCASTED_OPERATIONS = [
         wires=[0, 1],
     ),
     qml.QubitUnitary(1j * np.array([[[1, 0], [0, -1]], [[0, 1], [1, 0]]]), wires=0),
+    qml.PSWAP(np.array([0.123, -0.63, 0.235]), wires=["a", -1]),
     qml.DiagonalQubitUnitary(np.array([[1.0, 1.0j], [1.0j, 1.0j]]), wires=1),
 ]
 
@@ -2781,7 +2782,7 @@ class TestPauliRot:
         self, theta, pauli_word, compressed_pauli_word, wires, compressed_wires, tol
     ):
         """Test PauliRot matrix correctly accounts for identities."""
-        # pylint: disable=too-many-arguments
+        # pylint: disable=too-many-arguments, too-many-positional-arguments
 
         res = qml.PauliRot.compute_matrix(theta, pauli_word)
         expected = qml.math.expand_matrix(