Repository: PennyLaneAI/pennylane
Pull request #4745 (2023-10-27 08:00:29)

Support batching with `ArbitraryUnitary`
`ArbitraryUnitary` does not support batching yet.

This PR adds batching support.


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4745?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:
> Comparison is base [(`970d30c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/970d30c252c148f2bcf8ff14793f6ad349b55cc7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`5d2a73c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4745?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.

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


```diff
@@            Coverage Diff             @@
##           master    #4745      +/-   ##
==========================================
- Coverage   99.64%   99.64%   -0.01%     
==========================================
  Files         380      380              
  Lines       34242    33982     -260     
==========================================
- Hits        34121    33860     -261     
- Misses        121      122       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4745?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [...nnylane/templates/subroutines/arbitrary\_unitary.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4745?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdWJyb3V0aW5lcy9hcmJpdHJhcnlfdW5pdGFyeS5weQ==) | `100.00% <100.00%> (ø)` | |

... and [42 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4745/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8312e4d85..73865ec7b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -9,6 +9,9 @@
 * Updates to some relevant Pytests to enable its use as a suite of benchmarks.
   [(#4703)](https://github.com/PennyLaneAI/pennylane/pull/4703)
 
+* `qml.ArbitraryUnitary` now supports batching.
+  [(#4745)](https://github.com/PennyLaneAI/pennylane/pull/4745)
+
 <h3>Breaking changes 💔</h3>
 
 <h3>Deprecations 👋</h3>
@@ -21,4 +24,5 @@
 
 This release contains contributions from (in alphabetical order):
 
-Amintor Dusko
\ No newline at end of file
+Amintor Dusko,
+David Wierichs,
diff --git a/pennylane/templates/subroutines/arbitrary_unitary.py b/pennylane/templates/subroutines/arbitrary_unitary.py
index 2bb50b5fe..cb39f6348 100644
--- a/pennylane/templates/subroutines/arbitrary_unitary.py
+++ b/pennylane/templates/subroutines/arbitrary_unitary.py
@@ -93,20 +93,19 @@ class ArbitraryUnitary(Operation):
 
     num_wires = AnyWires
     grad_method = None
+    num_params = 1
+    ndim_params = (1,)
 
     def __init__(self, weights, wires, id=None):
         shape = qml.math.shape(weights)
-        if shape != (4 ** len(wires) - 1,):
+        dim = 4 ** len(wires) - 1
+        if len(shape) not in (1, 2) or shape[-1] != dim:
             raise ValueError(
-                f"Weights tensor must be of shape {(4 ** len(wires) - 1,)}; got {shape}."
+                f"Weights tensor must be of shape {(dim,)} or (batch_dim, {dim}); got {shape}."
             )
 
         super().__init__(weights, wires=wires, id=id)
 
-    @property
-    def num_params(self):
-        return 1
-
     @staticmethod
     def compute_decomposition(weights, wires):  # pylint: disable=arguments-differ
         r"""Representation of the operator as a product of other operators.
@@ -129,7 +128,7 @@ class ArbitraryUnitary(Operation):
         op_list = []
 
         for i, pauli_word in enumerate(_all_pauli_words_but_identity(len(wires))):
-            op_list.append(PauliRot(weights[i], pauli_word, wires=wires))
+            op_list.append(PauliRot(weights[..., i], pauli_word, wires=wires))
 
         return op_list
 


Test patch:
diff --git a/tests/templates/test_subroutines/test_arbitrary_unitary.py b/tests/templates/test_subroutines/test_arbitrary_unitary.py
index 7f6ac28cd..28d1e9a72 100644
--- a/tests/templates/test_subroutines/test_arbitrary_unitary.py
+++ b/tests/templates/test_subroutines/test_arbitrary_unitary.py
@@ -92,18 +92,21 @@ class TestHelpers:
         ],
     )
     def test_tuple_to_word(self, tuple, expected_word):
+        """Test that tuples of indices are correctly mapped to words."""
         assert _tuple_to_word(tuple) == expected_word
 
 
 class TestDecomposition:
     """Tests that the template defines the correct decomposition."""
 
-    def test_correct_gates_single_wire(self):
+    @pytest.mark.parametrize("batch_dim", [None, 1, 2])
+    def test_correct_gates_single_wire(self, batch_dim):
         """Test that the correct gates are applied on a single wire."""
-        weights = np.arange(3, dtype=float)
+        shape = (3,) if batch_dim is None else (batch_dim, 3)
+        weights = np.arange(np.prod(shape), dtype=float).reshape(shape)
 
         op = qml.ArbitraryUnitary(weights, wires=[0])
-        queue = op.expand().operations
+        queue = op.decomposition()
 
         for gate in queue:
             assert gate.name == "PauliRot"
@@ -112,15 +115,17 @@ class TestDecomposition:
         pauli_words = ["X", "Y", "Z"]
 
         for i, op in enumerate(queue):
-            assert op.data[0] == weights[i]
+            assert np.allclose(op.data[0], weights[..., i])
             assert op.hyperparameters["pauli_word"] == pauli_words[i]
 
-    def test_correct_gates_two_wires(self):
+    @pytest.mark.parametrize("batch_dim", [None, 1, 2])
+    def test_correct_gates_two_wires(self, batch_dim):
         """Test that the correct gates are applied on two wires."""
-        weights = np.arange(15, dtype=float)
+        shape = (15,) if batch_dim is None else (batch_dim, 15)
+        weights = np.arange(np.prod(shape), dtype=float).reshape(shape)
 
         op = qml.ArbitraryUnitary(weights, wires=[0, 1])
-        queue = op.expand().operations
+        queue = op.decomposition()
 
         for gate in queue:
             assert gate.name == "PauliRot"
@@ -145,7 +150,7 @@ class TestDecomposition:
         ]
 
         for i, op in enumerate(queue):
-            assert op.data[0] == weights[i]
+            assert np.allclose(op.data[0], weights[..., i])
             assert op.hyperparameters["pauli_word"] == pauli_words[i]
 
     def test_custom_wire_labels(self, tol):
@@ -175,7 +180,8 @@ class TestDecomposition:
 class TestInputs:
     """Test inputs and pre-processing."""
 
-    def test_exception_wrong_dim(self):
+    @pytest.mark.parametrize("shape", [(5,), (1, 2, 3), (15, 3)])
+    def test_exception_wrong_dim(self, shape):
         """Verifies that exception is raised if the
         number of dimensions of features is incorrect."""
         dev = qml.device("default.qubit", wires=2)
@@ -186,7 +192,7 @@ class TestInputs:
             return qml.expval(qml.PauliZ(0))
 
         with pytest.raises(ValueError, match="Weights tensor must be of shape"):
-            weights = np.array([0, 1])
+            weights = np.zeros(shape)
             circuit(weights)
 
     def test_id(self):