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:
Links to relevant pieces of source code:

* [`QubitUnitary`](https://github.com/PennyLaneAI/pennylane/blob/6633fcb91a7d83900fce876295df60594a241e48/pennylane/ops/qubit/matrix_ops.py#L28)
* [`zyz_decomposition`](https://github.com/PennyLaneAI/pennylane/blob/6633fcb91a7d83900fce876295df60594a241e48/pennylane/transforms/decompositions/single_qubit_unitary.py#L39)
* [`two_qubit_decomposition`](https://github.com/PennyLaneAI/pennylane/blob/6633fcb91a7d83900fce876295df60594a241e48/pennylane/transforms/decompositions/two_qubit_unitary.py#L499)

## The problem

Suppose have matrices with a batch dimension:

```pycon
>>> x = np.array([1.2, 2.3, 2.7])
>>> mats = qml.RX.compute_matrix(x)
>>> mats
tensor([[[0.82533561+0.j        , 0.        -0.56464247j],
         [0.        -0.56464247j, 0.82533561+0.j        ]],

        [[0.40848744+0.j        , 0.        -0.91276394j],
         [0.        -0.91276394j, 0.40848744+0.j        ]],

        [[0.21900669+0.j        , 0.        -0.97572336j],
         [0.        -0.97572336j, 0.21900669+0.j        ]]], requires_grad=True)
>>> mats.shape
(3, 2, 2)
``` 

We can construct a `QubitUnitary` operation using this matrix batch:

```pycon
>>> op = qml.QubitUnitary(mats, wires=0)
```

But trying to decompose it leads to a `DecompositionUndefinedError`:
```pycon
>>> op.decomposition()
DecompositionUndefinedError: The decomposition of QubitUnitary does not support broadcasting.
```

 If we didn't have a batch dimension, we could decompose the `QubitUnitary` into  `Rot` gate:

```pycon
>>> qml.QubitUnitary(mats[0, ...], 0).decomposition()
[Rot(tensor(1.57079633, requires_grad=True), tensor(1.2, requires_grad=True), tensor(-1.57079633, requires_grad=True), wires=[0])]
```

We would like to be able to decompose `QubitUnitary` when it has a batch dimension as well.

## Possible Solutions

There are (at least) two main ways to solve this problem.

### Solution 1

`QubitUnitary.decomposition` could:

* loop over the batch dimension,
* create a decomposition for each individual matrix
* combine the decompositions back together

### Solution 2

We could update the `zyz_decomposition` and/ or `two_qubit_decomposition` functions to work with batched matrices.  Batch support for these two functions could be implemented independently in separate PR's.  

Once `zyz_decomposition` (and/ or `two_qubit_decomposition`) supports batched matrices, we would need to update `QubitUnitary.compute_decomposition` as well.

### Additional Requirements

Any PR solving this problem would need tests and a changelog entry.  
Thanks @albi3ro, this is super useful and clear!
After reading your comment, it feels like solution 2 is the better of the two, what do you think? In particular, we would provide batching support to any other functions that may use `zyz_decomposition` and/or `two_qubit_decomposition` in the future, reducing the overhead for such future implementations.

I suppose whether there is a manual iteration over the batch dimension ultimately depends on whether the decomposition functions use lower-level functions that readily support broadcasting or not...

Has the priority of this feature changed?
@dwierichs updated and clarified for potential use as a technical assessment.



Resolving patch:
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -162,17 +162,21 @@ def compute_decomposition(U, wires):
         # 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
--- 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
--- a/tests/ops/qubit/test_matrix_ops.py
+++ b/tests/ops/qubit/test_matrix_ops.py
@@ -256,6 +256,15 @@ def test_qubit_unitary_jax_jit(self, U, num_wires):
                 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 @@ def test_qubit_unitary_decomposition(self, U, expected_gate, expected_params):
         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
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -42,6 +42,7 @@
     (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 @@
     ),
     (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 @@ def test_zyz_decomposition_torch(self, U, expected_gate, expected_params):
 
         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 @@ def test_zyz_decomposition_tf(self, U, expected_gate, expected_params):
 
         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 @@ def test_zyz_decomposition_jax(self, U, expected_gate, expected_params):
 
         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.