Add get samples function to InstructionToSignals for JAX-jit usage #149

to24toro · 2022-11-11T05:50:42Z

Summary

This PR adds get_samples to InstructionToSignals for JAX-jitting when using qiskit-pulse and removes the usage of get_waveform method of SymbolicPulse.

Details and comments

get_samples function gets the envelope expression form SymbolicPulse and calls sympy.lambdify with numerical backend specified by Array class. The lambdified function is lru cached for performance.

CLAassistant · 2022-11-11T05:50:46Z

All committers have signed the CLA.

DanPuzzuoli · 2022-11-14T18:09:47Z

Thanks for the PR @to24toro !

I think the code you've added to Dynamics looks good - however there is a core problem right now that prevents this from being useful. Your test:

def jit_func(amp, sigma):
    return get_samples(Gaussian(duration=5, amp=amp, sigma=sigma))

jit_samples = jax.jit(jit_func, static_argnums=(0, 1))(0.983, 2)
self.assertAllClose(jit_samples, self.gauss_get_waveform_samples, atol=1e-7, rtol=1e-7)

sets static_argnums=(0, 1), however, for this to be useful, this test should work without usingstatic_argnums. I.e. you should just be able to do:

jit_samples = jax.jit(jit_func)(0.983, 2)

If you try to do this now with the current main branch of terra, you get a JAX error, coming from the line:

...
File /opt/anaconda3/envs/devEnv310/lib/python3.10/site-packages/qiskit/pulse/library/symbolic_pulses.py:436, in SymbolicPulse.__init__(self, pulse_type, duration, parameters, name, limit_amplitude, envelope, constraints, valid_amp_conditions)
    431 # TODO remove this.
    432 #  This is due to convention in IBM Quantum backends where "amp" is treated as a
    433 #  special parameter that must be defined in the form [real, imaginary].
    434 #  this check must be removed because Qiskit pulse should be backend agnostic.
    435 if "amp" in parameters and not isinstance(parameters["amp"], ParameterExpression):
