Repository: PennyLaneAI/pennylane
Pull request #5838 (2024-06-11 18:20:24)

Fix bug where `split_non_commuting` erases trainability of observables
**Context:**
The post processing function of `split_non_commuting` incorrectly converts the interface of results and erases the trainability of `Hamiltonian` observables

**Description of the Change:**
Removes `convert_like` call.

**Related GitHub Issues:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5837
[sc-65649]

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://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5838?dropdown=coverage&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:
> Project coverage is 99.66%. Comparing base [(`d70f61a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/d70f61aed57497712f3c5c86033aa9747c7fb029?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`51dc029`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/51dc02919e985c027096881bef7b8ea699a27be8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


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


```diff
@@            Coverage Diff             @@
##           master    #5838      +/-   ##
==========================================
- Coverage   99.66%   99.66%   -0.01%     
==========================================
  Files         417      417              
  Lines       39978    39682     -296     
==========================================
- Hits        39846    39549     -297     
- Misses        132      133       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5838?dropdown=coverage&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 0a376848a..d4d3cb518 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -78,6 +78,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)
 
 * The qchem module has dedicated functions for calling `pyscf` and `openfermion` backends.
   [(#5553)](https://github.com/PennyLaneAI/pennylane/pull/5553)
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index ec8a5ba77..8c23061e8 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -693,7 +693,7 @@ def _sum_terms(res: ResultBatch, coeffs: List[float], offset: float, shape: Tupl
     if len(dot_products) == 0:
         return qml.math.ones(shape) * offset
     summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
-    return qml.math.convert_like(summed_dot_products + offset, res[0])
+    return summed_dot_products + offset
 
 
 def _mp_to_obs(mp: MeasurementProcess, tape: qml.tape.QuantumScript) -> qml.operation.Operator:


Test patch:
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index e95748a19..709a0ed8e 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -526,6 +526,11 @@ class TestUnits:
         """Tests that the tape is split correctly when containing non-Pauli observables"""
 
         obs_list = single_term_obs_list + non_pauli_obs
+
+        if not qml.operation.active_new_opmath():
+            non_pauli_obs = _convert_obs_to_legacy_opmath(non_pauli_obs)
+            obs_list = _convert_obs_to_legacy_opmath(obs_list)
+
         measurements = [
             qml.expval(c * o) for c, o in zip([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], obs_list)
         ]
@@ -956,9 +961,31 @@ class TestDifferentiability:
         assert qml.math.allclose(grad1, expected_grad_1)
         assert qml.math.allclose(grad2, expected_grad_2)
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_trainable_hamiltonian_autograd(self, grouping_strategy):
+        """Tests that measurements of trainable Hamiltonians are differentiable"""
+
+        import pennylane.numpy as pnp
+
+        dev = qml.device("default.qubit", wires=2, shots=50000)
+
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
+        @qml.qnode(dev)
+        def circuit(coeff1, coeff2):
+            qml.RX(np.pi / 4, wires=0)
+            qml.RY(np.pi / 4, wires=1)
+            return qml.expval(qml.Hamiltonian([coeff1, coeff2], [qml.Y(0) @ qml.Z(1), qml.X(1)]))
+
+        params = pnp.array(pnp.pi / 4), pnp.array(3 * pnp.pi / 4)
+        actual = qml.jacobian(circuit)(*params)
+
+        assert qml.math.allclose(actual, [-0.5, np.cos(np.pi / 4)], rtol=0.05)
+
     @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", [False, True])
     @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
-    def test_jax(self, grouping_strategy):
+    def test_jax(self, grouping_strategy, use_jit):
         """Tests that the output of ``split_non_commuting`` is differentiable with jax"""
 
         import jax
@@ -979,6 +1006,9 @@ class TestDifferentiability:
             qml.RY(phi, wires=1)
             return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
 
+        if use_jit:
+            circuit = jax.jit(circuit)
+
         def cost(theta, phi):
             res = circuit(theta, phi)
             return qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
@@ -996,44 +1026,30 @@ class TestDifferentiability:
         assert qml.math.allclose(grad2, expected_grad_2)
 
     @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", [False, True])
     @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
-    def test_jax_jit(self, grouping_strategy):
-        """Tests that the output of ``split_non_commuting`` is differentiable with jax and jit"""
+    def test_trainable_hamiltonian_jax(self, grouping_strategy, use_jit):
+        """Tests that measurements of trainable Hamiltonians are differentiable with jax"""
 
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=2, shots=50000)
 
-        obs_list = complex_obs_list
-        if not qml.operation.active_new_opmath():
-            obs_list = obs_list[:-1]  # exclude the identity term
-
-        @jax.jit
         @partial(split_non_commuting, grouping_strategy=grouping_strategy)
         @qml.qnode(dev)
-        def circuit(theta, phi):
-            qml.RX(theta, wires=0)
-            qml.RY(phi, wires=0)
-            qml.RX(theta, wires=1)
-            qml.RY(phi, wires=1)
-            return qml.probs(wires=[0, 1]), *[qml.expval(o) for o in obs_list]
+        def circuit(coeff1, coeff2):
+            qml.RX(np.pi / 4, wires=0)
+            qml.RY(np.pi / 4, wires=1)
+            return qml.expval(qml.Hamiltonian([coeff1, coeff2], [qml.Y(0) @ qml.Z(1), qml.X(1)]))
 
-        def cost(theta, phi):
-            res = circuit(theta, phi)
-            return qml.math.concatenate([res[0], qml.math.stack(res[1:])], axis=0)
+        if use_jit:
+            circuit = jax.jit(circuit)
 
-        params = jnp.array(jnp.pi / 4), jnp.array(3 * jnp.pi / 4)
-        grad1, grad2 = jax.jacobian(cost, argnums=[0, 1])(*params)
+        params = jnp.array(np.pi / 4), jnp.array(3 * np.pi / 4)
+        actual = jax.jacobian(circuit, argnums=[0, 1])(*params)
 
-        expected_grad_1 = expected_grad_param_0
-        expected_grad_2 = expected_grad_param_1
-        if not qml.operation.active_new_opmath():
-            expected_grad_1 = expected_grad_param_0[:-1]
-            expected_grad_2 = expected_grad_param_1[:-1]
-
-        assert qml.math.allclose(grad1, expected_grad_1)
-        assert qml.math.allclose(grad2, expected_grad_2)
+        assert qml.math.allclose(actual, [-0.5, np.cos(np.pi / 4)], rtol=0.05)
 
     @pytest.mark.torch
     @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
@@ -1074,10 +1090,32 @@ class TestDifferentiability:
         assert qml.math.allclose(grad1, expected_grad_1, atol=1e-5)
         assert qml.math.allclose(grad2, expected_grad_2, atol=1e-5)
 
+    @pytest.mark.torch
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_trainable_hamiltonian_torch(self, grouping_strategy):
+        """Tests that measurements of trainable Hamiltonians are differentiable with torch"""
+
+        import torch
+        from torch.autograd.functional import jacobian
+
+        dev = qml.device("default.qubit", wires=2, shots=50000)
+
+        @partial(split_non_commuting, grouping_strategy=grouping_strategy)
+        @qml.qnode(dev)
+        def circuit(coeff1, coeff2):
+            qml.RX(np.pi / 4, wires=0)
+            qml.RY(np.pi / 4, wires=1)
+            return qml.expval(qml.Hamiltonian([coeff1, coeff2], [qml.Y(0) @ qml.Z(1), qml.X(1)]))
+
+        params = torch.tensor(np.pi / 4), torch.tensor(3 * np.pi / 4)
+        actual = jacobian(circuit, params)
+
+        assert qml.math.allclose(actual, [-0.5, np.cos(np.pi / 4)], rtol=0.05)
+
     @pytest.mark.tf
     @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
     def test_tensorflow(self, grouping_strategy):
-        """Tests that the output of ``split_non_commuting`` is differentiable with torch"""
+        """Tests that the output of ``split_non_commuting`` is differentiable with tensorflow"""
 
         import tensorflow as tf
 
@@ -1111,3 +1149,27 @@ class TestDifferentiability:
 
         assert qml.math.allclose(grad1, expected_grad_1, atol=1e-5)
         assert qml.math.allclose(grad2, expected_grad_2, atol=1e-5)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("grouping_strategy", [None, "default", "qwc", "wires"])
+    def test_trainable_hamiltonian_tensorflow(self, grouping_strategy):
+        """Tests that measurements of trainable Hamiltonians are differentiable with tensorflow"""
+
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=2, shots=50000)
+
+        @qml.qnode(dev)
+        def circuit(coeff1, coeff2):
+            qml.RX(np.pi / 4, wires=0)
+            qml.RY(np.pi / 4, wires=1)
+            return qml.expval(qml.Hamiltonian([coeff1, coeff2], [qml.Y(0) @ qml.Z(1), qml.X(1)]))
+
+        params = tf.Variable(np.pi / 4), tf.Variable(3 * np.pi / 4)
+
+        with tf.GradientTape() as tape:
+            cost = split_non_commuting(circuit, grouping_strategy=grouping_strategy)(*params)
+
+        actual = tape.jacobian(cost, params)
+
+        assert qml.math.allclose(actual, [-0.5, np.cos(np.pi / 4)], rtol=0.05)