def test_simple_jacobian():
    'Test the use of jax.jaxrev'
    dev = qml.device('default.mixed', wires=2)

    @qml.qnode(dev, interface='jax', diff_method='parameter-shift')
    def circuit(weights):
        qml.RX(weights[0], wires=0)
        qml.RY(weights[1], wires=1)
        return qml.expval((qml.PauliZ(0) @ qml.PauliZ(1)))
    weights = jnp.array([0.1, 0.2])
    grads = jax.jacrev(circuit)(weights)
    assert ('DeviceArray' in grads.__repr__())
    assert (grads.shape == (2,))
    np.testing.assert_allclose(grads, np.array([(- 0.09784342), (- 0.19767685)]))
Test the use of jax.jaxrev
def test_simple_jacobian():
    
    dev = qml.device('default.mixed', wires=2)

    @qml.qnode(dev, interface='jax', diff_method='parameter-shift')
    def circuit(weights):
        qml.RX(weights[0], wires=0)
        qml.RY(weights[1], wires=1)
        return qml.expval((qml.PauliZ(0) @ qml.PauliZ(1)))
    weights = jnp.array([0.1, 0.2])
    grads = jax.jacrev(circuit)(weights)
    assert ('DeviceArray' in grads.__repr__())
    assert (grads.shape == (2,))
    np.testing.assert_allclose(grads, np.array([(- 0.09784342), (- 0.19767685)]))