@pytest.mark.parametrize('diff_method', ['parameter-shift', 'finite-diff'])
def test_differentiable_expand(diff_method):
    'Test that operation and nested tapes expansion\n    is differentiable'

    class U3(qml.U3):

        def expand(self):
            (theta, phi, lam) = self.data
            wires = self.wires
            with JacobianTape() as tape:
                qml.Rot(lam, theta, (- lam), wires=wires)
                qml.PhaseShift((phi + lam), wires=wires)
            return tape
    dev = qml.device('default.mixed', wires=1)
    a = jnp.array(0.1)
    p = jnp.array([0.1, 0.2, 0.3])

    @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))
    res = circuit(a, p)
    expected = (((np.cos(a) * np.cos(p[1])) * np.sin(p[0])) + (np.sin(a) * ((np.cos(p[2]) * np.sin(p[1])) + ((np.cos(p[0]) * np.cos(p[1])) * np.sin(p[2])))))
    tol = 1e-05
    assert np.allclose(res, expected, atol=tol, rtol=0)
    res = jax.grad(circuit, argnums=1)(a, p)
    expected = np.array([(np.cos(p[1]) * ((np.cos(a) * np.cos(p[0])) - ((np.sin(a) * np.sin(p[0])) * np.sin(p[2])))), (((np.cos(p[1]) * np.cos(p[2])) * np.sin(a)) - (np.sin(p[1]) * ((np.cos(a) * np.sin(p[0])) + ((np.cos(p[0]) * np.sin(a)) * np.sin(p[2]))))), (np.sin(a) * (((np.cos(p[0]) * np.cos(p[1])) * np.cos(p[2])) - (np.sin(p[1]) * np.sin(p[2]))))])
    assert np.allclose(res, expected, atol=tol, rtol=0)
Test that operation and nested tapes expansion
is differentiable
@pytest.mark.parametrize('diff_method', ['parameter-shift', 'finite-diff'])
def test_differentiable_expand(diff_method):
    'Test that operation and nested tapes expansion\n    is differentiable'

    class U3(qml.U3):

        def expand(self):
            (theta, phi, lam) = self.data
            wires = self.wires
            with JacobianTape() as tape:
                qml.Rot(lam, theta, (- lam), wires=wires)
                qml.PhaseShift((phi + lam), wires=wires)
            return tape
    dev = qml.device('default.mixed', wires=1)
    a = jnp.array(0.1)
    p = jnp.array([0.1, 0.2, 0.3])

    @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))
    res = circuit(a, p)
    expected = (((np.cos(a) * np.cos(p[1])) * np.sin(p[0])) + (np.sin(a) * ((np.cos(p[2]) * np.sin(p[1])) + ((np.cos(p[0]) * np.cos(p[1])) * np.sin(p[2])))))
    tol = 1e-05
    assert np.allclose(res, expected, atol=tol, rtol=0)
    res = jax.grad(circuit, argnums=1)(a, p)
    expected = np.array([(np.cos(p[1]) * ((np.cos(a) * np.cos(p[0])) - ((np.sin(a) * np.sin(p[0])) * np.sin(p[2])))), (((np.cos(p[1]) * np.cos(p[2])) * np.sin(a)) - (np.sin(p[1]) * ((np.cos(a) * np.sin(p[0])) + ((np.cos(p[0]) * np.sin(a)) * np.sin(p[2]))))), (np.sin(a) * (((np.cos(p[0]) * np.cos(p[1])) * np.cos(p[2])) - (np.sin(p[1]) * np.sin(p[2]))))])
    assert np.allclose(res, expected, atol=tol, rtol=0)