Repository: PennyLaneAI/pennylane
Pull request #1441 (2021-06-30 13:15:16)

Added return raw sample functionality to qml.sample when observable is not provided.
**Context:**

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


josh146: Thanks @Jaybsoni for opening this PR 🙂 Let us know if you have any questions (feel free to comment directly on the code within the 'files changed' tab), or otherwise, once this PR is ready for a review.

> Fix bug with plotting the resultant circuit !

Out of curiosity, what is the status of this bug? 
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1441](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (a983984) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/2623fcd74a5b5100d097babcebdd1a0bb0172d76?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2623fcd) will **decrease** coverage by `0.04%`.
> The diff coverage is `45.45%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1441      +/-   ##
==========================================
- Coverage   98.24%   98.19%   -0.05%     
==========================================
  Files         160      160              
  Lines       12016    12025       +9     
==========================================
+ Hits        11805    11808       +3     
- Misses        211      217       +6     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `97.58% <40.00%> (-1.02%)` | :arrow_down: |
| [pennylane/measure.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmUucHk=) | `88.88% <50.00%> (-2.98%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [2623fcd...a983984](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1441?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

Jaybsoni: Hi @josh146, 

Thanks for your patience, I got a bit busy and wasn't able to get as much work done as I would have liked these past few days! Here is an update for you: 

**Updates:** 
- Added unit tests for each of the files I have updated while implementing this feature 
    - Previously updated `measure.py` and `_qubit_device.py` so I added tests for each of these updates in `test_measure.py` and `test_qubit_device.py`  
- Fixed circuit visualization bug and added tests for the changes I made
    - Updated `representation_resolver.py` which can be found in the circuit_drawer module 
    - Added unit tests for these changes (in `test_representation_resolver.py`) 

**Here are some code snippets and context regarding the bug:**
The visualization bug was triggered when one would try to visualize a circuit which used the new `sample()` functionality
![circuit_drawer_visual bug](https://user-images.githubusercontent.com/21964864/124339093-cfdfab80-db79-11eb-9edd-9ff6ac8a44c0.png)

This bug was triggered because the object type passed into `representation_resolver.py` was a `MeasurementProcess` with `obs=None`. This would no longer trigger the initial if statement responsible for drawing the `sample` type. I added an additional sub-clause to account for this condition and represent it with the string representation: `raw`. Here is what the result looks like now: 
![circuit_drawer_visual update](https://user-images.githubusercontent.com/21964864/124339181-8f346200-db7a-11eb-95cd-0d22ec52f996.PNG)

**Here are some results from running the testing suite**:
At this point, I believe that the new feature has been fully integrated. In order to verify that it hasn't caused any other bugs in the library I ran the test suite. First, I ran the test suite on the master branch (which does not contain my updates) which is even with pennyLaneAI:master. I got the following results: 
![master_branch_test_suite](https://user-images.githubusercontent.com/21964864/124339243-30bbb380-db7b-11eb-943e-76c025159dc8.PNG)
![master_branch_test_suite2](https://user-images.githubusercontent.com/21964864/124339246-331e0d80-db7b-11eb-9e05-b76c5c84b4cd.PNG)

There seem to be some tests which fail or are skipped on the master branch itself. After looking through the error messages I have realized that most of these are due to me not having all of the optional dependencies (like dash["parallel"] or open fermion) installed in my environment. The numbers to keep track of are: ** 7 failed, 6656 passed, and 431 skipped ** 

Now here I run the test suite on the branch with my updated code: 
![raw_samples_branch_test_suite](https://user-images.githubusercontent.com/21964864/124339325-c8b99d00-db7b-11eb-867a-6d0913f683c1.PNG)
![raw_samples_branch_test_suite2](https://user-images.githubusercontent.com/21964864/124339326-cb1bf700-db7b-11eb-9956-02f1f3f77e97.PNG)

In this case we have ** 7 failed, 6665 passed, and 431 skipped **. The additional 9 tests which passed are the unit tests I added to test my updates. Since all of the error messages are identical to those resulting from running the test suite on the master branch, I am fairly confident that my feature implementation does not cause any other bugs. 

** To Do: **
- Update my local environment with optional dependencies and run the test suite again to make sure that there are no other errors! 
- Wait for any feedback from the Xanadu team on my implementation and updated code as required 

** Summary: ** 
I think the code is ready to be reviewed ! Hopefully I can integrate any changes you suggest and be ready to merge into master within the next few days. Feel free to provide feedback and let me know if you come across any other bugs (although I think we should be good as of now !). 

Thanks, 

josh146: Awesome, great to hear @Jaybsoni! I'll mark this as ready for review, and someone from the team will be on hand to go through on Monday.

Also, great job tracking down and fixing the circuit drawing issue - the circuit drawing code is definitely quite complex!
Jaybsoni: > Hi @Jaybsoni, thanks for this great PR! I thought you tackled it with a compact and clean approach:
> 
> * I like how you re-used existing objects and functions where appropriate
> * Nice attention to detail with both documentation, examples, and making sure the circuit drawer continues to work.
> 
> I have a couple of requested changes, but once addressed I think this is ready for merging:
> 
> * Don't forget to update the changelog, with (a) a description of the change, (b) an example codeblock, and (c) adding your name to the contributors section!
> * It looks like you will also need to run black (`black -l 100 pennylane tests`) to ensure the formatting CI passes.

Hi @josh146, 

I incorporated all of the changes you suggested, double checked that the test suite runs without errors, re-formatted the files using `black` and added information about the change to the ChangeLog. Thanks for all of your comments, and your help throughout this project! I had a lot of fun diving into Pennylane's source code! Let me know if there are any other finishing touches that I need to put on the project. Otherwise I think we are good to merge ! 

Thanks, 
Jaybsoni: On a side note: It seems that CodeFactor keeps failing, but it seems to be failing on some functions that I haven't changed. Not sure what is wrong here ... 
Jaybsoni: Hi @albi3ro, 
Thanks for the comments! 

> For the `circuit_1` example, the drawing does not work as expected.
> 
> `Sample[basis]` shows up when you explicitly set the wires, but it will not show up when you do not explicitly set the wires.

Nice job catching that display bug! I will look into it asap.

For your second comment:
> @qml.qnode(dev)
> def circ():
>     return qml.sample(wires=[])
> ```
> 
> This will return samples for all wires. Intuitively, I would expect that to return samples for no wires.

This might be a little bit complicated. Currently, `MeasurementProcess` instances can't have a None type for the wires. This can be seen in the code (line 54 of measure.py) as: `self._wires = wires or Wires([])`. Thus providing `wires = []` and `wires = None` is functionally identical in `qml.sample()`. I am a little bit hesitant to make a changes to this class since it is so fundamental and I haven't had a chance to explore the entire code base. That being said, I can offer two options: 


- Since sampling with `wires=[]` doesn't make physical sense, it might be more appropriate to just throw an error in `qml.sample` which tells the user that if they don't specify the observable then the wires must be non-empty or None. 

- I can try and augment `MeasurementProcess` such that it allows for None type wires, and then use that information to return None when `qml.sample(wires=[])` is called. I just have to make sure we won't throw any errors in any other functions which will have access to our object. 

Let me know which one makes more sense to implement in this case. If you have other suggestions I would love to hear them! 
Thanks, 

albi3ro: @Jaybsoni  Thanks for looking into that.  Given the trouble of fixing the `sample(wires=[)` problem, I think we can leave that be.  

Therefore, only the drawing problem needs to be fixed.
josh146: > Given the trouble of fixing the sample(wires=[) problem, I think we can leave that be.

Yep, I think this makes sense in the meantime. This way we can come back and fix/update the `MeasurementProcess` class in a separate PR :)
antalszava: Hi @Jaybsoni, cool addition! :slightly_smiling_face: :tada:  I got wondering if drawing `Sample[basis]` at the end of the circuit for having `qml.sample()` could somewhat lead to ambiguity regarding which basis it would be? Maybe it's me, just thought of leaving a comment up for consideration, wouldn't seem to be a blocker.
Jaybsoni: Hi all, 

Sorry for the delay, I got busy with work. 

This bug was a tricky one to pin down, especially since the plotting modules are very inter-connected. The root cause was in `circuit_graph.py` (lines 137 - 159), when initializing the attribute `self._grid`. This attribute represents the quantum circuit. Prior to my change, we would look through each operator or observable provided, iterate over the wires it acts upon and then add them accordingly to `self._grid`. This would fail when we call `qml.sample()` because implicitly `wires = qml.wires.Wires([])`. I added a simple catch to identify when a `qml.sample()` was called and no wires were specified. In this case the operation is replaced with a sample operation where all available wires are provided instead.  

Here is a snapshot of the result: 
![circuit_drawer_visual update group print](https://user-images.githubusercontent.com/21964864/126026567-de29b91a-1401-4d30-883a-c168e5b49211.PNG)

@albi3ro, I think everything should be good to go as of now! Let me know if there is anything I should change / fix about my implementation! 

Thanks, 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 021870159..c546152b6 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -138,6 +138,48 @@
 
 * Ising YY gate functionality added.
   [(#1358)](https://github.com/PennyLaneAI/pennylane/pull/1358)
+  
+* Added functionality to `qml.sample()` to extract samples from the basis states of
+  the device (currently only for qubit devices). Additionally, `wires` can be
+  specified to only return samples from those wires. 
+  [(#1441)](https://github.com/PennyLaneAI/pennylane/pull/1441)
+
+  ```python
+  dev = qml.device("default.qubit", wires=3, shots=5)
+
+  @qml.qnode(dev)
+  def circuit_1():
+      qml.Hadamard(wires=0)
+      qml.Hadamard(wires=1)
+      return qml.sample()
+
+  @qml.qnode(dev)
+  def circuit_2():
+      qml.Hadamard(wires=0)
+      qml.Hadamard(wires=1)
+      return qml.sample(wires=[0,2])    # no observable provided and wires specified 
+  ``` 
+
+  ```pycon
+  >>> print(circuit_1())
+  [[1, 0, 0],
+   [1, 1, 0],
+   [1, 0, 0],
+   [0, 0, 0],
+   [0, 1, 0]]
+
+  >>> print(circuit_2())
+  [[1, 0],
+   [1, 0],
+   [1, 0],
+   [0, 0],
+   [0, 0]]
+
+  >>> print(qml.draw(circuit_2))
+   0: ──H──╭┤ Sample[basis]
+   1: ──H──│┤
+   2: ──H──╰┤ Sample[basis]
+  ```
 
 <h3>Improvements</h3>
 
@@ -186,7 +228,8 @@
 
 This release contains contributions from (in alphabetical order):
 
-Olivia Di Matteo, Josh Izaac, Leonhard Kunczik, Romain Moyard, Ashish Panigrahi, Maria Schuld.
+Olivia Di Matteo, Josh Izaac, Leonhard Kunczik, Romain Moyard, Ashish Panigrahi, Maria Schuld,
+Jay Soni
 
 
 # Release 0.16.0 (current release)
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 5300cd088..77143d606 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -38,6 +38,8 @@ from pennylane import Device
 from pennylane.math import sum as qmlsum
 from pennylane.wires import Wires
 
+from pennylane.measure import MeasurementProcess
+
 
 class QubitDevice(Device):
     """Abstract base class for PennyLane qubit devices.
@@ -791,6 +793,16 @@ class QubitDevice(Device):
             # Process samples for observables with eigenvalues {1, -1}
             samples = 1 - 2 * self._samples[sample_slice, device_wires[0]]
 
+        elif isinstance(
+            observable, MeasurementProcess
+        ):  # if no observable was provided then return the raw samples
+            if (
+                len(observable.wires) != 0
+            ):  # if wires are provided, then we only return samples from those wires
+                samples = self._samples[sample_slice, np.array(device_wires)]
+            else:
+                samples = self._samples[sample_slice]
+
         else:
             # Replace the basis state in the computational basis with the correct eigenvalue.
             # Extract only the columns of the basis samples required based on ``wires``.
diff --git a/pennylane/circuit_drawer/representation_resolver.py b/pennylane/circuit_drawer/representation_resolver.py
index b8e7cb5fa..f4f51a15a 100644
--- a/pennylane/circuit_drawer/representation_resolver.py
+++ b/pennylane/circuit_drawer/representation_resolver.py
@@ -329,8 +329,11 @@ class RepresentationResolver:
         Returns:
             str: String representation of the Operator
         """
-        if isinstance(op, qml.measure.MeasurementProcess) and op.obs is not None:
-            op = op.obs
+        if isinstance(op, qml.measure.MeasurementProcess):
+            if op.obs is not None:
+                op = op.obs
+            else:
+                return "basis"  # when no observable is provided we perform a raw measurement
 
         if isinstance(op, qml.operation.Tensor):
             constituent_representations = [
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
index 38679a352..11c5a37b7 100644
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -23,6 +23,7 @@ import networkx as nx
 import pennylane as qml
 import numpy as np
 
+from pennylane.wires import Wires
 from .circuit_drawer import CHARSETS, CircuitDrawer
 
 
@@ -145,11 +146,17 @@ class CircuitGraph:
         for k, op in enumerate(queue):
             op.queue_idx = k  # store the queue index in the Operator
 
-            if hasattr(op, "return_type") and op.return_type is qml.operation.State:
-                # State measurements contain no wires by default, but wires are
-                # required for the circuit drawer, so we recreate the state
-                # measurement with all wires
-                op = qml.measure.MeasurementProcess(qml.operation.State, wires=wires)
+            if hasattr(op, "return_type"):
+                if op.return_type is qml.operation.State:
+                    # State measurements contain no wires by default, but wires are
+                    # required for the circuit drawer, so we recreate the state
+                    # measurement with all wires
+                    op = qml.measure.MeasurementProcess(qml.operation.State, wires=wires)
+
+                elif op.return_type is qml.operation.Sample and op.wires == Wires([]):
+                    # Sampling without specifying wires is treated as sampling all wires
+                    op = qml.measure.MeasurementProcess(qml.operation.Sample, wires=wires)
+
                 op.queue_idx = k
 
             for w in op.wires:
diff --git a/pennylane/measure.py b/pennylane/measure.py
index 343f15eb7..fa30a2f7a 100644
--- a/pennylane/measure.py
+++ b/pennylane/measure.py
@@ -273,16 +273,18 @@ def var(op):
     return MeasurementProcess(Variance, obs=op)
 
 
-def sample(op):
+def sample(op=None, wires=None):
     r"""Sample from the supplied observable, with the number of shots
     determined from the ``dev.shots`` attribute of the corresponding device.
+    If no observable is provided then basis state samples are returned directly
+    from the device.
 
     The samples are drawn from the eigenvalues :math:`\{\lambda_i\}` of the observable.
     The probability of drawing eigenvalue :math:`\lambda_i` is given by
     :math:`p(\lambda_i) = |\langle \xi_i | \psi \rangle|^2`, where :math:`| \xi_i \rangle`
     is the corresponding basis state from the observable's eigenbasis.
 
-    **Example:**
+    **Example**
 
     .. code-block:: python3
 
@@ -300,17 +302,53 @@ def sample(op):
     >>> circuit(0.5)
     array([ 1.,  1.,  1., -1.])
 
+    If no observable is provided, then the raw basis state samples obtained
+    from device are returned (e.g., for a qubit device, samples from the
+    computational device are returned). In this case, ``wires`` can be specified
+    so that sample results only include measurement results of the qubits of interest.
+
+    .. code-block:: python3
+
+        dev = qml.device("default.qubit", wires=2, shots=4)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.Hadamard(wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.sample()
+
+    Executing this QNode:
+
+    >>> circuit(0.5)
+    array([[0, 1],
+           [0, 0],
+           [1, 1],
+           [0, 0]])
+
+
     Args:
-        op (Observable): a quantum observable object
+        op (Observable or None): a quantum observable object
+        wires (Sequence[int] or int or None): the wires we wish to sample from, ONLY set wires if op is None
 
     Raises:
         QuantumFunctionError: `op` is not an instance of :class:`~.Observable`
+        ValueError: Cannot set wires if an observable is provided
     """
-    if not isinstance(op, Observable):
+    if not isinstance(op, Observable) and op is not None:  # None type is also allowed for op
         raise qml.QuantumFunctionError(
             "{} is not an observable: cannot be used with sample".format(op.name)
         )
 
+    if wires is not None:
+        if op is not None:
+            raise ValueError(
+                "Cannot specify the wires to sample if an observable is "
+                "provided. The wires to sample will be determined directly from the observable."
+            )
+
+        return MeasurementProcess(Sample, obs=op, wires=qml.wires.Wires(wires))
+
     return MeasurementProcess(Sample, obs=op)
 
 


Test patch:
diff --git a/tests/circuit_drawer/test_representation_resolver.py b/tests/circuit_drawer/test_representation_resolver.py
index c5e218685..987d6cd78 100644
--- a/tests/circuit_drawer/test_representation_resolver.py
+++ b/tests/circuit_drawer/test_representation_resolver.py
@@ -203,6 +203,8 @@ class TestRepresentationResolver:
             (qml.Toffoli(wires=[0, 2, 1]).inv(), 1, "X⁻¹"),
             (qml.Toffoli(wires=[0, 2, 1]).inv(), 0, "C"),
             (qml.Toffoli(wires=[0, 2, 1]).inv(), 2, "C"),
+            (qml.measure.sample(wires=[0, 1]), 0, "basis"),  # not providing an observable in
+            (qml.measure.sample(wires=[0, 1]), 1, "basis"),  # sample gets displayed as raw
         ],
     )
     def test_operator_representation_unicode(
@@ -343,6 +345,8 @@ class TestRepresentationResolver:
             (qml.Toffoli(wires=[0, 2, 1]).inv(), 1, "X^-1"),
             (qml.Toffoli(wires=[0, 2, 1]).inv(), 0, "C"),
             (qml.Toffoli(wires=[0, 2, 1]).inv(), 2, "C"),
+            (qml.measure.sample(wires=[0, 1]), 0, "basis"),  # not providing an observable in
+            (qml.measure.sample(wires=[0, 1]), 1, "basis"),  # sample gets displayed as raw
         ],
     )
     def test_operator_representation_ascii(self, ascii_representation_resolver, op, wire, target):
diff --git a/tests/circuit_graph/test_circuit_graph.py b/tests/circuit_graph/test_circuit_graph.py
index ea811689b..25b3579e3 100644
--- a/tests/circuit_graph/test_circuit_graph.py
+++ b/tests/circuit_graph/test_circuit_graph.py
@@ -93,6 +93,40 @@ def circuit_measure_multiple_with_max_twice():
     )
 
 
+def measurement_process_equality(mp1, mp2):
+    """Helper function to check for equality between two MeasurementProcess objects."""
+    same_return_type = mp1.return_type == mp2.return_type
+    same_obs = mp1.obs == mp2.obs
+    same_wires = mp1._wires == mp2._wires
+    same_eigs = mp1._eigvals == mp2._eigvals
+    same_name = mp1.name == mp2.name
+    same_data = mp1.data == mp2.data
+
+    if same_return_type and same_obs and same_wires and same_eigs and same_name and same_data:
+        return True
+
+    return False
+
+
+def lst_mp_equality(lst1, lst2):
+    """Wrapper function to check if two lists are identical when the
+    lists contain MeasurementProcess objects."""
+    assert len(lst1) == len(lst2)
+    lsts_same = True
+
+    for index in range(len(lst1)):
+        if not isinstance(lst1[index], type(lst2[index])):  # if they are different types
+            return False  # then return false
+
+        elif isinstance(lst1[index], qml.measure.MeasurementProcess):
+            lsts_same = lsts_same and measurement_process_equality(lst1[index], lst2[index])
+
+        else:
+            lsts_same = lsts_same and (lst1[index] == lst2[index])
+
+    return lsts_same
+
+
 class TestCircuitGraph:
     """Test conversion of queues to DAGs"""
 
@@ -246,3 +280,29 @@ class TestCircuitGraph:
         qnode()
         circuit = qnode.qtape.graph
         assert circuit.max_simultaneous_measurements == expected
+
+    def test_grid_when_sample_no_wires(self):
+        """A test to ensure the sample operation applies to all wires when
+        none are explicitly provided."""
+
+        ops = [qml.Hadamard(wires=0), qml.CNOT(wires=[0, 1])]
+        obs_no_wires = [qml.measure.sample(op=None, wires=None)]
+        obs_w_wires = [qml.measure.sample(op=None, wires=[0, 1, 2])]
+
+        circuit_no_wires = CircuitGraph(ops, obs_no_wires, wires=Wires([0, 1, 2]))
+        circuit_w_wires = CircuitGraph(ops, obs_w_wires, wires=Wires([0, 1, 2]))
+
+        sample_w_wires_op = qml.measure.sample(op=None, wires=[0, 1, 2])
+        expected_grid = {
+            0: [ops[0], ops[1], sample_w_wires_op],
+            1: [ops[1], sample_w_wires_op],
+            2: [sample_w_wires_op],
+        }
+
+        for key in range(3):
+            lst_w_wires = circuit_w_wires._grid[key]
+            lst_no_wires = circuit_no_wires._grid[key]
+            lst_expected = expected_grid[key]
+
+            assert lst_mp_equality(lst_no_wires, lst_w_wires)
+            assert lst_mp_equality(lst_no_wires, lst_expected)
diff --git a/tests/test_measure.py b/tests/test_measure.py
index 6c24bd864..9cf6d7190 100644
--- a/tests/test_measure.py
+++ b/tests/test_measure.py
@@ -255,6 +255,53 @@ class TestSample:
 
         circuit()
 
+    def test_providing_observable_and_wires(self):
+        """Test that a ValueError is raised if both an observable is provided and wires are specified"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            return qml.sample(qml.PauliZ(0), wires=[0, 1])
+
+        with pytest.raises(
+            ValueError,
+            match="Cannot specify the wires to sample if an observable is provided."
+            " The wires to sample will be determined directly from the observable.",
+        ):
+            res = circuit()
+
+    def test_providing_no_observable_and_no_wires(self):
+        """Test that we can provide no observable and no wires to sample function"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            res = qml.sample()
+            assert res.obs is None
+            assert res.wires == qml.wires.Wires([])
+            return res
+
+        circuit()
+
+    def test_providing_no_observable_and_wires(self):
+        """Test that we can provide no observable but specify wires to the sample function"""
+        wires = [0, 2]
+        wires_obj = qml.wires.Wires(wires)
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            res = qml.sample(wires=wires)
+
+            assert res.obs is None
+            assert res.wires == wires_obj
+            return res
+
+        circuit()
+
 
 @pytest.mark.parametrize(
     "stat_func,return_type", [(expval, Expectation), (var, Variance), (sample, Sample)]
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index cd9473a80..fcaf8d614 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -577,6 +577,34 @@ class TestSample:
 
         assert np.array_equal(res, np.array([-1, 1]))
 
+    def test_sample_with_no_observable_and_no_wires(
+        self, mock_qubit_device_with_original_statistics, tol
+    ):
+        """Test that when we sample a device without providing an observable or wires then it
+        will return the raw samples"""
+        obs = qml.measure.sample(op=None, wires=None)
+        dev = mock_qubit_device_with_original_statistics(wires=2)
+        generated_samples = np.array([[1, 0], [1, 1]])
+        dev._samples = generated_samples
+
+        res = dev.sample(obs)
+        assert np.array_equal(res, generated_samples)
+
+    def test_sample_with_no_observable_and_with_wires(
+        self, mock_qubit_device_with_original_statistics, tol
+    ):
+        """Test that when we sample a device without providing an observable but we specify
+        wires then it returns the generated samples for only those wires"""
+        obs = qml.measure.sample(op=None, wires=[1])
+        dev = mock_qubit_device_with_original_statistics(wires=2)
+        generated_samples = np.array([[1, 0], [1, 1]])
+        dev._samples = generated_samples
+
+        wire_samples = np.array([[0], [1]])
+        res = dev.sample(obs)
+
+        assert np.array_equal(res, wire_samples)
+
 
 class TestEstimateProb:
     """Test the estimate_probability method"""