[Bug Fix Release] raise warning if jax>0.6.2 is installed (#7949)

**Context:**

Jax released a [new
version](https://github.com/jax-ml/jax/releases/tag/jax-v0.7.0) with a
lot of breaking changes. Unfortunately our `v0.42` release is not
compatible with this version and so there are runtime errors.

ℹ️ : This PR was created by branching the v0.42.0 tag.

**Description of the Change:**

- The fix will check if the user has `jax` and if they do it will
perform a version check to ensure compatibility.

```python
>>> import pennylane as qml
```
```
~/Documents/pennylane/pennylane/__init__.py:29: RuntimeWarning: PennyLane is not yet compatible with JAX versions > 0.6.2. You have version 0.7.0 installed. Please downgrade JAX to <=0.6.2 to avoid runtime errors.
```

- 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 for some custom exceptions). This
change will not be merged into master.

**Benefits:** Better UI.

**Possible Drawbacks:** None.

[sc-96104]

---------

Co-authored-by: Mudit Pandey <[email protected]>
Co-authored-by: Lee James O'Riordan <[email protected]>

Author: 3 people

.github/workflows/interface-dependency-versions.yml

@@ -75,11 +75,11 @@ jobs:
 
       - name: Nightly Catalyst Version
         id: catalyst
-        run: echo "nightly=--index https://test.pypi.org/simple/ --prerelease=allow --upgrade-package PennyLane-Catalyst PennyLane-Catalyst" >> $GITHUB_OUTPUT
+        run: echo "nightly=--upgrade-package PennyLane-Catalyst PennyLane-Catalyst" >> $GITHUB_OUTPUT
 
       - name: PennyLane-Lightning Latest Version
         id: pennylane-lightning
-        run: echo "latest=--index https://test.pypi.org/simple/ --prerelease=allow --upgrade-package PennyLane-Lightning PennyLane-Lightning" >> $GITHUB_OUTPUT
+        run: echo "latest=--upgrade-package PennyLane-Lightning PennyLane-Lightning" >> $GITHUB_OUTPUT
 
     outputs:
       catalyst-jax-version: jax==${{ steps.catalyst-jax.outputs.version }} jaxlib==${{ steps.catalyst-jax.outputs.version }}

doc/development/release_notes.md

@@ -3,6 +3,8 @@ Release notes
 
 This page contains the release notes for PennyLane.
 
+.. mdinclude:: ../releases/changelog-0.42.1.md
+
 .. mdinclude:: ../releases/changelog-0.42.0.md
 
 .. mdinclude:: ../releases/changelog-0.41.1.md

doc/releases/changelog-0.42.0.md

@@ -1,6 +1,6 @@
 :orphan:
 
-# Release 0.42.0 (current release)
+# Release 0.42.0
 
 <h3>New features since last release</h3>
 

doc/releases/changelog-0.42.1.md

@@ -0,0 +1,14 @@
+:orphan:
+
+# Release 0.42.1 (current release)
+
+<h3>Bug fixes 🐛</h3>
+
+* A warning is raised if PennyLane is imported and a version of JAX greater than 0.6.2 is installed.
+  [(#7949)](https://github.com/PennyLaneAI/pennylane/pull/7949)
+
+<h3>Contributors ✍️</h3>
+
+This release contains contributions from (in alphabetical order):
+
+Andrija Paurevic

pennylane/__init__.py

@@ -186,6 +186,21 @@
 from pennylane.liealg import lie_closure, structure_constants, center
 import pennylane.qnn
 
+
+from importlib.metadata import version as _metadata_version
+from importlib.util import find_spec as _find_spec
+from packaging.version import Version as _Version
+
+if _find_spec("jax") is not None:
+    if (jax_version := _Version(_metadata_version("jax"))) > _Version("0.6.2"):  # pragma: no cover
+        warnings.warn(
+            "PennyLane is not yet compatible with JAX versions > 0.6.2. "
+            f"You have version {jax_version} installed. "
+            "Please downgrade JAX to 0.6.2 to avoid runtime errors using "
+            "python -m pip install jax~=0.6.0 jaxlib~=0.6.0",
+            RuntimeWarning,
+        )
+
 # Look for an existing configuration file
 default_config = Configuration("config.toml")
 

pennylane/_version.py

@@ -16,4 +16,4 @@
 Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.42.0"
+__version__ = "0.42.1"

pennylane/capture/base_interpreter.py

@@ -14,12 +14,15 @@
 """
 This submodule defines a strategy structure for defining custom plxpr interpreters
 """
-# pylint: disable=no-self-use
 from copy import copy
 from functools import partial, wraps
+
+# pylint: disable=no-self-use, wrong-import-position
+from importlib.metadata import version
 from typing import Callable, Optional, Sequence
 
 import jax
+from packaging.version import Version
 
 import pennylane as qml
 from pennylane import math
@@ -635,15 +638,22 @@ class FlattenedInterpreter(PlxprInterpreter):
     """
 
 
+jax_version = version("jax")
+if Version(jax_version) > Version("0.6.2"):  # pragma: no cover
+    from jax._src.pjit import jit_p as pjit_p
+else:  # pragma: no cover
+    from jax._src.pjit import pjit_p
+
+
 # pylint: disable=protected-access
-@FlattenedInterpreter.register_primitive(jax._src.pjit.pjit_p)
+@FlattenedInterpreter.register_primitive(pjit_p)
 def _(self, *invals, jaxpr, **params):
     if jax.config.jax_dynamic_shapes:
         # just evaluate it so it doesn't throw dynamic shape errors
         return copy(self).eval(jaxpr.jaxpr, jaxpr.consts, *invals)
 
-    subfuns, params = jax._src.pjit.pjit_p.get_bind_params({"jaxpr": jaxpr, **params})
-    return jax._src.pjit.pjit_p.bind(*subfuns, *invals, **params)
+    subfuns, params = pjit_p.get_bind_params({"jaxpr": jaxpr, **params})
+    return pjit_p.bind(*subfuns, *invals, **params)
 
 
 @FlattenedInterpreter.register_primitive(while_loop_prim)

pennylane/capture/make_plxpr.py

@@ -144,7 +144,7 @@ def fn(x):
     if not has_jax:  # pragma: no cover
         raise ImportError(
             "Module jax is required for the ``make_plxpr`` function. "
-            "You can install jax via: pip install jax"
+            "You can install jax via: pip install jax~=0.6.0"
         )
 
     if not qml.capture.enabled():

pennylane/devices/qubit/apply_operation.py

@@ -711,7 +711,7 @@ def _evolve_state_vector_under_parametrized_evolution(
     except ImportError as e:  # pragma: no cover
         raise ImportError(
             "Module jax is required for the ``ParametrizedEvolution`` class. "
-            "You can install jax via: pip install jax"
+            "You can install jax via: pip install jax~=0.6.0"
         ) from e
 
     if operation.data is None or operation.t is None:

pennylane/gradients/pulse_gradient.py

@@ -59,7 +59,7 @@ def _assert_has_jax(transform_name):
     if not has_jax:  # pragma: no cover
         raise ImportError(
             f"Module jax is required for the {transform_name} gradient transform. "
-            "You can install jax via: pip install jax jaxlib"
+            "You can install jax via: pip install jax~=0.6.0 jaxlib~=0.6.0"
         )