jit with sampling multi-wire observable

doc/releases/changelog-dev.md

@@ -236,6 +236,8 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `jax.jit` now works with `qml.sample` with a multi-wire observable.
+
 * We no longer perform unwanted dtype promotion in the `pauli_rep` of `SProd` instances when using tensorflow.
   [(#5246)](https://github.com/PennyLaneAI/pennylane/pull/5246)
 

pennylane/measurements/sample.py

@@ -198,7 +198,10 @@ 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:
+            len_wires = 1
+        else:
+            len_wires = len(self.wires) if len(self.wires) > 0 else len(device.wires)
 
         def _single_int_shape(shot_val, num_wires):
             # singleton dimensions, whether in shot val or num_wires are squeezed away

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)