class_docstr: Test working cases of stoch_pulse_grad.
class_docstr_tok: ["Test","working","cases","of","stoch_pulse_grad","."]
docstr: Test that the stochastic parameter-shift rule works with JITting.
docstr_tok: ["Test","that","the","stochastic","parameter-shift","rule","works","with","JITting","."]
code: def test_with_jit(
      self, dev_name, generator, exp_num_tapes, prefactor
  ):
      """Test that the stochastic parameter-shift rule works with JITting."""
      import jax
      import jax.numpy as jnp
  
      jax.config.update("jax_enable_x64", True)
      dev = qml.device(dev_name, wires=len(generator.wires))
      T = (0.2, 0.5)
      ham_single_q_const = qml.dot([qml.pulse.constant], [generator])
      meas = [qml.expval(qml.PauliZ(0))]
  
      def fun(params):
          """Create a pulse with the given parameters, build a tape from it, and
          differentiate it with stoch_pulse_grad."""
          op = qml.evolve(ham_single_q_const)(params, T)
          tape = qml.tape.QuantumScript([op], meas)
          tapes, fn = stoch_pulse_grad(tape)
          assert len(tapes) == exp_num_tapes
          res = fn(qml.execute(tapes, dev, "backprop"))
          return res
  
      params = [jnp.array(0.24)]
      # Effective rotation parameter
      p = params[0] * (T[1] - T[0]) * prefactor
      res = fun(params)
      assert qml.math.isclose(
          res, -2 * jnp.sin(2 * p) * (T[1] - T[0]) * prefactor
      )
      res_jit = jax.jit(fun)(params)
      assert qml.math.isclose(res, res_jit)
      jax.clear_caches()
  
imports: import warnings
  import copy
  import pytest
  import numpy
  import pennylane
  from pennylane.gradients.general_shift_rules import eigvals_to_frequencies, generate_shift_rule
  from pennylane.gradients.pulse_gradient import _parshift_and_integrate, _split_evol_ops, _split_evol_tape, stoch_pulse_grad
label_desc: Use image node_id 14 for calling the TestStochPulseGrad obj's underlying member method code with example usage: obj.test_with_jit(dev_name, generator, exp_num_tapes, prefactor) and returns: res