[BUG] Two-qubit Unitaries are not JIT-able when decomposing

Expected behavior

It would be nice to support the two-qubit unitary decomposition in JIT mode, which already works for the one-qubit decomposition. The issue is the use of if-statements that compare concrete values derived from function arguments.

Thus the following example succeeds:

import jax
import pennylane as qml

@jax.jit
@qml.qnode(qml.device("default.qubit", wires=2), interface="jax")
def f(U):
    for g in qml.QubitUnitary(U, wires=0, do_queue=False).decomposition():
        qml.apply(g)
    return qml.expval(qml.PauliZ(0))

f(jax.numpy.array([[1, 0], [0, 1]]))

While this one does not:

@jax.jit
@qml.qnode(qml.device("default.qubit", wires=2), interface="jax")
def f(U):
    for g in qml.QubitUnitary(U, wires=[0, 1], do_queue=False).decomposition():
        qml.apply(g)
    return qml.expval(qml.PauliZ(0))

f(jax.numpy.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]))

Actual behavior

The following tracing error is raised:

Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 
The error occurred while tracing the function f for jit. This concrete value was not available in Python because it depends on the value of the argument 'U'.

Additional information

Typically a way around this is first checking whether concrete values can be compared via qml.math.is_abstract in cases where the if-statement is merely an optimization. Alternatively, the output must be established via mathematical functions but this may not always be possible.

Related, but doesn't consider decomposition: #2247

Source code

No response

Tracebacks

No response

System information

Existing GitHub issues

[X] I have searched existing GitHub issues to make sure the issue does not already exist.

PennyLaneAI / pennylane