`qml.cond` does not trace the control flow op when used with an instantiated gate under capture

qml.capture.enable()

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=1))
def test(c: bool):
    qml.cond(c, qml.X(0))
    return qml.probs()

print(test(True), test(False))
print(test.mlir)

[0. 1.] [0. 1.]

  module @module_test {
    func.func public @test(%arg0: tensor<i1>) -> tensor<2xf64> attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
      %c0_i64 = arith.constant 0 : i64
      quantum.device shots(%c0_i64) ["/home/paul.wang/catalyst_new/catalyst/my_env/lib/python3.11/site-packages/pennylane_lightning/liblightning_qubit_catalyst.so", "LightningSimulator", "{'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
      %0 = quantum.alloc( 1) : !quantum.reg
      %1 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
      %out_qubits = quantum.custom "PauliX"() %1 : !quantum.bit
      %2 = quantum.insert %0[ 0], %out_qubits : !quantum.reg, !quantum.bit
      %3 = quantum.compbasis qreg %2 : !quantum.obs
      %4 = quantum.probs %3 : tensor<2xf64>
      quantum.dealloc %2 : !quantum.reg
      quantum.device_release
      return %4 : tensor<2xf64>
    }

Notice that the cond op is missing from the IR.

Note that this works if we provide just the qml.X class, instead of providing it with an instantiation:

qml.capture.enable()

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=1))
def test(c: bool):
    qml.cond(c, qml.X)(0)
    return qml.probs()

print(test(True), test(False))
print(test.mlir)

[0. 1.] [1. 0.]

    func.func public @test(%arg0: tensor<i1>) -> tensor<2xf64> attributes {diff_method = "parameter-shift", llvm.linkage = #llvm.linkage<internal>, qnode} {
      %c0_i64 = arith.constant 0 : i64
      quantum.device shots(%c0_i64) ["/home/paul.wang/catalyst_new/catalyst/my_env/lib/python3.11/site-packages/pennylane_lightning/liblightning_qubit_catalyst.so", "LightningSimulator", "{'mcmc': False, 'num_burnin': 0, 'kernel_name': None}"]
      %0 = quantum.alloc( 1) : !quantum.reg
      %extracted = tensor.extract %arg0[] : tensor<i1>
      %1 = scf.if %extracted -> (!quantum.reg) {
        %4 = quantum.extract %0[ 0] : !quantum.reg -> !quantum.bit
        %out_qubits = quantum.custom "PauliX"() %4 : !quantum.bit
        %5 = quantum.insert %0[ 0], %out_qubits : !quantum.reg, !quantum.bit
        scf.yield %5 : !quantum.reg
      } else {
        scf.yield %0 : !quantum.reg
      }
      %2 = quantum.compbasis qreg %1 : !quantum.obs
      %3 = quantum.probs %2 : tensor<2xf64>
      quantum.dealloc %1 : !quantum.reg
      quantum.device_release
      return %3 : tensor<2xf64>
    }

So this means one way to deal with this is to just disallow instantiated gates in cond in from_plxpr.

Start with the two qml.cond examples under qml.capture and compare their generated MLIR, then trace the from_plxpr handling of instantiated versus uninstantiated gates. Determine the intended behavior for instantiated gates, and verify completion by ensuring the conditional is represented in MLIR or the unsupported form is explicitly rejected.