--> 436     parameters["amp"] = complex(parameters["amp"])
...
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float64[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `complex` function. If trying to convert the data type of a value, try using `x.astype(complex)` or `jnp.array(x, complex)` instead.
The error occurred while tracing the function jit_func at /var/folders/v5/5t1xdchn2ws5l1nj9h02vm0m0000gn/T/ipykernel_3597/1345434633.py:1 for jit. This concrete value was not available in Python because it depends on the value of the argument 'amp'.

which is because in this case, parameters["amp"] is a JAX tracer, and calling complex(x) for x a JAX tracer type will cause this ConcretizationTypeError.

With this error being raised, there is unfortunately no benefit to having get_samples execute using JAX - the only benefit of JAX is being able to jit/grad/etc. If we can't truly jit a computation, then there is no reason to use JAX over numpy. Obviously, since this error is being raised by terra, it will be necessary to make changes to terra before this can work.

To move forward with this PR, I think it makes sense to figure out whatever changes are necessary in terra to make the above test case work with the static_argnums call kwarg removed, make those changes, then merge this PR to dynamics once those are released in terra. @nkanazawa1989 any comments on this thought process?

nkanazawa1989 · 2022-11-14T18:20:48Z

The special casing of amp must be removed from terra (only duration and amp are the non-float and unfortunately Python is not typed). Currently Qiskit/qiskit#9002 is trying to introduce (amp, float) representation to symbolic pulses, and this makes all parameters float-type except for duration. This approach allows us to remove builtin typecasting and we can eventually remove static_arguments.

Feel free to on hold this until above PR is merged, or you can merge this PR without JIT test and later add the the test in a separate (follow up) PR.

DanPuzzuoli · 2022-11-14T18:54:25Z

Okay cool thanks.

Feel free to on hold this until above PR is merged, or you can merge this PR without JIT test and later add the the test in a separate (follow up) PR.

I think it makes sense to hold on this PR, as I don't think the required functionality can truly be verified until terra is at a point where the static_argnums can be removed.

nkanazawa1989 · 2022-11-14T18:56:47Z

Fair enough. Let's merge Qiskit/qiskit#9002 first.

to24toro · 2022-11-25T09:05:33Z

test/dynamics/signals/test_pulse_to_signals.py

+        def jit_func(amp):
+            return get_samples(Constant(100, amp))
+
+        jit_samples = jax.jit(jit_func)(0.1)


We expected this function will pass without using static_argnums after merging Qiskit/qiskit#9002 .

DanPuzzuoli · 2022-12-01T16:37:03Z

Hey @to24toro , there are two more tests I think it'd be good to add:

Can you add a test case for computing gradients as well? You can take the same jit test, but change it so that at the end, instead of jit, you do jax.jit(jax.grad(jit_func))(0.1). You'll have to change the the function so that it returns a real scalar (as grad only works on real scalars), so you could just change what the function returns to be get_samples(instance).sum().real. Testing this combination of gradding/jitting is important because it can introduce a bit more complexity into what JAX is doing, and sometimes raise errors where jit/grad don't raise them on their own.
I'm also thinking we should add a test for Solver that compiles/differentiates a simulation of a pulse schedule with symbolic pulses. In the test/dynamics/solvers/test_solver_classes.py file there is a test class called TestPulseSimulationJAX. Can you add a test to this class that defines a function building a schedule with symbolic pulses, simulating it with a Solver, and returning the final state, and attempts to compile/differentiate it? This may be a little redundant with the other tests but it will be a good integration test.

One last point: the TestJaxBase class has two helper functions: jit_wrap and jit_grad_wrap, which can be used to jit a function, or to jit(grad(func)) a function func, where the latter takes the sum and the real part of the output. These may eb helpful/convenient.

to24toro · 2022-12-02T07:31:49Z

test/dynamics/signals/test_pulse_to_signals.py

+        self.jit_wrap(jit_func)(0.1)
+        self.jit_grad_wrap(jit_func)(0.1)


I tried to use jit_wrap and jit_grad_wrap.
Very useful.

to24toro · 2023-01-15T10:59:23Z

qiskit_dynamics/solvers/solver_classes.py

@@ -523,6 +524,8 @@ def solve(
            Array.default_backend() == "jax"
            and (method == "jax_odeint" or _is_diffrax_method(method))
            and all(isinstance(x, Schedule) for x in signals_list)
+            # check if jit transformation is already performed.
+            and not (isinstance(jnp.array(0), core.Tracer))


This is related with the issue #175

to24toro · 2023-01-16T12:20:23Z

I modified two points at 769e342 for JAX-jitting.

DanPuzzuoli · 2023-02-01T17:16:42Z

I've removed the "on hold" label, as you pointed out that terra has been sufficiently updated. Once the errors are resolved and I re-review we can merge this!

DanPuzzuoli

Aside from current test failing issue (which seems to just be a sympy issue), I have some comments about reorganizing the new tests a little.

Looks good though.

qiskit_dynamics/pulse/pulse_to_signals.py

DanPuzzuoli · 2023-02-01T17:40:52Z

test/dynamics/pulse/test_pulse_to_signals.py

        self.assertTrue(signals[0].carrier_freq == 5.0)
        self.assertTrue(signals[1].carrier_freq == 3.1)
        self.assertTrue(signals[2].carrier_freq == 0.0)
        self.assertTrue(signals[3].carrier_freq == 4.0)

+    def test_get_samples(self):


Rewrite these tests to still call the InstructionToSignals converter. get_samples and lru_cache_expr are both internal functions, so ideally we won't directly test them. There are exceptions to this in Dynamics, but in this case I don't think it's necessary. The inputs/outputs of these functions are directly fed from/to the converter, so I don't think anything is gained by directly testing these functions, as opposed to verifying the behaviour of the converter on a symbolic pulse.

Maybe change the name of this to test_SymbolicPulse.

improved and changed the name at e2c5148

DanPuzzuoli · 2023-02-01T17:42:58Z

test/dynamics/pulse/test_pulse_to_signals.py

+        )
+
+
+class TestJaxGetSamples(QiskitDynamicsTestCase, TestJaxBase):


Similar to the preceding comment, I don't think it's necessary to have a separate class for checking the behaviour of get_samples. These tests could be moved to the previous class, by rewriting them to verify that the converter returns the same thing, whether you pass in a Waveform or a SymbolicPulse.

The jit test is great and very important, but again, it can be moved to the previous class and rewritten to work with the converter.

removed get samples test at 72cd9ad

For this I mean to fully get rid of this test class, and to move the cases it's testing into the preceding test class. by calling InstructionToSignals directly.

The class you said moving to is TestPulseToSignals?
If I move to the class, we fail to pass the test because TestPulseToSignals is not a subclass of TestJaxBase.
On the other hand, if TestPulseToSignals inherits TestJaxBase, all the test in TestPulseToSignals will be skipped in python test and executed in only JAX test.

I understand now. Maybe just rename this class then to TestPulseToSignalsJAXTransformations or something along these lines.

rename and fix at 454f119 and 2814cf0

test/dynamics/solvers/test_solver_classes.py

Co-authored-by: Daniel Puzzuoli <dan.puzzuoli@gmail.com>

to24toro · 2023-02-02T09:54:15Z

My author and commit name were not correct. So I am sorry to have modify them and force-push to pass license/cla.

DanPuzzuoli

There are some conflicts to resolve, and some final tests calling get_samples to convert to calls directly to InstructionToSignals.

DanPuzzuoli · 2023-02-07T17:33:54Z

qiskit_dynamics/pulse/pulse_to_signals.py

-
-    See the :meth:`get_signals` method documentation for a detailed description of how pulse
-    schedules are interpreted and translated into :class:`.DiscreteSignal` objects.
+    schedule. The converter applies to instances of :class:`Schedule`. Instances of


I think something's happened here - this is undoing some changes I had made in a previous PR (that I don't think you are intentionally trying to change).

