Autograph exposing inconsistent conditional structures in qml.prod

[jsmallz333]

Related to #1824 , using PennyLane-Catalyst==0.11.0, PennyLane-qiskit==0.41.0, PennyLane_Lightning==0.41.0

The following:

@qml.prod
def ctrlThings():
    qml.Hadamard(wires=0)
    qml.ctrl(qml.Adder,control=[0],control_values=[1])(6,x_wires=[1,2])
    qml.PauliX(wires=0)

    return qml.state()

@qml.qjit(autograph=True, abstracted_axes=("n",)) 
@qml.qnode(qml.device("lightning.qubit", wires=3)) 
def ctrl_test():
    ctrlThings()
    return qml.state()
out=ctrl_test()

produces this TypeError error:

TypeError: Conditional requires a consistent return structure across all branches! Got PyTreeDef((*, CustomNode(Prod[()], [CustomNode(PauliX[(Wires([0]), ())], []), CustomNode(ControlledOp[(Wires([0]), (True,), Wires([]))], [CustomNode(Adder[(('k', 6), ('mod', 4), ('work_wires', Wires([])), ('x_wires', Wires([1, 2])))], [])])]))) and PyTreeDef((*, CustomNode(ControlledOp[(Wires([0]), (True,), Wires([]))], [CustomNode(Adder[(('k', 6), ('mod', 4), ('work_wires', Wires([])), ('x_wires', Wires([1, 2])))], [])]))).

This error is not raised when the autograph is set to false. Additionally, not using the qml.prod such as the following will not produce the error:

@qml.qjit(autograph=True, abstracted_axes=("n",)) 
@qml.qnode(qml.device("lightning.qubit", wires=3)) 
def ctrl_test():
    qml.Hadamard(wires=0)
    qml.ctrl(qml.Adder,control=[0],control_values=[1])(6,x_wires=[1,2])
    qml.PauliX(wires=0)
    return qml.state()
out=ctrl_test()

[dime10]

I believe this is another symptom of #1824, so I'll mark this as duplicate. We should have both fixed soon!

[jsmallz333]

Figured as much, thanks David!