Update qchem demos to be JAX or JAX/JIT compatible #1211

austingmhuang · 2024-09-05T18:32:55Z

Title:
Update qchem demos to be JAX or JAX/JIT compatible

Summary:
As we phase out autograd, we need to ensure that our current qchem demos can work with either vanilla numpy or jax. Demos that are JIT compatible were made to use JIT as well. For demos that don't have any differentiation, we opt to use vanilla numpy.

Notes: the qubit tapering demo is not JIT compatible due to 2 issues: 1) Wire ordering is not ascending (Not a bug, but might need attention) 2) qml.taper_operation() uses the Exp operator, which has a conditional statement (PennyLaneAI/pennylane#5993)

Relevant references:

Possible Drawbacks:

Related GitHub Issues:
[sc-69776]

github-actions · 2024-09-05T18:33:13Z

👋 Hey, looks like you've updated some demos!

🐘 Don't forget to update the dateOfLastModification in the associated metadata files so your changes are reflected in Glass Onion (search and recommendations).

Please hide this comment once the field(s) are updated. Thanks!

github-actions · 2024-10-01T18:05:54Z

Thank you for opening this pull request.

You can find the built site at this link.

Deployment Info:

Pull Request ID: 1211
Deployment SHA: 8f96c77e6e9445e15b29cec51b9198cb5cf39e69
(The Deployment SHA refers to the latest commit hash the docs were built from)

Note: It may take several minutes for updates to this pull request to be reflected on the deployed site.

[sc-73935] * Remove deprecated code from demos. * Update various `from pennylane import numpy as np` to `import numpy as np`.

demonstrations/tutorial_adaptive_circuits.py

**Summary:** Fixes a merge conflict between `master` and `dev` introduced by #1232 (due to divergent `dateOfLastModification`). To reproduce this PR: 1. Run `git checkout dev`. 2. Run `git checkout -b merge-master-into-dev`. 3. Run `git merge master`. 4. Accept all incoming changes for merge conflicts on `dateOfLastModification`. **Relevant GHA Workflow Runs:** * https://github.com/PennyLaneAI/qml/actions/runs/11241666164 * https://github.com/PennyLaneAI/qml/actions/runs/11253088983 --------- Co-authored-by: David Wierichs <[email protected]> Co-authored-by: Korbinian Kottmann <[email protected]> Co-authored-by: Ivana Kurečić <[email protected]>

- Branched off dev - Merged Master in via `--allow-unrelated-histories` Co-authored-by: Jack Brown <[email protected]>

soranjh

Thanks @austingmhuang, left my final set of comments. They are mainly about being consistent in using qml.grad and jax.config.update.

demonstrations/tutorial_adaptive_circuits.py

demonstrations/tutorial_chemical_reactions.py

soranjh · 2024-10-23T15:08:52Z

demonstrations/tutorial_classically_boosted_vqe.py

+import optax
+import jax
+
+jax.config.update("jax_enable_x64", True) # use double-precision numbers


We should make sure that there is no warning about this in any of the demos.

soranjh · 2024-10-23T15:11:22Z

demonstrations/tutorial_eqnn_force_field.py


 import pennylane as qml
 import numpy as np

 import jax
-
-jax.config.update("jax_platform_name", "cpu")
-jax.config.update("jax_enable_x64", True)


Why we need two jax config updates in this demo, while only used one in the others?

soranjh · 2024-10-23T15:16:26Z

demonstrations/tutorial_qubit_tapering.py

@@ -169,7 +169,7 @@

 H_tapered = qml.taper(H, generators, paulixops, paulix_sector)
 H_tapered_coeffs, H_tapered_ops = H_tapered.terms()
-H_tapered = qml.Hamiltonian(np.real(H_tapered_coeffs), H_tapered_ops)
+H_tapered = qml.Hamiltonian(jnp.real(jnp.array(H_tapered_coeffs)), H_tapered_ops)


So, the initial jax array gets changed to something else during the workflow? I was expecting no type change.

soranjh · 2024-10-23T15:18:17Z

demonstrations/tutorial_qubit_tapering.py

+
+def update_step(i, params, opt_state):
+    """Perform a single gradient update step"""
+    grads = catalyst.grad(tapered_circuit)(params)


Any reason for not using qml.grad here?

demonstrations/tutorial_adaptive_circuits.py

demonstrations/tutorial_chemical_reactions.py

demonstrations/tutorial_classically_boosted_vqe.py

obliviateandsurrender · 2024-10-24T13:51:41Z

demonstrations/tutorial_eqnn_force_field.py

@@ -134,7 +134,7 @@

 import jax

-jax.config.update("jax_platform_name", "cpu")
+jax.config.update('jax_platform_name', 'cpu')


Do we really need this? Shouldn't by default this run on the CPU?

It's in the original demo... I could delete it or just leave it be.

demonstrations/tutorial_eqnn_force_field.py

demonstrations/tutorial_givens_rotations.py

obliviateandsurrender · 2024-10-24T13:55:46Z

demonstrations/tutorial_how_to_quantum_just_in_time_compile_vqe_catalyst.py

@@ -44,7 +44,7 @@
 #

 import pennylane as qml
-from pennylane import numpy as np
+import numpy as np


We don't need jax numpy here?

No, this part of the demo is not JIT compiled so we don't necessarily need JAX.

obliviateandsurrender · 2024-10-24T13:57:15Z

demonstrations/tutorial_initial_state_preparation.py

@@ -159,7 +157,7 @@
 # Let's take this opportunity to create the Hartree-Fock initial state, to compare the
 # other states against it later on.

-from pennylane import numpy as np
+import numpy as np

 hf_primer = ([[3, 0, 0]], np.array([1.0]))
 wf_hf = import_state(hf_primer)


Does import_state work with jnp?

When I did my checks it did work with jnp as well. But we don't use jnp here since vanilla np works just as fine (might be even faster with vanilla np, not sure though)

I just thought if we could use jnp instead we wouldn't need to import numpy 🤔 makes things more uniform (only if it is possible)

I think it's preferred to only import jnp when it's necessary e.g. the VQE portion where we start doing jax.grad

For the beginning of the demo, i think we should stick to vanilla numpy since initial state preparation itself doesn't require jnp; I personally find it a little weird to use jnp there. I can delete this additional import numpy as np though since numpy is imported at the start of the demo 😅

Steps used: * Checkout new branch from `dev` * Merge `master` using `git merge master` * Publish branch and open PR --------- Co-authored-by: Jack Brown <[email protected]> Co-authored-by: Ivana Kurečić <[email protected]> Co-authored-by: bellekaplan <[email protected]>

actions-user and others added 2 commits September 5, 2024 15:46

Merge branch 'master' into dev

1a9d3ae

reset to master

697b8f9

actions-user and others added 17 commits September 9, 2024 10:34

Merge branch 'master' into dev

2bfd14f

Merge branch 'master' into dev

838b668

Merge branch 'master' into dev

5d0ad9b

Merge branch 'master' into dev

e2ed3dd

Merge branch 'master' into dev

7ffeaf3

Merge branch 'master' into dev

f3831d0

Merge branch 'master' into dev

1d8bb22

Merge branch 'master' into dev

9994cd5

Merge branch 'master' into dev

31148a6

Merge branch 'master' into dev

fb0658d

Merge branch 'master' into dev

a2d55d8

Merge branch 'master' into dev

2adb533

Merge branch 'master' into dev

1059369

Merge branch 'master' into use_jax_qchem_demos

89d9bbc

fix to basisstate

58edd08

Merge branch 'master' into dev

f22b7a7

Merge branch 'master' into use_jax_qchem_demos

f7617d5

github-actions bot deployed to preview October 1, 2024 18:00 View deployment

mudit2812 and others added 2 commits October 3, 2024 21:06

Remove code being deprecated in v0.39 (#1220)

3d02a26

[sc-73935] * Remove deprecated code from demos. * Update various `from pennylane import numpy as np` to `import numpy as np`.

Merge branch 'master' into dev

ea59037

josh146 reviewed Oct 6, 2024

View reviewed changes

demonstrations/tutorial_adaptive_circuits.py Outdated Show resolved Hide resolved

actions-user and others added 5 commits October 7, 2024 10:28

Merge branch 'master' into dev

624a1b3

Merge branch 'master' into dev

ac333d4

Merge branch 'master' into dev

c1f66d4

Fix nightly merge failure (#1239)

e5033f8

- Branched off dev - Merged Master in via `--allow-unrelated-histories` Co-authored-by: Jack Brown <[email protected]>

bug

9a500df

github-actions bot deployed to preview October 22, 2024 20:36 View deployment

fixes to tutorials

ad6bd88

github-actions bot had a problem deploying to preview October 22, 2024 21:33 Failure

austingmhuang requested a review from soranjh October 23, 2024 14:46

soranjh requested changes Oct 23, 2024

View reviewed changes

delete compute decomp

6dc060e

github-actions bot deployed to preview October 23, 2024 17:05 View deployment

austingmhuang changed the base branch from master to dev October 23, 2024 21:16

Merge branch 'dev' into use_jax_qchem_demos

7e38f27

github-actions bot deployed to preview October 23, 2024 21:29 View deployment

obliviateandsurrender requested changes Oct 24, 2024

View reviewed changes

mudit2812 and others added 12 commits October 24, 2024 14:45

Merge branch 'dev' into use_jax_qchem_demos

4e28171

small changes to use jnp to silence warnings

b542b9d

author names

864f9c7

fix

eca10a3

revert fix

59b8dcd

is this ok? delete author file

e78725c

undo

a0cae6f

fix

c48296d

fix

339c209

fix author txt files

6b55a79

fix

cbfd07e

github-actions bot deployed to preview October 24, 2024 18:07 View deployment

delete unnecessary import

8ebe814

github-actions bot deployed to preview October 24, 2024 18:25 View deployment

seems reasonable

8f96c77

github-actions bot deployed to preview October 24, 2024 18:53 View deployment

mudit2812 changed the base branch from dev to master November 6, 2024 19:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update qchem demos to be JAX or JAX/JIT compatible #1211

Update qchem demos to be JAX or JAX/JIT compatible #1211

austingmhuang commented Sep 5, 2024

github-actions bot commented Sep 5, 2024

github-actions bot commented Oct 1, 2024 •

edited

Loading

soranjh left a comment

soranjh Oct 23, 2024

soranjh Oct 23, 2024

soranjh Oct 23, 2024

soranjh Oct 23, 2024

obliviateandsurrender Oct 24, 2024

austingmhuang Oct 24, 2024

obliviateandsurrender Oct 24, 2024

austingmhuang Oct 24, 2024

obliviateandsurrender Oct 24, 2024

austingmhuang Oct 24, 2024

obliviateandsurrender Oct 24, 2024

austingmhuang Oct 24, 2024

Update qchem demos to be JAX or JAX/JIT compatible #1211

Are you sure you want to change the base?

Update qchem demos to be JAX or JAX/JIT compatible #1211

Conversation

austingmhuang commented Sep 5, 2024

github-actions bot commented Sep 5, 2024

github-actions bot commented Oct 1, 2024 • edited Loading

soranjh 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

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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Oct 1, 2024 •

edited

Loading