Merge branch 'master' into compact_pauli

Author: Jaybsoni

.github/workflows/documentation-tests.yml

@@ -55,8 +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
+          pip install sybil pytest "jax==0.7.1" "jaxlib==0.7.1" torch matplotlib pyzx
 
       - name: Print Dependencies
         run: | 

.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.7.0'
+        default: '0.7.1'
       catalyst_jax_version:
         description: The version of JAX to use for testing along with Catalyst
         required: false
         type: string
-        default: '0.7.0'
+        default: '0.7.1'
       torch_version:
         description: The version of PyTorch to use for testing
         required: false

.github/workflows/interface-unit-tests.yml

@@ -109,7 +109,7 @@ jobs:
       - name: Set pytest arguments for setting warnings level
         id: pytest_warning_flags
         env:
-          PYTEST_WARNING_ARGS: -W "${{ inputs.python_warning_level }}" --continue-on-collection-errors
+          PYTEST_WARNING_ARGS: -W "${{ inputs.python_warning_level }}" --continue-on-collection-errors "${{ inputs.pytest_additional_args }}"
         run: |
           if [ "${{ inputs.python_warning_level }}" != "default" ]; then
             echo "Setting pytest warning flags"

.gitignore

@@ -23,7 +23,6 @@ doc/code/api/*
 timer.dat
 tmp/*
 benchmark/revisions/
-venv
 *venv*/
 config.toml
 .envrc

doc/introduction/interfaces/jax.rst

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

doc/releases/changelog-dev.md

