Repository: PennyLaneAI/pennylane
Pull request #5923 (2024-07-01 13:02:31)

Support parameter broadcasting with `GlobalPhase`
**Context:**
`GlobalPhase` does not support broadcasting yet.

**Description of the Change:**
Introduce broadcasting to `GlobalPhase`, using the code from `PauliRot` with `set(pauli_word) == {"I"}`.
The latter falls back to the `GlobalPhase` implementation.

Also fixes a small bug where the global phase returned by `one_qubit_decomposition` gains a broadcasting dimension even if the input matrix does not have one.

**Benefits:**
Broadcasting support & Bug fix
Unlocks #4460

**Possible Drawbacks:**

**Related GitHub Issues:**
Implements #5815 
Fixes https://github.com/PennyLaneAI/pennylane/issues/5880

[sc-65316]

astralcai: Fixes https://github.com/PennyLaneAI/pennylane/issues/5880

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index cb77dfa61..2531f5273 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -43,6 +43,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `GlobalPhase` now supports parameter broadcasting.
+  [(#5923)](https://github.com/PennyLaneAI/pennylane/pull/5923)
+
 * `qml.devices.LegacyDeviceFacade` has been added to map the legacy devices to the new
   device interface.
   [(#5927)](https://github.com/PennyLaneAI/pennylane/pull/5927)
@@ -223,9 +226,13 @@
   [(#5974)](https://github.com/PennyLaneAI/pennylane/pull/5974)
 
 <h3>Bug fixes 🐛</h3>
+
 * Fix `jax.grad` + `jax.jit` not working for `AmplitudeEmbedding`, `StatePrep` and `MottonenStatePreparation`.
   [(#5620)](https://github.com/PennyLaneAI/pennylane/pull/5620) 
 
+* Fix a bug where the global phase returned by `one_qubit_decomposition` gained a broadcasting dimension.
+  [(#5923)](https://github.com/PennyLaneAI/pennylane/pull/5923)
+
 * Fixed a bug in `qml.SPSAOptimizer` that ignored keyword arguments in the objective function.
   [(#6027)](https://github.com/PennyLaneAI/pennylane/pull/6027)
 
@@ -275,4 +282,5 @@ Vincent Michaud-Rioux,
 Anurav Modak,
 Mudit Pandey,
 Erik Schultheis,
-nate stemen.
+nate stemen,
+David Wierichs,
\ No newline at end of file
diff --git a/pennylane/ops/identity.py b/pennylane/ops/identity.py
index a523c393f..b88aa2ce1 100644
--- a/pennylane/ops/identity.py
+++ b/pennylane/ops/identity.py
@@ -20,7 +20,13 @@ from functools import lru_cache
 from scipy import sparse
 
 import pennylane as qml
-from pennylane.operation import AllWires, AnyWires, CVObservable, Operation
+from pennylane.operation import (
+    AllWires,
+    AnyWires,
+    CVObservable,
+    Operation,
+    SparseMatrixUndefinedError,
+)
 
 
 class Identity(CVObservable, Operation):
@@ -227,6 +233,7 @@ class GlobalPhase(Operation):
 
     * Number of wires: All (the operation acts on all wires)
     * Number of parameters: 1
+    * Number of dimensions per parameter: (0,)
     * Gradient recipe: None
 
     Args:
@@ -282,11 +289,17 @@ class GlobalPhase(Operation):
 
     """
 
-    grad_method = None
-    num_params = 1
     num_wires = AllWires
     """int: Number of wires that the operator acts on."""
 
+    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."""
+
+    grad_method = None
+
     def __init__(self, phi, wires=None, id=None):
         super().__init__(phi, wires=[] if wires is None else wires, id=id)
 
@@ -313,7 +326,15 @@ class GlobalPhase(Operation):
         >>> qml.GlobalPhase.compute_eigvals(np.pi/2)
         array([6.123234e-17+1.j, 6.123234e-17+1.j])
         """
-        return qml.math.exp(-1j * phi) * qml.math.ones(2**n_wires)
+        exp = qml.math.exp(-1j * phi)
+        ones = qml.math.ones(2**n_wires, like=phi)
+        if qml.math.get_interface(phi) == "tensorflow":
+            ones = qml.math.cast_like(ones, 1j)
+
+        if qml.math.ndim(phi) == 0:
+            return exp * ones
+
+        return qml.math.tensordot(exp, ones, axes=0)
 
     @staticmethod
     def compute_matrix(phi, n_wires=1):  # pylint: disable=arguments-differ
@@ -333,15 +354,22 @@ class GlobalPhase(Operation):
                [0.        +0.j        , 0.70710678-0.70710678j]])
         """
         interface = qml.math.get_interface(phi)
+        eye = qml.math.eye(2**n_wires, like=phi)
+        exp = qml.math.exp(-1j * qml.math.cast(phi, complex))
         if interface == "tensorflow":
-            return qml.math.exp(-1j * qml.math.cast(phi, complex)) * qml.math.eye(int(2**n_wires))
-        return qml.math.exp(-1j * qml.math.cast(phi, complex)) * qml.math.eye(
-            int(2**n_wires), like=interface
-        )
+            eye = qml.math.cast_like(eye, 1j)
+        elif interface == "torch":
+            eye = eye.to(exp.device)
+
+        if qml.math.ndim(phi) == 0:
+            return exp * eye
+        return qml.math.tensordot(exp, eye, axes=0)
 
     @staticmethod
     def compute_sparse_matrix(phi, n_wires=1):  # pylint: disable=arguments-differ
-        return qml.math.exp(-1j * phi) * sparse.eye(int(2**n_wires), format="csr")
+        if qml.math.ndim(phi) > 0:
+            raise SparseMatrixUndefinedError("Sparse matrices do not support broadcasting")
+        return qml.math.exp(-1j * phi) * sparse.eye(2**n_wires, format="csr")
 
     @staticmethod
     def compute_diagonalizing_gates(
@@ -403,6 +431,9 @@ class GlobalPhase(Operation):
         """
         return []
 
+    def eigvals(self):
+        return self.compute_eigvals(self.data[0], n_wires=len(self.wires))
+
     def matrix(self, wire_order=None):
         n_wires = len(wire_order) if wire_order else len(self.wires)
         return self.compute_matrix(self.data[0], n_wires=n_wires)
diff --git a/pennylane/ops/op_math/decompositions/single_qubit_unitary.py b/pennylane/ops/op_math/decompositions/single_qubit_unitary.py
index 787639b79..eab40d549 100644
--- a/pennylane/ops/op_math/decompositions/single_qubit_unitary.py
+++ b/pennylane/ops/op_math/decompositions/single_qubit_unitary.py
@@ -146,6 +146,7 @@ def _rot_decomposition(U, wire, return_global_phase=False):
             angle = 2 * math.angle(U_det1[0, 1, 1]) % (4 * np.pi)
             operations = [qml.RZ(angle, wires=wire)]
             if return_global_phase:
+                alphas = math.squeeze(alphas)
                 operations.append(qml.GlobalPhase(-alphas))
             return operations
 
diff --git a/pennylane/ops/qubit/parametric_ops_multi_qubit.py b/pennylane/ops/qubit/parametric_ops_multi_qubit.py
index e29b1f029..f3dc77384 100644
--- a/pennylane/ops/qubit/parametric_ops_multi_qubit.py
+++ b/pennylane/ops/qubit/parametric_ops_multi_qubit.py
@@ -368,18 +368,7 @@ class PauliRot(Operation):
 
         # Simplest case is if the Pauli is the identity matrix
         if set(pauli_word) == {"I"}:
-            exp = qml.math.exp(-0.5j * theta)
-            iden = qml.math.eye(2 ** len(pauli_word), like=theta)
-            if qml.math.get_interface(theta) == "tensorflow":
-                iden = qml.math.cast_like(iden, 1j)
-            if qml.math.get_interface(theta) == "torch":
-                td = exp.device
-                iden = iden.to(td)
-
-            if qml.math.ndim(theta) == 0:
-                return exp * iden
-
-            return qml.math.stack([e * iden for e in exp])
+            return qml.GlobalPhase.compute_matrix(0.5 * theta, n_wires=len(pauli_word))
 
         # We first generate the matrix excluding the identity parts and expand it afterwards.
         # To this end, we have to store on which wires the non-identity parts act
@@ -445,15 +434,7 @@ class PauliRot(Operation):
 
         # Identity must be treated specially because its eigenvalues are all the same
         if set(pauli_word) == {"I"}:
-            exp = qml.math.exp(-0.5j * theta)
-            ones = qml.math.ones(2 ** len(pauli_word), like=theta)
-            if qml.math.get_interface(theta) == "tensorflow":
-                ones = qml.math.cast_like(ones, 1j)
-
-            if qml.math.ndim(theta) == 0:
-                return exp * ones
-
-            return qml.math.tensordot(exp, ones, axes=0)
+            return qml.GlobalPhase.compute_eigvals(0.5 * theta, n_wires=len(pauli_word))
 
         return MultiRZ.compute_eigvals(theta, len(pauli_word))
 


Test patch:
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
index 304006132..5d4775163 100644
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -269,14 +269,13 @@ class TestDecompositions:
         assert np.allclose(res[0].data[0], phi)
 
         decomposed_matrix = res[0].matrix()
-        global_phase = np.exp(-1j * phi / 2)[..., np.newaxis, np.newaxis]
+        global_phase = np.exp(1j * phi / 2)[..., np.newaxis, np.newaxis]
 
         assert res[1].name == "GlobalPhase"
-        assert np.allclose(qml.matrix(res[1]), np.exp(1j * phi / 2))
+        assert np.allclose(qml.matrix(res[1]), global_phase)
 
-        assert np.allclose(decomposed_matrix, global_phase * op.matrix(), atol=tol, rtol=0)
-        if qml.math.shape(phi) == ():  # GlobalPhase matrix doesn't support batching
-            assert np.allclose(op.matrix(), qml.prod(*res[::-1]).matrix())
+        assert np.allclose(decomposed_matrix, (1 / global_phase) * op.matrix(), atol=tol, rtol=0)
+        assert np.allclose(op.matrix(), qml.prod(*res[::-1]).matrix())
 
     def test_Rot_decomposition(self):
         """Test the decomposition of Rot."""
@@ -710,20 +709,42 @@ class TestMatrix:
         assert np.allclose(qml.PhaseShift.compute_matrix(phi), expected, atol=tol, rtol=0)
         assert np.allclose(qml.U1.compute_matrix(phi), expected, atol=tol, rtol=0)
 
-    def test_global_phase(self, tol):
+    @pytest.mark.parametrize("n_wires", [0, 1, 2])
+    def test_global_phase(self, tol, n_wires):
         """Test GlobalPhase matrix is correct"""
 
+        wires = list(range(n_wires))
+        eye = np.eye(2**n_wires)
+        eye2 = np.eye(2)
         # test identity for theta=0
-        assert np.allclose(qml.GlobalPhase.compute_matrix(0), np.identity(2), atol=tol, rtol=0)
         assert np.allclose(
-            qml.GlobalPhase(0).matrix(wire_order=[0]), np.identity(2), atol=tol, rtol=0
+            qml.GlobalPhase.compute_matrix(0, n_wires=n_wires), eye, atol=tol, rtol=0
         )
+        assert np.allclose(qml.GlobalPhase.compute_matrix(0), eye2, atol=tol, rtol=0)
+        assert np.allclose(qml.GlobalPhase(0).matrix(wire_order=wires), eye, atol=tol, rtol=0)
 
-        # test arbitrary phase shift
+        # test arbitrary global phase
         phi = 0.5432
-        expected = np.array([[qml.math.exp(-1j * phi), 0], [0, qml.math.exp(-1j * phi)]])
-        assert np.allclose(qml.GlobalPhase.compute_matrix(phi), expected, atol=tol, rtol=0)
-        assert np.allclose(qml.GlobalPhase(phi).matrix(wire_order=[0]), expected, atol=tol, rtol=0)
+        exp = np.exp(-1j * phi)
+        assert np.allclose(
+            qml.GlobalPhase.compute_matrix(phi, n_wires=n_wires), exp * eye, atol=tol, rtol=0
+        )
+        assert np.allclose(qml.GlobalPhase.compute_matrix(phi), exp * eye2, atol=tol, rtol=0)
+        assert np.allclose(
+            qml.GlobalPhase(phi).matrix(wire_order=wires), exp * eye, atol=tol, rtol=0
+        )
+
+        # test arbitrary broadcasted global phase with non-default n_wires=0
+        phi = np.array([0.5, 0.4, 0.3])
+        expected = np.tensordot(np.exp(-1j * phi), eye, axes=0)
+        expected2 = np.tensordot(np.exp(-1j * phi), eye2, axes=0)
+        assert np.allclose(
+            qml.GlobalPhase.compute_matrix(phi, n_wires=n_wires), expected, atol=tol, rtol=0
+        )
+        assert np.allclose(qml.GlobalPhase.compute_matrix(phi), expected2, atol=tol, rtol=0)
+        assert np.allclose(
+            qml.GlobalPhase(phi).matrix(wire_order=wires), expected, atol=tol, rtol=0
+        )
 
     def test_identity(self, tol):
         """Test Identity matrix is correct with no wires"""
@@ -1698,19 +1719,29 @@ class TestEigvals:
         )
         assert np.allclose(op.eigvals(), expected)
 
-    def test_global_phase_eigvals(self):
+    @pytest.mark.parametrize("n_wires", [0, 1, 2])
+    def test_global_phase_eigvals(self, n_wires):
         """Test GlobalPhase eigenvalues are correct"""
 
+        dim = 2**n_wires
         # test identity for theta=0
-        op = qml.GlobalPhase(0.0)
-        assert np.allclose(op.compute_eigvals(*op.parameters, **op.hyperparameters), np.ones(2))
-        assert np.allclose(op.eigvals(), np.ones(2))
+        phi = 0.0
+        op = qml.GlobalPhase(phi, wires=list(range(n_wires)))
+        assert np.allclose(op.compute_eigvals(phi, n_wires=n_wires), np.ones(dim))
+        assert np.allclose(op.eigvals(), np.ones(dim))
 
-        # test arbitrary phase shift
+        # test arbitrary global phase
         phi = 0.5432
-        op = qml.GlobalPhase(phi)
-        expected = np.array([np.exp(-1j * phi), np.exp(-1j * phi)])
-        assert np.allclose(op.compute_eigvals(*op.parameters, **op.hyperparameters), expected)
+        op = qml.GlobalPhase(phi, wires=list(range(n_wires)))
+        expected = np.array([np.exp(-1j * phi)] * dim)
+        assert np.allclose(op.compute_eigvals(phi, n_wires=n_wires), expected)
+        assert np.allclose(op.eigvals(), expected)
+
+        # test arbitrary broadcasted global phase
+        phi = np.array([0.5, 0.4, 0.3])
+        op = qml.GlobalPhase(phi, wires=list(range(n_wires)))
+        expected = np.array([np.exp(-1j * p) * np.ones(dim) for p in phi])
+        assert np.allclose(op.compute_eigvals(phi, n_wires=n_wires), expected)
         assert np.allclose(op.eigvals(), expected)
 
 
@@ -2790,10 +2821,8 @@ class TestPauliRot:
         qml.assert_equal(decomp_op, qml.GlobalPhase(theta / 2))
 
         op_matrices = op.matrix()
-        decomp_op_matrices = decomp_op.matrix().T
-        assert len(op_matrices) == len(decomp_op_matrices)
-        for op_matrix, decomp_phase in zip(op_matrices, decomp_op_matrices):
-            assert qml.math.allclose(op_matrix, decomp_phase * np.eye(4))
+        decomp_op_matrices = decomp_op.matrix(wire_order=[0, 1])
+        assert qml.math.allclose(op_matrices, decomp_op_matrices)
 
     @pytest.mark.parametrize("theta", [0.4, np.array([np.pi / 3, 0.1, -0.9])])
     def test_PauliRot_decomposition_ZZ(self, theta):
@@ -3922,7 +3951,7 @@ def test_diagonalization_static_global_phase():
 @pytest.mark.parametrize("phi", [0.123, np.pi / 4, 0])
 @pytest.mark.parametrize("n_wires", [0, 1, 2])
 def test_global_phase_compute_sparse_matrix(phi, n_wires):
-    """Test that compute_sparse_matrix"""
+    """Test compute_sparse_matrix"""
 
     sparse_matrix = qml.GlobalPhase.compute_sparse_matrix(phi, n_wires=n_wires)
     expected = np.exp(-1j * phi) * sparse.eye(int(2**n_wires), format="csr")
@@ -3930,6 +3959,15 @@ def test_global_phase_compute_sparse_matrix(phi, n_wires):
     assert np.allclose(sparse_matrix.todense(), expected.todense())
 
 
+@pytest.mark.parametrize("n_wires", [0, 1, 2])
+def test_global_phase_compute_sparse_matrix_broadcasted_raises(n_wires):
+    """Test that compute_sparse_matrix raises an error for broadcasted GlobalPhase"""
+
+    phi = np.array([0.123, np.pi / 4, 0])
+    with pytest.raises(qml.operation.SparseMatrixUndefinedError, match="broadcasting"):
+        _ = qml.GlobalPhase.compute_sparse_matrix(phi, n_wires=n_wires)
+
+
 def test_decomposition():
     """Test the decomposition of the GlobalPhase operation."""