Repository: PennyLaneAI/pennylane

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_

Discussion:
albi3ro: We currently have a PR up adding more documentation about broadcasting: https://github.com/PennyLaneAI/pennylane/pull/3356/files

Sorry I didn't link this earlier. Hope it helps.

For example of what we are going for, consider:

```pycon
>>> qml.IsingXX(np.array([1.2, 2.3, 3.4]), wires=(0,1)).decomposition()
[CNOT(wires=[0, 1]),
 RX(tensor([1.2, 2.3, 3.4], requires_grad=True), wires=[0]),
 CNOT(wires=[0, 1])]
```

Since the `IsingXX` gate started with a broadcasted parameter, it's decomposition ends up with a `RX` gate with a broadcasted parameter.

Unfortunately, the stated interface for [`Operator.decomposition`](https://docs.pennylane.ai/en/stable/code/api/pennylane.operation.Operator.html#pennylane.operation.Operator.decomposition) is that it returns `list[Operator]`. Deviating from that would require a breaking change to an important method for a  core data structure.

For now, we can make the assumption in `zyz_decompsoition` that if the input matrix is broadcasted, we will always return a `Rot` gate. That way we don't have to worry about choosing between `Rot` and `RZ`.
albi3ro: The end goal is to support the decomposition of `QubitUnitary` when the matrix input has a batch dimension and (at least) acts on one wire.  This could be accomplished through purely changing `QubitUnitary.decomposition`, but it could also be accomplished by changing `zyz_decomposition`.

I shared the `IsingXX` example to show what a decomposition with broadcasted parameters should look like.

Instead of:
```
[[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])]]
```

Something like:
```
>>> op.decomposition()
[Rot(tensor([0, 0, 0], requires_grad=True), tensor([1.2, 2.3, 2.7], requires_grad=True), tensor([0, 0, 0], requires_grad=True), wires=[0])]
```
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3477](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2196f2c) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/004aac284ad4785c89efb8de47dc7048518321a0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (004aac2) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #3477      +/-   ##
==========================================
+ Coverage   99.69%   99.71%   +0.01%     
==========================================
  Files         308      308              
  Lines       27164    27169       +5     
==========================================
+ Hits        27082    27092      +10     
+ Misses         82       77       -5     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/ops/qubit/matrix\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9tYXRyaXhfb3BzLnB5) | `100.00% <100.00%> (ø)` | |
| [.../transforms/decompositions/single\_qubit\_unitary.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVjb21wb3NpdGlvbnMvc2luZ2xlX3F1Yml0X3VuaXRhcnkucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `98.17% <0.00%> (+0.01%)` | :arrow_up: |
| [pennylane/transforms/qfunc\_transforms.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3477/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWZ1bmNfdHJhbnNmb3Jtcy5weQ==) | `100.00% <0.00%> (+6.94%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


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.