Repository: PennyLaneAI/pennylane

Issue:
[BUG] `QNSPSAOptimizer` with Autograd changes differentiable parameters type to vanilla NumPy array
### Expected behavior

Updating parameters with `opt = qml.QNSPSAOptimizer` should not change the type of the parameters.

### Actual behavior

The parameters change from `<class 'pennylane.numpy.tensor.tensor'>` to `<class 'numpy.ndarray'>` after an optimization step.

### Additional information

Stems from this forum post: https://discuss.pennylane.ai/t/qnspsa-error-using-a-dataset/4153/12
Related PR: https://github.com/PennyLaneAI/pennylane/pull/4920

### Source code

```shell
import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit")
@qml.qnode(dev)
def cost(params):
    qml.RX(params[0], wires=0)
    qml.CRY(params[1], wires=[0, 1])
    return qml.expval(qml.Z(0) @ qml.Z(1))

params = np.random.rand(2)
opt = qml.QNSPSAOptimizer(stepsize=5e-2)

for _ in range(3):
    print(type(params))
    params, loss = opt.step_and_cost(cost, params)
    #params = np.array(params) # this fixes the error
```


### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[17], line 16
     14 for _ in range(3):
     15     print(type(params))
---> 16     params, loss = opt.step_and_cost(cost, params)
     17     #params = np.array(params) # this fixes the error

File ~/.virtualenvs/pennylane-stable/lib/python3.11/site-packages/pennylane/optimize/qnspsa.py:185, in QNSPSAOptimizer.step_and_cost(self, cost, *args, **kwargs)
    172 def step_and_cost(self, cost, *args, **kwargs):
    173     r"""Update trainable parameters with one step of the optimizer and return
    174     the corresponding objective function value after the step.
    175 
   (...)
    183         function output prior to the step
    184     """
--> 185     params_next = self._step_core(cost, args, kwargs)
    187     if not self.blocking:
    188         loss_curr = cost(*args, **kwargs)

File ~/.virtualenvs/pennylane-stable/lib/python3.11/site-packages/pennylane/optimize/qnspsa.py:215, in QNSPSAOptimizer._step_core(self, cost, args, kwargs)
    213 grad_tapes, grad_dirs = self._get_spsa_grad_tapes(cost, args, kwargs)
    214 # metric_tapes contains 4 tapes for tensor estimation
--> 215 metric_tapes, tensor_dirs = self._get_tensor_tapes(cost, args, kwargs)
    216 all_grad_tapes += grad_tapes
    217 all_metric_tapes += metric_tapes

File ~/.virtualenvs/pennylane-stable/lib/python3.11/site-packages/pennylane/optimize/qnspsa.py:408, in QNSPSAOptimizer._get_tensor_tapes(self, cost, args, kwargs)
    406     args_list[2][index] = arg + self.finite_diff_step * (-dir1 + dir2)
    407     args_list[3][index] = arg - self.finite_diff_step * dir1
--> 408 dir_vecs = (np.concatenate(dir1_list), np.concatenate(dir2_list))
    409 tapes = [
    410     self._get_overlap_tape(cost, args, args_finite_diff, kwargs)
    411     for args_finite_diff in args_list
    412 ]
    414 return tapes, dir_vecs

ValueError: need at least one array to concatenate
```


### System information

```shell
Name: PennyLane
Version: 0.35.1
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/isaac/.virtualenvs/pennylane-stable/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-14.4-arm64-arm-64bit
Python version:          3.11.8
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.35.1)
- default.gaussian (PennyLane-0.35.1)
- default.mixed (PennyLane-0.35.1)
- default.qubit (PennyLane-0.35.1)
- default.qubit.autograd (PennyLane-0.35.1)
- default.qubit.jax (PennyLane-0.35.1)
- default.qubit.legacy (PennyLane-0.35.1)
- default.qubit.tf (PennyLane-0.35.1)
- default.qubit.torch (PennyLane-0.35.1)
- default.qutrit (PennyLane-0.35.1)
- null.qubit (PennyLane-0.35.1)
- lightning.qubit (PennyLane_Lightning-0.35.1)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/optimize/adaptive.py b/pennylane/optimize/adaptive.py
--- a/pennylane/optimize/adaptive.py
+++ b/pennylane/optimize/adaptive.py
@@ -79,12 +79,12 @@ class AdaptiveOptimizer:
     adaptive circuit for the :math:`\text{H}_3^+` cation.
 
     >>> import pennylane as qml
-    >>> from pennylane import numpy as np
+    >>> from pennylane import numpy as pnp
 
     The molecule is defined and the Hamiltonian is computed with:
 
     >>> symbols = ["H", "H", "H"]
