class_docstr: Integration tests for Rydberg system Hamiltonians.
docstr: Test that a Rydberg ensemble can be simulated within a
  jitted qnode when all coeffs are callable.
docstr_tok: ["Test","that","a","Rydberg","ensemble","can","be","simulated","within","a","jitted","qnode","when","all","coeffs","are","callable","."]
code: def test_jitted_qnode_all_coeffs_callable(self):
      """Test that a Rydberg ensemble can be simulated within a
      jitted qnode when all coeffs are callable."""
      import jax
      import jax.numpy as jnp
  
      H_drift = rydberg_interaction(
          register=atom_coordinates, wires=wires
      )
  
      def fa(p, t):
          return jnp.polyval(p, t)
  
      def fb(p, t):
          return p[0] * jnp.sin(p[1] * t)
  
      def fc(p, t):
          return p[0] * jnp.sin(t) + jnp.cos(p[1] * t)
  
      H_drive = rydberg_drive(
          amplitude=fa, phase=fb, detuning=fc, wires=1
      )
  
      dev = qml.device("default.qubit", wires=wires)
  
      ts = jnp.array([0.0, 3.0])
      H_obj = sum(qml.PauliZ(i) for i in range(2))
  
      @qml.qnode(dev, interface="jax")
      def qnode(params):
          qml.evolve(H_drift + H_drive)(params, ts)
          return qml.expval(H_obj)
  
      @jax.jit
      @qml.qnode(dev, interface="jax")
      def qnode_jit(params):
          qml.evolve(H_drift + H_drive)(params, ts)
          return qml.expval(H_obj)
  
      params = (
          jnp.ones(5),
          jnp.array([1.0, jnp.pi]),
          jnp.array([jnp.pi / 2, 0.5]),
      )
      res = qnode(params)
      res_jit = qnode_jit(params)
  
      assert isinstance(res, jax.Array)
      assert np.allclose(res, res_jit)
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane.pulse import rydberg_interaction, rydberg_drive
  from pennylane.pulse.hardware_hamiltonian import HardwareHamiltonian, HardwarePulse, AmplitudeAndPhase
  from pennylane.wires import Wires
  from pennylane.pulse.rydberg import RydbergSettings
label_desc: Use image node_id 3 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_jitted_qnode_all_coeffs_callable() and returns: jnp, unknown, unknown, qml, qml