Repository: PennyLaneAI/pennylane
Pull request #3966 (2023-04-04 19:13:47)

Rydberg Drive
Added `rydberg_drive` function to create drive Hamiltonians for neutral atom pulse programming.

Questions/notes for reviewers:
* Currently, the pulses for the `Hardwarehamiltonian` returned by `drive` (#3936) is overwritten by `rydberg_drive`. Should the `HardwarePulse` object creation be removed from `drive` and only be kept in `rydberg_drive`, and eventually `transmon_drive`?

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3966?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3966](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3966?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1548f4d) into [unify-transmon](https://codecov.io/gh/PennyLaneAI/pennylane/commit/17767d290f473ce8cd386c8a1f011a292de27a09?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (17767d2) will **decrease** coverage by `0.07%`.
> The diff coverage is `16.00%`.

```diff
@@                Coverage Diff                 @@
##           unify-transmon    #3966      +/-   ##
==================================================
- Coverage           99.73%   99.66%   -0.07%     
==================================================
  Files                 346      346              
  Lines               30246    30268      +22     
==================================================
+ Hits                30167    30168       +1     
- Misses                 79      100      +21     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3966?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/pulse/rydberg.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3966?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3B1bHNlL3J5ZGJlcmcucHk=) | `61.81% <8.69%> (-38.19%)` | :arrow_down: |
| [pennylane/pulse/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3966?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3B1bHNlL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

mudit2812: >I think the decision on the drive PR was that the HardwarePulse object on the drive function should be removed in this PR, and I think that's still missing and would be nice to add.

It's been removed 😄 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e9dd0bb21..fafd8d3d4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -24,10 +24,13 @@
     the Hamiltonian that describes the circuit QED interaction Hamiltonian of superconducting transmon systems.
   * A new user-facing `drive` function is added, which returns a `ParametrizedHamiltonian` (`HardwareHamiltonian`) containing
     the Hamiltonian of the interaction between a driving electro-magnetic field and a group of qubits.
+  * A new user-facing `rydberg_drive` function is added, which returns a `ParametrizedHamiltonian` (`HardwareHamiltonian`) containing
+    the Hamiltonian of the interaction between a driving laser field and a group of Rydberg atoms.
   [(#3749)](https://github.com/PennyLaneAI/pennylane/pull/3749)
   [(#3911)](https://github.com/PennyLaneAI/pennylane/pull/3911)
   [(#3930)](https://github.com/PennyLaneAI/pennylane/pull/3930)
   [(#3936)](https://github.com/PennyLaneAI/pennylane/pull/3936/)
+  [(#3966)](https://github.com/PennyLaneAI/pennylane/pull/3966)
 
 * Added `Operation.__truediv__` dunder method to be able to divide operators.
   [(#3749)](https://github.com/PennyLaneAI/pennylane/pull/3749)
diff --git a/pennylane/pulse/__init__.py b/pennylane/pulse/__init__.py
index 34829d6c4..383c39f55 100644
--- a/pennylane/pulse/__init__.py
+++ b/pennylane/pulse/__init__.py
@@ -69,6 +69,7 @@ Hardware Compatible Hamiltonians
 
     ~rydberg_interaction
     ~transmon_interaction
+    ~rydberg_drive
     ~drive
 
 
@@ -260,6 +261,6 @@ JIT-compiling is optional, and one can remove the decorator when only single exe
 from .convenience_functions import constant, pwc, pwc_from_function, rect
 from .parametrized_evolution import ParametrizedEvolution
 from .parametrized_hamiltonian import ParametrizedHamiltonian
-from .hardware_hamiltonian import HardwareHamiltonian, drive
-from .rydberg import rydberg_interaction
+from .hardware_hamiltonian import HardwareHamiltonian, HardwarePulse, drive
+from .rydberg import rydberg_interaction, rydberg_drive
 from .transmon import transmon_interaction
diff --git a/pennylane/pulse/hardware_hamiltonian.py b/pennylane/pulse/hardware_hamiltonian.py
index 164506594..cfb70f557 100644
--- a/pennylane/pulse/hardware_hamiltonian.py
+++ b/pennylane/pulse/hardware_hamiltonian.py
@@ -213,13 +213,7 @@ def drive(amplitude, phase, wires):
         >>> jax.grad(circuit)(params)
         [Array(1.75825695, dtype=float64)]
     """
-    if isinstance(wires, int):
-        wires = [wires]
-
-    if not callable(amplitude) and qml.math.isclose(amplitude, 0.0):
-        raise ValueError(
-            f"Expected non-zero value or callable for amplitude, but received amplitude={amplitude}. All terms are zero."
-        )
+    wires = Wires(wires)
 
     # TODO: use sigma+ and sigma- (not necessary as terms are the same, but for consistency)
     # We compute the `coeffs` and `observables` of the EM field
@@ -244,13 +238,12 @@ class HardwareHamiltonian(ParametrizedHamiltonian):
     but on top of that also contains attributes that store parameteres relevant for real hardware execution.
 
     .. warning::
+
         This class should NEVER be initialized directly! Please use the functions
         :func:`rydberg_interaction` and :func:`drive` instead.
 
-
     .. seealso:: :func:`rydberg_interaction`, :func:`drive`, :class:`ParametrizedHamiltonian`
 
-
     Args:
         coeffs (Union[float, callable]): coefficients of the Hamiltonian expression, which may be
             constants or parametrized functions. All functions passed as ``coeffs`` must have two
@@ -258,7 +251,6 @@ class HardwareHamiltonian(ParametrizedHamiltonian):
         observables (Iterable[Observable]): observables in the Hamiltonian expression, of same
             length as ``coeffs``
 
-
     Keyword Args:
         settings Union[RydbergSettings, TransmonSettings]: Dataclass containing the hardware specific settings. Default is ``None``.
         pulses (list[HardwarePulse]): list of ``HardwarePulse`` dataclasses containing the information about the
@@ -348,7 +340,6 @@ class HardwareHamiltonian(ParametrizedHamiltonian):
 class HardwarePulse:
     """Dataclass that contains the information of a single drive pulse. This class is used
     internally in PL to group into a single object all the data related to a single EM field.
-
     Args:
         amplitude (Union[float, Callable]): float or callable returning the amplitude of an EM
             field
@@ -420,7 +411,6 @@ class AmplitudeAndPhase:
 def _reorder_parameters(params, coeffs_parametrized):
     """Takes `params`, and reorganizes it based on whether the Hamiltonian has
     callable phase and/or callable amplitude.
-
     Consolidates phase and amplitude parameters in the case that both are callable,
     and duplicates phase and/or amplitude parameters if either are callables, since
     they will be passed to two operators in the Hamiltonian"""
diff --git a/pennylane/pulse/rydberg.py b/pennylane/pulse/rydberg.py
index 8cf691953..fafc54da3 100644
--- a/pennylane/pulse/rydberg.py
+++ b/pennylane/pulse/rydberg.py
@@ -17,7 +17,8 @@ from dataclasses import dataclass
 import numpy as np
 import pennylane as qml
 
-from pennylane.pulse import HardwareHamiltonian
+from pennylane.pulse import HardwareHamiltonian, HardwarePulse, drive
+from pennylane.wires import Wires
 
 
 def rydberg_interaction(
@@ -114,6 +115,152 @@ def rydberg_interaction(
     return HardwareHamiltonian(coeffs, observables, settings=settings)
 
 
+def rydberg_drive(amplitude, phase, detuning, wires):
+    r"""Returns a :class:`ParametrizedHamiltonian` representing the action of a driving laser field
+
+    .. math::
+
+        \frac{1}{2} \Omega(t) \sum_{q \in \text{wires}} (\cos(\phi(t))\sigma_q^x - \sin(\phi(t))\sigma_q^y) -
+        \delta(t) \sum_{q \in \text{wires}} \sigma_q^z
+
+    where :math:`\Omega`, :math:`\phi` and :math:`\delta` correspond to the amplitude, phase,
+    and detuning of the laser, :math:`i` correspond to the wire index, and :math:`\sigma^\alpha` for
+    :math:`\alpha = x,y,z` are the Pauli matrices. For hardware execution, time is expected to be in units
+    of :math:`\text{µs}`, and the frequency in units of :math:`\text{MHz}`. It is recommended to also follow
+    this convention for simulation, as it avoids numerical problems due to using very large and very small
+    numbers. This driving term can be combined with an interaction term to create a Hamiltonian describing a
+    driven Rydberg atom system. Multiple driving terms can be combined by summing them (see example).
+
+    Args:
+        amplitude (Union[float, Callable]): float or callable returning the amplitude (in MHz) of a
+            laser field
+        phase (Union[float, Callable]): float or callable returning the phase (in radians) of the laser field
+        detuning (Union[float, Callable]): float or callable returning the detuning (in MHz) of a
+            laser field
+        wires (Union[int, List[int]]): integer or list containing wire values for the Rydberg atoms that
+            the laser field acts on
+
+    Returns:
+        ParametrizedHamiltonian: a :class:`~.ParametrizedHamiltonian` representing the action of the laser field on the Rydberg atoms.
+
+    .. seealso::
+
+        :func:`~.rydberg_interaction`, :class:`~.ParametrizedHamiltonian`, :class:`~.ParametrizedEvolution`
+        and :func:`~.evolve`
+
+    **Example**
+
+    We create a Hamiltonian describing a laser acting on 4 wires (Rydberg atoms) with a fixed detuning and
+    phase, and a parametrized, time-dependent amplitude. The Hamiltonian includes an interaction term for
+    inter-atom interactions due to van der Waals forces, as well as the driving term for the laser driving
+    the atoms:
+
+    .. code-block:: python
+
+        atom_coordinates = [[0, 0], [0, 4], [4, 0], [4, 4]]
+        wires = [0, 1, 2, 3]
+        H_i = qml.pulse.rydberg_interaction(atom_coordinates, wires)
+
+        amplitude = lambda p, t: p * jnp.sin(jnp.pi * t)
+        phase = jnp.pi / 2
+        detuning = 3 * jnp.pi / 4
+        H_d = qml.pulse.rydberg_drive(amplitude, phase, detuning, wires)
+
+    >>> H_i
+    ParametrizedHamiltonian: terms=6
+    >>> H_d
+    ParametrizedHamiltonian: terms=3
+
+    The first two terms of the drive Hamiltonian ``H_d`` correspond to the first sum (the sine and cosine terms),
+    describing drive between the ground and excited states. The third term corresponds to the shift term
+    due to detuning from resonance. This drive term corresponds to a global drive that acts on all 4 wires of
+    the device.
+
+    The full Hamiltonian evolution and expectation value measurement can be executed in a ``QNode``:
+
+    .. code-block:: python
+
+        dev = qml.device("default.qubit.jax", wires=wires)
+        @qml.qnode(dev, interface="jax")
+        def circuit(params):
+            qml.evolve(H_i + H_d)(params, t=[0, 10])
+            return qml.expval(qml.PauliZ(0))
+
+    >>> params = [2.4]
+    >>> circuit(params)
+    Array(0.94301294, dtype=float64)
+    >>> jax.grad(circuit)(params)
+    [Array(0.59484969, dtype=float64)]
+
+    We can also create a Hamiltonian with local drives. The following circuit corresponds to the
+    evolution where additional local drives acting on wires ``0`` and ``1`` respectively are added to the
+    Hamiltonian:
+
+    .. code-block:: python
+
+        amplitude_local_0 = lambda p, t: p[0] * jnp.sin(2 * jnp.pi * t) + p[1]
+        phase_local_0 = jnp.pi / 4
+        detuning_local_0 = lambda p, t: p * jnp.exp(-0.25 * t)
+        H_local_0 = qml.pulse.rydberg_drive(amplitude_local_0, phase_local_0, detuning_local_0, [0])
+
+        amplitude_local_1 = lambda p, t: jnp.cos(jnp.pi * t) + p
+        phase_local_1 = jnp.pi
+        detuning_local_1 = lambda p, t: jnp.sin(jnp.pi * t) + p
+        H_local_1 = qml.pulse.rydberg_drive(amplitude_local_1, phase_local_1, detuning_local_1, [1])
+
+        H = H_i + H_d + H_local_0 + H_local_1
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def circuit_local(params):
+            qml.evolve(H)(params, t=[0, 10])
+            return qml.expval(qml.PauliZ(0))
+
+    >>> p_global = 2.4
+    >>> p_local_amp_0 = [1.3, -2.0]
+    >>> p_local_det_0 = -1.5
+    >>> p_local_amp_1 = -0.9
+    >>> p_local_det_1 = 3.1
+    >>> params = [p_global, p_local_amp_0, p_local_det_0, p_local_amp_1, p_local_det_1]
+    >>> circuit_local(params)
+    Array(0.6464017, dtype=float64)
+    >>> jax.grad(circuit_local)(params)
+    [Array(-1.47216703, dtype=float64),
+     [Array(0.25855072, dtype=float64, weak_type=True),
+      Array(0.1453378, dtype=float64, weak_type=True)],
+     Array(0.18329746, dtype=float64),
+     Array(0.05901987, dtype=float64),
+     Array(-0.23426886, dtype=float64)]
+    """
+    wires = Wires(wires)
+    trivial_detuning = not callable(detuning) and qml.math.isclose(detuning, 0.0)
+
+    if not callable(amplitude) and qml.math.isclose(amplitude, 0.0):
+        if trivial_detuning:
+            raise ValueError(
+                "Expected non-zero value for at least one of either amplitude or detuning, but "
+                f"received amplitude={amplitude} and detuning={detuning}. All terms are zero."
+            )
+
+        amplitude_term = HardwareHamiltonian([], [])
+
+    else:
+        amplitude_term = drive(amplitude, phase, wires)
+
+    detuning_obs, detuning_coeffs = [], []
+    if not trivial_detuning:
+        detuning_obs.append(-1.0 * sum(qml.PauliZ(wire) for wire in wires))
+        detuning_coeffs.append(detuning)
+
+    detuning_term = HardwareHamiltonian(detuning_coeffs, detuning_obs)
+    pulses = [HardwarePulse(amplitude, phase, detuning, wires)]
+
+    drive_term = amplitude_term + detuning_term
+    drive_term.pulses = pulses
+
+    return drive_term
+
+
 @dataclass
 class RydbergSettings:
     """Dataclass that contains the information of a Rydberg setup.


Test patch:
diff --git a/tests/pulse/test_hardware_hamiltonian.py b/tests/pulse/test_hardware_hamiltonian.py
index b87583d0f..f9b100cef 100644
--- a/tests/pulse/test_hardware_hamiltonian.py
+++ b/tests/pulse/test_hardware_hamiltonian.py
@@ -330,11 +330,6 @@ class TestDrive:
         # Hamiltonian is as expected
         assert qml.equal(Hd([0.5, -0.5], t=5), H_expected([0.5, -0.5], t=5))
 
-    def test_no_amplitude_no_detuning(self):
-        """Test that the correct error is raised if both amplitude and detuning are trivial."""
-        with pytest.raises(ValueError, match="Expected non-zero value or callable"):
-            _ = drive(0, np.pi, wires=[0])
-
 
 def callable_amp(p, t):
     return np.polyval(p, t)
diff --git a/tests/pulse/test_rydberg.py b/tests/pulse/test_rydberg.py
index b53933bde..36996500b 100644
--- a/tests/pulse/test_rydberg.py
+++ b/tests/pulse/test_rydberg.py
@@ -19,7 +19,12 @@ import numpy as np
 import pytest
 
 import pennylane as qml
-from pennylane.pulse import HardwareHamiltonian, rydberg_interaction
+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
@@ -82,6 +87,125 @@ class TestRydbergInteraction:
         assert qml.equal(H_res([], t=5), H_exp([], t=5))
 
 
+class TestRydbergDrive:
+    """Unit tests for the ``rydberg_drive`` function"""
+
+    def test_attributes_and_number_of_terms(self):
+        """Test that the attributes and the number of terms of the ``ParametrizedHamiltonian`` returned by
+        ``rydberg_drive`` are correct."""
+
+        Hd = rydberg_drive(amplitude=1, phase=2, detuning=3, wires=[1, 2])
+
+        assert isinstance(Hd, HardwareHamiltonian)
+        assert Hd.settings is None
+        assert Hd.wires == Wires([1, 2])
+        assert len(Hd.ops) == 3  # 2 amplitude/phase terms and one detuning term
+        assert Hd.pulses == [HardwarePulse(1, 2, 3, [1, 2])]
+
+    def test_multiple_local_drives(self):
+        """Test that adding multiple drive terms behaves as expected"""
+
+        def fa(p, t):
+            return np.sin(p * t)
+
+        def fb(p, t):
+            return np.cos(p * t)
+
+        H1 = rydberg_drive(amplitude=fa, phase=1, detuning=3, wires=[0, 3])
+        H2 = rydberg_drive(amplitude=1, phase=3, detuning=fb, wires=[1, 2])
+        Hd = H1 + H2
+
+        ops_expected = [
+            qml.Hamiltonian([-1, -1], [qml.PauliZ(0), qml.PauliZ(3)]),
+            qml.Hamiltonian([0.5, 0.5], [qml.PauliX(1), qml.PauliX(2)]),
+            qml.Hamiltonian([-0.5, -0.5], [qml.PauliY(1), qml.PauliY(2)]),
+            qml.Hamiltonian([0.5, 0.5], [qml.PauliX(0), qml.PauliX(3)]),
+            qml.Hamiltonian([-0.5, -0.5], [qml.PauliY(0), qml.PauliY(3)]),
+            qml.Hamiltonian([-1, -1], [qml.PauliZ(1), qml.PauliZ(2)]),
+        ]
+        coeffs_expected = [
+            3,
+            np.cos(3),
+            np.sin(3),
+            AmplitudeAndPhase(np.cos, fa, 1),
+            AmplitudeAndPhase(np.sin, fa, 1),
+            fb,
+        ]
+        H_expected = HardwareHamiltonian(coeffs_expected, ops_expected)
+
+        # structure of Hamiltonian is as expected
+        assert isinstance(Hd, HardwareHamiltonian)
+        assert Hd.wires == Wires([0, 3, 1, 2])
+        assert Hd.settings is None
+        assert len(Hd.ops) == 6  # 2 terms for amplitude/phase and one detuning for each drive
+
+        # coefficients are correct
+        # Callable coefficients are shifted to the end of the list.
+        assert Hd.coeffs[0:3] == [3, np.cos(3), np.sin(3)]
+        assert isinstance(Hd.coeffs[3], AmplitudeAndPhase)
+        assert isinstance(Hd.coeffs[4], AmplitudeAndPhase)
+        assert Hd.coeffs[5] is fb
+
+        # pulses were added correctly
+        assert len(Hd.pulses) == 2
+        assert Hd.pulses == H1.pulses + H2.pulses
+
+        # Hamiltonian is as expected
+        assert qml.equal(Hd([0.5, -0.5], t=5), H_expected([0.5, -0.5], t=5))
+
+    def test_no_amplitude(self):
+        """Test that when amplitude is not specified, the drive term is correctly defined."""
+
+        def f(p, t):
+            return np.cos(p * t)
+
+        Hd = rydberg_drive(amplitude=0, phase=1, detuning=f, wires=[0, 3])
+
+        ops_expected = [qml.Hamiltonian([-1, -1], [qml.PauliZ(0), qml.PauliZ(3)])]
+        coeffs_expected = [f]
+        H_expected = HardwareHamiltonian(coeffs_expected, ops_expected)
+
+        assert qml.equal(Hd([0.1], 10), H_expected([0.1], 10))
+        assert isinstance(Hd, HardwareHamiltonian)
+        assert Hd.wires == Wires([0, 3])
+        assert Hd.settings is None
+        assert len(Hd.coeffs) == 1
+        assert Hd.coeffs[0] is f
+        assert len(Hd.ops) == 1
+        assert qml.equal(Hd.ops[0], ops_expected[0])
+
+    def test_no_detuning(self):
+        """Test that when detuning not specified, the drive term is correctly defined."""
+
+        def f(p, t):
+            return np.cos(p * t)
+
+        Hd = rydberg_drive(amplitude=f, phase=1, detuning=0, wires=[0, 3])
+
+        ops_expected = [
+            qml.Hamiltonian([0.5, 0.5], [qml.PauliX(0), qml.PauliX(3)]),
+            qml.Hamiltonian([-0.5, -0.5], [qml.PauliY(0), qml.PauliY(3)]),
+        ]
+        coeffs_expected = [
+            AmplitudeAndPhase(np.cos, f, 1),
+            AmplitudeAndPhase(np.sin, f, 1),
+        ]
+        H_expected = HardwareHamiltonian(coeffs_expected, ops_expected)
+
+        assert qml.equal(Hd([0.1], 10), H_expected([0.1], 10))
+        assert isinstance(Hd, HardwareHamiltonian)
+        assert Hd.wires == Wires([0, 3])
+        assert Hd.settings is None
+        assert all(isinstance(coeff, AmplitudeAndPhase) for coeff in Hd.coeffs)
+        assert len(Hd.coeffs) == 2
+        assert all(qml.equal(op, op_expected) for op, op_expected in zip(Hd.ops, ops_expected))
+
+    def test_no_amplitude_no_detuning(self):
+        """Test that the correct error is raised if both amplitude and detuning are trivial."""
+        with pytest.raises(ValueError, match="Expected non-zero value for at least one of either"):
+            _ = rydberg_drive(0, np.pi, 0, wires=[0])
+
+
 # For rydberg settings test
 register0 = [[0.0, 1.0], [0.0, 2.0]]
 register1 = [[2.0, 0.3], [1.0, 4.0], [0.5, 0.4]]
@@ -129,3 +253,142 @@ class TestRydbergSettings:
         settings1 = RydbergSettings(register1)
         with pytest.raises(ValueError, match="Cannot add two"):
             res = settings0 + settings1
+
+
+class TestIntegration:
+    """Integration tests for Rydberg system Hamiltonians."""
+
+    @pytest.mark.jax
+    def test_jitted_qnode(self):
+        """Test that a Rydberg ensemble can be simulated within a jitted qnode."""
+        import jax
+        import jax.numpy as jnp
+
+        Hd = 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)
+
+        Ht = rydberg_drive(amplitude=fa, phase=0, detuning=fb, 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(Hd + Ht)(params, ts)
+            return qml.expval(H_obj)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def qnode_jit(params):
+            qml.evolve(Hd + Ht)(params, ts)
+            return qml.expval(H_obj)
+
+        params = (jnp.ones(5), jnp.array([1.0, jnp.pi]))
+        res = qnode(params)
+        res_jit = qnode_jit(params)
+
+        assert isinstance(res, jax.Array)
+        assert res == res_jit
+
+    @pytest.mark.jax
+    def test_jitted_qnode_multidrive(self):
+        """Test that a Rydberg ensemble with multiple drive terms can be
+        executed within a jitted qnode."""
+        import jax
+        import jax.numpy as jnp
+
+        Hd = 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)
+
+        def fd(p, t):
+            return p * jnp.cos(t)
+
+        H1 = rydberg_drive(amplitude=fa, phase=0, detuning=fb, wires=wires)
+        H2 = rydberg_drive(amplitude=fc, phase=3 * jnp.pi, detuning=0, wires=4)
+        H3 = rydberg_drive(amplitude=0, phase=0, detuning=fd, wires=[3, 0])
+
+        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(Hd + H1 + H2 + H3)(params, ts)
+            return qml.expval(H_obj)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def qnode_jit(params):
+            qml.evolve(Hd + H1 + H2 + H3)(params, ts)
+            return qml.expval(H_obj)
+
+        params = (
+            jnp.ones(5),
+            jnp.array([1.0, jnp.pi]),
+            jnp.array([jnp.pi / 2, 0.5]),
+            jnp.array(-0.5),
+        )
+        res = qnode(params)
+        res_jit = qnode_jit(params)
+
+        assert isinstance(res, jax.Array)
+        assert res == res_jit
+
+    @pytest.mark.jax
+    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 res == res_jit