Repository: PennyLaneAI/pennylane
Pull request #3915 (2023-03-15 21:34:52)

Various changes for lightning and/or arithmetic op support
**Context:**
These changes are needed to unblock PennyLaneAI/pennylane-lightning#424, along with a few others. I'll keep the ordering across sections below.

**Description of the Change:**
1. Do not run `sum_expand` on a Sum observable if a device supports Sum observables
3. CompositeOp now overrides `_check_batching` to use its own implementation
4. If a Sum has a `_pauli_rep`, then `Sum.is_hermitian` just checks that all coefficients in the Pauli basis are real
5. Support for batching is explicitly removed in Hamiltonian (it wouldn't have worked before, but I short-circuit by turning `Hamiltonian._check_batching` into a no-op)

**Benefits:**
1. Sums get native support when devices can support them
3. `CompositeOp._check_batching` is waaaaaay faster than using the base implementation, and it knows better anyway
4. It's equivalent, but faster. See performance improvements below.
5. Overriding `Hamiltonian._check_batching` isn't 100% needed, but it comes at no cost, and why not be explicit? I _think_ this speeds up computation too, but I can't remember in the mishmash of all these fixes.

**Possible Drawbacks:**
1. Maybe a device wants to use sum_expand? Idk
3. N/A
4. Maybe we don't want to lean on `_pauli_rep` too much? I think it's great whenever it's performant (it is).
5. N/A

**Related GitHub Issues:**
PennyLaneAI/pennylane-lightning#424

**Performance improvements for `Sum.is_hermitian`**
<details>
<summary>Prep code:</summary>

```python
import pennylane as qml
from pennylane import numpy as nnp

data = qml.data.load("qchem", molname=[
    "H2",  # 4 wires, 15 terms
    "H3+", # 6 wires, 66 terms
    "H4",  # 8 wires, 185 terms
    "H6",  # 12 wires, 919 terms
    "H2O", # 14 wires, 1086 terms
    "BH3", # 16 wires, 2497 terms
], basis="STO-3G", bondlength=0.86)
sum_ops = [qml.pauli.pauli_sentence(d.hamiltonian).operation() for d in data]

for op in sum_ops:
    print(f"\n{op.num_wires=}, {len(op)=}")
    %timeit all(s.is_hermitian for s in op)
    %timeit not any(qml.math.iscomplex(val) for val in op._pauli_rep.values())
```
</details>

<details>
<summary>Timings (old first, new second)</summary>

```
op.num_wires=4, len(op)=15
99.5 µs ± 525 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
45.6 µs ± 207 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

op.num_wires=6, len(op)=66
782 µs ± 4.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
199 µs ± 936 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

op.num_wires=8, len(op)=185
3.42 ms ± 15.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
559 µs ± 3.3 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

op.num_wires=12, len(op)=919
31.3 ms ± 89.5 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
2.77 ms ± 10.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

op.num_wires=14, len(op)=1086
45.6 ms ± 176 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
3.27 ms ± 5.45 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

op.num_wires=16, len(op)=2497
135 ms ± 216 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
7.56 ms ± 17.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```
</details>

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/3915?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3915](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (f91bb15) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/ad14bb2b1f084ff7b534df4fcfccb077d9bc510f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (ad14bb2) will **decrease** coverage by `0.01%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #3915      +/-   ##
==========================================
- Coverage   99.70%   99.70%   -0.01%     
==========================================
  Files         342      342              
  Lines       29856    29862       +6     
==========================================
+ Hits        29769    29774       +5     
- Misses         87       88       +1     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `99.68% <100.00%> (-0.32%)` | :arrow_down: |
| [pennylane/ops/op\_math/composite.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2NvbXBvc2l0ZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/op\_math/sum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3N1bS5weQ==) | `98.16% <100.00%> (ø)` | |
| [pennylane/ops/qubit/hamiltonian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9oYW1pbHRvbmlhbi5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/qubit/state\_preparation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3915?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9zdGF0ZV9wcmVwYXJhdGlvbi5weQ==) | `100.00% <100.00%> (ø)` | |

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

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)

timmysilv: I computed a bunch of timings (on this branch) with lightning, computing diagonalizing gates or not, doing sum expansion or not. In general, it seems that doing sum expansion is almost an order of magnitude faster, but there are some things to note:
1. Almost all of the extra time is in tape.get/set_parameters, and that needs to be improved anyway. If I time on another branch, results might be better
2. `sum_expand` is still broken, so perhaps some of the time saved comes from just not computing what's needed? lol. Comparing with Hamiltonian (which expands correctly), I'm sure this difference is negligible.
3. When not expanding, the first computation for a large Sum is extra-costly because it has to compute things. Subsequent computations use cached results

For small Hamiltonians on default.qubit, it's similar when you do or don't expand. Larger Hamiltonians (1k-2.5k term) are about 1.5x faster when not expanding.

On lightning.qubit, it is always faster to _not_ expand a Hamiltonian.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9f4921855..c625d0cca 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -99,6 +99,17 @@
 * Changed `Operator.num_wires` from an abstract value to `AnyWires`.
   [(#3919)](https://github.com/PennyLaneAI/pennylane/pull/3919)
 
+* Do not run `qml.transforms.sum_expand` in `Device.batch_transform` if the device supports Sum observables.
+  [(#3915)](https://github.com/PennyLaneAI/pennylane/pull/3915)
+
+* `CompositeOp` now overrides `Operator._check_batching`, providing a significant performance improvement.
+  `Hamiltonian` also overrides this method and does nothing, because it does not support batching.
+  [(#3915)](https://github.com/PennyLaneAI/pennylane/pull/3915)
+
+* If a `Sum` operator has a pre-computed Pauli representation, `is_hermitian` now checks that all coefficients
+  are real, providing a significant performance improvement.
+  [(#3915)](https://github.com/PennyLaneAI/pennylane/pull/3915)
+
 <h3>Breaking changes</h3>
 
 * Both JIT interfaces are not compatible with Jax `>0.4.3`, we raise an error for those versions.
diff --git a/pennylane/_device.py b/pennylane/_device.py
index ca72321d3..064011c3a 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -690,6 +690,7 @@ class Device(abc.ABC):
             will natively support all operations.
         """
         if self.custom_expand_fn is not None:
+            # pylint:disable=not-callable
             return self.custom_expand_fn(circuit, max_expansion=max_expansion)
 
         return self.default_expand_fn(circuit, max_expansion=max_expansion)
@@ -721,6 +722,7 @@ class Device(abc.ABC):
             to be applied to the list of evaluated circuit results.
         """
         supports_hamiltonian = self.supports_observable("Hamiltonian")
+        supports_sum = self.supports_observable("Sum")
         finite_shots = self.shots is not None
         grouping_known = all(
             obs.grouping_indices is not None
@@ -750,7 +752,7 @@ class Device(abc.ABC):
                 raise ValueError(
                     "Can only return the expectation of a single Hamiltonian observable"
                 ) from e
-        elif expval_sum_in_obs and not is_shadow:
+        elif expval_sum_in_obs and not is_shadow and not supports_sum:
             circuits, hamiltonian_fn = qml.transforms.sum_expand(circuit)
 
         elif (
diff --git a/pennylane/ops/op_math/composite.py b/pennylane/ops/op_math/composite.py
index 8d3fa972e..2653110b5 100644
--- a/pennylane/ops/op_math/composite.py
+++ b/pennylane/ops/op_math/composite.py
@@ -65,16 +65,16 @@ class CompositeOp(Operator):
         if do_queue:
             self.queue()
 
-        self._batch_size = self._check_and_compute_batch_size()
+        self._check_batching(None)  # unused param
 
-    def _check_and_compute_batch_size(self):
+    def _check_batching(self, _):
         batch_sizes = {op.batch_size for op in self if op.batch_size is not None}
         if len(batch_sizes) > 1:
             raise ValueError(
                 "Broadcasting was attempted but the broadcasted dimensions "
                 f"do not match: {batch_sizes}."
             )
-        return batch_sizes.pop() if batch_sizes else None
+        self._batch_size = batch_sizes.pop() if batch_sizes else None
 
     def __repr__(self):
         return f" {self._op_symbol} ".join(
diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py
index 85bb834fb..2a1288385 100644
--- a/pennylane/ops/op_math/sum.py
+++ b/pennylane/ops/op_math/sum.py
@@ -186,7 +186,11 @@ class Sum(CompositeOp):
     @property
     def is_hermitian(self):
         """If all of the terms in the sum are hermitian, then the Sum is hermitian."""
-        return all(s.is_hermitian for s in self)
+        return (
+            all(s.is_hermitian for s in self)
+            if self._pauli_rep is None
+            else not any(qml.math.iscomplex(val) for val in self._pauli_rep.values())
+        )
 
     def terms(self):
         r"""Representation of the operator as a linear combination of other operators.
diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index 06a1d6de0..bfb44b07e 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -163,6 +163,8 @@ class Hamiltonian(Observable):
 
     num_wires = qml.operation.AnyWires
     grad_method = "A"  # supports analytic gradients
+    batch_size = None
+    ndim_params = None  # could be (0,) * len(coeffs), but it is not needed. Define at class-level
 
     def __init__(
         self,
@@ -216,6 +218,9 @@ class Hamiltonian(Observable):
         # while H.coeffs is the original tensor
         super().__init__(*coeffs_flat, wires=self._wires, id=id, do_queue=do_queue)
 
+    def _check_batching(self, params):
+        """Override for Hamiltonian, batching is not yet supported."""
+
     def label(self, decimals=None, base_label=None, cache=None):
         decimals = None if (len(self.parameters) > 3) else decimals
         return super().label(decimals=decimals, base_label=base_label or "𝓗", cache=cache)
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index 46cdf2bbc..c4ee19c72 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -61,6 +61,9 @@ class BasisState(StatePrep):
     num_params = 1
     """int: Number of trainable parameters that the operator depends on."""
 
+    ndim_params = (1,)
+    """int: Number of dimensions per trainable parameter of the operator."""
+
     @staticmethod
     def compute_decomposition(n, wires):
         r"""Representation of the operator as a product of other operators (static method). :


Test patch:
diff --git a/tests/test_device.py b/tests/test_device.py
index b4b4f0234..b1af30e4c 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -983,35 +983,63 @@ class TestBatchExecution:
 class TestGrouping:
     """Tests for the use_grouping option for devices."""
 
+    class SomeDevice(qml.Device):
+        name = ""
+        short_name = ""
+        pennylane_requires = ""
+        version = ""
+        author = ""
+        operations = ""
+        observables = ""
+        apply = lambda *args, **kwargs: 0
+        expval = lambda *args, **kwargs: 0
+        reset = lambda *args, **kwargs: 0
+        supports_observable = lambda *args, **kwargs: True
+
     @pytest.mark.parametrize("use_grouping", (True, False))
-    def test_batch_transform_checks_use_grouping_property(self, use_grouping):
+    def test_batch_transform_checks_use_grouping_property(self, use_grouping, mocker):
         """If the device specifies `use_grouping=False`, the batch transform
         method won't expand the hamiltonian when the measured hamiltonian has
         grouping indices.
         """
 
-        class TestDevice(qml.Device):
-            name = ""
-            short_name = ""
-            pennylane_requires = ""
-            version = ""
-            author = ""
-            operations = ""
-            observables = ""
-            apply = lambda *args, **kwargs: 0
-            expval = lambda *args, **kwargs: 0
-            reset = lambda *args, **kwargs: 0
-            supports_observable = lambda *args, **kwargs: True
-
         H = qml.Hamiltonian([1.0, 1.0], [qml.PauliX(0), qml.PauliY(0)], grouping_type="qwc")
         qs = qml.tape.QuantumScript(measurements=[qml.expval(H)])
+        spy = mocker.spy(qml.transforms, "hamiltonian_expand")
 
-        dev = TestDevice()
+        dev = self.SomeDevice()
         dev.shots = None
         dev.use_grouping = use_grouping
-        new_qscripts, post_proc_fn = dev.batch_transform(qs)
+        new_qscripts, _ = dev.batch_transform(qs)
 
         if use_grouping:
             assert len(new_qscripts) == 2
+            spy.assert_called_once()
         else:
             assert len(new_qscripts) == 1
+            spy.assert_not_called()
+
+    def test_batch_transform_does_not_expand_supported_sum(self, mocker):
+        """Tests that batch_transform does not expand Sums if they are supported."""
+        H = qml.sum(qml.PauliX(0), qml.PauliY(0))
+        qs = qml.tape.QuantumScript(measurements=[qml.expval(H)])
+        spy = mocker.spy(qml.transforms, "sum_expand")
+
+        dev = self.SomeDevice()
+        new_qscripts, _ = dev.batch_transform(qs)
+
+        assert len(new_qscripts) == 1
+        spy.assert_not_called()
+
+    def test_batch_transform_expands_not_supported_sums(self, mocker):
+        """Tests that batch_transform expand Sums if they are not supported."""
+        H = qml.sum(qml.PauliX(0), qml.PauliY(0))
+        qs = qml.tape.QuantumScript(measurements=[qml.expval(H)])
+        spy = mocker.spy(qml.transforms, "sum_expand")
+
+        dev = self.SomeDevice()
+        dev.supports_observable = lambda *args, **kwargs: False
+        new_qscripts, _ = dev.batch_transform(qs)
+
+        assert len(new_qscripts) == 2
+        spy.assert_called()