Controlled subclasses no longer bind base class primitives on initialization

[mudit2812]

The constructors of subclasses of Controlled (CNOT, CRX, etc.) construct the base operator so that the parent Controlled class can be initialized correctly. However, in doing so, the primitive of the base operator is also binded. We don't want this, because calling constructors is part of the concrete implementation, and binding primitives can create abstract values (when using make_jaxpr, etc.), which we don't want.

Changes:

• Update all currently existing subclasses of Controlled to use type.__call__(base_op_class, ...) instead of base_op_class(...) to construct the base operator. This skips the call to op._primitive.bind and just directly creates the class we want to construct.

Downsides:

• This will need to be done for all subclasses of Controlled, so it is an additional detail that devs have to keep in mind when adding new controlled ops.

[sc-79769]

[albi3ro]

Do we have time to write some tests?

One way to test if this works if to do:

@DefaultQubitInterpreter(num_wires=3)
def f():
    qml.CNOT((0,1))
    return qml.state()

jax.jit(f)()

or something like that.

[mudit2812]

@albi3ro Thanks for the suggestion. I was stuck on what tests should be added 😄 . I'd like to localize the tests further, so I'm going to add a test that uses PlxprInterpreter instead of DefaultQubitInterpreter, and jax.make_jaxpr instead of jax.jit.

[albi3ro]

@albi3ro Thanks for the suggestion. I was stuck on what tests should be added 😄 . I'd like to localize the tests further, so I'm going to add a test that uses PlxprInterpreter instead of DefaultQubitInterpreter, and jax.make_jaxpr instead of jax.jit.

Even something like:

@qml.capture.PlxprInterpreter()
def f():
    qml.CNOT((0,1))
    qml.RX(0.1, (0,1))
    qml.Toffoli((0,1,2))

jax.make_jaxpr(f)()

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.64%. Comparing base (413fb32) to head (f65a59a).
Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6672      +/-   ##
==========================================
- Coverage   99.66%   99.64%   -0.02%     
==========================================
  Files         467      467              
  Lines       44095    44119      +24     
==========================================
+ Hits        43946    43963      +17     
- Misses        149      156       +7     

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

[lillian542]

I ask this without detailed information about the PlxprInterpreter - would there be a way to have a version of the test in assert_valid and check that the base isn't being bound for all Controlled operations?

No worries if that's tricky to do, just wondering if we can add an easy check to make sure we remember to do this with future operations.

[albi3ro]

I ask this without detailed information about the PlxprInterpreter - would there be a way to have a version of the test in assert_valid and check that the base isn't being bound for all Controlled operations?

No worries if that's tricky to do, just wondering if we can add an easy check to make sure we remember to do this with future operations.

Maybe:

def _validate(op):
      qml.capture.enable()
      try:
            @qml.capture.PlxprInterpreter()
            def f():
                op._unflatten(*op._flatten())
                
            jax.make_jaxpr(f)()
    finally:
        qml.capture.disable()

[mudit2812]

@albi3ro @lillian542 mind if we make the assert_valid addition a backlog story and tackle it there?

[albi3ro]

@albi3ro @lillian542 mind if we make the assert_valid addition a backlog story and tackle it there?

I could actually add this in #6343 in the convert_to_tape jaxpr-> tape function.

[lillian542]

LGTM then!