class_docstr: Test that stoch_pulse_grad is differentiable.
class_docstr_tok: ["Test","that","stoch_pulse_grad","is","differentiable","."]
docstr: Test that stoch_pulse_grad is differentiable with JAX.
docstr_tok: ["Test","that","stoch_pulse_grad","is","differentiable","with","JAX","."]
code: def test_jax(self, dev_name):
      """Test that stoch_pulse_grad is differentiable with JAX."""
      import jax
      import jax.numpy as jnp
  
      jax.config.update("jax_enable_x64", True)
      dev = qml.device(dev_name, wires=1)
      T = 0.5
      ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
  
      def fun(params):
          op = qml.evolve(ham_single_q_const)(params, T)
          tape = qml.tape.QuantumScript(
              [op], [qml.expval(qml.PauliZ(0))]
          )
          tape.trainable_params = [0]
          tapes, fn = stoch_pulse_grad(tape)
          return fn(qml.execute(tapes, dev, "backprop"))
  
      params = [jnp.array(0.4)]
      p = params[0] * T
      exp_grad = -2 * jnp.sin(2 * p) * T
      grad = fun(params)
      assert qml.math.isclose(grad, exp_grad)
  
      exp_diff_of_grad = -4 * jnp.cos(2 * p) * T**2
      diff_of_grad = jax.grad(fun)(params)
      assert qml.math.isclose(diff_of_grad, exp_diff_of_grad)
  
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 1 for calling the TestStochPulseGradDiff obj's underlying member method code with example usage: obj.test_jax(dev_name) and returns: fn