Repository: PennyLaneAI/pennylane

Issue:
[BUG] `Sample` does not work with jitting when a subset of wires are used
### Expected behavior

`qml.sample` works without errors when a circuit with parameters is jitted and a subset of the circuit wires are specified in the call to `qml.sample`.

### Actual behavior

An exception is raised due to a mismatch between exepected and actual output shape. Note that when the circuit does not have any parameters, no error is raised and the output is correct.

### Additional information

_No response_

### Source code

```shell
>>> import pennylane as qml
>>> import jax
>>> import jax.numpy as jnp
>>> jax.config.update("jax_enable_x64", True)
>>> dev = qml.device("default.qubit", wires=3, shots=5)
>>> @qml.qnode(dev, interface="jax")
... def circuit(p):
...     qml.broadcast(qml.Hadamard, pattern="single", wires=dev.wires)
...     qml.RX(p, wires=0)
...     return qml.sample(wires=0)
... 
>>> jitted_qnode = jax.jit(circuit)
>>> jitted_qnode(jnp.array(0.123))
```


### Tracebacks

```shell
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/api.py", line 695, in cache_miss
    out_flat = call_bind_continuation(execute(*args_flat))
  File "/opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
    return func(*args, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py", line 2108, in __call__
    out_bufs = self.xla_executable.execute_sharded_on_local_devices(
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: Incorrect output shape for return value 0: Expected: (1, 5, 3), Actual: (1, 5, 1)

At:
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/interpreters/mlir.py(1796): _wrapped_callback
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py(2108): __call__
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/profiler.py(314): wrapper
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/api.py(695): cache_miss
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/traceback_util.py(163): reraise_with_filtered_traceback
  <stdin>(1): <module>

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: Incorrect output shape for return value 0: Expected: (1, 5, 3), Actual: (1, 5, 1)

At:
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/interpreters/mlir.py(1796): _wrapped_callback
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/interpreters/pxla.py(2108): __call__
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/profiler.py(314): wrapper
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/api.py(695): cache_miss
  /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages/jax/_src/traceback_util.py(163): reraise_with_filtered_traceback
  <stdin>(1): <module>
```


### System information

```shell
Name: PennyLane
Version: 0.30.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /opt/homebrew/Caskroom/miniforge/base/envs/qml/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Cirq, PennyLane-Lightning, PennyLane-qiskit

Platform info:           macOS-13.2.1-arm64-arm-64bit
Python version:          3.9.15
Numpy version:           1.23.5
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.30.0.dev0)
- default.mixed (PennyLane-0.30.0.dev0)
- default.qubit (PennyLane-0.30.0.dev0)
- default.qubit.autograd (PennyLane-0.30.0.dev0)
- default.qubit.jax (PennyLane-0.30.0.dev0)
- default.qubit.tf (PennyLane-0.30.0.dev0)
- default.qubit.torch (PennyLane-0.30.0.dev0)
- default.qutrit (PennyLane-0.30.0.dev0)
- null.qubit (PennyLane-0.30.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.28.1)
- cirq.mixedsimulator (PennyLane-Cirq-0.28.0)
- cirq.pasqal (PennyLane-Cirq-0.28.0)
- cirq.qsim (PennyLane-Cirq-0.28.0)
- cirq.qsimh (PennyLane-Cirq-0.28.0)
- cirq.simulator (PennyLane-Cirq-0.28.0)
- qiskit.aer (PennyLane-qiskit-0.28.0)
- qiskit.basicaer (PennyLane-qiskit-0.28.0)
- qiskit.ibmq (PennyLane-qiskit-0.28.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.28.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.28.0)
```


### 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
@@ -169,7 +169,7 @@ def shape(self, device=None):
             return tuple(
                 (shot_val,) if shot_val != 1 else tuple() for shot_val in device._raw_shot_sequence
             )
-        len_wires = len(device.wires)
+        len_wires = len(self.wires) if len(self.wires) > 0 else len(device.wires)
         return (1, device.shots) if self.obs is not None else (1, device.shots, len_wires)
 
     def _shape_new(self, device=None):
@@ -178,19 +178,22 @@ def _shape_new(self, device=None):
                 "The device argument is required to obtain the shape of the measurement "
                 f"{self.__class__.__name__}."
             )
