Repository: PennyLaneAI/pennylane
Pull request #1883 (2021-11-09 22:00:40)

Added Batching to BasicEntangerLayers
### Before submitting

Please complete the following checklist when submitting a PR:

- [ ] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [ ] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [ ] Ensure that the test suite passes, by running `make test`.

- [ ] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [ ] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1883](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (60238ef) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/95dfeda9a8be010b11c9e77b3a1daedcbc35b951?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (95dfeda) will **decrease** coverage by `0.00%`.
> The diff coverage is `75.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1883      +/-   ##
==========================================
- Coverage   98.98%   98.97%   -0.01%     
==========================================
  Files         218      218              
  Lines       16453    16455       +2     
==========================================
+ Hits        16286    16287       +1     
- Misses        167      168       +1     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/templates/layers/basic\_entangler.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9sYXllcnMvYmFzaWNfZW50YW5nbGVyLnB5) | `97.14% <75.00%> (-2.86%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [95dfeda...60238ef](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1883?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

antalszava: [sc-11222]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1cc442f51..b056245d9 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -343,6 +343,8 @@
 
 * AngleEmbedding now supports `batch_params` decorator. [(#1812)](https://github.com/PennyLaneAI/pennylane/pull/1812)
 
+* BasicEntanglerLayers now supports `batch_params` decorator. [(#1883)](https://github.com/PennyLaneAI/pennylane/pull/1883)
+
 * MottonenStatePreparation now supports `batch_params` decorator. [(#1893)](https://github.com/PennyLaneAI/pennylane/pull/1893)
 
 * CircuitDrawer now supports a `max_length` argument to help prevent text overflows when printing circuits to the CLI. [#1841](https://github.com/PennyLaneAI/pennylane/pull/1841)
@@ -442,5 +444,5 @@
 
 This release contains contributions from (in alphabetical order):
 
-Guillermo Alonso-Linaje, Benjamin Cordier, Olivia Di Matteo, David Ittah, Josh Izaac, Jalani Kanem, Ankit Khandelwal, Shumpei Kobayashi,
+Guillermo Alonso-Linaje, Samuel Banning, Benjamin Cordier, Olivia Di Matteo, David Ittah, Josh Izaac, Jalani Kanem, Ankit Khandelwal, Shumpei Kobayashi,
 Robert Lang, Christina Lee, Cedric Lin, Alejandro Montanez, Romain Moyard, Maria Schuld, Jay Soni, David Wierichs
diff --git a/pennylane/templates/layers/basic_entangler.py b/pennylane/templates/layers/basic_entangler.py
index 1090e5c36..a8acce473 100644
--- a/pennylane/templates/layers/basic_entangler.py
+++ b/pennylane/templates/layers/basic_entangler.py
@@ -125,14 +125,23 @@ class BasicEntanglerLayers(Operation):
 
     def __init__(self, weights, wires=None, rotation=None, do_queue=True, id=None):
 
+        # convert weights to numpy array if weights is list otherwise keep unchanged
+        interface = qml.math.get_interface(weights)
+        weights = qml.math.asarray(weights, like=interface)
+
         self.rotation = rotation or qml.RX
 
         shape = qml.math.shape(weights)
-        if len(shape) != 2:
-            raise ValueError(f"Weights tensor must be 2-dimensional; got shape {shape}")
-        if shape[1] != len(wires):
+        if not (len(shape) == 3 or len(shape) == 2):  # 3 is when batching, 2 is no batching
+            raise ValueError(
+                f"Weights tensor must be 2-dimensional "
+                f"or 3-dimensional if batching; got shape {shape}"
+            )
+
+        if shape[-1] != len(wires):
+            # index with -1 since we may or may not have batching in first dimension
             raise ValueError(
-                f"Weights tensor must have second dimension of length {len(wires)}; got {shape[1]}"
+                f"Weights tensor must have last dimension of length {len(wires)}; got {shape[-1]}"
             )
 
         super().__init__(weights, wires=wires, do_queue=do_queue, id=id)
@@ -145,15 +154,15 @@ class BasicEntanglerLayers(Operation):
 
         weights = self.parameters[0]
 
-        # first dimension of the weights tensor determines
+        # first dimension of the weights tensor (second when batching) determines
         # the number of layers
-        repeat = qml.math.shape(weights)[0]
+        repeat = qml.math.shape(weights)[-2]
 
         with qml.tape.QuantumTape() as tape:
 
             for layer in range(repeat):
                 for i in range(len(self.wires)):
-                    self.rotation(weights[layer][i], wires=self.wires[i : i + 1])
+                    self.rotation(weights[..., layer, i], wires=self.wires[i : i + 1])
 
                 if len(self.wires) == 2:
                     qml.CNOT(wires=self.wires)


Test patch:
diff --git a/tests/templates/test_layers/test_basic_entangler.py b/tests/templates/test_layers/test_basic_entangler.py
index bec3f38c5..c6908256a 100644
--- a/tests/templates/test_layers/test_basic_entangler.py
+++ b/tests/templates/test_layers/test_basic_entangler.py
@@ -121,7 +121,7 @@ class TestInputs:
         with pytest.raises(ValueError, match="Weights tensor must be 2-dimensional"):
             circuit([1, 0])
 
-        with pytest.raises(ValueError, match="Weights tensor must have second dimension of length"):
+        with pytest.raises(ValueError, match="Weights tensor must have last dimension of length"):
             circuit([[1, 0], [1, 0]])
 
     def test_id(self):
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
index 7e070c86b..daa722587 100644
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -45,6 +45,26 @@ def test_simple_circuit(mocker):
     assert len(spy.call_args[0][0]) == batch_size
 
 
+def test_basic_entangler_layers(mocker):
+    """Test that batching works for BasicEngtanglerLayers"""
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(weights):
+        qml.templates.BasicEntanglerLayers(weights, wires=[0, 1])
+        qml.RY(0.2, wires=1)
+        return qml.probs(wires=[0, 1])
+
+    batch_size = 5
+    weights = np.random.random((batch_size, 2, 2))
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(weights)
+    assert res.shape == (batch_size, 4)
+    assert len(spy.call_args[0][0]) == batch_size
+
+
 def test_angle_embedding(mocker):
     """Test that batching works for AngleEmbedding"""
     dev = qml.device("default.qubit", wires=3)