diff --git a/frontend/catalyst/from_plxpr/from_plxpr.py b/frontend/catalyst/from_plxpr/from_plxpr.py index f3b0366f57..8e66973d0c 100644 --- a/frontend/catalyst/from_plxpr/from_plxpr.py +++ b/frontend/catalyst/from_plxpr/from_plxpr.py @@ -18,12 +18,15 @@ # pylint: disable=protected-access import warnings +from contextlib import nullcontext from copy import copy from functools import partial from typing import Callable import jax import pennylane as qp +from jax._src import config as jax_config +from jax.core import ensure_compile_time_eval from jax.extend.core import ClosedJaxpr, Jaxpr from pennylane.capture import PlxprInterpreter, qnode_prim from pennylane.capture.primitives import transform_prim @@ -312,23 +315,26 @@ def handle_qnode( ) def calling_convention(*args): - device_init_p.bind( - shots, - auto_qubit_management=(device.wires is None), - **_get_device_kwargs(device), - ) + with jax_config.eager_constant_folding(False): + device_init_p.bind( + shots, + auto_qubit_management=(device.wires is None), + **_get_device_kwargs(device), + ) # https://github.com/PennyLaneAI/pennylane/pull/9248 assert not is_dynamic_wires( device.wires ), "plxpr does not support dynamic number of wires on the device yet" - qreg = qref_alloc_p.bind(static_num_qubits=len(device.wires)) + with jax_config.eager_constant_folding(False): + qreg = qref_alloc_p.bind(static_num_qubits=len(device.wires)) self.init_qreg = qreg converter = PLxPRToQuantumJaxprInterpreter(device, shots, self.init_qreg, {}) retvals = converter(closed_jaxpr, *args) qref_dealloc_p.bind(self.init_qreg) - device_release_p.bind() + with jax_config.eager_constant_folding(False): + device_release_p.bind() return retvals if self.requires_decompose_lowering and graph_succeeded: @@ -467,6 +473,7 @@ def trace_from_pennylane( static_argnums, abstracted_axes, skip_preprocess=False, + const_eval=True, debug_info=None, ): """Capture the JAX program representation (JAXPR) of the wrapped function, using @@ -507,9 +514,11 @@ def trace_from_pennylane( # Therefore we need to coordinate them manually fn.static_argnums = static_argnums + const_eval_ctx_manager = ensure_compile_time_eval() if const_eval else nullcontext() + with transient_jax_config( {"jax_dynamic_shapes": True, "jax_use_shardy_partitioner": False} - ), Patcher(*get_jax_patches()): + ), Patcher(*get_jax_patches()), const_eval_ctx_manager: make_jaxpr_kwargs = { "static_argnums": static_argnums, @@ -533,8 +542,8 @@ def wrapper(*inner_args, **inner_kwargs): jaxpr = from_plxpr(plxpr, skip_preprocess=skip_preprocess)( *abstract_shapes, *flat_inputs ) - - return _dummy_hop.bind(jaxpr=jaxpr, out_type=out_type, out_treedef=out_treedef) + with jax_config.eager_constant_folding(False): + return _dummy_hop.bind(jaxpr=jaxpr, out_type=out_type, out_treedef=out_treedef) nested_jaxpr = jax.make_jaxpr(wrapper, **make_jaxpr_kwargs)(*args, **kwargs) jaxpr = nested_jaxpr.eqns[0].params["jaxpr"] diff --git a/frontend/catalyst/jit.py b/frontend/catalyst/jit.py index 8d137a8268..c8d5ddaad3 100644 --- a/frontend/catalyst/jit.py +++ b/frontend/catalyst/jit.py @@ -96,6 +96,7 @@ def qjit( dialect_plugins=None, capture="global", skip_preprocess=False, + const_eval=True, ): # pylint: disable=too-many-arguments,unused-argument """A just-in-time decorator for PennyLane and JAX programs using Catalyst. @@ -808,6 +809,7 @@ def capture(self, args, **kwargs): static_argnums=static_argnums, abstracted_axes=abstracted_axes, skip_preprocess=self.compile_options.skip_preprocess, + const_eval=self.compile_options.const_eval, debug_info=dbg, ) return jaxpr, out_type, out_treedef, full_sig diff --git a/frontend/catalyst/pipelines.py b/frontend/catalyst/pipelines.py index 241b37a42d..585d0759e2 100644 --- a/frontend/catalyst/pipelines.py +++ b/frontend/catalyst/pipelines.py @@ -170,6 +170,7 @@ class CompileOptions: dialect_plugins: Optional[Set[Path]] = None capture: bool | Literal["global"] = "global" skip_preprocess: bool = False + const_eval: bool = True def __post_init__(self): # Convert keep_intermediate to Enum