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

[paul0403]

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.

[paul0403]

If we were to raise an error, we can say something like

   catalyst.cond() must take in callables, but got f{instantiated_gate}. See [documentation].