[Bug Fix Release] raise warning if jax>0.6.2 is installed

.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"
         )
 
 

pennylane/labs/dla/variational_kak.py

@@ -222,7 +222,7 @@ def Kc(theta_opt):
 
     if not has_jax:  # pragma: no cover
         raise ImportError(
-            "jax and optax are required for variational_kak_adj. You can install them with pip install jax jaxlib optax."
+            "jax and optax are required for variational_kak_adj. You can install them with pip install jax~=0.6.0 jaxlib~=0.6.0 optax."
         )  # pragma: no cover
     if verbose >= 1 and not has_plt:  # pragma: no cover
         print(
@@ -392,7 +392,7 @@ def cost(x):
 
     if not has_jax:  # pragma: no cover
         raise ImportError(
-            "jax and optax are required for run_opt. You can install them with pip install jax jaxlib optax."
+            "jax and optax are required for run_opt. You can install them with pip install jax~=0.6.0 jaxlib~=0.6.0 optax."
         )  # pragma: no cover
 
     if optimizer is None:

pennylane/pulse/convenience_functions.py

@@ -176,7 +176,7 @@ def f(p, t):
     if not has_jax:
         raise ImportError(
             "Module jax is required for any pulse-related convenience function. "
-            "You can install jax via: pip install jax==0.4.10 jaxlib==0.4.10"
+            "You can install jax via: pip install jax~=0.6.0 jaxlib~=0.6.0"
         )
     if windows is not None:
         is_nested = any(hasattr(w, "__len__") for w in windows)
@@ -291,7 +291,7 @@ def wrapped(p, t):
     if not has_jax:
         raise ImportError(
             "Module jax is required for any pulse-related convenience function. "
-            "You can install jax via: pip install jax==0.4.3 jaxlib==0.4.3"
+            "You can install jax via: pip install jax~=0.6.0 jaxlib~=0.6.0"
         )
 
     if isinstance(timespan, (tuple, list)):
@@ -365,7 +365,7 @@ def fn(params, t):
     if not has_jax:
         raise ImportError(
             "Module jax is required for any pulse-related convenience function. "
-            "You can install jax via: pip install jax==0.4.3 jaxlib==0.4.3"
+            "You can install jax via: pip install jax~=0.6.0 jaxlib~=0.6.0"
         )
 
     if isinstance(timespan, tuple):

pennylane/pulse/parametrized_evolution.py

@@ -418,7 +418,7 @@ def __call__(
         if not has_jax:
             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"
             )
         # Need to cast all elements inside params to `jnp.arrays` to make sure they are not cast
         # to `np.arrays` inside `Operator.__init__`
@@ -511,7 +511,7 @@ def matrix(self, wire_order=None):
         if not has_jax:
             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"
             )
         if not self.has_matrix:
             raise ValueError(

pennylane/qchem/factorization.py

@@ -77,7 +77,7 @@ def factorize(
     .. note::
 
         Packages JAX and Optax are required when performing CDF with ``compressed=True``.
-        Install them using ``pip install jax optax``.
+        Install them using ``pip install jax~=0.6.0 optax``.
 
     Args:
         two_electron (array[array[float]]): Two-electron integral tensor in the molecular orbital
@@ -262,7 +262,7 @@ def factorize(
         if not has_jax_optax:
             raise ImportError(
                 "Jax and Optax libraries are required for optimizing the factors. Install them via "
-                "pip install jax optax"
+                "pip install jax~=0.6.0 optax"
             )  # pragma: no cover
 
         norm_order = {None: None, "L1": 1, "L2": 2}.get(regularization, "LX")