Repository: PennyLaneAI/pennylane
Pull request #3912 (2023-03-15 20:00:54)

fast-path for getting all parameters from a tape
**Context:**
`tape.get_parameters(trainable_only=False)` was very slow for nested data (like a large composite operator!). This will speed it up. Below I compare the timing of `get_parameters` for a tape with just one medium-size Hamiltonian (185 terms, I also checked and got a 1000x speedup with a 1k-term Hamiltonian).

Before this change:
![image](https://user-images.githubusercontent.com/23283972/225427745-939ea6e3-a51f-44b5-9b46-9d7a87ed44ce.png)

With this change:
![image](https://user-images.githubusercontent.com/23283972/225427201-8ddb3dd7-7edc-4127-bf6e-e585eaeb235c.png)

**Description of the Change:**
When `trainable_only=False`, don't use `tape._par_info` to get all parameters. Instead just return a flat combination of op.data for every operator (and measurement, if requested) as a list.

**Benefits:**
Massive speedup

**Possible Drawbacks:**
Slightly uglier code. I can move it into two functions if preferred for legibility.

**Related GitHub Issues:**
While proposing #3910, I realized that this is the true essence of the problem, and would like to fix it in this minimal PR instead.

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.
timmysilv: this problem is fundamentally rooted in the fact that the data of composite operators is too complex in shape. I'm closing this, and instead going to look into improving that shape.

That said, the Torch unit tests are failing because while a tape is unwrapped, it gets passed to `tape_expand` which adds a rotation gate to the tape's `_ops`. Perhaps it shouldn't do that, or handle it better, because when it exits the unwrap context, the shape of the tape's parameters has changed (or in the case of master, it hasn't but it should) assuming the rotation gate is parametric (and it usually is). This is an aside, although it's part of why I'm giving up on this :D
codecov[bot]: ## [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3912](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (b703f1b) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c0fb9b851fcdf004b943f475cbeae47a86b7d7db?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c0fb9b8) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3912   +/-   ##
=======================================
  Coverage   99.57%   99.57%           
=======================================
  Files         349      349           
  Lines       30425    30442   +17     
=======================================
+ Hits        30296    30313   +17     
  Misses        129      129           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/hadamard\_gradient.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9oYWRhbWFyZF9ncmFkaWVudC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/parameter\_shift\_cv.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9wYXJhbWV0ZXJfc2hpZnRfY3YucHk=) | `99.61% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/tape/qscript.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvcXNjcmlwdC5weQ==) | `93.67% <100.00%> (+0.11%)` | :arrow_up: |
| [pennylane/tape/tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvdGFwZS5weQ==) | `99.14% <100.00%> (+0.03%)` | :arrow_up: |

... and [2 files with indirect coverage changes](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3912/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)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7440fa33a..e08c6dacd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -204,6 +204,11 @@
 * Execution code has been updated to use the new `qml.transforms.convert_to_numpy_parameters` instead of `qml.tape.Unwrap`.
   [(#3989)](https://github.com/PennyLaneAI/pennylane/pull/3989)
 
+* Converted a sub-routine of `expand_tape` into `qml.tape.tape.rotations_and_diagonal_measurements`,
+  a helper function that computes rotations and diagonal measurements for a tape with measurements
+  with overlapping wires.
+  [(#3912)](https://github.com/PennyLaneAI/pennylane/pull/3912)
+
 <h3>Breaking changes 💔</h3>
 
 * Both JIT interfaces are not compatible with JAX `>0.4.3`, we raise an error for those versions.
@@ -234,6 +239,11 @@
   [(#3958)](https://github.com/PennyLaneAI/pennylane/pull/3958)
   [(#3983)](https://github.com/PennyLaneAI/pennylane/pull/3983)
 
+* `qml.tape.tape.expand_tape` (and consequentially `QuantumScript.expand`) no longer updates the inputted tape
+  with rotations and diagonal measurements. Note that the newly expanded tape that is returned will still have
+  the rotations and diagonal measurements.
+  [(#3912)](https://github.com/PennyLaneAI/pennylane/pull/3912)
+
 <h3>Deprecations 👋</h3>
 
 <h3>Documentation 📝</h3>
diff --git a/pennylane/gradients/hadamard_gradient.py b/pennylane/gradients/hadamard_gradient.py
index 209e8ba5b..bb1a50825 100644
--- a/pennylane/gradients/hadamard_gradient.py
+++ b/pennylane/gradients/hadamard_gradient.py
@@ -304,8 +304,10 @@ def _expval_hadamard_grad(tape, argnum, aux_wire):
 
             new_tape = qml.tape.QuantumScript(ops=ops, measurements=measurements)
 
-            new_tape.expand()
+            _rotations, _measurements = qml.tape.tape.rotations_and_diagonal_measurements(new_tape)
             # pylint: disable=protected-access
+            new_tape._ops = new_tape._ops + _rotations
+            new_tape._measurements = _measurements
             new_tape._update()
 
             num_tape += 1
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index d8c7c1b69..784975ecd 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -634,16 +634,25 @@ class QuantumScript:
         >>> qscript.get_parameters(trainable_only=False)
         [0.432, 0.543, 0.133]
         """
-        params = []
-        iterator = self.trainable_params if trainable_only else range(len(self._par_info))
-
-        for p_idx in iterator:
-            op = self._par_info[p_idx]["op"]
-            if operations_only and hasattr(op, "return_type"):
-                continue
-
-            op_idx = self._par_info[p_idx]["p_idx"]
-            params.append(op.data[op_idx])
+        if trainable_only:
+            params = []
+            for p_idx in self.trainable_params:
+                op = self._par_info[p_idx]["op"]
+                if operations_only and hasattr(op, "return_type"):
+                    continue
+
+                op_idx = self._par_info[p_idx]["p_idx"]
+                params.append(op.data[op_idx])
+            return params
+
+        # If trainable_only=False, return all parameters
+        # This is faster than the above and should be used when indexing `_par_info` is not needed
+        params = [d for op in self.operations for d in op.data]
+        if operations_only:
+            return params
+        for m in self.measurements:
+            if m.obs is not None:
+                params.extend(m.obs.data)
         return params
 
     def set_parameters(self, params, trainable_only=True):
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index 3872b924d..4d4774a72 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -77,6 +77,35 @@ def _validate_computational_basis_sampling(measurements):
                 raise qml.QuantumFunctionError(_err_msg_for_some_meas_not_qwc(measurements))
 
 
+def rotations_and_diagonal_measurements(tape):
+    """Compute the rotations for overlapping observables, and return them along with the diagonalized observables."""
+    if not tape._obs_sharing_wires:
+        return [], tape._measurements
+
+    with QueuingManager.stop_recording():  # stop recording operations to active context when computing qwc groupings
+        try:
+            rotations, diag_obs = qml.pauli.diagonalize_qwc_pauli_words(tape._obs_sharing_wires)
+        except (TypeError, ValueError) as e:
+            if any(isinstance(m, (ProbabilityMP, SampleMP, CountsMP)) for m in tape.measurements):
+                raise qml.QuantumFunctionError(
+                    "Only observables that are qubit-wise commuting "
+                    "Pauli words can be returned on the same wire.\n"
+                    "Try removing all probability, sample and counts measurements "
+                    "this will allow for splitting of execution and separate measurements "
+                    "for each non-commuting observable."
+                ) from e
+
+            raise qml.QuantumFunctionError(_err_msg_for_some_meas_not_qwc(tape.measurements)) from e
+
+        measurements = copy.copy(tape._measurements)
+
+        for o, i in zip(diag_obs, tape._obs_sharing_wires_id):
+            new_m = tape.measurements[i].__class__(obs=o)
+            measurements[i] = new_m
+
+    return rotations, measurements
+
+
 # TODO: move this function to its own file and rename
 def expand_tape(tape, depth=1, stop_at=None, expand_measurements=False):
     """Expand all objects in a tape to a specific depth.
@@ -149,37 +178,11 @@ def expand_tape(tape, depth=1, stop_at=None, expand_measurements=False):
     if tape.samples_computational_basis and len(tape.measurements) > 1:
         _validate_computational_basis_sampling(tape.measurements)
 
-    if tape._obs_sharing_wires:
-        with QueuingManager.stop_recording():  # stop recording operations to active context when computing qwc groupings
-            try:
-                rotations, diag_obs = qml.pauli.diagonalize_qwc_pauli_words(tape._obs_sharing_wires)
-                tape._obs_sharing_wires = diag_obs
-            except (TypeError, ValueError) as e:
-                if any(
-                    isinstance(m, (ProbabilityMP, SampleMP, CountsMP)) for m in tape.measurements
-                ):
-                    raise qml.QuantumFunctionError(
-                        "Only observables that are qubit-wise commuting "
-                        "Pauli words can be returned on the same wire.\n"
-                        "Try removing all probability, sample and counts measurements "
-                        "this will allow for splitting of execution and separate measurements "
-                        "for each non-commuting observable."
-                    ) from e
-
-                raise qml.QuantumFunctionError(
-                    _err_msg_for_some_meas_not_qwc(tape.measurements)
-                ) from e
-
-            tape._ops.extend(rotations)
-
-            for o, i in zip(diag_obs, tape._obs_sharing_wires_id):
-                new_m = tape.measurements[i].__class__(obs=o)
-                tape._measurements[i] = new_m
-
+    diagonalizing_gates, diagonal_measurements = rotations_and_diagonal_measurements(tape)
     for queue, new_queue in [
         (tape._prep, new_prep),
-        (tape._ops, new_ops),
-        (tape._measurements, new_measurements),
+        (tape._ops + diagonalizing_gates, new_ops),
+        (diagonal_measurements, new_measurements),
     ]:
         for obj in queue:
             stop = stop_at(obj)


Test patch:
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index badaa318f..cf76ed570 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1198,6 +1198,28 @@ class TestExpand:
             ]
         )
 
+    def test_expand_does_not_affect_original_tape(self):
+        """Test that expand_tape does not modify the inputted tape while creating a new one."""
+        ops = [qml.RX(1.1, 0)]
+        measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(0))]
+        tape = qml.tape.QuantumTape(ops, measurements)
+        expanded = tape.expand()
+
+        assert len(tape.operations) == 1
+        assert qml.equal(tape.operations[0], ops[0])
+        assert len(tape._obs_sharing_wires) == 2
+        assert all(qml.equal(obs, qml.PauliX(0)) for obs in tape._obs_sharing_wires)
+        assert qml.equal(tape.measurements[0], qml.expval(qml.PauliX(0)))
+        assert qml.equal(tape.measurements[1], qml.expval(qml.PauliX(0)))
+
+        assert len(expanded.operations) == 2
+        assert qml.equal(expanded.operations[0], ops[0])
+        assert qml.equal(expanded.operations[1], qml.RY(-np.pi / 2, 0))  # new rotation
+        assert len(expanded._obs_sharing_wires) == 2
+        assert all(qml.equal(obs, qml.PauliZ(0)) for obs in expanded._obs_sharing_wires)
+        assert qml.equal(expanded.measurements[0], qml.expval(qml.PauliZ(0)))
+        assert qml.equal(expanded.measurements[1], qml.expval(qml.PauliZ(0)))
+
     def test_is_sampled_reserved_after_expansion(self, monkeypatch, mocker):
         """Test that the is_sampled property is correctly set when tape
         expansion happens."""