Repository: PennyLaneAI/pennylane

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.


Discussion:

Resolving patch:
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 @@ def shape(self, device, shots):
                 "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)
 


Test patch:
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
--- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
@@ -908,7 +908,7 @@ def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interfa
         if diff_method == "adjoint":
             pytest.skip("Adjoint warns with finite shots")
 
-        @qnode(
+        @qml.qnode(
             dev,
             diff_method=diff_method,
             interface=interface,
@@ -918,7 +918,7 @@ def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interfa
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
-            return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1))
+            return qml.sample(qml.Z(0)), qml.sample(qml.s_prod(2, qml.X(0) @ qml.Y(1)))
 
         res = jax.jit(circuit, static_argnames="shots")(shots=10)