Plxpr can capture operations

[albi3ro]

HIGHLY EXPERIMENTAL FEATURE

Context:

To improve integration with catalyst, we want a way to be able to natively capture the creation of quantum operations into an intermediate representation. This is one of the early PR's for this experimental push.

Description of the Change:

This PR adds a PLXPR metaclass that Operator uses. This allows for the capture of all Operator classes into jaxpr.

qml.capture.enable_plxpr()

def qfunc(a):
    qml.X(0)
    qml.IsingXX(a, wires=range(2))
    qml.ctrl(qml.adjoint(qml.X(0)), 1)

    0.5 * qml.X(0) @ qml.Y(1) + qml.Z(2)
    
jaxpr = jax.make_jaxpr(qfunc)(0)
jaxpr

{ lambda ; a:i32[]. let
    _:AbstractOperator() = PauliX[n_wires=1] 0
    _:AbstractOperator() = IsingXX[n_wires=2] a 0 1
    b:AbstractOperator() = PauliX[n_wires=1] 0
    c:AbstractOperator() = Adjoint b
    _:AbstractOperator() = Controlled[control_values=None work_wires=None] c 1
    d:AbstractOperator() = PauliX[n_wires=1] 0
    e:AbstractOperator() = SProd[id=None] 0.5 d
    f:AbstractOperator() = PauliY[n_wires=1] 1
    g:AbstractOperator() = Prod[id=None] e f
    h:AbstractOperator() = PauliZ[n_wires=1] 2
    _:AbstractOperator() = Sum[grouping_type=None id=None method=rlf] g h
  in () }

We can also return the jaxpr to normal pennylane qfunc behaviour via jax.core.eval_jaxpr:

with qml.queuing.AnnotatedQueue() as q:
    jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 0.1)
q.queue

[X(0),
 IsingXX(0.1, wires=[0, 1]),
 Controlled(Adjoint(X(0)), control_wires=[1]),
 (0.5 * X(0)) @ Y(1) + Z(2)]

Benefits:

Possible Drawbacks:

• Metaprogramming in python is an edge skill, and often not the best way to solve a problem. Messing around with things like this can often have unintended consequences down the line.

• With PLXPR, wires will be restricted to be jax-tracable friendly labels.

Related GitHub Issues:

[sc-61199]

[dwierichs]

Some comments :)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.67%. Comparing base (9c9b6ba) to head (34133fe).
Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5511      +/-   ##
==========================================
- Coverage   99.68%   99.67%   -0.01%     
==========================================
  Files         412      414       +2     
  Lines       38666    38522     -144     
==========================================
- Hits        38545    38398     -147     
- Misses        121      124       +3     

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

[dwierichs]

This is in really good shape now :) Thanks for the continued effort!

[dime10]

Love it 😍 The simplicity of it is really nice especially. I just have a few questions that came up while reading through it:

[dime10]

🙌