We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Sampling an observable with more than one wire is not compatible with JAX-JIT
### Expected behavior

That I can JIT a circuit that returns `qml.sample(obs)`, where `obs` is some observable with more than one wire.

### Actual behavior

A shape mismatch causes the circuit to fail if executing with JAX-JIT

### Additional information

_No response_

### Source code

```shell
import jax

jax.config.update("jax_enable_x64", True)

dev = qml.device("default.qubit", wires=2, shots=100)

@qml.qnode(dev, interface="jax")
def circuit(x):
    qml.RX(x, wires=0)
    return qml.sample(qml.PauliX(0) @ qml.PauliY(1))

results = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))
```


### Tracebacks

```shell
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Cell In[131], line 12
      9     qml.RX(x, wires=0)
     10     return qml.sample(qml.PauliX(0) @ qml.PauliY(1))
---> 12 results = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))

    [... skipping hidden 10 frame]

File /opt/homebrew/Caskroom/miniforge/base/envs/pennylane/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:1229, in ExecuteReplicated.__call__(self, *args)
   1224   self._handle_token_bufs(
   1225       results.disassemble_prefix_into_single_device_arrays(
   1226           len(self.ordered_effects)),
   1227       results.consume_token())
   1228 else:
-> 1229   results = self.xla_executable.execute_sharded(input_bufs)
   1230 if dispatch.needs_check_special():
   1231   out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: Incorrect output shape for return value 0: Expected: (100, 2), Actual: (100,)
```


### System information

