@pytest.mark.parametrize('dev_name,diff_method', [('default.mixed', 'finite-diff'), ('default.qubit.autograd', 'parameter-shift')])
def test_transform(dev_name, diff_method, monkeypatch, tol):
    'Test an example transform'
    monkeypatch.setattr(qml.operation.Operation, 'do_check_domain', False)
    dev = qml.device(dev_name, wires=1)

    @qnode(dev, interface='jax', diff_method=diff_method)
    def circuit(weights):
        op1 = qml.RY(weights[0], wires=0)
        op2 = qml.RX(weights[1], wires=0)
        return qml.expval(qml.PauliZ(wires=0))
    weights = np.array([0.32, 0.543])
    a = np.array(0.5)

    def loss(weights, a):
        new_circuit = qtransform(circuit, a)
        res = new_circuit(weights)
        res2 = circuit(jnp.sin(weights))
        return (res + res2)
    res = loss(weights, a)
    grad = jax.grad(loss, argnums=[0, 1])(weights, a)
    assert (len(grad) == 2)
    assert (grad[0].shape == weights.shape)
    assert (grad[1].shape == a.shape)
    tol = 1e-05
    assert np.allclose(res, 1.8244501889992706, atol=tol, rtol=0)
    assert np.allclose(grad[0], [(- 0.26610258), (- 0.47053553)], atol=tol, rtol=0)
    assert np.allclose(grad[1], 0.06486032, atol=tol, rtol=0)
Test an example transform
@pytest.mark.parametrize('dev_name,diff_method', [('default.mixed', 'finite-diff'), ('default.qubit.autograd', 'parameter-shift')])
def test_transform(dev_name, diff_method, monkeypatch, tol):
    
    monkeypatch.setattr(qml.operation.Operation, 'do_check_domain', False)
    dev = qml.device(dev_name, wires=1)

    @qnode(dev, interface='jax', diff_method=diff_method)
    def circuit(weights):
        op1 = qml.RY(weights[0], wires=0)
        op2 = qml.RX(weights[1], wires=0)
        return qml.expval(qml.PauliZ(wires=0))
    weights = np.array([0.32, 0.543])
    a = np.array(0.5)

    def loss(weights, a):
        new_circuit = qtransform(circuit, a)
        res = new_circuit(weights)
        res2 = circuit(jnp.sin(weights))
        return (res + res2)
    res = loss(weights, a)
    grad = jax.grad(loss, argnums=[0, 1])(weights, a)
    assert (len(grad) == 2)
    assert (grad[0].shape == weights.shape)
    assert (grad[1].shape == a.shape)
    tol = 1e-05
    assert np.allclose(res, 1.8244501889992706, atol=tol, rtol=0)
    assert np.allclose(grad[0], [(- 0.26610258), (- 0.47053553)], atol=tol, rtol=0)
    assert np.allclose(grad[1], 0.06486032, atol=tol, rtol=0)