Raise RuntimeWarning if jax > 0.4.28 is installed

doc/releases/changelog-dev.md

@@ -6,6 +6,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* A `RuntimeWarning` is now raised by `qml.QNode` and `qml.execute` if executing JAX workflows and the installed version of JAX
+  is greater than `0.4.28`.
+  [(#6864)](https://github.com/PennyLaneAI/pennylane/pull/6864)
+
 * `QNode` objects now have an `update` method that allows for re-configuring settings like `diff_method`, `mcm_method`, and more. This allows for easier on-the-fly adjustments to workflows. Any arguments not specified will retain their original value.
   [(#6803)](https://github.com/PennyLaneAI/pennylane/pull/6803)
 
@@ -105,4 +109,5 @@ Diksha Dhawan,
 Pietropaolo Frisoni,
 Marcus Gisslén,
 Christina Lee,
+Mudit Pandey,
 Andrija Paurevic

pennylane/math/interface_utils.py

@@ -15,9 +15,13 @@
 
 import warnings
 from enum import Enum
+from importlib.metadata import version
+from importlib.util import find_spec
 from typing import Literal, Union
+from warnings import warn
 
 import autoray as ar
+from packaging.version import Version
 
 
 class Interface(Enum):
@@ -212,23 +216,50 @@ def get_deep_interface(value):
     return _get_interface_of_single_tensor(itr)
 
 
-def get_canonical_interface_name(user_input: InterfaceLike) -> Interface:
+def _check_supported_jax() -> None:
+    """Checks if the installed version of JAX is supported. If an unsupported version of
+    JAX is installed, a ``RuntimeWarning`` is raised."""
+    if not find_spec("jax"):
+        return
+
+    jax_version = version("jax")
+    if Version(jax_version) > Version("0.4.28"):  # pragma: no cover
+        warn(
+            "PennyLane is currently not compatible with versions of JAX > 0.4.28. "
+            f"You have version {jax_version} installed.",
+            RuntimeWarning,
+        )
+
+
+def get_canonical_interface_name(
+    user_input: InterfaceLike, _validate_jax_version=False
+) -> Interface:
     """Helper function to get the canonical interface name.
 
     Args:
-        interface (str, Interface): reference interface
+        interface (str, Interface): Reference interface
+        _validate_jax_version (bool): Whether we should check if a supported version of
+            JAX is installed. If ``True``, a ``RuntimeWarning`` is raised if the canonical
+            interface is found to be JAX. ``False`` by default.
 
     Raises:
-        ValueError: key does not exist in the interface map
+        ValueError: Key does not exist in the interface map
 
     Returns:
-        Interface: canonical interface
+        Interface: Canonical interface
     """
 
     if user_input in SUPPORTED_INTERFACE_NAMES:
+        if _validate_jax_version and user_input in (Interface.JAX, Interface.JAX_JIT):
+            _check_supported_jax()
+
         return user_input
     try:
-        return INTERFACE_MAP[user_input]
+        out = INTERFACE_MAP[user_input]
+
+        if out in (Interface.JAX, Interface.JAX_JIT):
+            _check_supported_jax()
+        return out
     except KeyError as exc:
         raise ValueError(
             f"Unknown interface {user_input}. Interface must be one of {SUPPORTED_INTERFACE_NAMES}."

pennylane/workflow/qnode.py

@@ -551,7 +551,7 @@ def __init__(
         # input arguments
         self.func = func
         self.device = device
-        self._interface = get_canonical_interface_name(interface)
+        self._interface = get_canonical_interface_name(interface, _validate_jax_version=True)
         self.diff_method = diff_method
         mcm_config = qml.devices.MCMConfig(mcm_method=mcm_method, postselect_mode=postselect_mode)
         cache = (max_diff > 1) if cache == "auto" else cache

pennylane/workflow/resolution.py

@@ -97,7 +97,7 @@ def _resolve_interface(interface: Union[str, Interface], tapes: QuantumScriptBat
         Interface: resolved interface
     """
 
-    interface = get_canonical_interface_name(interface)
+    interface = get_canonical_interface_name(interface, _validate_jax_version=True)
 
     if interface == Interface.AUTO:
         params = []
@@ -106,7 +106,7 @@ def _resolve_interface(interface: Union[str, Interface], tapes: QuantumScriptBat
         interface = get_interface(*params)
         if interface != Interface.NUMPY:
             try:
-                interface = get_canonical_interface_name(interface)
+                interface = get_canonical_interface_name(interface, _validate_jax_version=True)
             except ValueError:
                 interface = Interface.NUMPY
     if interface == Interface.TF and _use_tensorflow_autograph():