Remove "experimental_capture" kwarg from "qjit"

[rauletorresc]

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)

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.

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:

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:

        @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:

        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]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.88%. Comparing base (dc51613) to head (1f4cc72).
Report is 192 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1657      +/-   ##
==========================================
- Coverage   96.89%   96.88%   -0.01%     
==========================================
  Files          80       80              
  Lines        8879     8871       -8     
  Branches      843      841       -2     
==========================================
- Hits         8603     8595       -8     
  Misses        222      222              
  Partials       54       54              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dime10]

Nice, much cleaner this way!

Is the integration documented anywhere in Catalyst I wonder? 🤔 @josh146 what do you think?

[josh146]

Is the integration documented anywhere in Catalyst I wonder?

Only in two places as far as I am aware:

• In the catalyst.qjit docstring
• In the PennyLane docs: https://github.com/PennyLaneAI/pennylane/blob/master/doc/news/program_capture_sharp_bits.rst?plain=1#L403

[rauletorresc]

Is the integration documented anywhere in Catalyst I wonder?

Only in two places as far as I am aware:

• In the catalyst.qjit docstring
• In the PennyLane docs: https://github.com/PennyLaneAI/pennylane/blob/master/doc/news/program_capture_sharp_bits.rst?plain=1#L403

I'll update this accordingly

[isaacdevlugt]

Is the integration documented anywhere in Catalyst I wonder?

Only in two places as far as I am aware:

* In the `catalyst.qjit` docstring

* In the PennyLane docs: https://github.com/PennyLaneAI/pennylane/blob/master/doc/news/program_capture_sharp_bits.rst?plain=1#L403

And in the sharp bits page!

[isaacdevlugt]

Is the integration documented anywhere in Catalyst I wonder?

Only in two places as far as I am aware:

* In the `catalyst.qjit` docstring

* In the PennyLane docs: https://github.com/PennyLaneAI/pennylane/blob/master/doc/news/program_capture_sharp_bits.rst?plain=1#L403

And in the sharp bits page!

Oh lol @josh146 I didn't read the link 🤦 hahaha

[albi3ro]

👍 Nice to have a consistent way of entering this ecosystem now.

[isaacdevlugt]

Thanks Raul! This is really nice to have 😎