Raising a warning if latest JAX is installed

[PietropaoloFrisoni]

Context: We need a bugfix release for PennyLane 0.41.

The reason is that if a user installs PennyLane, then latext JAX, and then imports PennyLane (in this sequence), the following error is raised:

AttributeError: jax.core.Primitive was removed in JAX v0.6.0. Use jax.extend.core.Primitive instead, and see https://docs.jax.dev/en/latest/jax.extend.html for details..

and there are no further details.
Therefore, we need to raise an informative warning at some point that tells the user that PL is not yet compatible with the latest JAX version.

Notice that this branch was created from the tag of the latest PennyLane release, as the PL release guidance recommends.

Description of the Change:

• This is not a change implemented in this PR, but notice that __version__ is set to "0.41.1" on this branch. This change was implemented before pushing this branch to GitHub, as the PL release guidance recommends.

• We need to install the stable version of pennylane-lightning and pennylane-catalyst instead of the latest, because the latest version of these packages is now incompatible with this branch (they are using the pennylane.exceptions module, which was introduced after the 0.41 release). Obviously, this change will not be merged into master.

• The actual warning message. This triggered pylint failures in the entire file, and I had to stick:

# pylint: disable=wrong-import-position, wrong-import-order, import-outside-toplevel

at the very top of the file. This is certainly not ideal, but this is only temporary before PL is compatible with the latest JAX version (WIP to be completed before the next release).
Notice that we might also raise this warning somewhere else, catching the first import jax inside the capture folder that causes the issue, instead of the init file in PennyLane. There was a brief discussion about this with @mlxd. Happy to change if needed, so please let me know what you think.

• Downgraded sphinx. The reason is that I was getting so many sphinx failures.

I noticed that sphinx was updated recently in this PR (@rashidnhm or @runora95), in which PennyLaneAI/[email protected] has been replaced with PennyLaneAI/sphinx-action@master in .github/workflows/docs.yml.

This branch was using PennyLaneAI/sphinx-action@master, and it was created from the 0.41 tag, which therefore was also using PennyLaneAI/sphinx-action@master (you can check by downloading the source code from GitHub).
Therefore, in between the 0.41 release and this PR, we must have changed PennyLaneAI/sphinx-action@master to PennyLaneAI/[email protected] at some point.

Maybe @rashidnhm or @runora95 have some insight about what and why this happened. Thanks!

Obviously, this change (that is, downgrading sphinx) will not be merged into master.

• Replaced an outdated link in qnn/__init__.py. The reason is that I was getting an annoying sphinx failure:

/github/workspace/pennylane/qnn/__init__.py:docstring of pennylane.qnn:6: WARNING: unknown document: demos/tutorial_qnn_module_tf

And I noticed that the link to this demo in the stable doc version is broken (most probably because the demo has been moved/renamed on March 20th, after the 0.41 release). I replaced the link with a working one in the only way I know.

Benefits: An informative warning is raised when importing pennylane with the latest JAX installed.

Possible Drawbacks: Strictly speaking, none as this is a bugfix release.

Related GitHub Issues: None.

Related Shortcut Story: [sc-90355]

Further details: I tested this in a virtual conda environment.

If JAX is installed with version <= 0.4.28 or if it is not installed at all, we get the following output when importing pennylane:

~/Desktop/repos/pennylane warning_latest_jax > python -c "import pennylane"                                                                                                       py PLPostRelease 19:15:18

~/Desktop/repos/pennylane warning_latest_jax >                                                                                                                                    py PLPostRelease 19:15:31

If JAX is installed with version > 0.4.28, we get the following output when importing pennylane:

~/Desktop/repos/pennylane warning_latest_jax > python -c "import pennylane"                                                                                                    6s py PLPostRelease 19:08:26
/home/pietropaolo.frisoni/Desktop/repos/pennylane/pennylane/__init__.py:35: RuntimeWarning: PennyLane is not yet compatible with JAX versions > 0.4.28. You have JAX 0.6.0 installed. Please downgrade JAX to <=0.4.28 to avoid runtime errors. 
  warnings.warn(
Traceback (most recent call last):
 (...)

I encourage you to try as well : )

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Please upload report for BASE (v0.41.1-rc0@fe87033). Learn more about missing BASE report.

Additional details and impacted files

@@              Coverage Diff               @@
##             v0.41.1-rc0    #7369   +/-   ##
==============================================
  Coverage               ?   99.64%           
==============================================
  Files                  ?      513           
  Lines                  ?    50191           
  Branches               ?        0           
==============================================
  Hits                   ?    50015           
  Misses                 ?      176           
  Partials               ?        0           

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[PietropaoloFrisoni]

If the v0.41.1-rc0 branch gets merged back into master

Once the rc branch is merged to master

Thanks @rashidnhm for raising this point.

I don't think we should merge this branch into master, precisely for the reasons that you mentioned (also, there would be several merge conflicts). If needed, it's much faster for me to open a PR directly in master to copy the warning message implemented here.
For the record, we are indeed working to make PL compatible with the latest JAX (and this will be completed before PL 0.42), so in principle we wouldn't even need this warning in PL 0.42.

[EmilianoG-byte]

Hey there! So what is the current solution if I don't want to install from this branch but just from pip? Should I uninstall jax and install a version <= 0.4.28

Also, hi :D. Long time :)

[PietropaoloFrisoni]

Hi @EmilianoG-byte! Long time indeed : )

The change in this PR is now part of PL 0.41.1. That is, the current stable version that you can install with pip (that is, pip install pennylane). If you don't want to downgrade to JAX 0.4.28, you can use the latest PennyLane (installed from the master branch), which is compatible with JAX 0.6.0, but be careful that this is just a dev version so far.
The future PL 0.42 version should be released in about a month and will be compatible with JAX 0.6.0. I hope that helps!

Good luck with all your current and future endeavours!

[EmilianoG-byte]

Thank you for the answer @PietropaoloFrisoni :D. I will downgrade to 0.4.28 for now then and wait for the next release!

Also best of luck to you as well 😄🌞!