Clean up Device.batch_transform to use split_non_commuting (#5828)

**Context:**
#5729 introduced the
unified `split_non_commuting`. Now the legacy device can use
`split_non_commuting` in all scenarios.

**Description of the Change:**
Cleans up `batch_transform` to use `split_non_commuting`

**Benefits:**
Cleaner code, prepares for the deprecation of `hamiltonian_expand` and
`sum_expand`.

**Related GitHub Issues:**
[sc-61253]

doc/releases/changelog-dev.md

@@ -115,6 +115,7 @@
 * `qml.transforms.split_non_commuting` can now handle circuits containing measurements of multi-term observables.
  [(#5729)](https://github.com/PennyLaneAI/pennylane/pull/5729)
  [(#5853)](https://github.com/PennyLaneAI/pennylane/pull/5838)
+ [(#5828)](https://github.com/PennyLaneAI/pennylane/pull/5828)
  [(#5869)](https://github.com/PennyLaneAI/pennylane/pull/5869)
 
 * The qchem module has dedicated functions for calling `pyscf` and `openfermion` backends. The

pennylane/_device.py

@@ -26,15 +26,11 @@
 
 import pennylane as qml
 from pennylane.measurements import (
- CountsMP,
  Expectation,
- ExpectationMP,
  MeasurementProcess,
  MidMeasureMP,
  Probability,
- ProbabilityMP,
  Sample,
- SampleMP,
  ShadowExpvalMP,
  State,
  Variance,
@@ -738,79 +734,92 @@ def batch_transform(self, circuit: QuantumTape):
  the sequence of circuits to be executed, and a post-processing function
  to be applied to the list of evaluated circuit results.
  """
- supports_hamiltonian = self.supports_observable("Hamiltonian")
- supports_sum = self.supports_observable("Sum")
+
+ def null_postprocess(results):
+ return results[0]
+
  finite_shots = self.shots is not None
- grouping_known = all(
- obs.grouping_indices is not None
- for obs in circuit.observables
- if isinstance(obs, (Hamiltonian, LinearCombination))
+ has_shadow = any(isinstance(m, ShadowExpvalMP) for m in circuit.measurements)
+ is_analytic_or_shadow = not finite_shots or has_shadow
+ all_obs_usable = self._all_multi_term_obs_supported(circuit)
+ exists_multi_term_obs = any(
+ isinstance(m.obs, (Hamiltonian, Sum, Prod, SProd)) for m in circuit.measurements
  )
- # device property present in braket plugin
- use_grouping = getattr(self, "use_grouping", True)
-
- hamiltonian_in_obs = any(
- isinstance(obs, (Hamiltonian, LinearCombination)) for obs in circuit.observables
+ has_overlapping_wires = len(circuit.obs_sharing_wires) > 0
+ single_hamiltonian = len(circuit.measurements) == 1 and isinstance(
+ circuit.measurements[0].obs, (Hamiltonian, Sum)
  )
-
- expval_sum_or_prod_in_obs = any(
- isinstance(m.obs, (Sum, Prod, SProd)) and isinstance(m, ExpectationMP)
- for m in circuit.measurements
+ single_hamiltonian_with_grouping_known = (
+ single_hamiltonian and circuit.measurements[0].obs.grouping_indices is not None
  )
 
- is_shadow = any(isinstance(m, ShadowExpvalMP) for m in circuit.measurements)
+ if not getattr(self, "use_grouping", True) and single_hamiltonian and all_obs_usable:
+ # Special logic for the braket plugin
+ circuits = [circuit]
+ processing_fn = null_postprocess
 
- hamiltonian_unusable = not supports_hamiltonian or (finite_shots and not is_shadow)
+ elif not exists_multi_term_obs and not has_overlapping_wires:
+ circuits = [circuit]
+ processing_fn = null_postprocess
 
- if hamiltonian_in_obs and (hamiltonian_unusable or (use_grouping and grouping_known)):
- # If the observable contains a Hamiltonian and the device does not
- # support Hamiltonians, or if the simulation uses finite shots, or
- # if the Hamiltonian explicitly specifies an observable grouping,
- # split tape into multiple tapes of diagonalizable known observables.
- try:
- circuits, hamiltonian_fn = qml.transforms.hamiltonian_expand(circuit, group=False)
- except ValueError:
- circuits, hamiltonian_fn = qml.transforms.sum_expand(circuit)
+ elif is_analytic_or_shadow and all_obs_usable and not has_overlapping_wires:
+ circuits = [circuit]
+ processing_fn = null_postprocess
 
- elif expval_sum_or_prod_in_obs and not is_shadow and not supports_sum:
- circuits, hamiltonian_fn = qml.transforms.sum_expand(circuit)
+ elif single_hamiltonian_with_grouping_known:
 
- elif (
- len(circuit.obs_sharing_wires) > 0
- and not hamiltonian_in_obs
- and all(
- not isinstance(m, (SampleMP, ProbabilityMP, CountsMP)) for m in circuit.measurements
- )
- ):
- # Check for case of non-commuting terms and that there are no Hamiltonians
- # TODO: allow for Hamiltonians in list of observables as well.
- circuits, hamiltonian_fn = qml.transforms.split_non_commuting(circuit)
+ # Use qwc grouping if the circuit contains a single measurement of a
+ # Hamiltonian/Sum with grouping indices already calculated.
+ circuits, processing_fn = qml.transforms.split_non_commuting(circuit, "qwc")
 
- else:
- # otherwise, return the output of an identity transform
- circuits = [circuit]
+ elif any(isinstance(m.obs, (Hamiltonian, LinearCombination)) for m in circuit.measurements):
 
- def hamiltonian_fn(res):
- return res[0]
+ # Otherwise, use wire-based grouping if the circuit contains a Hamiltonian
+ # that is potentially very large.
+ circuits, processing_fn = qml.transforms.split_non_commuting(circuit, "wires")
 
- # Check whether the circuit was broadcasted (then the Hamiltonian-expanded
- # ones will be as well) and whether broadcasting is supported
+ else:
+ circuits, processing_fn = qml.transforms.split_non_commuting(circuit)
+
+ # Check whether the circuit was broadcasted and whether broadcasting is supported
  if circuit.batch_size is None or self.capabilities().get("supports_broadcasting"):
  # If the circuit wasn't broadcasted or broadcasting is supported, no action required
- return circuits, hamiltonian_fn
+ return circuits, processing_fn
 
  # Expand each of the broadcasted Hamiltonian-expanded circuits
  expanded_tapes, expanded_fn = qml.transforms.broadcast_expand(circuits)
 
  # Chain the postprocessing functions of the broadcasted-tape expansions and the Hamiltonian
  # expansion. Note that the application order is reversed compared to the expansion order,
- # i.e. while we first applied `hamiltonian_expand` to the tape, we need to process the
+ # i.e. while we first applied `split_non_commuting` to the tape, we need to process the
  # results from the broadcast expansion first.
  def total_processing(results):
- return hamiltonian_fn(expanded_fn(results))
+ return processing_fn(expanded_fn(results))
 
  return expanded_tapes, total_processing
 
+ def _all_multi_term_obs_supported(self, circuit):
+ """Check whether all multi-term observables in the circuit are supported."""
+
+ for mp in circuit.measurements:
+
+ if mp.obs is None:
+ # Some measurements are not observable based.
+ continue
+
+ if mp.obs.name == "LinearCombination" and not self.supports_observable("Hamiltonian"):
+ return False
+
+ if mp.obs.name in (
+ "Hamiltonian",
+ "Sum",
+ "Prod",
+ "SProd",
+ ) and not self.supports_observable(mp.obs.name):
+ return False
+
+ return True
+
  @property
  def op_queue(self):
  """The operation queue to be applied.

pennylane/devices/qubit/sampling.py

@@ -103,7 +103,30 @@ def _get_num_executions_for_expval_H(obs):
  indices = obs.grouping_indices
  if indices:
  return len(indices)
- return sum(int(not isinstance(o, qml.Identity)) for o in obs.terms()[1])
+ return _get_num_wire_groups_for_expval_H(obs)
+
+
+def _get_num_wire_groups_for_expval_H(obs):
+ _, obs_list = obs.terms()
+ wires_list = []
+ added_obs = []
+ num_groups = 0
+ for o in obs_list:
+ if o in added_obs:
+ continue
+ if isinstance(o, qml.Identity):
+ continue
+ added = False
+ for wires in wires_list:
+ if len(qml.wires.Wires.shared_wires([wires, o.wires])) == 0:
+ added_obs.append(o)
+ added = True
+ break
+ if not added:
+ added_obs.append(o)
+ wires_list.append(o.wires)
+ num_groups += 1
+ return num_groups
 
 
 def _get_num_executions_for_sum(obs):

tests/conftest.py

@@ -228,6 +228,12 @@ def new_opmath_only():
  pytest.skip("This feature only works with new opmath enabled")
 
 
+@pytest.fixture
+def legacy_opmath_only():
+ if qml.operation.active_new_opmath():
+ pytest.skip("This test exclusively tests legacy opmath")
+
+
 #######################################################################
 
 try:

tests/devices/default_qubit/test_default_qubit_tracking.py

@@ -195,40 +195,28 @@ def circuit_3(y):
 
 shot_testing_combos = [
  # expval combinations
- ([qml.expval(qml.PauliX(0))], 1, 10),
- ([qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0))], 2, 20),
+ ([qml.expval(qml.X(0))], 1, 10),
+ ([qml.expval(qml.X(0)), qml.expval(qml.Y(0))], 2, 20),
  # Hamiltonian test cases
- ([qml.expval(qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(1)]))], 2, 20),
- (
- [qml.expval(qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(1)], grouping_type="qwc"))],
- 1,
- 10,
- ),
- (
- [qml.expval(qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliY(0)], grouping_type="qwc"))],
- 2,
- 20,
- ),
+ ([qml.expval(qml.Hamiltonian([1, 0.5, 1], [qml.X(0), qml.Y(0), qml.X(1)]))], 2, 20),
+ ([qml.expval(qml.Hamiltonian([1, 1], [qml.X(0), qml.X(1)], grouping_type="qwc"))], 1, 10),
+ ([qml.expval(qml.Hamiltonian([1, 1], [qml.X(0), qml.Y(0)], grouping_type="qwc"))], 2, 20),
  # op arithmetic test cases
- ([qml.expval(qml.sum(qml.PauliX(0), qml.PauliY(0)))], 2, 20),
- ([qml.expval(qml.sum(qml.PauliX(0), qml.PauliX(0) @ qml.PauliX(1)))], 1, 10),
- ([qml.expval(qml.sum(qml.PauliX(0), qml.Hadamard(0)))], 2, 20),
- (
- [qml.expval(qml.sum(qml.PauliX(0), qml.PauliY(1) @ qml.PauliX(1), grouping_type="qwc"))],
- 1,
- 10,
- ),
+ ([qml.expval(qml.sum(qml.X(0), qml.Y(0)))], 2, 20),
+ ([qml.expval(qml.sum(qml.X(0), qml.X(0) @ qml.X(1)))], 1, 10),
+ ([qml.expval(qml.sum(qml.X(0), qml.Hadamard(0)))], 2, 20),
+ ([qml.expval(qml.sum(qml.X(0), qml.Y(1) @ qml.X(1), grouping_type="qwc"))], 1, 10),
  (
  [
- qml.expval(qml.prod(qml.PauliX(0), qml.PauliX(1))),
- qml.expval(qml.prod(qml.PauliX(1), qml.PauliX(2))),
+ qml.expval(qml.prod(qml.X(0), qml.X(1))),
+ qml.expval(qml.prod(qml.X(1), qml.X(2))),
  ],
  1,
  10,
  ),
  # computational basis measurements
  ([qml.probs(wires=(0, 1)), qml.sample(wires=(0, 1))], 1, 10),
- ([qml.probs(wires=(0, 1)), qml.sample(wires=(0, 1)), qml.expval(qml.PauliX(0))], 2, 20),
+ ([qml.probs(wires=(0, 1)), qml.sample(wires=(0, 1)), qml.expval(qml.X(0))], 2, 20),
  # classical shadows
  ([qml.shadow_expval(H0)], 10, 10),
  ([qml.shadow_expval(H0), qml.probs(wires=(0, 1))], 11, 20),

tests/devices/qutrit_mixed/test_qutrit_mixed_tracking.py

@@ -157,7 +157,7 @@ def circuit_3(y):
  ([qml.expval(qml.GellMann(0, 1))], 1, 10),
  ([qml.expval(qml.GellMann(0, 1)), qml.expval(qml.GellMann(0, 2))], 2, 20),
  # Hamiltonian test cases
- ([qml.expval(qml.Hamiltonian([1, 1], [qml.GellMann(0, 1), qml.GellMann(1, 5)]))], 2, 20),
+ ([qml.expval(qml.Hamiltonian([1, 1], [qml.GellMann(0, 1), qml.GellMann(1, 5)]))], 1, 10),
  # op arithmetic test cases
  ([qml.expval(qml.sum(qml.GellMann(0, 1), qml.GellMann(1, 4)))], 2, 20),
  (

tests/interfaces/test_autograd.py

@@ -182,7 +182,6 @@ class TestBatchTransformExecution:
  """Tests to ensure batch transforms can be correctly executed
  via qml.execute and map_batch_transform"""
 
- @pytest.mark.usefixtures("use_new_opmath")
  def test_no_batch_transform(self, mocker):
  """Test that batch transforms can be disabled and enabled"""
  dev = qml.device("default.qubit.legacy", wires=2, shots=100000)
@@ -200,47 +199,20 @@ def test_no_batch_transform(self, mocker):
  tape = qml.tape.QuantumScript.from_queue(q)
  spy = mocker.spy(dev, "batch_transform")
 
- res = qml.execute([tape], dev, None, device_batch_transform=False)
- assert np.allclose(res[0], np.cos(y), atol=0.1)
-
- spy.assert_not_called()
-
- res = qml.execute([tape], dev, None, device_batch_transform=True)
- spy.assert_called()
-
- assert isinstance(res[0], np.ndarray)
- assert res[0].shape == ()
- assert np.allclose(res[0], np.cos(y), atol=0.1)
-
- @pytest.mark.usefixtures("use_legacy_opmath")
- def test_no_batch_transform_legacy_opmath(self, mocker):
- """Test functionality to enable and disable"""
- dev = qml.device("default.qubit.legacy", wires=2, shots=100000)
-
- H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
- x = 0.6
- y = 0.2
-
- with qml.queuing.AnnotatedQueue() as q:
- qml.RX(x, wires=0)
- qml.RY(y, wires=1)
- qml.CNOT(wires=[0, 1])
- qml.expval(H)
-
- tape = qml.tape.QuantumScript.from_queue(q)
- spy = mocker.spy(dev, "batch_transform")
-
- with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"):
+ if not qml.operation.active_new_opmath():
+ with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"):
+ _ = qml.execute([tape], dev, None, device_batch_transform=False)
+ else:
  res = qml.execute([tape], dev, None, device_batch_transform=False)
+ assert np.allclose(res[0], np.cos(y), atol=0.1)
 
  spy.assert_not_called()
 
  res = qml.execute([tape], dev, None, device_batch_transform=True)
  spy.assert_called()
 
- assert isinstance(res[0], np.ndarray)
- assert res[0].shape == ()
- assert np.allclose(res[0], np.cos(y), atol=0.1)
+ assert qml.math.shape(res[0]) == ()
+ assert np.allclose(res[0], np.cos(y), rtol=0.05)
 
  def test_batch_transform_dynamic_shots(self):
  """Tests that the batch transform considers the number of shots for the execution, not those
@@ -462,10 +434,10 @@ def f(x):
  assert qml.math.allclose(out, expected)
 
  def test_single_backward_pass_split_hamiltonian(self):
- """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift
- derivatives are requested for a a tape that the device split into batches."""
+ """Tests that the backward pass is one single batch, not a bunch of batches, when parameter
+ shift derivatives are requested for a tape that the device split into batches."""
 
- dev = qml.device("default.qubit.legacy", wires=2)
+ dev = qml.device("default.qubit.legacy", wires=2, shots=50000)
 
  H = qml.Hamiltonian([1, 1], [qml.PauliY(0), qml.PauliZ(0)], grouping_type="qwc")
 
@@ -480,7 +452,7 @@ def f(x):
  assert dev.tracker.totals["batches"] == 2
  assert dev.tracker.history["batch_len"] == [2, 4]
 
- assert qml.math.allclose(out, -np.cos(x) - np.sin(x))
+ assert qml.math.allclose(out, -np.cos(x) - np.sin(x), atol=0.05)
 
 
 execute_kwargs_integration = [

tests/interfaces/test_autograd_qnode.py

@@ -1666,7 +1666,7 @@ def test_hamiltonian_expansion_finite_shots(
  gradient_kwargs = {"h": 0.05}
 
  dev = qml.device(dev_name, wires=3, shots=50000)
- spy = mocker.spy(qml.transforms, "hamiltonian_expand")
+ spy = mocker.spy(qml.transforms, "split_non_commuting")
  obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
 
  @qnode(

tests/interfaces/test_jax_jit_qnode.py

@@ -1540,7 +1540,7 @@ def test_hamiltonian_expansion_analytic(
  tol = TOL_FOR_SPSA
 
  dev = qml.device(dev_name, wires=3, shots=None)
- spy = mocker.spy(qml.transforms, "hamiltonian_expand")
+ spy = mocker.spy(qml.transforms, "split_non_commuting")
  obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
 
  @qnode(
@@ -1609,7 +1609,7 @@ def test_hamiltonian_expansion_finite_shots(
  tol = TOL_FOR_SPSA
 
  dev = qml.device(dev_name, wires=3, shots=50000)
- spy = mocker.spy(qml.transforms, "hamiltonian_expand")
+ spy = mocker.spy(qml.transforms, "split_non_commuting")
  obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
 
  @qnode(