I am so sorry to confuse you. As you say, it is not my intention. I restored to that of main.

DanPuzzuoli · 2023-02-07T17:34:24Z

qiskit_dynamics/pulse/pulse_to_signals.py

-            dt: Length of the samples. This is required by the converter as pulse schedule are
-                specified in units of dt and typically do not carry the value of dt with them.
-            carriers: A dict of analog carrier frequencies. The keys are the names of the channels
+            dt: Length of the samples. This is required by the converter as pulse


These two code blocks are similarly changing some documentation I had changed previously. Maybe some issue with a previous attempt at merging main into this branch?

DanPuzzuoli · 2023-02-07T17:36:00Z

qiskit_dynamics/pulse/pulse_to_signals.py

-                :class:`~qiskit.pulse.ScheduleBlock` must first be converted to
-                :class:`~qiskit.pulse.Schedule` using the
-                :func:`~qiskit.pulse.transforms.block_to_schedule` function in Qiskit Pulse.
+                :class:`ScheduleBlock` must first be converted to :class:`Schedule` using the


Similarly here, this is undoing some changes I had made.

DanPuzzuoli · 2023-02-07T17:37:56Z

test/dynamics/pulse/test_pulse_to_signals.py

+        )
+
+
+class TestJaxGetSamples(QiskitDynamicsTestCase, TestJaxBase):


For this I mean to fully get rid of this test class, and to move the cases it's testing into the preceding test class. by calling InstructionToSignals directly.

DanPuzzuoli · 2023-02-07T17:38:32Z

test/dynamics/pulse/test_pulse_to_signals.py

+                envelope=envelope_expr,
+                valid_amp_conditions=valid_amp_conditions_expr,
+            )
+            return get_samples(instance)


E.g. this test should use InstructionToSignals instead of get_samples directly.

DanPuzzuoli · 2023-02-07T17:38:49Z

test/dynamics/pulse/test_pulse_to_signals.py

+        jit_samples = jax.jit(jit_func_get_samples)(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):


This is a great test, that could be moved to the preceding class.

…sformations

DanPuzzuoli

Looks good! One final documentation comment.

DanPuzzuoli · 2023-02-15T18:20:46Z

qiskit_dynamics/pulse/pulse_to_signals.py

-
    The converter can be initialized with the optional arguments ``carriers`` and ``channels``. When
    ``channels`` is given, only the signals specified by name in ``channels`` are returned. The
    ``carriers`` dictionary specifies the analog carrier frequency of each channel. Here, the keys
    are the channel name, e.g. ``d12`` for drive channel number ``12``, and the values are the
    corresponding frequency. If a channel is not present in ``carriers`` it is assumed that the
    analog carrier frequency is zero.
-


My last question is: did you intentionally remove these spaces? Deleting these causes this whole block of text to be a single paragraph in the docs. I prefer these as separate paragraphs, as each is highlighting something different. Unless there is a specific reason for this, can you put the spaces back in.

Sorry. I didn't take into account the paragraph change. I returned back at c3338c5.

to24toro commented Nov 25, 2022

View reviewed changes

to24toro force-pushed the jax_implementation_to_pulse branch from ad0e3aa to b82c152 Compare December 2, 2022 07:29

to24toro commented Dec 2, 2022

View reviewed changes

DanPuzzuoli added the On hold label Dec 6, 2022

to24toro force-pushed the jax_implementation_to_pulse branch from b82c152 to f36cb66 Compare December 10, 2022 03:53

DanPuzzuoli mentioned this pull request Jan 6, 2023

Solver.solve auto-jitting behaviour for pulse schedules breaks when JAX transformations are called from outside #175

Closed

to24toro changed the title ~~Add get samples function to InstructionToSignals~~ Add get samples function to InstructionToSignals for JAX-jit usage Jan 15, 2023

to24toro commented Jan 15, 2023

View reviewed changes

DanPuzzuoli removed the On hold label Feb 1, 2023

DanPuzzuoli requested changes Feb 1, 2023

View reviewed changes

to24toro force-pushed the jax_implementation_to_pulse branch from 3f214c5 to 03592c1 Compare February 2, 2023 09:04

to24toro and others added 11 commits February 2, 2023 18:36

add sampling pulse for jax jit

b5e2fad

add a release note

7e2c039

comment change of tolerance

9c87d08

change test_jit_get_samples

8adbc79

add jax_grad test

6ad5815

Fix pulse to signal converter (qiskit-community#164)

c309427

Co-authored-by: Daniel Puzzuoli <dan.puzzuoli@gmail.com>

wrap signal samples as Array

5713e13

format

9a3acaa

modify test_pulse_to_signals

3c6a78a

add test_jit_solve_with_internal_jit

ccb9bdc

impl jax-jit for pulse simulation

42d4b8c

to24toro added 6 commits February 2, 2023 18:48

format by black

a1637e0

add type hint

ef502d3

add test_SymbolicPulse

a73d315

modify test_SymbolicPulse

32ad385

modify test_SymbolicPulse

8839c78

add test_pulse_types_combination_with_jax

bf66412

to24toro force-pushed the jax_implementation_to_pulse branch from 03592c1 to bf66412 Compare February 2, 2023 09:50

remove redefined fucntion

9bc4972

DanPuzzuoli reviewed Feb 7, 2023

View reviewed changes

to24toro added 2 commits February 8, 2023 10:49

restore previous comments

39a1e90

add removed sentence

aa10fbe

DanPuzzuoli added this to the pulse sim v1 milestone Feb 9, 2023

to24toro and others added 4 commits February 15, 2023 08:00

rename class name from TestJaxGetSamples to TestPulseToSignalsJAXTran…

454f119

…sformations

fix codes about TestPulseToSignalsJAXTransformations

2814cf0

Merge branch 'main' into jax_implementation_to_pulse

4044ebd

modify importing and naming of pulse function

f8b055e

DanPuzzuoli requested changes Feb 15, 2023

View reviewed changes

to24toro and others added 3 commits February 16, 2023 06:10

insert spaces in InstructionToSignals docs

c3338c5

Merge branch 'main' into jax_implementation_to_pulse

5a5d2ac

Merge branch 'main' into jax_implementation_to_pulse

c56c9ca

DanPuzzuoli self-requested a review February 15, 2023 22:16

DanPuzzuoli approved these changes Feb 15, 2023

View reviewed changes

DanPuzzuoli merged commit 85128b7 into qiskit-community:main Feb 16, 2023

to24toro mentioned this pull request Feb 27, 2023

Userguide and tutorial JAX transformation of Pulse Scheduler #194

Merged

wshanks mentioned this pull request Jun 10, 2023

sympy missing from dependencies #230

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add get samples function to InstructionToSignals for JAX-jit usage #149

Add get samples function to InstructionToSignals for JAX-jit usage #149

to24toro commented Nov 11, 2022 •

edited

Loading

CLAassistant commented Nov 11, 2022 •

edited

Loading

DanPuzzuoli commented Nov 14, 2022 •

edited

Loading

nkanazawa1989 commented Nov 14, 2022

DanPuzzuoli commented Nov 14, 2022

nkanazawa1989 commented Nov 14, 2022

to24toro Nov 25, 2022 •

edited

Loading

DanPuzzuoli commented Dec 1, 2022

to24toro Dec 2, 2022 •

edited

Loading

to24toro Jan 15, 2023

to24toro commented Jan 16, 2023

DanPuzzuoli commented Feb 1, 2023

DanPuzzuoli left a comment

DanPuzzuoli Feb 1, 2023

to24toro Feb 2, 2023

DanPuzzuoli Feb 1, 2023

to24toro Feb 2, 2023 •

edited

Loading

DanPuzzuoli Feb 7, 2023

to24toro Feb 8, 2023

DanPuzzuoli Feb 14, 2023

to24toro Feb 14, 2023

to24toro commented Feb 2, 2023 •

edited

Loading

DanPuzzuoli left a comment

DanPuzzuoli Feb 7, 2023

to24toro Feb 8, 2023 •

edited

Loading

DanPuzzuoli Feb 7, 2023

DanPuzzuoli Feb 7, 2023

DanPuzzuoli Feb 7, 2023

DanPuzzuoli Feb 7, 2023

DanPuzzuoli Feb 7, 2023

DanPuzzuoli left a comment

DanPuzzuoli Feb 15, 2023

to24toro Feb 15, 2023

		self.jit_wrap(jit_func)(0.1)
		self.jit_grad_wrap(jit_func)(0.1)

		)


		class TestJaxGetSamples(QiskitDynamicsTestCase, TestJaxBase):

Add get samples function to InstructionToSignals for JAX-jit usage #149

Add get samples function to InstructionToSignals for JAX-jit usage #149

Conversation

to24toro commented Nov 11, 2022 • edited Loading

Summary

Details and comments

CLAassistant commented Nov 11, 2022 • edited Loading

DanPuzzuoli commented Nov 14, 2022 • edited Loading

nkanazawa1989 commented Nov 14, 2022

DanPuzzuoli commented Nov 14, 2022

nkanazawa1989 commented Nov 14, 2022

to24toro Nov 25, 2022 • edited Loading

Choose a reason for hiding this comment

DanPuzzuoli commented Dec 1, 2022

to24toro Dec 2, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro commented Jan 16, 2023

DanPuzzuoli commented Feb 1, 2023

DanPuzzuoli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro Feb 2, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro commented Feb 2, 2023 • edited Loading

DanPuzzuoli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro Feb 8, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DanPuzzuoli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro commented Nov 11, 2022 •

edited

Loading

CLAassistant commented Nov 11, 2022 •

edited

Loading

DanPuzzuoli commented Nov 14, 2022 •

edited

Loading

to24toro Nov 25, 2022 •

edited

Loading

to24toro Dec 2, 2022 •

edited

Loading

to24toro Feb 2, 2023 •

edited

Loading

to24toro commented Feb 2, 2023 •

edited

Loading

to24toro Feb 8, 2023 •

edited

Loading