-    >>> geometry = np.array([[0.01076341, 0.04449877, 0.0],
+    >>> geometry = pnp.array([[0.01076341, 0.04449877, 0.0],
     ...                      [0.98729513, 1.63059094, 0.0],
     ...                      [1.87262415, -0.00815842, 0.0]], requires_grad=False)
     >>> H, qubits = qml.qchem.molecular_hamiltonian(symbols, geometry, charge = 1)
diff --git a/pennylane/optimize/qnspsa.py b/pennylane/optimize/qnspsa.py
--- a/pennylane/optimize/qnspsa.py
+++ b/pennylane/optimize/qnspsa.py
@@ -14,7 +14,7 @@
 """Quantum natural SPSA optimizer"""
 import warnings
 from scipy.linalg import sqrtm
-import numpy as np
+from pennylane import numpy as pnp
 import pennylane as qml
 
 
@@ -92,7 +92,8 @@ class QNSPSAOptimizer:
     Once constructed, the qnode can be passed directly to the ``step`` or ``step_and_cost``
     function of the optimizer.
 
-    >>> params = np.random.rand(2)
+    >>> from pennylane import numpy as pnp
+    >>> params = pnp.random.rand(2)
     >>> opt = QNSPSAOptimizer(stepsize=5e-2)
     >>> for i in range(51):
     >>>     params, loss = opt.step_and_cost(cost, params)
@@ -140,8 +141,8 @@ def __init__(
         self.k = 1
         self.resamplings = resamplings
         self.blocking = blocking
-        self.last_n_steps = np.zeros(history_length)
-        self.rng = np.random.default_rng(seed)
+        self.last_n_steps = pnp.zeros(history_length)
+        self.rng = pnp.random.default_rng(seed)
 
     def step(self, cost, *args, **kwargs):
         """Update trainable arguments with one step of the optimizer.
@@ -156,7 +157,7 @@ def step(self, cost, *args, **kwargs):
             kwargs : variable length of keyword arguments for the qnode
 
         Returns:
-            np.array: the new variable values after step-wise update :math:`x^{(t+1)}`
+            pnp.array: the new variable values after step-wise update :math:`x^{(t+1)}`
         """
         if self.blocking:
             warnings.warn(
@@ -201,7 +202,7 @@ def _step_core(self, cost, args, kwargs):
             kwargs : variable length of keyword arguments for the qnode
 
         Returns:
-            np.array: the new variable values :math:`x^{(t+1)}` before the blocking condition
+            pnp.array: the new variable values :math:`x^{(t+1)}` before the blocking condition
             is applied.
         """
         all_grad_tapes = []
@@ -236,7 +237,7 @@ def _step_core(self, cost, args, kwargs):
             self._post_process_grad(raw_results[2 * i : 2 * i + 2], all_grad_dirs[i])
             for i in range(self.resamplings)
         ]
-        grads = np.array(grads)
+        grads = pnp.array(grads)
         metric_tensors = [
             self._post_process_tensor(
                 raw_results[2 * self.resamplings + 4 * i : 2 * self.resamplings + 4 * i + 4],
@@ -244,9 +245,9 @@ def _step_core(self, cost, args, kwargs):
             )
             for i in range(self.resamplings)
         ]
-        metric_tensors = np.array(metric_tensors)
-        grad_avg = np.mean(grads, axis=0)
-        tensor_avg = np.mean(metric_tensors, axis=0)
+        metric_tensors = pnp.array(metric_tensors)
+        grad_avg = pnp.mean(grads, axis=0)
+        tensor_avg = pnp.mean(metric_tensors, axis=0)
 
         self._update_tensor(tensor_avg)
         params_next = self._get_next_params(args, grad_avg)
@@ -283,7 +284,7 @@ def _post_process_tensor(self, tensor_raw_results, tensor_dirs):
             been concatenated
 
         Returns:
-            np.array: estimated Fubini-Study metric tensor
+            pnp.array: estimated Fubini-Study metric tensor
         """
         tensor_raw_results = [result.squeeze() for result in tensor_raw_results]
         # For each element of tensor_raw_results, the first dimension is the measured probability in
@@ -297,8 +298,8 @@ def _post_process_tensor(self, tensor_raw_results, tensor_dirs):
         )
         return (
             -(
-                np.tensordot(tensor_dirs[0], tensor_dirs[1], axes=0)
-                + np.tensordot(tensor_dirs[1], tensor_dirs[0], axes=0)
+                pnp.tensordot(tensor_dirs[0], tensor_dirs[1], axes=0)
+                + pnp.tensordot(tensor_dirs[1], tensor_dirs[0], axes=0)
             )
             * tensor_finite_diff
             / (8 * self.finite_diff_step**2)
@@ -319,12 +320,12 @@ def _get_next_params(self, args, gradient):
 
         # params_vec and grad_vec group multiple inputs into the same vector to solve the
         # linear equation
-        params_vec = np.concatenate([param.reshape(-1) for param in params])
-        grad_vec = np.concatenate([grad.reshape(-1) for grad in gradient])
+        params_vec = pnp.concatenate([param.reshape(-1) for param in params])
+        grad_vec = pnp.concatenate([grad.reshape(-1) for grad in gradient])
 
-        new_params_vec = np.linalg.solve(
+        new_params_vec = pnp.linalg.solve(
             self.metric_tensor,
-            (-self.stepsize * grad_vec + np.matmul(self.metric_tensor, params_vec)),
+            (-self.stepsize * grad_vec + pnp.matmul(self.metric_tensor, params_vec)),
         )
         # reshape single-vector new_params_vec into new_params, to match the input params
         params_split_indices = []
@@ -332,7 +333,7 @@ def _get_next_params(self, args, gradient):
         for param in params:
             tmp += param.size
             params_split_indices.append(tmp)
-        new_params = np.split(new_params_vec, params_split_indices)
+        new_params = pnp.split(new_params_vec, params_split_indices)
         new_params_reshaped = [new_params[i].reshape(params[i].shape) for i in range(len(params))]
 
         next_args = []
@@ -375,12 +376,12 @@ def _get_spsa_grad_tapes(self, cost, args, kwargs):
     def _update_tensor(self, tensor_raw):
         def get_tensor_moving_avg(metric_tensor):
             if self.metric_tensor is None:
-                self.metric_tensor = np.identity(metric_tensor.shape[0])
+                self.metric_tensor = pnp.identity(metric_tensor.shape[0])
             return self.k / (self.k + 1) * self.metric_tensor + 1 / (self.k + 1) * metric_tensor
 
         def regularize_tensor(metric_tensor):
-            tensor_reg = np.real(sqrtm(np.matmul(metric_tensor, metric_tensor)))
-            return (tensor_reg + self.reg * np.identity(metric_tensor.shape[0])) / (1 + self.reg)
+            tensor_reg = pnp.real(sqrtm(pnp.matmul(metric_tensor, metric_tensor)))
+            return (tensor_reg + self.reg * pnp.identity(metric_tensor.shape[0])) / (1 + self.reg)
 
         tensor_avg = get_tensor_moving_avg(tensor_raw)
         tensor_regularized = regularize_tensor(tensor_avg)
@@ -405,7 +406,7 @@ def _get_tensor_tapes(self, cost, args, kwargs):
             args_list[1][index] = arg + self.finite_diff_step * dir1
             args_list[2][index] = arg + self.finite_diff_step * (-dir1 + dir2)
             args_list[3][index] = arg - self.finite_diff_step * dir1
-        dir_vecs = (np.concatenate(dir1_list), np.concatenate(dir2_list))
+        dir_vecs = (pnp.concatenate(dir1_list), pnp.concatenate(dir2_list))
         tapes = [
             self._get_overlap_tape(cost, args, args_finite_diff, kwargs)
             for args_finite_diff in args_list


Test patch:
diff --git a/tests/optimize/test_adaptive.py b/tests/optimize/test_adaptive.py
--- a/tests/optimize/test_adaptive.py
+++ b/tests/optimize/test_adaptive.py
@@ -118,7 +118,7 @@ def test_step_and_cost_nodrain(circuit, energy_ref, pool):
     """Test that step_and_cost function returns the correct results when drain_pool is False."""
     opt = qml.AdaptiveOptimizer()
     for _ in range(4):
-        circuit, energy, __ = opt.step_and_cost(circuit, pool, drain_pool=False)
+        circuit, energy, _ = opt.step_and_cost(circuit, pool, drain_pool=False)
 
     circuit()
     selected_excitations = [op.wires for op in circuit.tape.operations[1:]]
diff --git a/tests/optimize/test_qnspsa.py b/tests/optimize/test_qnspsa.py
--- a/tests/optimize/test_qnspsa.py
+++ b/tests/optimize/test_qnspsa.py
@@ -435,3 +435,26 @@ def cost(params):
     opt = qml.QNSPSAOptimizer(stepsize=5e-2)
     assert opt.step_and_cost(cost, params)  # just checking it runs without error
     assert not qml.RandomLayers.has_adjoint
+
+
+def test_workflow_integration():
+    """Test that the optimizer can optimize a workflow."""
+
+    num_qubits = 2
+    dev = qml.device("default.qubit", wires=num_qubits)
+
+    @qml.qnode(dev)
+    def cost(params):
+        qml.RX(params[0], wires=0)
+        qml.CRY(params[1], wires=[0, 1])
+        return qml.expval(qml.Z(0) @ qml.Z(1))
+
+    params = qml.numpy.array([0.5, 0.5], requires_grad=True)
+    opt = qml.optimize.QNSPSAOptimizer(stepsize=5e-2)
+    for _ in range(51):
+        params, loss = opt.step_and_cost(cost, params)
+
+    assert qml.math.allclose(loss, -1, atol=1e-3)
+    # compare sine of params and target params as could converge to params + 2* np.pi
+    target_params = np.array([np.pi, 0])
+    assert qml.math.allclose(np.sin(params), np.sin(target_params), atol=1e-2)