[Frontend] Enable Jax dynamic API and use `abstracted_axes` in qjit #366

erick-xanadu · 2023-11-16T17:48:07Z

Context: Enable Jax Dynamic API and add abstracted_axes top-level argument

Description of the Change:

Enable dynamic_shapes option in JAX.
Automatically deduce implicit parameters when calling the compiled function
Remove implicit parameters from wrapped function arguments.

Benefits:

Jax dynamic API allows to compile tensors with variable dimension sizes.
jax.numpy.ones/zeros/empy etc could now be used in the compiled programs.

Drawbacks:

Accessing dynamically-shaped tensors inside Catalyst control-flow primitives is not yet supported
Abstracted axes of more than 3 dimensions appear to not work on some jax.lax functions.

Related GitHub Issues:

This PR contains a workaround for the following Jax issue

Do not pass AbstractValues to get_aval. google/jax#18579

Also related: Please do not remove

Avoid querying named_shape to support DShapedArrays. google/jax#18505

[sc-47630]
[sc-47631]

frontend/catalyst/compilation_pipelines.py

Co-authored-by: Josh Izaac <josh146@gmail.com>

frontend/test/pytest/test_jax_dynamic_api.py

dime10

🚀

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

frontend/test/pytest/test_jax_dynamic_api.py

frontend/catalyst/utils/jax_extras.py

josh146

documentation looks great @erick-xanadu @grwlf!

doc/changelog.md

frontend/catalyst/compilation_pipelines.py

Co-authored-by: Josh Izaac <josh146@gmail.com>

erick-xanadu added 20 commits November 13, 2023 11:27

Set jax_dynamic_shapes to True.

e17bc4c

Add abstracted_axes to compiled options.

fa14d02

Avoid recomputing runtime_signature.

002a4ed

Adds very basic test.

6ad7d5a

Call make_jaxpr with abstracted_axes

2a6247c

Pass arguments from Python to compiled function.

a7a2d42

Changed deduce_avals.

612356f

Get aval from from tracer.

918a6fd

Pass correct arguments to eval_jaxpr.

9af0d48

Multiple.

397d612

Fixed named_shape in the caller.

9697443

Outline function to get implicit and explicit flat args.

acf400b

Remove print statements from test.

d861aa7

Move get_implicit_and_explicit_flat_args to jax_extras.

c09178e

Move imports to top-level.

5686702

Remove unused variables.

2cf2af8

Remove dead code.

82ca008

Remove unused variables.

2b23715

Use function.

22af452

Remove comments in deduce_avals.

7dc5ca2

erick-xanadu force-pushed the eochoa/2023-11-16/abstracted-axes branch 3 times, most recently from bda393e to bd41005 Compare November 17, 2023 20:27

Style.

1737f9d

erick-xanadu force-pushed the eochoa/2023-11-16/abstracted-axes branch from bd41005 to 1737f9d Compare November 17, 2023 20:29

Adds documentation.

89c6bbb

erick-xanadu changed the title ~~Use abstracted_axes in qjit~~ [Frontend] Use abstracted_axes in qjit Nov 17, 2023

erick-xanadu and others added 3 commits November 20, 2023 16:16

Can perform computation.

afafce0

Do not return implicit returns.

6c383d9

Do not use shaped types in quantum tracing

09a244b

josh146 reviewed Nov 29, 2023

View reviewed changes

frontend/catalyst/compilation_pipelines.py Outdated Show resolved Hide resolved

josh146 reviewed Nov 29, 2023

View reviewed changes

frontend/catalyst/compilation_pipelines.py Outdated Show resolved Hide resolved

josh146 reviewed Nov 29, 2023

View reviewed changes

frontend/catalyst/compilation_pipelines.py Outdated Show resolved Hide resolved

josh146 reviewed Nov 29, 2023

View reviewed changes

frontend/catalyst/compilation_pipelines.py Outdated Show resolved Hide resolved

erick-xanadu and others added 9 commits November 29, 2023 09:10

Apply suggestions from code review

ba57ebc

Co-authored-by: Josh Izaac <josh146@gmail.com>

Style

87c50f5

Pylint

ba9e6cb

Add test to avoid recompilation when using abstracted_axes.

bb32710

Adds documentation for recompilation.

bd46085

Special case.

d9f798c

Patching JAX.

f492e8c

Import order.

956afac

pylint.

4f8da0c

dime10 reviewed Nov 30, 2023

View reviewed changes

dime10 approved these changes Nov 30, 2023

View reviewed changes

erick-xanadu and others added 3 commits December 1, 2023 09:08

Apply suggestions from code review

06233e7

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

Merge branch 'main' into eochoa/2023-12-1/abstracted-axes

4d4985b

black

9bf23aa

erick-xanadu commented Dec 1, 2023

View reviewed changes

frontend/test/pytest/test_jax_dynamic_api.py Show resolved Hide resolved

erick-xanadu commented Dec 1, 2023

View reviewed changes

frontend/catalyst/utils/jax_extras.py Show resolved Hide resolved

erick-xanadu and others added 5 commits December 1, 2023 14:09

Small test

f3fe778

Address review suggestions: comment for the shape-conversion test

d3df4de

Fix test.

54a0229

Black.

e83553d

Code factor.

f65eebf

erick-xanadu requested a review from josh146 December 4, 2023 16:06

josh146 approved these changes Dec 5, 2023

View reviewed changes

doc/changelog.md Outdated Show resolved Hide resolved

doc/changelog.md Outdated Show resolved Hide resolved

frontend/catalyst/compilation_pipelines.py Show resolved Hide resolved

Apply suggestions from code review

23bd286

Co-authored-by: Josh Izaac <josh146@gmail.com>

erick-xanadu merged commit b0f9a3f into main Dec 5, 2023
19 checks passed

erick-xanadu deleted the eochoa/2023-11-16/abstracted-axes branch December 5, 2023 13:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Frontend] Enable Jax dynamic API and use `abstracted_axes` in qjit #366

[Frontend] Enable Jax dynamic API and use `abstracted_axes` in qjit #366

erick-xanadu commented Nov 16, 2023 •

edited

Loading

dime10 left a comment

josh146 left a comment

[Frontend] Enable Jax dynamic API and use abstracted_axes in qjit #366

[Frontend] Enable Jax dynamic API and use abstracted_axes in qjit #366

Conversation

erick-xanadu commented Nov 16, 2023 • edited Loading

dime10 left a comment

Choose a reason for hiding this comment

josh146 left a comment

Choose a reason for hiding this comment

[Frontend] Enable Jax dynamic API and use `abstracted_axes` in qjit #366

[Frontend] Enable Jax dynamic API and use `abstracted_axes` in qjit #366

erick-xanadu commented Nov 16, 2023 •

edited

Loading