Repository: PennyLaneAI/pennylane
Pull request #4992 (2024-01-04 17:19:25)

update snapshot; little docs changes
**Context:**
Reviewed my PRs, found some stuffs.

**Description of the Change:**
- fix code markdown in deprecations file
- add examples to `Shots` docs on creating a new Shots instance by scaling an old one
- bugfix: `Snapshot` wrongly flattened batched states, so I handle it accordingly (for `default.qubit`, at least)
- bugfix: `Snapshot` assumed the current device state is complex, but `default.qubit` does not, so I always cast to complex in `apply_operation` when getting the state. Fixed the associated docs as well.

**Benefits:**
Improved docs, `Snapshot` is fixed up

**Possible Drawbacks:**
I am actually changing something (`Snapshot`) right before a release.

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/4992?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:
> :exclamation: No coverage uploaded for pull request base (`v0.34.0-rc0@db4b762`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).


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


```diff
@@              Coverage Diff               @@
##             v0.34.0-rc0    #4992   +/-   ##
==============================================
  Coverage               ?   99.67%           
==============================================
  Files                  ?      392           
  Lines                  ?    35291           
  Branches               ?        0           
==============================================
  Hits                   ?    35175           
  Misses                 ?      116           
  Partials               ?        0           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4992?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/development/deprecations.rst b/doc/development/deprecations.rst
index c96e68538..1f2bbfd05 100644
--- a/doc/development/deprecations.rst
+++ b/doc/development/deprecations.rst
@@ -9,13 +9,13 @@ deprecations are listed below.
 Pending deprecations
 --------------------
 
-* `qml.transforms.one_qubit_decomposition` and `qml.transforms.two_qubit_decomposition` are deprecated. Instead,
-  you should use `qml.ops.one_qubit_decomposition` and `qml.ops.two_qubit_decomposition` .
+* ``qml.transforms.one_qubit_decomposition`` and ``qml.transforms.two_qubit_decomposition`` are deprecated. Instead,
+  you should use ``qml.ops.one_qubit_decomposition`` and ``qml.ops.two_qubit_decomposition`` .
 
   - Deprecated in v0.34
   - Will be removed in v0.35
 
-* `Observable.return_type` is deprecated. Instead, you should inspect the type
+* ``Observable.return_type`` is deprecated. Instead, you should inspect the type
   of the surrounding measurement process.
 
   - Deprecated in v0.34
diff --git a/doc/introduction/inspecting_circuits.rst b/doc/introduction/inspecting_circuits.rst
index 5b916f657..50eac348a 100644
--- a/doc/introduction/inspecting_circuits.rst
+++ b/doc/introduction/inspecting_circuits.rst
@@ -136,7 +136,7 @@ During normal execution, the snapshots are ignored:
 
     @qml.qnode(dev, interface=None)
     def circuit():
-        qml.Snapshot(measurement=qml.expval(qml.PauliZ(0))
+        qml.Snapshot(measurement=qml.expval(qml.PauliZ(0)))
         qml.Hadamard(wires=0)
         qml.Snapshot("very_important_state")
         qml.CNOT(wires=[0, 1])
@@ -149,8 +149,8 @@ results.
 
 >>> qml.snapshots(circuit)()
 {0: 1.0,
-'very_important_state': array([0.70710678, 0., 0.70710678, 0.]),
-2: array([0.70710678, 0., 0., 0.70710678]),
+'very_important_state': array([0.707+0.j, 0.+0.j, 0.707+0.j, 0.+0.j]),
+2: array([0.707+0.j, 0.+0.j, 0.+0.j, 0.707+0.j]),
 'execution_results': 0.0}
 
 Graph representation
diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 3a7249dc9..feb0639ab 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -81,8 +81,8 @@ def snapshots(qnode):
 
     >>> qml.snapshots(circuit)()
     {0: 1.0,
-    'very_important_state': array([0.70710678, 0.        , 0.70710678, 0.        ]),
-    2: array([0.70710678, 0.        , 0.        , 0.70710678]),
+    'very_important_state': array([0.70710678+0.j, 0.        +0.j, 0.70710678+0.j, 0.        +0.j]),
+    2: array([0.70710678+0.j, 0.        +0.j, 0.        +0.j, 0.70710678+0.j]),
     'execution_results': 0.0}
     """
 
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 69913d5b6..cd6e4949e 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -365,7 +365,8 @@ def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debu
         if measurement:
             snapshot = qml.devices.qubit.measure(measurement, state)
         else:
-            snapshot = math.flatten(state)
+            flat_shape = (math.shape(state)[0], -1) if is_state_batched else (-1,)
+            snapshot = math.cast(math.reshape(state, flat_shape), complex)
         if op.tag:
             debugger.snapshots[op.tag] = snapshot
         else:
diff --git a/pennylane/measurements/shots.py b/pennylane/measurements/shots.py
index f5220f77e..448f796e7 100644
--- a/pennylane/measurements/shots.py
+++ b/pennylane/measurements/shots.py
@@ -114,6 +114,13 @@ class Shots:
     >>> shots.total_shots, shots.shot_vector
     (1210, (ShotCopies(10 shots x 1), ShotCopies(100 shots x 4), ShotCopies(200 shots x 4)))
 
+    Example constructing a Shots instance by multiplying an existing one by an int or float:
+
+    >>> Shots(100) * 2
+    Shots(total_shots=200, shot_vector=(ShotCopies(200 shots x 1),))
+    >>> Shots([7, (100, 2)]) * 1.5
+    Shots(total_shots=310, shot_vector=(ShotCopies(10 shots x 1), ShotCopies(150 shots x 2)))
+
     One should also note that specifying a single tuple of length 2 is considered two different
     shot values, and *not* a tuple-pair representing shots and copies to avoid special behaviour
     depending on the iterable type:


Test patch:
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 18929bc1b..9135ec0a1 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -624,6 +624,17 @@ class TestSnapshot:
         assert debugger.snapshots[0].shape == ()
         assert debugger.snapshots[0] == qml.devices.qubit.measure(measurement, initial_state)
 
+    def test_batched_state(self, ml_framework):
+        """Test that batched states create batched snapshots."""
+        initial_state = qml.math.asarray([[1.0, 0.0], [0.0, 0.1]], like=ml_framework)
+        debugger = self.Debugger()
+        new_state = apply_operation(
+            qml.Snapshot(), initial_state, is_state_batched=True, debugger=debugger
+        )
+        assert new_state.shape == initial_state.shape
+        assert set(debugger.snapshots) == {0}
+        assert np.array_equal(debugger.snapshots[0], initial_state)
+
 
 @pytest.mark.parametrize("method", methods)
 class TestRXCalcGrad: