[Bugfix] `qml.BasisRotation` is not jax compatible #6019

PabloAMC · 2024-07-22T12:12:21Z

Before submitting

Please complete the following checklist when submitting a PR:

All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to the
test directory!
All new functions and code must be clearly commented and documented.
If you do make documentation changes, make sure that the docs build and
render correctly by running make docs.
Ensure that the test suite passes, by running make test.
Add a new entry to the doc/releases/changelog-dev.md file, summarizing the
change, and including a link back to the PR.
The PennyLane source code conforms to
PEP8 standards.
We check all of our code against Pylint.
To lint modified files, simply pip install pylint, and then
run pylint pennylane/path/to/file.py.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

Context: qml.BasisRotation is was not jit-compatible.

Description of the Change: We modified all the numpy arrays to jax.numpy and ensure the tests were passing

Benefits: The basis rotation is jittable and thus qjit compatible.

Possible Drawbacks: jax numpy may be slower than basis numpy.

Related GitHub Issues: #6004

…//github.com/PennyLaneAI/pennylane into 6004-bug-basisrotation-is-not-jax-compatible

codecov · 2024-07-23T05:55:51Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.58%. Comparing base (d0344b0) to head (def0979).
Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6019      +/-   ##
==========================================
- Coverage   99.59%   99.58%   -0.01%     
==========================================
  Files         443      443              
  Lines       42255    42274      +19     
==========================================
+ Hits        42082    42097      +15     
- Misses        173      177       +4

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

…//github.com/PennyLaneAI/pennylane into 6004-bug-basisrotation-is-not-jax-compatible

doc/releases/changelog-dev.md

Co-authored-by: Utkarsh <[email protected]>

pennylane/qchem/givens_decomposition.py

josh146 · 2024-07-26T02:36:37Z

Thanks @obliviateandsurrender and @PabloAMC!

@obliviateandsurrender something odd I notice: using this PR, the following works:

dev = qml.device("lightning.qubit", wires=4)

@jax.jit
@qml.qnode(dev)
def f(U):
    for i in range(4):
        qml.Hadamard(i)
    qml.BasisRotation.compute_decomposition(unitary_matrix=U, wires=[0, 1, 2, 3], check=False)
    return qml.expval(qml.PauliZ(0))

from scipy.stats import unitary_group
U = jnp.array(unitary_group.rvs(4))
f(U)

But it no longer works if you remove the compute_decomposition (I still get an error message ValueError: Converting a JAX array to a NumPy array not supported when using the JAX JIT.).

obliviateandsurrender · 2024-07-26T03:36:40Z

Hey @josh146! I am not sure if it is a problem with the template cause it works with default.qubit.
I suspect the error is coming from the convert_to_numpy_parameters transform being used in the lightning.qubit for under _execute_jvp function. I think the error comes from interface being recognized as jax instead of jax-jit for lightning.qubit. Manually setting the interface like below works for me -

dev = qml.device("lightning.qubit", wires=4)

@jax.jit
@qml.qnode(dev, interface="jax-jit")
def f(U):
    for i in range(4):
        qml.Hadamard(i)
    qml.BasisRotation(unitary_matrix=U, wires=[0, 1, 2, 3], check=False)
    return qml.expval(qml.PauliZ(0))
U = jnp.array(unitary_group.rvs(4))

>>> f(U)
Array(-0.02033669, dtype=float64)

pennylane/qchem/givens_decomposition.py

Co-authored-by: Josh Izaac <[email protected]>

soranjh

Thanks @PabloAMC, left some comments. The failing test can be fixed by running black on givens_decomposition.py:

black -l 100 givens_decomposition.py

doc/releases/changelog-dev.md

soranjh · 2024-08-02T20:32:55Z

pennylane/qchem/givens_decomposition.py

