code: 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.legacy", 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, Shots(samples))
          == (samples, 2)
          if samples != 1
          else (2,)
      )
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane.measurements import MeasurementShapeError, Sample, Shots, MeasurementValue
  from pennylane.operation import Operator
label_desc: Use image node_id 2 for calling a global function with example usage: test_jitting_with_sampling_on_subset_of_wires(samples) and returns: qml