How to use the pennylane.device function in PennyLane

def test_invalid_qubit_unitary(self):
        """Test that an invalid unitary operation is not allowed"""
        dev = qml.device("forest.wavefunction", wires=3)

        def circuit(Umat):
            """Test QNode"""
            qml.QubitUnitary(Umat, wires=[0, 1])
            return qml.expval(qml.PauliZ(0))

        circuit1 = qml.QNode(circuit, dev)
        with pytest.raises(ValueError, match="must be a square matrix"):
            circuit1(np.array([[0, 1]]))

        circuit1 = qml.QNode(circuit, dev)
        with pytest.raises(ValueError, match="must be unitary"):
            circuit1(np.array([[1, 1], [1, 1]]))

        circuit1 = qml.QNode(circuit, dev)
        with pytest.raises(ValueError, match=r"must be 2\^Nx2\^N"):

def test_qubit_unitary(self, shots, qvm, compiler):
        """Test that an arbitrary unitary operation works"""
        dev1 = qml.device("forest.qvm", device="3q-qvm", shots=shots)
        dev2 = qml.device("forest.qvm", device="9q-square-qvm", shots=shots)

        def circuit():
            """Reference QNode"""
            qml.Hadamard(wires=0)
            qml.CNOT(wires=[0, 1])
            qml.QubitUnitary(U2, wires=[0, 1])
            return qml.expval(qml.PauliZ(0))

        circuit1 = qml.QNode(circuit, dev1)
        circuit2 = qml.QNode(circuit, dev2)

        out_state = U2 @ np.array([1, 0, 0, 1]) / np.sqrt(2)
        obs = np.kron(np.array([[1, 0], [0, -1]]), I)

        self.assertAllAlmostEqual(
            circuit1(), np.vdot(out_state, obs @ out_state), delta=3 / np.sqrt(shots)

def test_nonzero_shots(self):
        """Test that the fock plugin provides correct result for high shot number"""
        shots = 10 ** 2
        dev = qml.device("strawberryfields.fock", wires=1, cutoff_dim=10, shots=shots)

        @qml.qnode(dev)
        def circuit(x):
            qml.Displacement(x, 0, wires=0)
            return qml.expval(qml.NumberOperator(0))

        x = 1

        runs = []
        for _ in range(100):
            runs.append(circuit(x))

        expected_var = np.sqrt(1 / shots)
        assert np.allclose(np.mean(runs), x, atol=expected_var)

def test_CVOperation_with_heisenberg_and_no_parshift(self, name, tol):
        """An integration test for Gaussian CV gates that have a Heisenberg representation
        but cannot be differentiated using the parameter-shift method themselves
        (for example, they may accept no parameters, or have no gradient recipe).

        Tests that the parameter-shift method can still be used with other gates in the circuit.
        """
        gaussian_dev = qml.device("default.gaussian", wires=2)

        cls = getattr(qml.ops, name)
        if cls.supports_heisenberg and (not cls.supports_parameter_shift):
            U = np.array([[0.51310276+0.81702166j, 0.13649626+0.22487759j],
                          [0.26300233+0.00556194j, -0.96414101-0.03508489j]])

            if cls.num_wires <= 0:
                w = list(range(2))
            else:
                w = list(range(cls.num_wires))

            def circuit(x):
                qml.Displacement(x, 0, wires=0)

                if cls.par_domain == "A":
                    cls(U, wires=w)

def test_paulix_pauliy(self, theta, phi, varphi, tol):
        """Test that a tensor product involving PauliX and PauliY works correctly"""
        dev = qml.device("default.qubit", wires=3)
        dev.reset()

        dev.apply("RX", wires=[0], par=[theta])
        dev.apply("RX", wires=[1], par=[phi])
        dev.apply("RX", wires=[2], par=[varphi])
        dev.apply("CNOT", wires=[0, 1], par=[])
        dev.apply("CNOT", wires=[1, 2], par=[])

        res = dev.expval(["PauliX", "PauliY"], [[0], [2]], [[], [], []])
        expected = np.sin(theta) * np.sin(phi) * np.sin(varphi)

        assert np.allclose(res, expected, atol=tol, rtol=0)

def test_qnode_gradient_repeated_gate_parameters(self, tol):
        """Tests that repeated use of a free parameter in a
        multi-parameter gate yield correct gradients."""
        par = [0.8, 1.3]

        def qf(x, y):
            qml.RX(np.pi / 4, wires=[0])
            qml.Rot(y, x, 2 * x, wires=[0])
            return qml.expval(qml.PauliX(0))

        dev = qml.device("default.qubit", wires=1)
        q = qml.QNode(qf, dev)
        grad_A = q.jacobian(par, method="A")
        grad_F = q.jacobian(par, method="F")

        # the different methods agree
        assert np.allclose(grad_A, grad_F, atol=tol, rtol=0)

##############################################################################
# To estimate the overlap of the ground state with the post-selected state, one could
# directly make use of the measurement samples. However, since we want to optimize the cost
# function, it is useful to express everything in terms of expectation values through
# Bayes' theorem:
#
# .. math::
#   |\langle b | \Psi \rangle|^2=
#   P( \mathrm{sys}=\mathrm{ground}\,|\, \mathrm{anc} = \mathrm{ground}) =
#   P( \mathrm{all}=\mathrm{ground})/P( \mathrm{anc}=\mathrm{ground})
#
# To evaluate the two probabilities appearing on the right hand side of the previous equation
# we initialize a ``default.qubit`` device and we define two different ``qnode`` circuits.

dev = qml.device("default.qubit", wires=tot_qubits)

@qml.qnode(dev)
def global_ground(weights):
    # Circuit gates
    full_circuit(weights)
    # Projector on the global ground state
    P = np.zeros((2 ** tot_qubits, 2 ** tot_qubits))
    P[0, 0] = 1.0
    return qml.expval(qml.Hermitian(P, wires=range(tot_qubits)))

@qml.qnode(dev)
def ancilla_ground(weights):
    # Circuit gates
    full_circuit(weights)
    # Projector on the ground state of the ancillary system
    P_anc = np.zeros((2 ** m, 2 ** m))

# Hilbert space, we minimize the cost
# :math:`C = 1 - \frac{1}{2}D_{\mathrm{hs}}(A, B)`.
#
# To set up the "quantum part" of the cost function in PennyLane, we have
# to create a quantum node. Here, the quantum node is simulated on
# PennyLane's ``'default.qubit'`` backend.
#
# .. note:: One could also connect the
#           quantum node to a hardware backend to find out if the noise of a
#           physical implementation still allows us to train the embedding.
#

n_features = 2
n_qubits = 2 * n_features + 1

dev = qml.device("default.qubit", wires=n_qubits)


######################################################################
# We use a SWAP test to measure the overlap
# :math:`|\langle \psi | \phi \rangle|^2` between two quantum feature
# states :math:`|\psi\rangle` and :math:`|\phi\rangle`, prepared by a
# ``QAOAEmbedding`` with weights ``q_weights``.
#


@qml.qnode(dev)
def swap_test(q_weights, x1, x2):

    # load the two inputs into two different registers
    QAOAEmbedding(features=x1, weights=q_weights, wires=[1, 2])
    QAOAEmbedding(features=x2, weights=q_weights, wires=[3, 4])

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.ticker import LinearLocator, FormatStrFormatter
from matplotlib import cm


import pennylane as qml
from pennylane import numpy as np
from pennylane._optimize import GradientDescentOptimizer

dev = qml.device('strawberryfields.fock', wires=2, cutoff_dim=10)

@qml.qfunc(dev)
def func(x, y):

    qml.FockState(1, [0])
    qml.Beamsplitter(x, y, [0, 1])

    return qml.expectation.Fock(1)


fig = plt.figure(figsize = (5, 3))
ax = fig.gca(projection='3d')

# Landscape.
X = np.arange(-3.1, 3.1, 0.1)
Y = np.arange(-3.1, 3.1, 0.1)

# Finally, the ensemble model chooses the QPU which is most confident about its prediction
# (i.e., the class with the highest overall probability over all QPUs) and uses that to make a
# prediction.
#
# .. figure:: /demonstrations/ensemble_multi_qpu/ensemble_diagram.png
#    :width: 50%
#    :align: center
#
# Quantum nodes
# ^^^^^^^^^^^^^
#
# We begin by defining the two quantum devices and the circuits to be run on them.

n_wires = 4

dev0 = qml.device("forest.qvm", device="4q-qvm")
dev1 = qml.device("qiskit.aer", wires=4)
devs = [dev0, dev1]

##############################################################################
# .. note::
#    If you have access to Rigetti hardware, you can swap out ``forest.qvm`` for ``forest.qpu``
#    and specify the hardware device to run on. Users with access to the IBM Q Experience can
#    swap ``qiskit.aer`` for ``qiskit.ibmq`` and specify their chosen backend (see `here
#    `__).
#
# .. warning::
#    Rigetti's QVM and Quil Compiler services must be running for this tutorial to execute. They
#    can be installed by consulting the `Rigetti documentation
#    `__ or, for users with Docker, by running:
#
#    .. code-block:: bash