Remove "experimental_capture" kwarg from "qjit" (#1657)

**Context:** When dealing with transforms, the `experimental_capture`
keyword from `qjit` fails at capturing them because they are defined
before the capture functionality has been activated. To circumvent this,
program capture has to be enabled manually beforehand, which defies the
purpose of the mentioned keyword. We propose using only the PL program
capture enabling/disabling mechanism across the whole ecosystem to
prevent such cases to happen.

**Description of the Change:** Removed the `experimental_capture`
keyword from its `qjit` function in favor of a
  unified program capture behavior. 
  [(#1657)](#1657)

Program capture has to be enabled before the definition of the function
to be qjitted.
  
For AOT compilation, program capture can be disabled right after the
qjit usage and before execution.
  
  ```python
  import pennylane as qml
  from catalyst import qjit

  dev = qml.device("lightning.qubit", wires=1)

  qml.capture.enable()

  @qjit()
  @qml.qnode(dev)
  def circuit(x: float):
      qml.Hadamard(0)
      qml.CNOT([0, 1])
      return qml.expval(qml.Z(0))

  qml.capture.disable()

  circuit(0.1)
  ```

But for JIT compilation, program capture cannot be disabled before
execution,
  otherwise the capture will not take place:

  ```python
  import pennylane as qml
  from catalyst import qjit

  dev = qml.device("lightning.qubit", wires=1)

  qml.capture.enable()

  @qjit()
  @qml.qnode(dev)
  def circuit(x):
      qml.Hadamard(0)
      qml.CNOT([0, 1])
      return qml.expval(qml.Z(0))

  circuit(0.1)

  qml.capture.disable()
  ```

**Benefits:** 
- Unified program capture behavior.
- No longer necessary to use a wrapper function around a circuit with
transforms:

Before:
```python
        @qjit(experimental_capture=True)
        def wrapper():
           @qml.transforms.merge_amplitude_embedding
           @qml.qnode(qml.device(backend, wires=2))
           def captured_circuit():
              qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
              qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
              return qml.expval(qml.PauliZ(0))
```

After:

```python
        qml.capture.enable()

        @qjit()
        @qml.transforms.merge_amplitude_embedding
        @qml.qnode(qml.device(backend, wires=2))
        def captured_circuit():
            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
            return qml.expval(qml.PauliZ(0))
```


**Possible Drawbacks:** Users might find it difficult to understand
where exactly program capture should be enabled/disabled.

[sc-89121]

---------

Co-authored-by: Isaac De Vlugt <34751083+isaacdevlugt@users.noreply.github.com>

Author: rauletorresc

doc/releases/changelog-dev.md

@@ -15,6 +15,33 @@
 
 <h3>Breaking changes 💔</h3>
 
+* Catalyst has removed the `experimental_capture` keyword from the `qjit` decorator in favour of
+  unified behaviour with PennyLane. 
+  [(#1657)](https://github.com/PennyLaneAI/catalyst/pull/1657)
+
+  Instead of enabling program capture with Catalyst via `qjit(experimental_capture=True)`, program capture
+  can be enabled via the global toggle `qml.capture.enable()`:
+  
+  ```python
+  import pennylane as qml
+  from catalyst import qjit
+
+  dev = qml.device("lightning.qubit", wires=2)
+
+  qml.capture.enable()
+
+  @qjit
+  @qml.qnode(dev)
+  def circuit(x):
+      qml.Hadamard(0)
+      qml.CNOT([0, 1])
+      return qml.expval(qml.Z(0))
+
+  circuit(0.1)
+  ```
+
+  Disabling program capture can be done with `qml.capture.disable()`.
+
 <h3>Deprecations 👋</h3>
 
 <h3>Bug fixes 🐛</h3>

frontend/catalyst/from_plxpr.py

@@ -24,7 +24,7 @@
 import pennylane as qml
 from jax.extend.linear_util import wrap_init
 from jax.interpreters.partial_eval import convert_constvars_jaxpr
-from pennylane.capture import PlxprInterpreter, disable, enable, enabled, qnode_prim
+from pennylane.capture import PlxprInterpreter, qnode_prim
 from pennylane.capture.expand_transforms import ExpandTransformsInterpreter
 from pennylane.capture.primitives import cond_prim as plxpr_cond_prim
 from pennylane.capture.primitives import for_loop_prim as plxpr_for_loop_prim
@@ -739,18 +739,9 @@ def trace_from_pennylane(fn, static_argnums, abstracted_axes, sig, kwargs):
             "abstracted_axes": abstracted_axes,
         }
 
-        if enabled():
-            capture_on = True
-        else:
-            capture_on = False
-            enable()
-
         args = sig
-        try:
-            plxpr, out_type, out_treedef = make_jaxpr2(fn, **make_jaxpr_kwargs)(*args, **kwargs)
-            jaxpr = from_plxpr(plxpr)(*args, **kwargs)
-        finally:
-            if not capture_on:
-                disable()
+
+        plxpr, out_type, out_treedef = make_jaxpr2(fn, **make_jaxpr_kwargs)(*args, **kwargs)
+        jaxpr = from_plxpr(plxpr)(*args, **kwargs)
 
     return jaxpr, out_type, out_treedef, sig

frontend/catalyst/jit.py

@@ -91,7 +91,6 @@ def qjit(
     abstracted_axes=None,
     disable_assertions=False,
     seed=None,
-    experimental_capture=False,
     circuit_transform_pipeline=None,
     pass_plugins=None,
     dialect_plugins=None,
@@ -151,9 +150,6 @@ def qjit(
             :func:`qml.sample() <pennylane.sample>`, :func:`qml.counts() <pennylane.counts>`,
             :func:`qml.probs() <pennylane.probs>`, :func:`qml.expval() <pennylane.expval>`,
             :func:`qml.var() <pennylane.var>`.
-        experimental_capture (bool): If set to ``True``, the qjit decorator
-            will use PennyLane's experimental program capture capabilities
-            to capture the decorated function for compilation.
         circuit_transform_pipeline (Optional[dict[str, dict[str, str]]]):
             A dictionary that specifies the quantum circuit transformation pass pipeline order,
             and optionally arguments for each pass in the pipeline. Keys of this dictionary
@@ -715,7 +711,7 @@ def capture(self, args, **kwargs):
         dynamic_sig = get_abstract_signature(dynamic_args)
         full_sig = merge_static_args(dynamic_sig, args, static_argnums)
 
-        if self.compile_options.experimental_capture:
+        if qml.capture.enabled():
             return trace_from_pennylane(
                 self.user_function, static_argnums, abstracted_axes, full_sig, kwargs
             )

frontend/catalyst/pipelines.py

@@ -68,9 +68,6 @@ class CompileOptions:
         disable_assertions (Optional[bool]): disables all assertions. Default is ``False``.
         seed (Optional[int]) : the seed for random operations in a qjit call.
             Default is None.
-        experimental_capture (bool): If set to ``True``,
-            use PennyLane's experimental program capture capabilities
-            to capture the function for compilation.
         circuit_transform_pipeline (Optional[dict[str, dict[str, str]]]):
             A dictionary that specifies the quantum circuit transformation pass pipeline order,
             and optionally arguments for each pass in the pipeline.
@@ -94,7 +91,6 @@ class CompileOptions:
     checkpoint_stage: Optional[str] = ""
     disable_assertions: Optional[bool] = False
     seed: Optional[int] = None
-    experimental_capture: Optional[bool] = False
     circuit_transform_pipeline: Optional[dict[str, dict[str, str]]] = None
     pass_plugins: Optional[Set[Path]] = None
     dialect_plugins: Optional[Set[Path]] = None