class_docstr: Test integration of pwc functions with the pulse module.
class_docstr_tok: ["Test","integration","of","pwc","functions","with","the","pulse","module","."]
docstr: Test that the evolution of a parametrized hamiltonian defined with a pwc function can executed on
  a QNode using jax-jit, and the results don't differ from execution without jitting.
docstr_tok: ["Test","that","the","evolution","of","a","parametrized","hamiltonian","defined","with","a","pwc","function","can","executed","on","a","QNode","using","jax-jit",",","and","the","results","do","n't","differ","from","execution","without","jitting","."]
code: def test_qnode_pwc_jit(self):
      """Test that the evolution of a parametrized hamiltonian defined with a pwc function can executed on
      a QNode using jax-jit, and the results don't differ from execution without jitting.
      """
      import jax
  
      f1 = qml.pulse.pwc((1, 6))
      f2 = qml.pulse.pwc((0.5, 3))
      H = f1 * qml.PauliX(0) + f2 * qml.PauliY(1)
  
      t = (0, 4)
  
      dev = qml.device("default.qubit", wires=2)
  
      @qml.qnode(dev, interface="jax")
      def circuit(params):
          qml.evolve(H)(params=params, t=t)
          return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
  
      @jax.jit
      @qml.qnode(dev, interface="jax")
      def jitted_circuit(params):
          qml.evolve(H)(params=params, t=t)
          return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
  
      constants = np.linspace(0, 9, 10)
      params = [constants, constants]
  
      assert qml.math.allclose(
          jitted_circuit(params), circuit(params), atol=5e-3
      )
      assert qml.math.allclose(
          jax.grad(jitted_circuit)(params),
          jax.grad(circuit)(params),
          atol=5e-3,
      )
  
imports: import inspect
  from functools import partial
  import numpy
  import pytest
  import pennylane
  from pennylane.pulse.parametrized_hamiltonian import ParametrizedHamiltonian
label_desc: Use image node_id 4 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_qnode_pwc_jit() and returns: qml, qml