Bump jax api hashability (#8701)

**Context:**
Subset of #8525 focusing on
hashability, as well as the jax API updates

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**
[sc-102157]

---------

Co-authored-by: Christina Lee <[email protected]>
Co-authored-by: Mudit Pandey <[email protected]>

Author: 3 people

.github/workflows/documentation-tests.yml

@@ -55,6 +55,7 @@ jobs:
         run: |
           pip install --upgrade --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ pennylane-catalyst pennylane-lightning
           pip install -e .
+          # TODO: use 0.7.0 after updating all the documentation 
           pip install sybil pytest "jax~=0.6.0" "jaxlib~=0.6.0" torch matplotlib pyzx
 
       - name: Print Dependencies

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

@@ -16,12 +16,12 @@ on:
         description: The version of JAX to use for testing
         required: false
         type: string
-        default: '0.6.2'
+        default: '0.7.0'
       catalyst_jax_version:
         description: The version of JAX to use for testing along with Catalyst
         required: false
         type: string
-        default: '0.6.2'
+        default: '0.7.0'
       torch_version:
         description: The version of PyTorch to use for testing
         required: false

.gitignore

@@ -24,6 +24,7 @@ timer.dat
 tmp/*
 benchmark/revisions/
 venv
+*venv*/
 config.toml
 .envrc
 qml_debug.log

doc/introduction/interfaces/jax.rst

@@ -10,7 +10,7 @@ JAX interface
 
     .. code-block:: bash
 
-        pip install jax~=0.6.0 jaxlib~=0.6.0
+        pip install jax==0.7.0 jaxlib==0.7.0
     
     You can then import PennyLane and JAX as follows:
 

doc/requirements.txt

@@ -3,8 +3,8 @@ appdirs
 autograd
 autoray
 cachetools
-jax==0.6.0
-jaxlib==0.6.0
+jax==0.7.0
+jaxlib==0.7.0
 mistune==0.8.4
 m2r2
 # TODO: Remove once galois becomes compatible with latest numpy 

pennylane/__init__.py

@@ -204,12 +204,12 @@
 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
+    if (jax_version := _Version(_metadata_version("jax"))) > _Version("0.7.0"):  # pragma: no cover
         warnings.warn(
-            "PennyLane is not yet compatible with JAX versions > 0.6.2. "
+            "PennyLane is not yet compatible with JAX versions > 0.7.0. "
             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",
+            "Please downgrade JAX to 0.7.0 to avoid runtime errors using "
+            "python -m pip install jax==0.7.0 jaxlib==0.7.0",
             RuntimeWarning,
         )
 

pennylane/_grad.py

@@ -86,7 +86,9 @@ def _grad_abstract(*args, argnums, jaxpr, n_consts, method, h, scalar_out, fn):
 
 
 def _shape(shape, dtype, weak_type=False):
-    if jax.config.jax_dynamic_shapes and any(not isinstance(s, int) for s in shape):
+    if jax.config.jax_dynamic_shapes and any(
+        not isinstance(s, int) for s in shape
+    ):  # pragma: no cover
         return jax.core.DShapedArray(shape, dtype, weak_type=weak_type)
     return jax.core.ShapedArray(shape, dtype, weak_type=weak_type)
 

pennylane/capture/base_interpreter.py

@@ -82,7 +82,7 @@ def _fill_in_shape_with_dyn_shape(dyn_shape: tuple["jax.core.Tracer"], shape: tu
     for s in shape:
         if s is not None:
             new_shape.append(s)
-        else:
+        else:  # pragma: no cover
             # pull from iterable of dynamic shapes
             next_s = next(dyn_shape_iter)
             if not qml.math.is_abstract(next_s):
@@ -496,9 +496,10 @@ def handle_for_loop(
     self, start, stop, step, *args, jaxpr_body_fn, consts_slice, args_slice, abstract_shapes_slice
 ):
     """Handle a for loop primitive."""
-    consts = args[consts_slice]
-    init_state = args[args_slice]
-    abstract_shapes = args[abstract_shapes_slice]
+    # Convert tuples back to slices (tuples are used for JAX 0.7.0 hashability)
+    consts = args[slice(*consts_slice)]
+    init_state = args[slice(*args_slice)]
+    abstract_shapes = args[slice(*abstract_shapes_slice)]
     new_jaxpr_body_fn = jaxpr_to_jaxpr(
         copy(self), jaxpr_body_fn, consts, *abstract_shapes, start, *init_state
     )
@@ -523,6 +524,10 @@ def handle_for_loop(
 @PlxprInterpreter.register_primitive(cond_prim)
 def handle_cond(self, *invals, jaxpr_branches, consts_slices, args_slice):
     """Handle a cond primitive."""
+    # Convert tuples back to slices (tuples are used for JAX 0.7.0 hashability)
+    args_slice = slice(*args_slice)
+    consts_slices = [slice(*s) for s in consts_slices]
+
     args = invals[args_slice]
 
     new_jaxprs = []
@@ -560,6 +565,11 @@ def handle_while_loop(
     args_slice,
 ):
     """Handle a while loop primitive."""
+    # Convert tuples back to slices (tuples are used for JAX 0.7.0 hashability)
+    body_slice = slice(*body_slice)
+    cond_slice = slice(*cond_slice)
+    args_slice = slice(*args_slice)
+
     consts_body = invals[body_slice]
     consts_cond = invals[cond_slice]
     init_state = invals[args_slice]
@@ -654,6 +664,11 @@ def flatten_while_loop(
     args_slice,
 ):
     """Handle the while loop by a flattened python strategy."""
+    # Convert tuples back to slices (tuples are used for JAX 0.7.0 hashability)
+    body_slice = slice(*body_slice)
+    cond_slice = slice(*cond_slice)
+    args_slice = slice(*args_slice)
+
     consts_body = invals[body_slice]
     consts_cond = invals[cond_slice]
     init_state = invals[args_slice]
@@ -671,6 +686,10 @@ def flatten_while_loop(
 @FlattenedInterpreter.register_primitive(cond_prim)
 def flattened_cond(self, *invals, jaxpr_branches, consts_slices, args_slice):
     """Handle the cond primitive by a flattened python strategy."""
+    # Convert tuples back to slices (tuples are used for JAX 0.7.0 hashability)
+    args_slice = slice(*args_slice)
+    consts_slices = [slice(*s) for s in consts_slices]
+
     n_branches = len(jaxpr_branches)
     conditions = invals[:n_branches]
     args = invals[args_slice]
@@ -694,6 +713,11 @@ def flattened_for(
     self, start, stop, step, *invals, jaxpr_body_fn, consts_slice, args_slice, abstract_shapes_slice
 ):
     """Handle the for loop by a flattened python strategy."""
+    # Convert tuples back to slices (tuples are used for JAX 0.7.0 hashability)
+    consts_slice = slice(*consts_slice)
+    args_slice = slice(*args_slice)
+    abstract_shapes_slice = slice(*abstract_shapes_slice)
+
     consts = invals[consts_slice]
     init_state = invals[args_slice]
     abstract_shapes = invals[abstract_shapes_slice]

pennylane/capture/custom_primitives.py

@@ -16,6 +16,7 @@
 """
 
 from enum import Enum
+from typing import Any
 
 from jax.extend.core import Primitive
 
@@ -30,10 +31,32 @@ class PrimitiveType(Enum):
     TRANSFORM = "transform"
 
 
+def _make_hashable(obj: Any) -> Any:
+    """Convert potentially unhashable objects to hashable equivalents for JAX 0.7.0+.
+
+    JAX 0.7.0 requires all primitive parameters to be hashable. This helper converts
+    common unhashable types (list, dict, slice) to hashable tuples.
+
+    Args:
+        obj: Object to potentially convert to hashable form
+
+    Returns:
+        Hashable version of the object
+    """
+    if isinstance(obj, slice):
+        return (obj.start, obj.stop, obj.step)
+    if isinstance(obj, list):
+        return tuple(_make_hashable(item) for item in obj)
+    if isinstance(obj, dict):
+        return tuple((k, _make_hashable(v)) for k, v in obj.items())
+
+    return obj
+
+
 # pylint: disable=abstract-method,too-few-public-methods
 class QmlPrimitive(Primitive):
     """A subclass for JAX's Primitive that differentiates between different
-    classes of primitives."""
+    classes of primitives and automatically makes parameters hashable for JAX 0.7.0+."""
 
     _prim_type: PrimitiveType = PrimitiveType.DEFAULT
 
@@ -47,3 +70,13 @@ def prim_type(self):
     def prim_type(self, value: str | PrimitiveType):
         """Setter for QmlPrimitive.prim_type."""
         self._prim_type = PrimitiveType(value)
+
+    def bind(self, *args, **params):
+        """Bind with automatic parameter hashability conversion for JAX 0.7.0+.
+
+        Overrides the parent bind method to automatically convert unhashable parameters
+        (like lists, dicts, and slices) to hashable tuples, which is required by JAX 0.7.0+.
+        """
+        # Convert all parameters to hashable forms
+        hashable_params = {k: _make_hashable(v) for k, v in params.items()}
+        return super().bind(*args, **hashable_params)

pennylane/capture/dynamic_shapes.py

@@ -19,6 +19,7 @@
 has_jax = True
 try:
     import jax
+    from jax._src.interpreters.partial_eval import TracingEqn
     from jax.interpreters import partial_eval as pe
 except ImportError:  # pragma: no cover
     has_jax = False  # pragma: no cover
@@ -47,7 +48,7 @@ def _get_shape_for_array(x, abstract_shapes: list, previous_ints: list) -> dict:
         return {}
 
     abstract_axes = {}
-    for i, s in enumerate(getattr(x, "shape", ())):
+    for i, s in enumerate(getattr(x, "shape", ())):  # pragma: no cover
         if not isinstance(s, int):  #  if not int, then abstract
             found = False
             # check if the shape tracer is one we have already encountered
@@ -137,8 +138,8 @@ def f(n):
     if not any(abstracted_axes):
         return None, ()
 
-    abstracted_axes = jax.tree_util.tree_unflatten(structure, abstracted_axes)
-    return abstracted_axes, abstract_shapes
+    abstracted_axes = jax.tree_util.tree_unflatten(structure, abstracted_axes)  # pragma: no cover
+    return abstracted_axes, abstract_shapes  # pragma: no cover
 
 
 def register_custom_staging_rule(
@@ -164,7 +165,14 @@ def register_custom_staging_rule(
     # see https://github.com/jax-ml/jax/blob/9e62994bce7c7fcbb2f6a50c9ef89526cd2c2be6/jax/_src/lax/lax.py#L3538
     # and https://github.com/jax-ml/jax/blob/9e62994bce7c7fcbb2f6a50c9ef89526cd2c2be6/jax/_src/lax/lax.py#L208
     # for reference to how jax is handling staging rules for dynamic shapes in v0.4.28
-    # see also capture/intro_to_dynamic_shapes.md
+    # JAX 0.6.2 to 0.7.0 introduced breaking changes in custom staging rules for dynamic shapes:
+    # 1. DynamicJaxprTracer constructor now requires the var as 3rd argument (previously created internally)
+    # 2. TracingEqn must be used instead of JaxprEqn for trace.frame.add_eqn
+    #
+    # This implementation creates vars first using trace.frame.newvar() before constructing
+    # DynamicJaxprTracer instances, fixing dynamic shape support that was broken in JAX 0.7.0.
+    # See pennylane/capture/jax_patches.py for related fixes to JAX's own staging rules.
+    # See also capture/intro_to_dynamic_shapes.md for dynamic shapes documentation.
 
     def _tracer_and_outvar(
         jaxpr_trace: pe.DynamicJaxprTrace,
@@ -176,15 +184,18 @@ def _tracer_and_outvar(
         Returned vars are cached in env for use in future shapes
         """
         if not hasattr(outvar.aval, "shape"):
-            out_tracer = pe.DynamicJaxprTracer(jaxpr_trace, outvar.aval, None)
-            return out_tracer, jaxpr_trace.makevar(out_tracer)
+            # JAX 0.7.0: Create variable first, then pass to DynamicJaxprTracer
+            new_var = jaxpr_trace.frame.newvar(outvar.aval)
+            out_tracer = pe.DynamicJaxprTracer(jaxpr_trace, outvar.aval, new_var)
+            return out_tracer, new_var
         new_shape = [s if isinstance(s, int) else env[s] for s in outvar.aval.shape]
         if all(isinstance(s, int) for s in outvar.aval.shape):
             new_aval = jax.core.ShapedArray(tuple(new_shape), outvar.aval.dtype)
-        else:
+        else:  # pragma: no cover
             new_aval = jax.core.DShapedArray(tuple(new_shape), outvar.aval.dtype)
-        out_tracer = pe.DynamicJaxprTracer(jaxpr_trace, new_aval, None)
-        new_var = jaxpr_trace.makevar(out_tracer)
+        # JAX 0.7.0: Create variable first, then pass to DynamicJaxprTracer
+        new_var = jaxpr_trace.frame.newvar(new_aval)
+        out_tracer = pe.DynamicJaxprTracer(jaxpr_trace, new_aval, new_var)
 
         if not isinstance(outvar, jax.extend.core.Literal):
             env[outvar] = new_var
@@ -211,15 +222,26 @@ def custom_staging_rule(
         else:
             out_tracers, returned_vars = (), ()
 
-        invars = [jaxpr_trace.getvar(x) for x in tracers]
+        # JAX 0.7.0: Use t.val to get var from tracer, and TracingEqn for frame.add_eqn
+        invars = [t.val for t in tracers]
         eqn = jax.core.new_jaxpr_eqn(
             invars,
             returned_vars,
             primitive,
             params,
             jax.core.no_effects,
+            source_info,
+        )
+        tracing_eqn = TracingEqn(
+            list(tracers),
+            returned_vars,
+            primitive,
+            params,
+            eqn.effects,
+            source_info,
+            eqn.ctx,
         )
-        jaxpr_trace.frame.add_eqn(eqn)
+        jaxpr_trace.frame.add_eqn(tracing_eqn)
         return out_tracers
 
     pe.custom_staging_rules[primitive] = custom_staging_rule