-        if device.shot_vector is not None:
-            if self.obs is None:
-                return tuple(
-                    (shot_val, len(device.wires)) if shot_val != 1 else (len(device.wires),)
-                    for shot_val in device._raw_shot_sequence
-                )
-            return tuple(
-                (shot_val,) if shot_val != 1 else tuple() for shot_val in device._raw_shot_sequence
-            )
-        if self.obs is None:
-            len_wires = len(device.wires)
-            return (device.shots, len_wires) if device.shots != 1 else (len_wires,)
-        return (device.shots,) if device.shots != 1 else ()
+        num_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
+            inner_shape = []
+            if shot_val != 1:
+                inner_shape.append(shot_val)
+            if num_wires != 1:
+                inner_shape.append(num_wires)
+            return tuple(inner_shape)
+
+        if device.shot_vector is None:
+            return _single_int_shape(device.shots, num_wires)
+        return tuple(
+            _single_int_shape(shot_val, num_wires) for shot_val in device._raw_shot_sequence
+        )
 
     def process_samples(
         self,


Test patch:
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -45,9 +45,9 @@ def custom_measurement_process(device, spy):
 class TestSample:
     """Tests for the sample function"""
 
-    def test_sample_dimension(self, mocker):
+    @pytest.mark.parametrize("n_sample", (1, 10))
+    def test_sample_dimension(self, mocker, n_sample):
         """Test that the sample function outputs samples of the right size"""
-        n_sample = 10
 
         dev = qml.device("default.qubit", wires=2, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
@@ -57,9 +57,11 @@ def circuit():
             qml.RX(0.54, wires=0)
             return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1))
 
-        sample = circuit()
+        output = circuit()
 
-        assert np.array_equal(sample.shape, (2, n_sample))
+        assert np.array_equal(output.shape, (2, n_sample))
+        assert circuit._qfunc_output[0].shape(dev) == (1, n_sample)
+        assert circuit._qfunc_output[1].shape(dev) == (1, n_sample)
 
         custom_measurement_process(dev, spy)
 
@@ -81,6 +83,7 @@ def circuit():
 
         assert len(result) == 3
         assert np.array_equal(result[0].shape, (n_sample,))
+        assert circuit._qfunc_output[0].shape(dev) == (1, n_sample)
         assert isinstance(result[1], np.ndarray)
         assert isinstance(result[2], np.ndarray)
 
@@ -96,13 +99,13 @@ def test_single_wire_sample(self, mocker):
         @qml.qnode(dev)
         def circuit():
             qml.RX(0.54, wires=0)
-
             return qml.sample(qml.PauliZ(0))
 
         result = circuit()
 
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result.shape, (n_sample,))
+        assert circuit._qfunc_output.shape(dev) == (1, n_sample)
 
         custom_measurement_process(dev, spy)
 
@@ -120,6 +123,10 @@ def circuit():
 
         result = circuit()
 
+        assert circuit._qfunc_output[0].shape(dev) == (1, n_sample)
+        assert circuit._qfunc_output[1].shape(dev) == (1, n_sample)
+        assert circuit._qfunc_output[2].shape(dev) == (1, n_sample)
+
         # If all the dimensions are equal the result will end up to be a proper rectangular array
         assert isinstance(result, np.ndarray)
         assert np.array_equal(result.shape, (3, n_sample))
@@ -319,6 +326,13 @@ def test_shape(self, obs):
         expected = (1, shots) if obs is not None else (1, shots, 3)
         assert res.shape(dev) == expected
 
+    @pytest.mark.parametrize("n_samples", (1, 10))
+    def test_shape_wires(self, n_samples):
+        """Test that the shape is correct when wires are provided."""
+        dev = qml.device("default.qubit", wires=3, shots=n_samples)
+        mp = qml.sample(wires=(0, 1))
+        assert mp.shape(dev) == (1, n_samples, 2)
+
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
@@ -331,8 +345,8 @@ def test_shape_shot_vector(self, obs):
         expected = ((), (2,), (3,))
         assert res.shape(dev) == expected
 
-    def test_shape_shot_vector_no_obs(self):
-        """Test that the shape is correct with the shot vector and no observable too."""
+    def test_shape_shot_vector_obs(self):
+        """Test that the shape is correct with the shot vector and a observable too."""
         shot_vec = (2, 2)
         dev = qml.device("default.qubit", wires=3, shots=shot_vec)
 
@@ -374,3 +388,27 @@ class DummyOp(Operator):
 
         with pytest.raises(EigvalsUndefinedError, match="Cannot compute samples of"):
             qml.sample(op=DummyOp(0)).process_samples(samples=np.array([[1, 0]]), wire_order=[0])
+
+
+@pytest.mark.jax
+@pytest.mark.parametrize("samples", (1, 10))
+def test_jitting_with_sampling_on_subset_of_wires(samples):
+    """Test case covering bug in Issue #3904.  Sampling should be jit-able
+    when sampling occurs on a subset of wires. The bug was occuring due an improperly
+    set shape method."""
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+
+    dev = qml.device("default.qubit", wires=3, shots=samples)
+
+    @qml.qnode(dev, interface="jax")
+    def circuit(x):
+        qml.RX(x, wires=0)
+        return qml.sample(wires=(0, 1))
+
+    results = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))
+
+    expected = (2,) if samples == 1 else (samples, 2)
+    assert results.shape == expected
+    assert circuit._qfunc_output.shape(dev) == (1, samples, 2)