Merge remote-tracking branch 'Dan/fix-symbolic-pulse-jax' into fix/pu…

…lse_to_signal

.github/workflows/docs.yml

@@ -15,7 +15,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: '3.8'
+        python-version: '3.10'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

.github/workflows/release.yml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/setup-python@v4
         name: Install Python
         with:
-          python-version: '3.8'
+          python-version: '3.10'
       - name: Install Deps
         run: pip install -U wheel
       - name: Build Artifacts

test/dynamics/pulse/test_pulse_to_signals.py

@@ -25,10 +25,15 @@
 
 from qiskit_ibm_runtime.fake_provider import FakeQuito
 
+try:
+    from jax import jit
+except ImportError:
+    pass
+
 from qiskit_dynamics.pulse import InstructionToSignals
 from qiskit_dynamics.signals import DiscreteSignal
 
-from ..common import QiskitDynamicsTestCase, TestJaxBase
+from ..common import QiskitDynamicsTestCase, JAXTestBase
 
 
 class TestPulseToSignals(QiskitDynamicsTestCase):
@@ -358,7 +363,7 @@ def test_barrier_instructions(self):
         self.assertAllClose(sigs[1].samples, np.array([0.0, 0.0, 0.0, -0.5, -0.5, -0.5]))
 
 
-class TestPulseToSignalsJAXTransformations(QiskitDynamicsTestCase, TestJaxBase):
+class TestPulseToSignalsJAXTransformations(JAXTestBase):
     """Tests InstructionToSignals class by using Jax."""
 
     def setUp(self):
@@ -393,9 +398,9 @@ def jit_func_instruction_to_signals(amp):
             converter = InstructionToSignals(self._dt, carriers={"d0": 5})
             return converter.get_signals(schedule)[0].samples
 
-        self.jit_wrap(jit_func_instruction_to_signals)(0.1)
-        self.jit_grad_wrap(jit_func_instruction_to_signals)(0.1)
-        jit_samples = self.jit_wrap(jit_func_instruction_to_signals)(0.1)
+        jit(jit_func_instruction_to_signals)(0.1)
+        self.jit_grad(jit_func_instruction_to_signals)(0.1)
+        jit_samples = jit(jit_func_instruction_to_signals)(0.1)
         self.assertAllClose(jit_samples, self.constant_get_waveform_samples, atol=1e-7, rtol=1e-7)
 
     def test_pulse_types_combination_with_jax(self):
@@ -430,8 +435,8 @@ def jit_func_symbolic_pulse(amp):
             converter = InstructionToSignals(self._dt, carriers={"d0": 5})
             return converter.get_signals(schedule)[0].samples
 
-        self.jit_wrap(jit_func_symbolic_pulse)(0.1)
-        self.jit_grad_wrap(jit_func_symbolic_pulse)(0.1)
+        jit(jit_func_symbolic_pulse)(0.1)
+        self.jit_grad(jit_func_symbolic_pulse)(0.1)
 
 
 @ddt