```shell
pl-dev
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pennylane/measurements/sample.py`
Content:
```
1 # Copyright 2018-2021 Xanadu Quantum Technologies Inc.
2 
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """
15 This module contains the qml.sample measurement.
16 """
17 import functools
18 import warnings
19 from typing import Sequence, Tuple, Optional, Union
20 
21 import pennylane as qml
22 from pennylane.operation import Operator
23 from pennylane.wires import Wires
24 
25 from .measurements import MeasurementShapeError, Sample, SampleMeasurement
26 from .mid_measure import MeasurementValue
27 
28 
29 def sample(op: Optional[Union[Operator, MeasurementValue]] = None, wires=None) -> "SampleMP":
30     r"""Sample from the supplied observable, with the number of shots
31     determined from the ``dev.shots`` attribute of the corresponding device,
32     returning raw samples. If no observable is provided then basis state samples are returned
33     directly from the device.
34 
35     Note that the output shape of this measurement process depends on the shots
36     specified on the device.
37 
38     Args:
39         op (Observable or MeasurementValue): a quantum observable object. To get samples
40             for mid-circuit measurements, ``op`` should be a``MeasurementValue``.
41         wires (Sequence[int] or int or None): the wires we wish to sample from; ONLY set wires if
42             op is ``None``
43 
44     Returns:
45         SampleMP: Measurement process instance
46 
47     Raises:
48         ValueError: Cannot set wires if an observable is provided
49 
50     The samples are drawn from the eigenvalues :math:`\{\lambda_i\}` of the observable.
51     The probability of drawing eigenvalue :math:`\lambda_i` is given by
52     :math:`p(\lambda_i) = |\langle \xi_i | \psi \rangle|^2`, where :math:`| \xi_i \rangle`
53     is the corresponding basis state from the observable's eigenbasis.
54 
55     .. note::
56 
57         QNodes that return samples cannot, in general, be differentiated, since the derivative
58         with respect to a sample --- a stochastic process --- is ill-defined. An alternative
59         approach would be to use single-shot expectation values. For example, instead of this:
60 
61         .. code-block:: python
62 
63             dev = qml.device("default.qubit", shots=10)
64 
65             @qml.qnode(dev, diff_method="parameter-shift")
66             def circuit(angle):
67                 qml.RX(angle, wires=0)
68                 return qml.sample(qml.PauliX(0))
69 
70             angle = qml.numpy.array(0.1)
71             res = qml.jacobian(circuit)(angle)
72 
73         Consider using :func:`~pennylane.expval` and a sequence of single shots, like this:
74 
75         .. code-block:: python
76 
77             dev = qml.device("default.qubit", shots=[(1, 10)])
78 
79             @qml.qnode(dev, diff_method="parameter-shift")
80             def circuit(angle):
81                 qml.RX(angle, wires=0)
82                 return qml.expval(qml.PauliX(0))
83 
84             def cost(angle):
85                 return qml.math.hstack(circuit(angle))
86 
87             angle = qml.numpy.array(0.1)
88             res = qml.jacobian(cost)(angle)
89 
90     **Example**
91 
92     .. code-block:: python3
93 
94         dev = qml.device("default.qubit", wires=2, shots=4)
95 
96         @qml.qnode(dev)
97         def circuit(x):
98             qml.RX(x, wires=0)
99             qml.Hadamard(wires=1)
100             qml.CNOT(wires=[0, 1])
101             return qml.sample(qml.Y(0))
102 
103     Executing this QNode:
104 
105     >>> circuit(0.5)
106     array([ 1.,  1.,  1., -1.])
107 
108     If no observable is provided, then the raw basis state samples obtained
109     from device are returned (e.g., for a qubit device, samples from the
110     computational device are returned). In this case, ``wires`` can be specified
111     so that sample results only include measurement results of the qubits of interest.
112 
113     .. code-block:: python3
114 
115         dev = qml.device("default.qubit", wires=2, shots=4)
116 
117         @qml.qnode(dev)
118         def circuit(x):
119             qml.RX(x, wires=0)
120             qml.Hadamard(wires=1)
121             qml.CNOT(wires=[0, 1])
122             return qml.sample()
123 
124     Executing this QNode:
125 
126     >>> circuit(0.5)
127     array([[0, 1],
128            [0, 0],
129            [1, 1],
130            [0, 0]])
131 
132     """
133     if isinstance(op, MeasurementValue):
134         return SampleMP(obs=op)
135 
136     if isinstance(op, Sequence):
137         if not all(isinstance(o, MeasurementValue) and len(o.measurements) == 1 for o in op):
138             raise qml.QuantumFunctionError(
139                 "Only sequences of single MeasurementValues can be passed with the op argument. "
140                 "MeasurementValues manipulated using arithmetic operators cannot be used when "
141                 "collecting statistics for a sequence of mid-circuit measurements."
142             )
143 
144         return SampleMP(obs=op)
145 
146     if op is not None and not op.is_hermitian:  # None type is also allowed for op
147         warnings.warn(f"{op.name} might not be hermitian.")
148 
149     if wires is not None:
150         if op is not None:
151             raise ValueError(
152                 "Cannot specify the wires to sample if an observable is "
153                 "provided. The wires to sample will be determined directly from the observable."
154             )
155         wires = Wires(wires)
156 
157     return SampleMP(obs=op, wires=wires)
158 
159 
160 class SampleMP(SampleMeasurement):
161     """Measurement process that returns the samples of a given observable. If no observable is
162     provided then basis state samples are returned directly from the device.
163 
164     Please refer to :func:`sample` for detailed documentation.
165 
166     Args:
167         obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured
168             as part of the measurement process. Not all measurement processes require observables
169             (for example ``Probability``); this argument is optional.
170         wires (.Wires): The wires the measurement process applies to.
171             This can only be specified if an observable was not provided.
172         eigvals (array): A flat array representing the eigenvalues of the measurement.
173             This can only be specified if an observable was not provided.
174         id (str): custom label given to a measurement instance, can be useful for some applications
175             where the instance has to be identified
176     """
177 
178     @property
179     def return_type(self):
180         return Sample
181 
182     @property
183     @functools.lru_cache()
184     def numeric_type(self):
185         # Note: we only assume an integer numeric type if the observable is a
186         # built-in observable with integer eigenvalues or a tensor product thereof
187         if self.obs is None:
188             # Computational basis samples
189             return int
190         int_eigval_obs = {qml.X, qml.Y, qml.Z, qml.Hadamard, qml.Identity}
191         tensor_terms = self.obs.obs if hasattr(self.obs, "obs") else [self.obs]
192         every_term_standard = all(o.__class__ in int_eigval_obs for o in tensor_terms)
193         return int if every_term_standard else float
194 
195     def shape(self, device, shots):
196         if not shots:
197             raise MeasurementShapeError(
198                 "Shots are required to obtain the shape of the measurement "
199                 f"{self.__class__.__name__}."
200             )
201         len_wires = len(self.wires) if len(self.wires) > 0 else len(device.wires)
202 
203         def _single_int_shape(shot_val, num_wires):
204             # singleton dimensions, whether in shot val or num_wires are squeezed away
205             inner_shape = []
206             if shot_val != 1:
207                 inner_shape.append(shot_val)
208             if num_wires != 1:
209                 inner_shape.append(num_wires)
210             return tuple(inner_shape)
211 
212         if not shots.has_partitioned_shots:
213             return _single_int_shape(shots.total_shots, len_wires)
214 
215         shape = []
216         for s in shots.shot_vector:
217             for _ in range(s.copies):
218                 shape.append(_single_int_shape(s.shots, len_wires))
219 
220         return tuple(shape)
221 
222     def process_samples(
223         self,
224         samples: Sequence[complex],
225         wire_order: Wires,
226         shot_range: Tuple[int] = None,
227         bin_size: int = None,
228     ):
229         wire_map = dict(zip(wire_order, range(len(wire_order))))
230         mapped_wires = [wire_map[w] for w in self.wires]
231         name = self.obs.name if self.obs is not None else None
232         # Select the samples from samples that correspond to ``shot_range`` if provided
233         if shot_range is not None:
234             # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last
235             # colon (:) is required because shots is the second-to-last axis and the
236             # Ellipsis (...) otherwise would take up broadcasting and shots axes.
237             samples = samples[..., slice(*shot_range), :]
238 
239         if mapped_wires:
240             # if wires are provided, then we only return samples from those wires
241             samples = samples[..., mapped_wires]
242 
243         num_wires = samples.shape[-1]  # wires is the last dimension
244 
245         # If we're sampling wires or a list of mid-circuit measurements
246         if self.obs is None and not isinstance(self.mv, MeasurementValue):
247             # if no observable was provided then return the raw samples
248             return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)
249 
250         # If we're sampling observables
251         if str(name) in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
252             # Process samples for observables with eigenvalues {1, -1}
253             samples = 1 - 2 * qml.math.squeeze(samples, axis=-1)
254         else:
255             # Replace the basis state in the computational basis with the correct eigenvalue.
256             # Extract only the columns of the basis samples required based on ``wires``.
257             powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
258             indices = samples @ powers_of_two
259             indices = qml.math.array(indices)  # Add np.array here for Jax support.
260             try:
261                 # This also covers statistics for mid-circuit measurements manipulated using
262                 # arithmetic operators
263                 samples = self.eigvals()[indices]
264             except qml.operation.EigvalsUndefinedError as e:
265                 # if observable has no info on eigenvalues, we cannot return this measurement
266                 raise qml.operation.EigvalsUndefinedError(
267                     f"Cannot compute samples of {self.obs.name}."
268                 ) from e
269 
270         return samples if bin_size is None else samples.reshape((bin_size, -1))
271
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -198,24 +198,28 @@
                 "Shots are required to obtain the shape of the measurement "
                 f"{self.__class__.__name__}."
             )
-        len_wires = len(self.wires) if len(self.wires) > 0 else len(device.wires)
+        if self.obs:
+            num_values_per_shot = 1  # one single eigenvalue
+        else:
+            # one value per wire
+            num_values_per_shot = len(self.wires) if len(self.wires) > 0 else len(device.wires)
 
-        def _single_int_shape(shot_val, num_wires):
+        def _single_int_shape(shot_val, num_values):
             # singleton dimensions, whether in shot val or num_wires are squeezed away
             inner_shape = []
             if shot_val != 1:
                 inner_shape.append(shot_val)
-            if num_wires != 1:
-                inner_shape.append(num_wires)
+            if num_values != 1:
+                inner_shape.append(num_values)
             return tuple(inner_shape)
 
         if not shots.has_partitioned_shots:
-            return _single_int_shape(shots.total_shots, len_wires)
+            return _single_int_shape(shots.total_shots, num_values_per_shot)
 
         shape = []
         for s in shots.shot_vector:
             for _ in range(s.copies):
-                shape.append(_single_int_shape(s.shots, len_wires))
+                shape.append(_single_int_shape(s.shots, num_values_per_shot))
 
         return tuple(shape)

{"golden_diff": "diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py\n--- a/pennylane/measurements/sample.py\n+++ b/pennylane/measurements/sample.py\n@@ -198,24 +198,28 @@\n                 \"Shots are required to obtain the shape of the measurement \"\n                 f\"{self.__class__.__name__}.\"\n             )\n-        len_wires = len(self.wires) if len(self.wires) > 0 else len(device.wires)\n+        if self.obs:\n+            num_values_per_shot = 1  # one single eigenvalue\n+        else:\n+            # one value per wire\n+            num_values_per_shot = len(self.wires) if len(self.wires) > 0 else len(device.wires)\n \n-        def _single_int_shape(shot_val, num_wires):\n+        def _single_int_shape(shot_val, num_values):\n             # singleton dimensions, whether in shot val or num_wires are squeezed away\n             inner_shape = []\n             if shot_val != 1:\n                 inner_shape.append(shot_val)\n-            if num_wires != 1:\n-                inner_shape.append(num_wires)\n+            if num_values != 1:\n+                inner_shape.append(num_values)\n             return tuple(inner_shape)\n \n         if not shots.has_partitioned_shots:\n-            return _single_int_shape(shots.total_shots, len_wires)\n+            return _single_int_shape(shots.total_shots, num_values_per_shot)\n \n         shape = []\n         for s in shots.shot_vector:\n             for _ in range(s.copies):\n-                shape.append(_single_int_shape(s.shots, len_wires))\n+                shape.append(_single_int_shape(s.shots, num_values_per_shot))\n \n         return tuple(shape)\n", "issue": "[BUG] Sampling an observable with more than one wire is not compatible with JAX-JIT\n### Expected behavior\n\nThat I can JIT a circuit that returns `qml.sample(obs)`, where `obs` is some observable with more than one wire.\n\n### Actual behavior\n\nA shape mismatch causes the circuit to fail if executing with JAX-JIT\n\n### Additional information\n\n_No response_\n\n### Source code\n\n```shell\nimport jax\r\n\r\njax.config.update(\"jax_enable_x64\", True)\r\n\r\ndev = qml.device(\"default.qubit\", wires=2, shots=100)\r\n\r\n@qml.qnode(dev, interface=\"jax\")\r\ndef circuit(x):\r\n    qml.RX(x, wires=0)\r\n    return qml.sample(qml.PauliX(0) @ qml.PauliY(1))\r\n\r\nresults = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))\n```\n\n\n### Tracebacks\n\n```shell\n---------------------------------------------------------------------------\r\nXlaRuntimeError                           Traceback (most recent call last)\r\nCell In[131], line 12\r\n      9     qml.RX(x, wires=0)\r\n     10     return qml.sample(qml.PauliX(0) @ qml.PauliY(1))\r\n---> 12 results = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))\r\n\r\n    [... skipping hidden 10 frame]\r\n\r\nFile /opt/homebrew/Caskroom/miniforge/base/envs/pennylane/lib/python3.10/site-packages/jax/_src/interpreters/pxla.py:1229, in ExecuteReplicated.__call__(self, *args)\r\n   1224   self._handle_token_bufs(\r\n   1225       results.disassemble_prefix_into_single_device_arrays(\r\n   1226           len(self.ordered_effects)),\r\n   1227       results.consume_token())\r\n   1228 else:\r\n-> 1229   results = self.xla_executable.execute_sharded(input_bufs)\r\n   1230 if dispatch.needs_check_special():\r\n   1231   out_arrays = results.disassemble_into_single_device_arrays()\r\n\r\nXlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: Incorrect output shape for return value 0: Expected: (100, 2), Actual: (100,)\n```\n\n\n### System information\n\n```shell\npl-dev\n```\n\n\n### Existing GitHub issues\n\n- [X] I have searched existing GitHub issues to make sure the issue does not already exist.\n", "before_files": [{"content": "# Copyright 2018-2021 Xanadu Quantum Technologies Inc.\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThis module contains the qml.sample measurement.\n\"\"\"\nimport functools\nimport warnings\nfrom typing import Sequence, Tuple, Optional, Union\n\nimport pennylane as qml\nfrom pennylane.operation import Operator\nfrom pennylane.wires import Wires\n\nfrom .measurements import MeasurementShapeError, Sample, SampleMeasurement\nfrom .mid_measure import MeasurementValue\n\n\ndef sample(op: Optional[Union[Operator, MeasurementValue]] = None, wires=None) -> \"SampleMP\":\n    r\"\"\"Sample from the supplied observable, with the number of shots\n    determined from the ``dev.shots`` attribute of the corresponding device,\n    returning raw samples. If no observable is provided then basis state samples are returned\n    directly from the device.\n\n    Note that the output shape of this measurement process depends on the shots\n    specified on the device.\n\n    Args:\n        op (Observable or MeasurementValue): a quantum observable object. To get samples\n            for mid-circuit measurements, ``op`` should be a``MeasurementValue``.\n        wires (Sequence[int] or int or None): the wires we wish to sample from; ONLY set wires if\n            op is ``None``\n\n    Returns:\n        SampleMP: Measurement process instance\n\n    Raises:\n        ValueError: Cannot set wires if an observable is provided\n\n    The samples are drawn from the eigenvalues :math:`\\{\\lambda_i\\}` of the observable.\n    The probability of drawing eigenvalue :math:`\\lambda_i` is given by\n    :math:`p(\\lambda_i) = |\\langle \\xi_i | \\psi \\rangle|^2`, where :math:`| \\xi_i \\rangle`\n    is the corresponding basis state from the observable's eigenbasis.\n\n    .. note::\n\n        QNodes that return samples cannot, in general, be differentiated, since the derivative\n        with respect to a sample --- a stochastic process --- is ill-defined. An alternative\n        approach would be to use single-shot expectation values. For example, instead of this:\n\n        .. code-block:: python\n\n            dev = qml.device(\"default.qubit\", shots=10)\n\n            @qml.qnode(dev, diff_method=\"parameter-shift\")\n            def circuit(angle):\n                qml.RX(angle, wires=0)\n                return qml.sample(qml.PauliX(0))\n\n            angle = qml.numpy.array(0.1)\n            res = qml.jacobian(circuit)(angle)\n\n        Consider using :func:`~pennylane.expval` and a sequence of single shots, like this:\n\n        .. code-block:: python\n\n            dev = qml.device(\"default.qubit\", shots=[(1, 10)])\n\n            @qml.qnode(dev, diff_method=\"parameter-shift\")\n            def circuit(angle):\n                qml.RX(angle, wires=0)\n                return qml.expval(qml.PauliX(0))\n\n            def cost(angle):\n                return qml.math.hstack(circuit(angle))\n\n            angle = qml.numpy.array(0.1)\n            res = qml.jacobian(cost)(angle)\n\n    **Example**\n\n    .. code-block:: python3\n\n        dev = qml.device(\"default.qubit\", wires=2, shots=4)\n\n        @qml.qnode(dev)\n        def circuit(x):\n            qml.RX(x, wires=0)\n            qml.Hadamard(wires=1)\n            qml.CNOT(wires=[0, 1])\n            return qml.sample(qml.Y(0))\n\n    Executing this QNode:\n\n    >>> circuit(0.5)\n    array([ 1.,  1.,  1., -1.])\n\n    If no observable is provided, then the raw basis state samples obtained\n    from device are returned (e.g., for a qubit device, samples from the\n    computational device are returned). In this case, ``wires`` can be specified\n    so that sample results only include measurement results of the qubits of interest.\n\n    .. code-block:: python3\n\n        dev = qml.device(\"default.qubit\", wires=2, shots=4)\n\n        @qml.qnode(dev)\n        def circuit(x):\n            qml.RX(x, wires=0)\n            qml.Hadamard(wires=1)\n            qml.CNOT(wires=[0, 1])\n            return qml.sample()\n\n    Executing this QNode:\n\n    >>> circuit(0.5)\n    array([[0, 1],\n           [0, 0],\n           [1, 1],\n           [0, 0]])\n\n    \"\"\"\n    if isinstance(op, MeasurementValue):\n        return SampleMP(obs=op)\n\n    if isinstance(op, Sequence):\n        if not all(isinstance(o, MeasurementValue) and len(o.measurements) == 1 for o in op):\n            raise qml.QuantumFunctionError(\n                \"Only sequences of single MeasurementValues can be passed with the op argument. \"\n                \"MeasurementValues manipulated using arithmetic operators cannot be used when \"\n                \"collecting statistics for a sequence of mid-circuit measurements.\"\n            )\n\n        return SampleMP(obs=op)\n\n    if op is not None and not op.is_hermitian:  # None type is also allowed for op\n        warnings.warn(f\"{op.name} might not be hermitian.\")\n\n    if wires is not None:\n        if op is not None:\n            raise ValueError(\n                \"Cannot specify the wires to sample if an observable is \"\n                \"provided. The wires to sample will be determined directly from the observable.\"\n            )\n        wires = Wires(wires)\n\n    return SampleMP(obs=op, wires=wires)\n\n\nclass SampleMP(SampleMeasurement):\n    \"\"\"Measurement process that returns the samples of a given observable. If no observable is\n    provided then basis state samples are returned directly from the device.\n\n    Please refer to :func:`sample` for detailed documentation.\n\n    Args:\n        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured\n            as part of the measurement process. Not all measurement processes require observables\n            (for example ``Probability``); this argument is optional.\n        wires (.Wires): The wires the measurement process applies to.\n            This can only be specified if an observable was not provided.\n        eigvals (array): A flat array representing the eigenvalues of the measurement.\n            This can only be specified if an observable was not provided.\n        id (str): custom label given to a measurement instance, can be useful for some applications\n            where the instance has to be identified\n    \"\"\"\n\n    @property\n    def return_type(self):\n        return Sample\n\n    @property\n    @functools.lru_cache()\n    def numeric_type(self):\n        # Note: we only assume an integer numeric type if the observable is a\n        # built-in observable with integer eigenvalues or a tensor product thereof\n        if self.obs is None:\n            # Computational basis samples\n            return int\n        int_eigval_obs = {qml.X, qml.Y, qml.Z, qml.Hadamard, qml.Identity}\n        tensor_terms = self.obs.obs if hasattr(self.obs, \"obs\") else [self.obs]\n        every_term_standard = all(o.__class__ in int_eigval_obs for o in tensor_terms)\n        return int if every_term_standard else float\n\n    def shape(self, device, shots):\n        if not shots:\n            raise MeasurementShapeError(\n                \"Shots are required to obtain the shape of the measurement \"\n                f\"{self.__class__.__name__}.\"\n            )\n        len_wires = len(self.wires) if len(self.wires) > 0 else len(device.wires)\n\n        def _single_int_shape(shot_val, num_wires):\n            # singleton dimensions, whether in shot val or num_wires are squeezed away\n            inner_shape = []\n            if shot_val != 1:\n                inner_shape.append(shot_val)\n            if num_wires != 1:\n                inner_shape.append(num_wires)\n            return tuple(inner_shape)\n\n        if not shots.has_partitioned_shots:\n            return _single_int_shape(shots.total_shots, len_wires)\n\n        shape = []\n        for s in shots.shot_vector:\n            for _ in range(s.copies):\n                shape.append(_single_int_shape(s.shots, len_wires))\n\n        return tuple(shape)\n\n    def process_samples(\n        self,\n        samples: Sequence[complex],\n        wire_order: Wires,\n        shot_range: Tuple[int] = None,\n        bin_size: int = None,\n    ):\n        wire_map = dict(zip(wire_order, range(len(wire_order))))\n        mapped_wires = [wire_map[w] for w in self.wires]\n        name = self.obs.name if self.obs is not None else None\n        # Select the samples from samples that correspond to ``shot_range`` if provided\n        if shot_range is not None:\n            # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last\n            # colon (:) is required because shots is the second-to-last axis and the\n            # Ellipsis (...) otherwise would take up broadcasting and shots axes.\n            samples = samples[..., slice(*shot_range), :]\n\n        if mapped_wires:\n            # if wires are provided, then we only return samples from those wires\n            samples = samples[..., mapped_wires]\n\n        num_wires = samples.shape[-1]  # wires is the last dimension\n\n        # If we're sampling wires or a list of mid-circuit measurements\n        if self.obs is None and not isinstance(self.mv, MeasurementValue):\n            # if no observable was provided then return the raw samples\n            return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)\n\n        # If we're sampling observables\n        if str(name) in {\"PauliX\", \"PauliY\", \"PauliZ\", \"Hadamard\"}:\n            # Process samples for observables with eigenvalues {1, -1}\n            samples = 1 - 2 * qml.math.squeeze(samples, axis=-1)\n        else:\n            # Replace the basis state in the computational basis with the correct eigenvalue.\n            # Extract only the columns of the basis samples required based on ``wires``.\n            powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]\n            indices = samples @ powers_of_two\n            indices = qml.math.array(indices)  # Add np.array here for Jax support.\n            try:\n                # This also covers statistics for mid-circuit measurements manipulated using\n                # arithmetic operators\n                samples = self.eigvals()[indices]\n            except qml.operation.EigvalsUndefinedError as e:\n                # if observable has no info on eigenvalues, we cannot return this measurement\n                raise qml.operation.EigvalsUndefinedError(\n                    f\"Cannot compute samples of {self.obs.name}.\"\n                ) from e\n\n        return samples if bin_size is None else samples.reshape((bin_size, -1))\n", "path": "pennylane/measurements/sample.py"}], "after_files": [{"content": "# Copyright 2018-2021 Xanadu Quantum Technologies Inc.\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThis module contains the qml.sample measurement.\n\"\"\"\nimport functools\nimport warnings\nfrom typing import Sequence, Tuple, Optional, Union\n\nimport pennylane as qml\nfrom pennylane.operation import Operator\nfrom pennylane.wires import Wires\n\nfrom .measurements import MeasurementShapeError, Sample, SampleMeasurement\nfrom .mid_measure import MeasurementValue\n\n\ndef sample(op: Optional[Union[Operator, MeasurementValue]] = None, wires=None) -> \"SampleMP\":\n    r\"\"\"Sample from the supplied observable, with the number of shots\n    determined from the ``dev.shots`` attribute of the corresponding device,\n    returning raw samples. If no observable is provided then basis state samples are returned\n    directly from the device.\n\n    Note that the output shape of this measurement process depends on the shots\n    specified on the device.\n\n    Args:\n        op (Observable or MeasurementValue): a quantum observable object. To get samples\n            for mid-circuit measurements, ``op`` should be a``MeasurementValue``.\n        wires (Sequence[int] or int or None): the wires we wish to sample from; ONLY set wires if\n            op is ``None``\n\n    Returns:\n        SampleMP: Measurement process instance\n\n    Raises:\n        ValueError: Cannot set wires if an observable is provided\n\n    The samples are drawn from the eigenvalues :math:`\\{\\lambda_i\\}` of the observable.\n    The probability of drawing eigenvalue :math:`\\lambda_i` is given by\n    :math:`p(\\lambda_i) = |\\langle \\xi_i | \\psi \\rangle|^2`, where :math:`| \\xi_i \\rangle`\n    is the corresponding basis state from the observable's eigenbasis.\n\n    .. note::\n\n        QNodes that return samples cannot, in general, be differentiated, since the derivative\n        with respect to a sample --- a stochastic process --- is ill-defined. An alternative\n        approach would be to use single-shot expectation values. For example, instead of this:\n\n        .. code-block:: python\n\n            dev = qml.device(\"default.qubit\", shots=10)\n\n            @qml.qnode(dev, diff_method=\"parameter-shift\")\n            def circuit(angle):\n                qml.RX(angle, wires=0)\n                return qml.sample(qml.PauliX(0))\n\n            angle = qml.numpy.array(0.1)\n            res = qml.jacobian(circuit)(angle)\n\n        Consider using :func:`~pennylane.expval` and a sequence of single shots, like this:\n\n        .. code-block:: python\n\n            dev = qml.device(\"default.qubit\", shots=[(1, 10)])\n\n            @qml.qnode(dev, diff_method=\"parameter-shift\")\n            def circuit(angle):\n                qml.RX(angle, wires=0)\n                return qml.expval(qml.PauliX(0))\n\n            def cost(angle):\n                return qml.math.hstack(circuit(angle))\n\n            angle = qml.numpy.array(0.1)\n            res = qml.jacobian(cost)(angle)\n\n    **Example**\n\n    .. code-block:: python3\n\n        dev = qml.device(\"default.qubit\", wires=2, shots=4)\n\n        @qml.qnode(dev)\n        def circuit(x):\n            qml.RX(x, wires=0)\n            qml.Hadamard(wires=1)\n            qml.CNOT(wires=[0, 1])\n            return qml.sample(qml.Y(0))\n\n    Executing this QNode:\n\n    >>> circuit(0.5)\n    array([ 1.,  1.,  1., -1.])\n\n    If no observable is provided, then the raw basis state samples obtained\n    from device are returned (e.g., for a qubit device, samples from the\n    computational device are returned). In this case, ``wires`` can be specified\n    so that sample results only include measurement results of the qubits of interest.\n\n    .. code-block:: python3\n\n        dev = qml.device(\"default.qubit\", wires=2, shots=4)\n\n        @qml.qnode(dev)\n        def circuit(x):\n            qml.RX(x, wires=0)\n            qml.Hadamard(wires=1)\n            qml.CNOT(wires=[0, 1])\n            return qml.sample()\n\n    Executing this QNode:\n\n    >>> circuit(0.5)\n    array([[0, 1],\n           [0, 0],\n           [1, 1],\n           [0, 0]])\n\n    \"\"\"\n    if isinstance(op, MeasurementValue):\n        return SampleMP(obs=op)\n\n    if isinstance(op, Sequence):\n        if not all(isinstance(o, MeasurementValue) and len(o.measurements) == 1 for o in op):\n            raise qml.QuantumFunctionError(\n                \"Only sequences of single MeasurementValues can be passed with the op argument. \"\n                \"MeasurementValues manipulated using arithmetic operators cannot be used when \"\n                \"collecting statistics for a sequence of mid-circuit measurements.\"\n            )\n\n        return SampleMP(obs=op)\n\n    if op is not None and not op.is_hermitian:  # None type is also allowed for op\n        warnings.warn(f\"{op.name} might not be hermitian.\")\n\n    if wires is not None:\n        if op is not None:\n            raise ValueError(\n                \"Cannot specify the wires to sample if an observable is \"\n                \"provided. The wires to sample will be determined directly from the observable.\"\n            )\n        wires = Wires(wires)\n\n    return SampleMP(obs=op, wires=wires)\n\n\nclass SampleMP(SampleMeasurement):\n    \"\"\"Measurement process that returns the samples of a given observable. If no observable is\n    provided then basis state samples are returned directly from the device.\n\n    Please refer to :func:`sample` for detailed documentation.\n\n    Args:\n        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured\n            as part of the measurement process. Not all measurement processes require observables\n            (for example ``Probability``); this argument is optional.\n        wires (.Wires): The wires the measurement process applies to.\n            This can only be specified if an observable was not provided.\n        eigvals (array): A flat array representing the eigenvalues of the measurement.\n            This can only be specified if an observable was not provided.\n        id (str): custom label given to a measurement instance, can be useful for some applications\n            where the instance has to be identified\n    \"\"\"\n\n    @property\n    def return_type(self):\n        return Sample\n\n    @property\n    @functools.lru_cache()\n    def numeric_type(self):\n        # Note: we only assume an integer numeric type if the observable is a\n        # built-in observable with integer eigenvalues or a tensor product thereof\n        if self.obs is None:\n            # Computational basis samples\n            return int\n        int_eigval_obs = {qml.X, qml.Y, qml.Z, qml.Hadamard, qml.Identity}\n        tensor_terms = self.obs.obs if hasattr(self.obs, \"obs\") else [self.obs]\n        every_term_standard = all(o.__class__ in int_eigval_obs for o in tensor_terms)\n        return int if every_term_standard else float\n\n    def shape(self, device, shots):\n        if not shots:\n            raise MeasurementShapeError(\n                \"Shots are required to obtain the shape of the measurement \"\n                f\"{self.__class__.__name__}.\"\n            )\n        if self.obs:\n            num_values_per_shot = 1  # one single eigenvalue\n        else:\n            # one value per wire\n            num_values_per_shot = len(self.wires) if len(self.wires) > 0 else len(device.wires)\n\n        def _single_int_shape(shot_val, num_values):\n            # singleton dimensions, whether in shot val or num_wires are squeezed away\n            inner_shape = []\n            if shot_val != 1:\n                inner_shape.append(shot_val)\n            if num_values != 1:\n                inner_shape.append(num_values)\n            return tuple(inner_shape)\n\n        if not shots.has_partitioned_shots:\n            return _single_int_shape(shots.total_shots, num_values_per_shot)\n\n        shape = []\n        for s in shots.shot_vector:\n            for _ in range(s.copies):\n                shape.append(_single_int_shape(s.shots, num_values_per_shot))\n\n        return tuple(shape)\n\n    def process_samples(\n        self,\n        samples: Sequence[complex],\n        wire_order: Wires,\n        shot_range: Tuple[int] = None,\n        bin_size: int = None,\n    ):\n        wire_map = dict(zip(wire_order, range(len(wire_order))))\n        mapped_wires = [wire_map[w] for w in self.wires]\n        name = self.obs.name if self.obs is not None else None\n        # Select the samples from samples that correspond to ``shot_range`` if provided\n        if shot_range is not None:\n            # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last\n            # colon (:) is required because shots is the second-to-last axis and the\n            # Ellipsis (...) otherwise would take up broadcasting and shots axes.\n            samples = samples[..., slice(*shot_range), :]\n\n        if mapped_wires:\n            # if wires are provided, then we only return samples from those wires\n            samples = samples[..., mapped_wires]\n\n        num_wires = samples.shape[-1]  # wires is the last dimension\n\n        # If we're sampling wires or a list of mid-circuit measurements\n        if self.obs is None and not isinstance(self.mv, MeasurementValue):\n            # if no observable was provided then return the raw samples\n            return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)\n\n        # If we're sampling observables\n        if str(name) in {\"PauliX\", \"PauliY\", \"PauliZ\", \"Hadamard\"}:\n            # Process samples for observables with eigenvalues {1, -1}\n            samples = 1 - 2 * qml.math.squeeze(samples, axis=-1)\n        else:\n            # Replace the basis state in the computational basis with the correct eigenvalue.\n            # Extract only the columns of the basis samples required based on ``wires``.\n            powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]\n            indices = samples @ powers_of_two\n            indices = qml.math.array(indices)  # Add np.array here for Jax support.\n            try:\n                # This also covers statistics for mid-circuit measurements manipulated using\n                # arithmetic operators\n                samples = self.eigvals()[indices]\n            except qml.operation.EigvalsUndefinedError as e:\n                # if observable has no info on eigenvalues, we cannot return this measurement\n                raise qml.operation.EigvalsUndefinedError(\n                    f\"Cannot compute samples of {self.obs.name}.\"\n                ) from e\n\n        return samples if bin_size is None else samples.reshape((bin_size, -1))\n", "path": "pennylane/measurements/sample.py"}]}