Repository: PennyLaneAI/pennylane (python)

Issue:
Implement decomposition of `QubitUnitary` with parameter broadcasting
### Feature details

When using a broadcasted unitary matrix in `QubitUnitary` (possible as of #2609), the `decomposition` is not defined yet.
This can be implemented but may not be entirely straightforward.

### Implementation

To implement this feature:
1. The broadcasted matrix of shape `(batch_size, dim, dim)` has to be split into `batch_size` matrices (trivial)
2. The correct decomposition method (`zyz_decomposition` for 1 qubit/`two_qubit_decomposition` for 2 qubits) has to be applied to each of the matrices from step 1. (for example list comprehension)
3. The data of the operations in the resulting decompositions has to be read out and unified (may be tricky depending on how homogeneous the decompositions obtained in step 2. are),
4. operations containing the broadcasted data from step 3 have to be collected into a single decomposition list (might also be a bit tricky).

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

_No response_

Pull request description:
Add support for broadcast `QubitUnitary` in `zyz_decomposition`
**Context:**
This is one of two PRs (the second being #3494) adding support for broadcast parameters to the unitary decomposition routine of the `QubitUnitary` operator.

**Description of the Change:**
- Users now can retrieve individual decompositions for each item in a broadcast single-qubit `QubitUnitary`. This used to raise an error, but now is allowed. An example usage is shown below:
```python
>>> mats = qml.RY.compute_matrix(np.array([1.2, 2.3, 2.7]))
>>> op = qml.QubitUnitary(mats, wires=0)
>>> decompositions = op.decomposition()
>>> pprint(decompositions)
[[Rot(tensor(-2.98266051e-128, requires_grad=True), tensor(1.2, requires_grad=True), tensor(5.59403882e-129, requires_grad=True), wires=[0])],
 [Rot(tensor(-3.5436292e-128, requires_grad=True), tensor(2.3, requires_grad=True), tensor(-1.35248216e-128, requires_grad=True), wires=[0])],
 [Rot(tensor(-5.59095128e-128, requires_grad=True), tensor(2.7, requires_grad=True), tensor(-3.54118996e-128, requires_grad=True), wires=[0])]]
```
- The approach relies on taking advantage of vectorized operations to maximize performance in comparison to a simple loop-and-gather approach.
- The already-existent behavior for non-batched input is still maintained 

**Benefits:**
- Performance gain for users who previously had to calculate decompositions individually and later on merge the results.

**Possible Drawbacks:**
The current implementation doesn't distinguish between the following two usages.
```python
# Usage 1
>>> mats = qml.RY.compute_matrix(np.array([1.2]))
>>> op = qml.QubitUnitary(mats, wires=0)
>>> decompositions = op.decomposition()
# Usage 2
>>> mats = qml.RY.compute_matrix(1.2)
>>> op = qml.QubitUnitary(mats, wires=0)
>>> decompositions = op.decomposition()
```
with the output in both bases having the form `[Rot(tensor(-2.98266051e-128, requires_grad=True), tensor(1.2, requires_grad=True), tensor(5.59403882e-129, requires_grad=True), wires=[0])]` (a single-item list). This can be a source of confusion to the users. A more consistent output may be a list of a list containing the decomposition for Usage 1, and a minor change of the code would be needed if others agree this is a better output. 

**Related GitHub Issues:**
Resolves #2636 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0df4dd3a6..5fe8e4fbd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -303,6 +303,11 @@
   the optional `use_grouping` attribute.
   [(#3456)](https://github.com/PennyLaneAI/pennylane/pull/3456)
 
+* Updated `zyz_decomposition` function such that it now supports broadcast operators. This
+  means that single-qubit `QubitUnitary` operators, instantiated from a batch of unitaries,
+  can now be decomposed.
+  [(#3477)](https://github.com/PennyLaneAI/pennylane/pull/3477)
+
 * Reduce usage of `MeasurementProcess.return_type`. Use `isinstance` checks instead.
   [(#3399)](https://github.com/PennyLaneAI/pennylane/pull/3399)
 
@@ -638,6 +643,7 @@ This release contains contributions from (in alphabetical order):
 Juan Miguel Arrazola
 Utkarsh Azad
 Astral Cai
+Ahmed Darwish
 Isaac De Vlugt
 Pieter Eendebak
 Lillian M. A. Frederiksen
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index cd4b1ae2a..864b66596 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -162,17 +162,21 @@ class QubitUnitary(Operation):
         # Decomposes arbitrary single-qubit unitaries as Rot gates (RZ - RY - RZ format),
         # or a single RZ for diagonal matrices.
         shape = qml.math.shape(U)
-        if shape == (2, 2):
+
+        is_batched = len(shape) == 3
+        shape_without_batch_dim = shape[1:] if is_batched else shape
+
+        if shape_without_batch_dim == (2, 2):
             return qml.transforms.decompositions.zyz_decomposition(U, Wires(wires)[0])
 
-        if shape == (4, 4):
-            return qml.transforms.two_qubit_decomposition(U, Wires(wires))
+        if shape_without_batch_dim == (4, 4):
+            # TODO[dwierichs]: Implement decomposition of broadcasted unitary
+            if is_batched:
+                raise DecompositionUndefinedError(
+                    "The decomposition of a two-qubit QubitUnitary does not support broadcasting."
+                )
 
-        # TODO[dwierichs]: Implement decomposition of broadcasted unitary
-        if len(shape) == 3:
-            raise DecompositionUndefinedError(
-                "The decomposition of QubitUnitary does not support broadcasting."
-            )
+            return qml.transforms.two_qubit_decomposition(U, Wires(wires))
 
         return super(QubitUnitary, QubitUnitary).compute_decomposition(U, wires=wires)
 
diff --git a/pennylane/transforms/decompositions/single_qubit_unitary.py b/pennylane/transforms/decompositions/single_qubit_unitary.py
index 07d875a6e..ece75c1b4 100644
--- a/pennylane/transforms/decompositions/single_qubit_unitary.py
+++ b/pennylane/transforms/decompositions/single_qubit_unitary.py
@@ -29,11 +29,11 @@ def _convert_to_su2(U):
         array[complex]: A :math:`2 \times 2` matrix in :math:`SU(2)` that is
         equivalent to U up to a global phase.
     """
-    # Compute the determinant
-    det = U[0, 0] * U[1, 1] - U[0, 1] * U[1, 0]
+    # Compute the determinants
+    dets = math.linalg.det(U)
 
-    exp_angle = -1j * math.cast_like(math.angle(det), 1j) / 2
-    return U * math.exp(exp_angle)
+    exp_angles = -1j * math.cast_like(math.angle(dets), 1j) / 2
+    return math.cast_like(U, dets) * math.exp(exp_angles)[:, None, None]
 
 
 def zyz_decomposition(U, wire):
@@ -76,26 +76,39 @@ def zyz_decomposition(U, wire):
     >>> decomp
     [Rot(-0.24209529417800013, 1.14938178234275, 1.7330581433950871, wires=[0])]
     """
-    U = _convert_to_su2(U)
 
-    # If the value of U is not abstract, we can include a conditional statement
-    # that will check if the off-diagonal elements are 0; if so, just use one RZ
-    if not math.is_abstract(U):
-        if math.allclose(U[0, 1], 0.0):
-            return [qml.RZ(2 * math.angle(U[1, 1]), wires=wire)]
+    # Cast to batched format for more consistent code
+    U = math.expand_dims(U, axis=0) if len(U.shape) == 2 else U
 
-    # Derive theta from the off-diagonal element. Clip to ensure valid arcsin input
-    element = math.clip(math.abs(U[0, 1]), 0, 1)
-    theta = 2 * math.arcsin(element)
+    U = _convert_to_su2(U)
 
-    # Compute phi and omega from the angles of the top row; use atan2 to keep
-    # the angle within -np.pi and np.pi, and add very small values to avoid the
-    # undefined case of 0/0. We add a smaller value to the imaginary part than
-    # the real part because it is imag / real in the definition of atan2.
-    angle_U00 = math.arctan2(math.imag(U[0, 0]) + 1e-128, math.real(U[0, 0]) + 1e-64)
-    angle_U10 = math.arctan2(math.imag(U[1, 0]) + 1e-128, math.real(U[1, 0]) + 1e-64)
+    # If U is only one unitary and its value is not abstract, we can include a conditional
+    # statement that will check if the off-diagonal elements are 0; if so, just use one RZ
+    if len(U) == 1 and not math.is_abstract(U[0]):
+        if math.allclose(U[0, 0, 1], 0.0):
+            return [qml.RZ(2 * math.angle(U[0, 1, 1]), wires=wire)]
 
-    phi = -angle_U10 - angle_U00
-    omega = angle_U10 - angle_U00
+    # For batched U or single U with non-zero off-diagonal, compute the
+    # Rot operator decomposition instead
+    off_diagonal_elements = math.clip(math.abs(U[:, 0, 1]), 0, 1)
+    thetas = 2 * math.arcsin(off_diagonal_elements)
 
-    return [qml.Rot(phi, theta, omega, wires=wire)]
+    # Compute phi and omega from the angles of the top row; use atan2 to keep
+    # the angle within -np.pi and np.pi, and add very small value to the real
+    # part to avoid division by zero.
+    epsilon = 1e-64
+    angles_U00 = math.arctan2(
+        math.imag(U[:, 0, 0]),
+        math.real(U[:, 0, 0]) + epsilon,
+    )
+    angles_U10 = math.arctan2(
+        math.imag(U[:, 1, 0]),
+        math.real(U[:, 1, 0]) + epsilon,
+    )
+
+    phis = -angles_U10 - angles_U00
+    omegas = angles_U10 - angles_U00
+
+    phis, thetas, omegas = map(math.squeeze, [phis, thetas, omegas])
+
+    return [qml.Rot(phis, thetas, omegas, wires=wire)]


Test patch:
diff --git a/tests/ops/qubit/test_matrix_ops.py b/tests/ops/qubit/test_matrix_ops.py
index 711801ae2..5a0b3a512 100644
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -256,6 +256,15 @@ class TestQubitUnitary:
                 qml.Rot,
                 [-1.0, 2.0, -3.0],
             ),
+            # An instance of a broadcast unitary
+            (
+                np.exp(1j * 0.02)
+                * qml.Rot(
+                    np.array([1.2, 2.3]), np.array([0.12, 0.5]), np.array([0.98, 0.567]), wires=0
+                ).matrix(),
+                qml.Rot,
+                [[1.2, 2.3], [0.12, 0.5], [0.98, 0.567]],
+            ),
         ],
     )
     def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
@@ -269,16 +278,14 @@ class TestQubitUnitary:
         assert isinstance(decomp2[0], expected_gate)
         assert np.allclose(decomp2[0].parameters, expected_params, atol=1e-7)
 
-    def test_error_qubit_unitary_decomposition_broadcasted(self):
+    def test_broadcasted_two_qubit_qubit_unitary_decomposition_raises_error(self):
         """Tests that broadcasted QubitUnitary decompositions are not supported."""
-        U = np.array(
-            [[0.98877108 + 0.0j, 0.0 - 0.14943813j], [0.0 - 0.14943813j, 0.98877108 + 0.0j]]
-        )
-        U = np.tensordot([1j, -1.0, (1 + 1j) / np.sqrt(2)], U, axes=0)
+        U = qml.IsingYY.compute_matrix(np.array([1.2, 2.3, 3.4]))
+
         with pytest.raises(DecompositionUndefinedError, match="QubitUnitary does not support"):
-            qml.QubitUnitary.compute_decomposition(U, wires=0)
+            qml.QubitUnitary.compute_decomposition(U, wires=[0, 1])
         with pytest.raises(DecompositionUndefinedError, match="QubitUnitary does not support"):
-            qml.QubitUnitary(U, wires=0).decomposition()
+            qml.QubitUnitary(U, wires=[0, 1]).decomposition()
 
     def test_qubit_unitary_decomposition_multiqubit_invalid(self):
         """Test that QubitUnitary is not decomposed for more than two qubits."""
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index 6d936b9df..8f815facb 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -42,6 +42,7 @@ single_qubit_decomps = [
     (qml.RZ(0.3, wires=0).matrix(), qml.RZ, [0.3]),
     (qml.RZ(-0.5, wires=0).matrix(), qml.RZ, [-0.5]),
     # Next set of gates are non-diagonal and decomposed as Rots
+    (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix(), qml.Rot, [0.2, 0.5, -0.3]),
     (
         np.array(
             [
@@ -54,35 +55,57 @@ single_qubit_decomps = [
     ),
     (H, qml.Rot, [np.pi, np.pi / 2, 0.0]),
     (X, qml.Rot, [np.pi / 2, np.pi, -np.pi / 2]),
-    (qml.Rot(0.2, 0.5, -0.3, wires=0).matrix(), qml.Rot, [0.2, 0.5, -0.3]),
     (
         np.exp(1j * 0.02) * qml.Rot(-1.0, 2.0, -3.0, wires=0).matrix(),
         qml.Rot,
         [-1.0, 2.0, -3.0],
     ),
+    # Add two instances of broadcasted unitaries, one coming from RZ and another from Rot
+    (
+        qml.QubitUnitary(qml.RZ.compute_matrix(np.array([np.pi, np.pi / 2])), wires=0).matrix(),
+        qml.Rot,
+        [[np.pi / 2, np.pi / 4], [0.0, 0.0], [np.pi / 2, np.pi / 4]],
+    ),
+    (
+        qml.Rot(np.array([1.2, 2.3]), np.array([1.2, 2.3]), np.array([1.2, 2.3]), wires=0).matrix(),
+        qml.Rot,
+        [[1.2, 2.3], [1.2, 2.3], [1.2, 2.3]],
+    ),
 ]
 
 
 class TestQubitUnitaryZYZDecomposition:
     """Test that the decompositions are correct."""
 
-    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
-    def test_zyz_decomposition(self, U, expected_gate, expected_params):
-        """Test that a one-qubit matrix in isolation is correctly decomposed."""
-        obtained_gates = zyz_decomposition(U, Wires("a"))
-
+    def _run_assertions(self, U, expected_gate, expected_params, obtained_gates):
         assert len(obtained_gates) == 1
         assert isinstance(obtained_gates[0], expected_gate)
         assert obtained_gates[0].wires == Wires("a")
 
-        assert qml.math.allclose(obtained_gates[0].parameters, expected_params, atol=1e-7)
+        assert qml.math.allclose(
+            qml.math.unwrap(obtained_gates[0].parameters), expected_params, atol=1e-7
+        )
 
         if obtained_gates[0].num_params == 1:
             obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).matrix()
         else:
             obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).matrix()
 
-        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
+        if len(obtained_mat.shape) == 2:
+            U = [U]
+            obtained_mat = [obtained_mat]
+
+        assert all(
+            check_matrix_equivalence(curr_obtained_mat, curr_U, atol=1e-7)
+            for curr_obtained_mat, curr_U in zip(obtained_mat, qml.math.unwrap(U))
+        )
+
+    @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
+    def test_zyz_decomposition(self, U, expected_gate, expected_params):
+        """Test that a one-qubit matrix in isolation is correctly decomposed."""
+        obtained_gates = zyz_decomposition(U, Wires("a"))
+
+        self._run_assertions(U, expected_gate, expected_params, obtained_gates)
 
     @pytest.mark.torch
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
@@ -94,20 +117,7 @@ class TestQubitUnitaryZYZDecomposition:
 
         obtained_gates = zyz_decomposition(U, wire="a")
 
-        assert len(obtained_gates) == 1
-        assert isinstance(obtained_gates[0], expected_gate)
-        assert obtained_gates[0].wires == Wires("a")
-
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[0].parameters), expected_params, atol=1e-7
-        )
-
-        if obtained_gates[0].num_params == 1:
-            obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).matrix()
-        else:
-            obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).matrix()
-
-        assert check_matrix_equivalence(obtained_mat, qml.math.unwrap(U), atol=1e-7)
+        self._run_assertions(U, expected_gate, expected_params, obtained_gates)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
@@ -119,19 +129,7 @@ class TestQubitUnitaryZYZDecomposition:
 
         obtained_gates = zyz_decomposition(U, wire="a")
 
-        assert len(obtained_gates) == 1
-        assert isinstance(obtained_gates[0], expected_gate)
-        assert obtained_gates[0].wires == Wires("a")
-
-        assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
-
-        if obtained_gates[0].num_params == 1:
-            # With TF and RZ, need to cast since can't just unwrap
-            obtained_mat = qml.RZ(obtained_gates[0].parameters[0].numpy(), wires=0).matrix()
-        else:
-            obtained_mat = qml.Rot(*qml.math.unwrap(obtained_gates[0].parameters), wires=0).matrix()
-
-        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
+        self._run_assertions(U, expected_gate, expected_params, obtained_gates)
 
     @pytest.mark.jax
     @pytest.mark.parametrize("U,expected_gate,expected_params", single_qubit_decomps)
@@ -149,20 +147,7 @@ class TestQubitUnitaryZYZDecomposition:
 
         obtained_gates = zyz_decomposition(U, wire="a")
 
-        assert len(obtained_gates) == 1
-        assert isinstance(obtained_gates[0], expected_gate)
-        assert obtained_gates[0].wires == Wires("a")
-
-        assert qml.math.allclose(
-            qml.math.unwrap(obtained_gates[0].parameters), expected_params, atol=1e-7
-        )
-
-        if obtained_gates[0].num_params == 1:
-            obtained_mat = qml.RZ(obtained_gates[0].parameters[0], wires=0).matrix()
-        else:
-            obtained_mat = qml.Rot(*obtained_gates[0].parameters, wires=0).matrix()
-
-        assert check_matrix_equivalence(obtained_mat, U, atol=1e-7)
+        self._run_assertions(U, expected_gate, expected_params, obtained_gates)
 
 
 # Randomly generated set (scipy.unitary_group) of five U(4) operations.