@qnode(dev, diff_method=diff_method, interface='jax')
def circuit(a, p):
    qml.RX(a, wires=0)
    U3(p[0], p[1], p[2], wires=0)
    return qml.expval(qml.PauliX(0))
@qnode(dev, diff_method=diff_method, interface='jax')
def circuit(a, p):
    qml.RX(a, wires=0)
    U3(p[0], p[1], p[2], wires=0)
    return qml.expval(qml.PauliX(0))