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:
albi3ro: > I have a question: I see that here `pennylane.numpy` is used in all cases instead of vanilla `numpy`. Is this fine or should we differentiate based on whether trainability is needed for each particular case?

@astralcai  Something we could potentially look into. Some of the variables we are creating might not actually end up being trainable, and could just being intermediaries.  But the optimizers are autograd-only, so we don't have to worry about mixing ml-framework types. I'm not as concerned about using autograd numpy as I would be in other parts of the code base.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d1f17ae1d..621b094b7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -291,6 +291,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* The `QNSPSAOptimizer` now properly handles differentiable parameters, resulting in being able to use it for more than one optimization step.
+  [(#5439)](https://github.com/PennyLaneAI/pennylane/pull/5439)
+
 * The `QNode` interface now resets if an error occurs during execution.
   [(#5449)](https://github.com/PennyLaneAI/pennylane/pull/5449)
 
diff --git a/pennylane/optimize/adaptive.py b/pennylane/optimize/adaptive.py
index 649ae0655..2b2dcd3c4 100644
--- 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
index ba6f1d82a..be183d285 100644
--- 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 @@ class QNSPSAOptimizer:
         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 @@ class QNSPSAOptimizer:
             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 @@ class QNSPSAOptimizer:
             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 @@ class QNSPSAOptimizer:
             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 @@ class QNSPSAOptimizer:
             )
             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 @@ class QNSPSAOptimizer:
             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 @@ class QNSPSAOptimizer:
         )
         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 @@ class QNSPSAOptimizer:
 
         # 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 @@ class QNSPSAOptimizer:
         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 @@ class QNSPSAOptimizer:
     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 @@ class QNSPSAOptimizer:
             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
index b117b3492..74fa75507 100644
--- 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
index 132ba2ae4..40c8bb37a 100644
--- a/tests/optimize/test_qnspsa.py
+++ b/tests/optimize/test_qnspsa.py
@@ -435,3 +435,26 @@ def test_template_no_adjoint():
     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)