class_docstr: Test that pulse_odegen is differentiable, i.e. that computing
  the derivative with pulse_odegen is differentiable a second time,
  yielding the Hessian.
class_docstr_tok: ["Test","that","pulse_odegen","is","differentiable",",","i.e",".","that","computing","the","derivative","with","pulse_odegen","is","differentiable","a","second","time",",","yielding","the","Hessian","."]
docstr: Test that pulse_odegen is differentiable,
  allowing to compute the Hessian, with JAX..
docstr_tok: ["Test","that","pulse_odegen","is","differentiable",",","allowing","to","compute","the","Hessian",",","with","JAX",".."]
code: def test_jax(self, dev_name):
      """Test that pulse_odegen is differentiable,
      allowing to compute the Hessian, 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 * Y(0)
  
      def fun(params):
          op = qml.evolve(ham_single_q_const)(params, T)
          tape = qml.tape.QuantumScript([op], [qml.expval(Z(0))])
          tape.trainable_params = [0]
          _tapes, fn = pulse_odegen(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 copy
  import pytest
  import numpy
  import pennylane
  from pennylane import numpy
  from pennylane.ops.qubit.special_unitary import pauli_basis_matrices, pauli_basis_strings
  from pennylane.math import expand_matrix
  from pennylane.gradients.pulse_gradient_odegen import _generate_tapes_and_coeffs, pulse_odegen, _insert_op, _nonzero_coeffs_and_words, _one_parameter_generators, _one_parameter_paulirot_coeffs, _parshift_and_contract
label_desc: Use image node_id 1 for calling the TestPulseOdegenDiff obj's underlying member method code with example usage: obj.test_jax(dev_name) and returns: fn