Error when using Control flow + QNode + Dynamic shapes

[dime10]

The following program raises an error:

from catalyst import *
import jax.numpy as jnp
import pennylane as qml

@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(n):
    a = jnp.zeros((n,), dtype=int)
    return a

@qjit
def main(n: int):

    def f():
        circuit(n)

    cond(n != 0)(f)()  # fails

    return circuit(n)  # succeeds

This occurs only when the three elements are used in tandem:

• control flow calling a subroutine
• the subroutine is a qnode
• the qnode returns a dynamically-shaped array