Fix a few docs examples in qnn, fermi and math (#5763)

**Context:**

**Description of the Change:**
As title says.
Also reverts a few invalid changes from #5596 in `qml.kernels` docs.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com>

doc/code/qml_fermi.rst

@@ -51,18 +51,18 @@ The fermionic operators can be mapped to the qubit basis by using the
 Fermi sentences.
 
 >>> qml.jordan_wigner(qml.FermiA(1))
-(0.5*(PauliZ(wires=[0]) @ PauliX(wires=[1])))
-+ (0.5j*(PauliZ(wires=[0]) @ PauliY(wires=[1])))
+0.5 * (Z(0) @ X(1)) + 0.5j * (Z(0) @ Y(1))
 
 >>> qml.jordan_wigner(qml.FermiC(1) * qml.FermiA(1))
-((0.5+0j)*(Identity(wires=[1])))
-+ ((-0.5+0j)*(PauliZ(wires=[1])))
+(0.5+0j) * I(1) + (-0.5+0j) * Z(1)
 
 >>> f = 0.5 * qml.FermiC(1) * qml.FermiA(1) + 0.75 * qml.FermiC(2) * qml.FermiA(2)
 >>> qml.jordan_wigner(f)
-((0.625+0j)*(Identity(wires=[1])))
-+ ((-0.25+0j)*(PauliZ(wires=[1])))
-+ ((-0.375+0j)*(PauliZ(wires=[2])))
+(
+    (0.625+0j) * I(1)
+  + (-0.25+0j) * Z(1)
+  + (-0.375+0j) * Z(2)
+)
 
 FermiWord and FermiSentence
 ---------------------------

pennylane/fermi/conversion.py

@@ -211,9 +211,9 @@ def parity_transform(
     >>> parity_transform(w, n=6)
     (
         -0.25j * Y(0)
-    + (-0.25+0j) * (X(0) @ Z(1))
-    + (0.25+0j) * X(0)
-    + 0.25j * (Y(0) @ Z(1))
+      + (-0.25+0j) * (X(0) @ Z(1))
+      + (0.25+0j) * X(0)
+      + 0.25j * (Y(0) @ Z(1))
     )
 
     >>> parity_transform(w, n=6, ps=True)
@@ -374,10 +374,10 @@ def bravyi_kitaev(
     >>> w = qml.fermi.from_string('0+ 1-')
     >>> bravyi_kitaev(w, n=6)
     (
-    -0.25j * Y(0)
-    + (-0.25+0j) * (X(0) @ Z(1))
-    + (0.25+0j) * X(0)
-    + 0.25j * (Y(0) @ Z(1))
+        -0.25j * Y(0)
+      + (-0.25+0j) * (X(0) @ Z(1))
+      + (0.25+0j) * X(0)
+      + 0.25j * (Y(0) @ Z(1))
     )
 
     >>> bravyi_kitaev(w, n=6, ps=True)

pennylane/kernels/cost_functions.py

@@ -68,7 +68,7 @@ def polarity(
 
     .. code-block :: python
 
-        dev = qml.device('default.qubit')
+        dev = qml.device('default.qubit', wires=2)
         @qml.qnode(dev)
         def circuit(x1, x2):
             qml.templates.AngleEmbedding(x1, wires=dev.wires)
@@ -144,7 +144,7 @@ def target_alignment(
 
     .. code-block :: python
 
-        dev = qml.device('default.qubit')
+        dev = qml.device('default.qubit', wires=2)
         @qml.qnode(dev)
         def circuit(x1, x2):
             qml.templates.AngleEmbedding(x1, wires=dev.wires)

pennylane/kernels/utils.py

@@ -39,7 +39,7 @@ def square_kernel_matrix(X, kernel, assume_normalized_kernel=False):
 
     .. code-block :: python
 
-        dev = qml.device('default.qubit')
+        dev = qml.device('default.qubit', wires=2)
         @qml.qnode(dev)
         def circuit(x1, x2):
             qml.templates.AngleEmbedding(x1, wires=dev.wires)
@@ -103,7 +103,7 @@ def kernel_matrix(X1, X2, kernel):
 
     .. code-block :: python
 
-        dev = qml.device('default.qubit')
+        dev = qml.device('default.qubit', wires=2)
         @qml.qnode(dev)
         def circuit(x1, x2):
             qml.templates.AngleEmbedding(x1, wires=dev.wires)

pennylane/math/multi_dispatch.py

@@ -187,7 +187,7 @@ def block_diag(values, like=None):
     >>> t = [
     ...     np.array([[1, 2], [3, 4]]),
     ...     torch.tensor([[1, 2, 3], [-1, -6, -3]]),
-    ...     torch.tensor(5)
+    ...     torch.tensor([[5]])
     ... ]
     >>> qml.math.block_diag(t)
     tensor([[ 1,  2,  0,  0,  0,  0],
@@ -225,8 +225,9 @@ def concatenate(values, axis=0, like=None):
     >>> y = tf.Variable([0.1, 0.2, 0.3])
     >>> z = np.array([5., 8., 101.])
     >>> concatenate([x, y, z])
-    <tf.Tensor: shape=(3, 3), dtype=float32, numpy=
-    array([6.00e-01, 1.00e-01, 6.00e-01, 1.00e-01, 2.00e-01, 3.00e-01, 5.00e+00, 8.00e+00, 1.01e+02], dtype=float32)>
+    <tf.Tensor: shape=(9,), dtype=float32, numpy=
+    array([6.00e-01, 1.00e-01, 6.00e-01, 1.00e-01, 2.00e-01, 3.00e-01,
+           5.00e+00, 8.00e+00, 1.01e+02], dtype=float32)>
     """
 
     if like == "torch":
@@ -415,11 +416,12 @@ def get_trainable_indices(values, like=None):
 
     **Example**
 
+    >>> from pennylane import numpy as pnp
     >>> def cost_fn(params):
     ...     print("Trainable:", qml.math.get_trainable_indices(params))
     ...     return np.sum(np.sin(params[0] * params[1]))
-    >>> values = [np.array([0.1, 0.2], requires_grad=True),
-    ... np.array([0.5, 0.2], requires_grad=False)]
+    >>> values = [pnp.array([0.1, 0.2], requires_grad=True),
+    ... pnp.array([0.5, 0.2], requires_grad=False)]
     >>> cost_fn(values)
     Trainable: {0}
     tensor(0.0899685, requires_grad=True)
@@ -453,7 +455,7 @@ def ones_like(tensor, dtype=None):
 
     >>> x = torch.tensor([1., 2.])
     >>> ones_like(x)
-    tensor([1, 1])
+    tensor([1., 1.])
     >>> y = tf.Variable([[0], [5]])
     >>> ones_like(y, dtype=np.complex128)
     <tf.Tensor: shape=(2, 1), dtype=complex128, numpy=
@@ -563,7 +565,7 @@ def einsum(indices, *operands, like=None, optimize=None):
 
 
 def where(condition, x=None, y=None):
-    """Returns elements chosen from x or y depending on a boolean tensor condition,
+    r"""Returns elements chosen from x or y depending on a boolean tensor condition,
     or the indices of entries satisfying the condition.
 
     The input tensors ``condition``, ``x``, and ``y`` must all be broadcastable to the same shape.
@@ -594,12 +596,10 @@ def where(condition, x=None, y=None):
         The output format for ``x=None`` and ``y=None`` follows the respective
         interface and differs between TensorFlow and all other interfaces:
         For TensorFlow, the output is a tensor with shape
-        ``(num_true, len(condition.shape))`` where ``num_true`` is the number
+        ``(len(condition.shape), num_true)`` where ``num_true`` is the number
         of entries in ``condition`` that are ``True`` .
-        The entry at position ``(i, j)`` is the ``j`` th entry of the ``i`` th
-        index.
         For all other interfaces, the output is a tuple of tensor-like objects,
-        with the ``j`` th object indicating the ``j`` th entries of all indices.
+        with the ``j``\ th object indicating the ``j``\ th entries of all indices.
         Also see the examples below.
 
     **Example with single argument**
@@ -614,13 +614,10 @@ def where(condition, x=None, y=None):
     ``(2, 4)`` . For TensorFlow, on the other hand:
 
     >>> math.where(tf.constant(a) < 1)
-    tf.Tensor(
-    [[0 0]
-     [0 1]
-     [1 1]
-     [1 2]], shape=(4, 2), dtype=int64)
+    <tf.Tensor: shape=(2, 4), dtype=int64, numpy=
+    array([[0, 0, 1, 1],
+           [0, 1, 1, 2]])>
 
-    As we can see, the dimensions are swapped and the output is a single Tensor.
     Note that the number of dimensions of the output does *not* depend on the input
     shape, it is always two-dimensional.
 
@@ -767,7 +764,7 @@ def unwrap(values, max_depth=None):
     ...     return np.sum(np.sin(params))
     >>> params = np.array([0.1, 0.2, 0.3])
     >>> grad = autograd.grad(cost_fn)(params)
-    Unwrapped: [(0.1, <class 'float'>), (0.2, <class 'float'>), (0.3, <class 'float'>)]
+    Unwrapped: [(0.1, <class 'numpy.float64'>), (0.2, <class 'numpy.float64'>), (0.3, <class 'numpy.float64'>)]
     >>> print(grad)
     [0.99500417 0.98006658 0.95533649]
     """

pennylane/math/quantum.py

@@ -77,10 +77,11 @@ def circuit(weights):
     We can now compute the covariance matrix:
 
     >>> shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=3)
-    >>> weights = np.random.random(shape, requires_grad=True)
+    >>> weights = pnp.random.random(shape, requires_grad=True)
     >>> cov = qml.math.cov_matrix(circuit(weights), obs_list)
     >>> cov
-    tensor([[0.9275379 , 0.05233832], [0.05233832, 0.99335545]], requires_grad=True)
+    tensor([[0.98125435, 0.4905541 ],
+            [0.4905541 , 0.99920878]], requires_grad=True)
 
     Autodifferentiation is fully supported using all interfaces.
     Here we use autograd:
@@ -204,7 +205,7 @@ def reduce_dm(density_matrix, indices, check_state=False, c_dtype="complex128"):
      [0.+0.j 0.+0.j]]
 
     >>> z = tf.Variable([[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]], dtype=tf.complex128)
-    >>> reduce_dm(x, indices=[1])
+    >>> reduce_dm(z, indices=[1])
     tf.Tensor(
     [[1.+0.j 0.+0.j]
      [0.+0.j 0.+0.j]], shape=(2, 2), dtype=complex128)
@@ -268,24 +269,31 @@ def partial_trace(matrix, indices, c_dtype="complex128"):
 
     >>> x = np.array([[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]])
     >>> partial_trace(x, indices=[0])
-    array([[1, 0], [0, 0]])
+    array([[1.+0.j, 0.+0.j],
+           [0.+0.j, 0.+0.j]])
 
     We can also pass a batch of matrices ``x`` to the function and return the partial trace of each matrix with respect to each matrix's 0th index.
 
     >>> x = np.array([
-        [[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]],
-        [[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
-    ])
+    ... [[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]],
+    ... [[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
+    ... ])
     >>> partial_trace(x, indices=[0])
-    array([[[1, 0], [0, 0]], [[0, 0], [0, 1]]])
+    array([[[1.+0.j, 0.+0.j],
+            [0.+0.j, 0.+0.j]],
+           [[0.+0.j, 0.+0.j],
+            [0.+0.j, 1.+0.j]]])
 
     The partial trace can also be computed with respect to multiple indices within different frameworks such as TensorFlow.
 
     >>> x = tf.Variable([[[1, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]],
     ... [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]]], dtype=tf.complex128)
     >>> partial_trace(x, indices=[1])
     <tf.Tensor: shape=(2, 2, 2), dtype=complex128, numpy=
-    array([[[1.+0.j, 0.+0.j], [0.+0.j, 0.+0.j]], [[1.+0.j, 0.+0.j], [0.+0.j, 0.+0.j]]])>
+    array([[[1.+0.j, 0.+0.j],
+            [0.+0.j, 0.+0.j]],
+           [[0.+0.j, 0.+0.j],
+            [0.+0.j, 1.+0.j]]])>
 
     """
     # Autograd does not support same indices sum in backprop, and tensorflow
@@ -771,6 +779,7 @@ def vn_entanglement_entropy(
     The entanglement entropy between subsystems for a state vector can be returned as follows:
 
     >>> x = np.array([0, -1, 1, 0]) / np.sqrt(2)
+    >>> x = qml.math.dm_from_state_vector(x)
     >>> qml.math.vn_entanglement_entropy(x, indices0=[0], indices1=[1])
     0.6931471805599453
 
@@ -934,12 +943,12 @@ def relative_entropy(state0, state1, base=None, check_state=False, c_dtype="comp
     >>> rho = np.array([[0.3, 0], [0, 0.7]])
     >>> sigma = np.array([[0.5, 0], [0, 0.5]])
     >>> qml.math.relative_entropy(rho, sigma)
-    tensor(0.08228288, requires_grad=True)
+    0.08228288
 
     It is also possible to change the log base:
 
     >>> qml.math.relative_entropy(rho, sigma, base=2)
-    tensor(0.1187091, requires_grad=True)
+    0.1187091
 
     .. seealso:: :func:`pennylane.qinfo.transforms.relative_entropy`
     """

pennylane/math/utils.py

@@ -301,7 +301,7 @@ def get_deep_interface(value):
 
     >>> qml.math.asarray(x, like=qml.math.get_deep_interface(x))
     Array([[1, 2],
-       [3, 4]], dtype=int64)
+           [3, 4]], dtype=int64)
 
     """
     itr = value

pennylane/qnn/torch.py

@@ -186,7 +186,7 @@ def qnode(inputs, weights_0, weights_1, weights_2, weight_3, weight_4):
 
             init_method = {
                 "weights_0": torch.nn.init.normal_,
-                "weights_1": torch.nn.init.uniform,
+                "weights_1": torch.nn.init.uniform_,
                 "weights_2": torch.tensor([1., 2., 3.]),
                 "weight_3": torch.tensor(1.),  # scalar when shape is not an iterable and is <= 1
                 "weight_4": torch.tensor([1.]),
@@ -261,7 +261,7 @@ def qnode(inputs, weights_0, weights_1, weights_2, weight_3, weight_4):
             def qnode(inputs, weights):
                 qml.templates.AngleEmbedding(inputs, wires=range(n_qubits))
                 qml.templates.StronglyEntanglingLayers(weights, wires=range(n_qubits))
-                return qml.expval(qml.Z(0)), qml.expval(qml.Z(1))
+                return [qml.expval(qml.Z(0)), qml.expval(qml.Z(1))]
 
             weight_shapes = {"weights": (3, n_qubits, 3)}