jit with sampling multi-wire observable

doc/releases/changelog-dev.md

@@ -255,6 +255,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `jax.jit` now works with `qml.sample` with a multi-wire observable.
+  [(#5422)](https://github.com/PennyLaneAI/pennylane/pull/5422)
+
 * `qml.qinfo.quantum_fisher` now works with non-`default.qubit` devices.
   [(#5423)](https://github.com/PennyLaneAI/pennylane/pull/5423)
 

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)
 

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)