@@ -147,43 +147,68 @@ def givens_decomposition(U):
                            # Update U = T(N+j-i-1, N+j-i) @ U

    """
+    interface = qml.math.get_deep_interface(unitary)
+    unitary = qml.math.copy(unitary) if interface == "jax" else qml.math.toarray(unitary).copy()


qml.math.copy is not general enough for all interfaces?

No, it is not. 😅

soranjh · 2024-08-02T20:38:17Z

pennylane/qchem/givens_decomposition.py

+    if interface == "jax":
+        for i in range(1, N):
+            if i % 2:
+                for j in range(0, i):
+                    indices = [i - j - 1, i - j]
+                    grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
+                    unitary = unitary.at[:, indices].set(unitary[:, indices] @ grot_mat.T, indices_are_sorted=True, unique_indices=True)
+                    right_givens.append((qml.math.conj(grot_mat), indices))
+            else:
+                for j in range(1, i + 1):
+                    indices = [N + j - i - 2, N + j - i - 1]
+                    grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
+                    unitary = unitary.at[indices, :].set(grot_mat @ unitary[indices, :], indices_are_sorted=True, unique_indices=True)
+                    left_givens.append((grot_mat, indices))
+    else:
+        for i in range(1, N):
+            if i % 2:
+                for j in range(0, i):
+                    indices = [i - j - 1, i - j]
+                    grot_mat = _givens_matrix(*unitary[N - j - 1, indices].T, left=True)
+                    unitary[:, indices] = unitary[:, indices] @ grot_mat.T
+                    right_givens.append((qml.math.conj(grot_mat), indices))
+            else:
+                for j in range(1, i + 1):
+                    indices = [N + j - i - 2, N + j - i - 1]
+                    grot_mat = _givens_matrix(*unitary[indices, j - 1], left=False)
+                    unitary[indices] = grot_mat @ unitary[indices]
+                    left_givens.append((grot_mat, indices))


This could be come more compact. There is a lot of repetition and the interface check can be just put before the lines that do not work with jax.

Also, is it possible to be general enough such that there is no need for checking the interface and writting interface-specific code blocks?

Also, is it possible to be general enough such that there is no need for checking the interface and writing interface-specific code blocks?

Nope. The best I could do was to separate out the interface checking to a separate function altogether. I was trying to prevent it since each if condition check might lead to a tiny overhead but at least this way codebase becomes more manageable and modular.

soranjh · 2024-08-02T20:39:39Z

pennylane/qchem/givens_decomposition.py

+        if interface == "jax":
+            unitary = unitary.at[i, i].set(qml.math.diag(nphase_mat)[0])
+            unitary = unitary.at[j, j].set(qml.math.diag(nphase_mat)[1])
+        else:
+            unitary[i, i], unitary[j, j] = qml.math.diag(nphase_mat)
+        nleft_givens.append((qml.math.conj(givens_mat), (i, j)))


Nothing in qml.math can make this block general enough to work with any interface?

.at and .set is unique to JAX array since those are immutable and do not support item assignments. I did try making it work via qml.math but couldn't.

So the closest thing I could find in qml.math was this, but it doesn't completely capture our purpose and might cause performance degradation due to jax array re-initializing. So, I have separated the index setting out to a separate private method.

soranjh · 2024-08-02T20:41:34Z

pennylane/templates/subroutines/basis_rotation.py

-            if not np.allclose(umat @ umat.conj().T, np.eye(M, dtype=complex), atol=1e-4):
-                raise ValueError("The provided transformation matrix should be unitary.")
+            umat = qml.math.copy(unitary_matrix)
+            if not qml.math.is_abstract(unitary_matrix):


Why this check is needed?

This prevents the below unitary check, which would be incompatible in a JIT-environment.

soranjh · 2024-08-02T20:45:13Z

tests/templates/test_subroutines/test_basis_rotation.py

+        assert jnp.allclose(res, res2, atol=tol, rtol=0)
+        assert qml.math.allclose(res, res3, atol=tol, rtol=0)


Is this for checking both jax and autograd at the same time? Seems to me that assert qml.math.allclose should be enough?

yes, qml.math.allclose should be sufficient for both. I use jnp.allclose in the first one as I know both res and res2 must be JAX-arrays.

doc/releases/changelog-dev.md

soranjh · 2024-08-02T20:54:51Z

pennylane/templates/subroutines/basis_rotation.py

Not a blocker for this PR, but it seems that BasisRotation is not included in the Templates doc here.

It is present 🤔

Co-authored-by: soranjh <[email protected]>

…//github.com/PennyLaneAI/pennylane into 6004-bug-basisrotation-is-not-jax-compatible

PabloAMC added 4 commits July 19, 2024 19:27

JAX compatible givens decompositions. Unit tests passing.

76abf63

Converting to jax.numpy the basis rotation

c17cf83

Update basis_rotation.py

cb663d3

Update changelog-dev.md

cb4d1f7

PabloAMC linked an issue Jul 22, 2024 that may be closed by this pull request

[BUG] BasisRotation is not JAX compatible #6004

Open

PabloAMC and others added 21 commits July 22, 2024 15:12

Update givens_decomposition.py

11a4b27

Reverting from jnp. to qml.math.

d811833

Remove import numpy

7edc75a

Unit test for jax jitting

d47d601

Removing whitespaces

ca95f7a

qml.math.copy and (not qml.math.is_abstract(unitary_matrix))

eb77772

Update basis_rotation.py

7d6c223

branching due to .at

38e3e68

Branching functions with/without jax backend

07816df

Update givens_decomposition.py

cd76203

Update test_basis_rotation.py

75e4f1e

Update givens_decomposition.py

8bfba22

Update givens_decomposition.py

34d4257

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

379b630

Update test_basis_rotation.py

96028a2

Update basis_rotation.py

bc34c6f

Merge branch '6004-bug-basisrotation-is-not-jax-compatible' of https:…

2ccc9ab

…//github.com/PennyLaneAI/pennylane into 6004-bug-basisrotation-is-not-jax-compatible

Update changelog-dev.md

57bfc8a

Update changelog-dev.md

650379c

Checks for abstract arrays

a563754

make jax-jit friendly

e33bf03

PabloAMC added 3 commits July 23, 2024 11:19

Update basis_rotation.py

43686b9

Merge branch '6004-bug-basisrotation-is-not-jax-compatible' of https:…

2b107c1

…//github.com/PennyLaneAI/pennylane into 6004-bug-basisrotation-is-not-jax-compatible

Update givens_decomposition.py

75edef1

obliviateandsurrender requested review from obliviateandsurrender and soranjh July 23, 2024 16:50

obliviateandsurrender reviewed Jul 23, 2024

View reviewed changes

doc/releases/changelog-dev.md Outdated Show resolved Hide resolved

PabloAMC and others added 2 commits July 24, 2024 17:20

Update doc/releases/changelog-dev.md

9010de1

Co-authored-by: Utkarsh <[email protected]>

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

5541c89

trbromley reviewed Jul 25, 2024

View reviewed changes

pennylane/qchem/givens_decomposition.py Outdated Show resolved Hide resolved

pennylane/qchem/givens_decomposition.py Outdated Show resolved Hide resolved

pennylane/qchem/givens_decomposition.py Outdated Show resolved Hide resolved

PabloAMC and others added 5 commits July 29, 2024 13:19

Updating docstrings

62c4c23

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

9fb4d26

attemp warning fix

8b4eef5

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

545be3e

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

eec5149

josh146 reviewed Jul 31, 2024

View reviewed changes

pennylane/qchem/givens_decomposition.py Outdated Show resolved Hide resolved

josh146 reviewed Jul 31, 2024

View reviewed changes

pennylane/qchem/givens_decomposition.py Outdated Show resolved Hide resolved

PabloAMC and others added 5 commits July 31, 2024 16:56

Update pennylane/qchem/givens_decomposition.py

4064fa6

Co-authored-by: Josh Izaac <[email protected]>

Update pennylane/qchem/givens_decomposition.py

be3df11

Co-authored-by: Josh Izaac <[email protected]>

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

a512ae3

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

7914f53

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

8ebfebe

soranjh requested changes Aug 2, 2024

View reviewed changes

PabloAMC and others added 5 commits August 3, 2024 00:14

Update givens_decomposition.py

7da2f88

Update changelog-dev.md

db61a73

Update doc/releases/changelog-dev.md

1c752a1

Co-authored-by: soranjh <[email protected]>

code cleanup

200481d

Merge branch '6004-bug-basisrotation-is-not-jax-compatible' of https:…

af011ac

…//github.com/PennyLaneAI/pennylane into 6004-bug-basisrotation-is-not-jax-compatible

obliviateandsurrender requested a review from soranjh August 5, 2024 17:06

Update qdrift.py

24d13c7

Alex-Preciado changed the title ~~6004 bug basisrotation is not jax compatible~~ [Bugfix] qml.BasisRotation is not jax compatible Aug 12, 2024

Merge branch 'master' into 6004-bug-basisrotation-is-not-jax-compatible

def0979

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Bugfix] `qml.BasisRotation` is not jax compatible #6019

[Bugfix] `qml.BasisRotation` is not jax compatible #6019

PabloAMC commented Jul 22, 2024 •

edited by obliviateandsurrender

Loading

codecov bot commented Jul 23, 2024 •

edited

Loading

josh146 commented Jul 26, 2024

obliviateandsurrender commented Jul 26, 2024 •

edited

Loading

soranjh left a comment

soranjh Aug 2, 2024

obliviateandsurrender Aug 2, 2024

soranjh Aug 2, 2024

obliviateandsurrender Aug 5, 2024

soranjh Aug 2, 2024

obliviateandsurrender Aug 2, 2024

obliviateandsurrender Aug 5, 2024

soranjh Aug 2, 2024

obliviateandsurrender Aug 2, 2024

soranjh Aug 2, 2024

obliviateandsurrender Aug 5, 2024

soranjh Aug 2, 2024

obliviateandsurrender Aug 5, 2024

		assert jnp.allclose(res, res2, atol=tol, rtol=0)
		assert qml.math.allclose(res, res3, atol=tol, rtol=0)

[Bugfix] qml.BasisRotation is not jax compatible #6019

Are you sure you want to change the base?

[Bugfix] qml.BasisRotation is not jax compatible #6019

Conversation

PabloAMC commented Jul 22, 2024 • edited by obliviateandsurrender Loading

Before submitting

codecov bot commented Jul 23, 2024 • edited Loading

Codecov Report

josh146 commented Jul 26, 2024

obliviateandsurrender commented Jul 26, 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

Choose a reason for hiding this comment

[Bugfix] `qml.BasisRotation` is not jax compatible #6019

[Bugfix] `qml.BasisRotation` is not jax compatible #6019

PabloAMC commented Jul 22, 2024 •

edited by obliviateandsurrender

Loading

codecov bot commented Jul 23, 2024 •

edited

Loading

obliviateandsurrender commented Jul 26, 2024 •

edited

Loading