@@ -21,6 +21,12 @@
   that produces a set of gate names to be used as the target gate set in decompositions.
   [(#8522)](https://github.com/PennyLaneAI/pennylane/pull/8522)
 
+* The :class:`~pennylane.decomposition.DecompositionGraph` now tracks the minimum number of
+  dynamic wire allocations required to solve the circuit, and provides a `minimize_work_wires`
+  option that enables the graph to select the best decomposition rules while minimizing the
+  number of additional allocations of work wires.
+  [(#8729)](https://github.com/PennyLaneAI/pennylane/pull/8729)
+
 <h4>Pauli product measurements</h4>
 
 * Added a :func:`~pennylane.ops.pauli_measure` that takes a Pauli product measurement.
@@ -31,6 +37,16 @@
 
 <h3>Improvements 🛠</h3>
 
+* Added a new decomposition, `_decompose_2_cnots`, for the two-qubit decomposition for `QubitUnitary`.
+  It supports the analytical decomposition a two-qubit unitary known to require exactly 2 CNOTs.
+  [(#8666)](https://github.com/PennyLaneAI/pennylane/issues/8666)
+  
+* Arithmetic dunder methods (`__add__`, `__mul__`, `__rmul__`) have been added to 
+  :class:`~.transforms.core.TransformDispatcher`, :class:`~.transforms.core.TransformContainer`, 
+  and :class:`~.transforms.core.TransformProgram` to enable intuitive composition of transform 
+  programs using `+` and `*` operators.
+  [(#8703)](https://github.com/PennyLaneAI/pennylane/pull/8703)
+
 * Quantum compilation passes in MLIR and XDSL can now be applied using the core PennyLane transform
   infrastructure, instead of using Catalyst-specific tools. This is made possible by a new argument in
   :func:`~pennylane.transform` and `~.TransformDispatcher` called ``pass_name``, which accepts a string
@@ -360,6 +376,15 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* Bump `jax` version to `0.7.1` for `capture` module.
+  [(#8715)](https://github.com/PennyLaneAI/pennylane/pull/8715)
+
+* Bump `jax` version to `0.7.0` for `capture` module.
+  [(#8701)](https://github.com/PennyLaneAI/pennylane/pull/8701)
+
+* Improve error handling when using PennyLane's experimental program capture functionality with an incompatible JAX version. 
+  [(#8723)](https://github.com/PennyLaneAI/pennylane/pull/8723)
+
 * Bump `autoray` package version to `0.8.2`.
   [(#8674)](https://github.com/PennyLaneAI/pennylane/pull/8674)
 
@@ -515,6 +540,9 @@ A warning message has been added to :doc:`Building a plugin <../development/plug
 
 <h3>Bug fixes 🐛</h3>
 
+* Update `interface-unit-tests.yml` to use its input parameter `pytest_additional_args` when running pytest.
+  [(#8705)](https://github.com/PennyLaneAI/pennylane/pull/8705)
+
 * Fixes a bug where in `resolve_work_wire_type` we incorrectly returned a value of `zeroed` if `both work_wires` 
   and `base_work_wires` were empty, causing an incorrect work wire type.
   [(#8718)](https://github.com/PennyLaneAI/pennylane/pull/8718)
@@ -564,6 +592,10 @@ A warning message has been added to :doc:`Building a plugin <../development/plug
 * Fixes a bug where :func:`~.change_op_basis` cannot be captured when the `uncompute_op` is left out.
   [(#8695)](https://github.com/PennyLaneAI/pennylane/pull/8695)
 
+* Fixes a bug where decomposition rules are sometimes incorrectly disregarded by the `DecompositionGraph` when a higher level
+  decomposition rule uses dynamically allocated work wires.
+  [(#8725)](https://github.com/PennyLaneAI/pennylane/pull/8725)
+
 * Fixes a bug where :class:`~.ops.ChangeOpBasis` is not correctly reconstructed using `qml.pytrees.unflatten(*qml.pytrees.flatten(op))`
   [(#8721)](https://github.com/PennyLaneAI/pennylane/issues/8721)
 
@@ -587,6 +619,7 @@ Mudit Pandey,
 Shuli Shu,
 Jay Soni,
 nate stemen,
+Theodoros Trochatos,
 David Wierichs,
 Hongsheng Zheng,
 Zinan Zhou

doc/requirements.txt

@@ -3,8 +3,8 @@ appdirs
 autograd
 autoray
 cachetools
-jax==0.7.0
-jaxlib==0.7.0
+jax==0.7.1
+jaxlib==0.7.1
 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.7.0"):  # pragma: no cover
+    if (jax_version := _Version(_metadata_version("jax"))) > _Version("0.7.1"):  # pragma: no cover
         warnings.warn(
-            "PennyLane is not yet compatible with JAX versions > 0.7.0. "
+            "PennyLane is not yet compatible with JAX versions > 0.7.1. "
             f"You have version {jax_version} installed. "
-            "Please downgrade JAX to 0.7.0 to avoid runtime errors using "
-            "python -m pip install jax==0.7.0 jaxlib==0.7.0",
+            "Please downgrade JAX to 0.7.1 to avoid runtime errors using "
+            "python -m pip install jax==0.7.1 jaxlib==0.7.1",
             RuntimeWarning,
         )
 

pennylane/_version.py

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

pennylane/capture/dynamic_shapes.py

@@ -14,15 +14,16 @@
 """
 Contains a utility for handling inputs with dynamically shaped arrays.
 """
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 
 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
+
+
+except ImportError as e:  # pragma: no cover
+    has_jax = False
 
 
 def _get_shape_for_array(x, abstract_shapes: list, previous_ints: list) -> dict:
@@ -123,7 +124,10 @@ def f(n):
 
     """
     if not has_jax:  # pragma: no cover
-        raise ImportError("jax must be installed to use determine_abstracted_axes")
+        raise ImportError(
+            "JAX == 0.7.0 must be installed to use determine_abstracted_axes. "
+            "Install with: pip install jax==0.7.0 jaxlib==0.7.0 "
+        )
     if not jax.config.jax_dynamic_shapes:
         return None, ()
 
@@ -175,18 +179,20 @@ def register_custom_staging_rule(
     # See also capture/intro_to_dynamic_shapes.md for dynamic shapes documentation.
 
     def _tracer_and_outvar(
-        jaxpr_trace: pe.DynamicJaxprTrace,
+        jaxpr_trace,
         outvar: jax.extend.core.Var,
         env: dict[jax.extend.core.Var, jax.extend.core.Var],
-    ) -> tuple[pe.DynamicJaxprTracer, jax.extend.core.Var]:
+    ):
         """
         Create a new tracer and return var from the true branch outvar.
         Returned vars are cached in env for use in future shapes
         """
         if not hasattr(outvar.aval, "shape"):
             # 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)
+            out_tracer = jax.interpreters.partial_eval.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):
@@ -195,15 +201,15 @@ def _tracer_and_outvar(
             new_aval = jax.core.DShapedArray(tuple(new_shape), outvar.aval.dtype)
         # 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)
+        out_tracer = jax.interpreters.partial_eval.DynamicJaxprTracer(
+            jaxpr_trace, new_aval, new_var
+        )
 
         if not isinstance(outvar, jax.extend.core.Literal):
             env[outvar] = new_var
         return out_tracer, new_var
 
-    def custom_staging_rule(
-        jaxpr_trace: pe.DynamicJaxprTrace, source_info, *tracers: pe.DynamicJaxprTracer, **params
-    ) -> Sequence[pe.DynamicJaxprTracer] | pe.DynamicJaxprTracer:
+    def custom_staging_rule(jaxpr_trace, source_info, *tracers, **params):
         """
         Add new jaxpr equation to the jaxpr_trace and return new tracers.
         """
@@ -244,4 +250,4 @@ def custom_staging_rule(
         jaxpr_trace.frame.add_eqn(tracing_eqn)
         return out_tracers
 
-    pe.custom_staging_rules[primitive] = custom_staging_rule
+    jax.interpreters.partial_eval.custom_staging_rules[primitive] = custom_staging_rule