[BUG] qml.ShotAdaptiveOptimizer doesn’t work with qiskit.aer

[CatalinaAlbornoz]

Describe the bug
qml.ShotAdaptiveOptimizer doesn’t work with qiskit.aer

To Reproduce
Running this code

from pennylane import numpy as np
import pennylane as qml

coeffs = [2, 4, -1, 5, 2]
obs = [
  qml.X(1),
  qml.Z(1),
  qml.X(0) @ qml.X(1),
  qml.Y(0) @ qml.Y(1),
  qml.Z(0) @ qml.Z(1)
]
H = qml.Hamiltonian(coeffs, obs)
dev = qml.device('qiskit.aer', wires=2,shots=100)

@qml.qnode(dev)
def cost(weights):
    qml.StronglyEntanglingLayers(weights, wires=range(2))
    return qml.expval(H)

shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=2)
params = np.random.random(shape)
opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="weighted_random_sampling")
for i in range(60):
   params = opt.step(cost, params)
   print(f"Step {i}: cost = {cost(params):.2f}, shots_used = {opt.total_shots_used}")

Throws this error

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
.venv/adaptiveShot_pennylane_example.ipynb Cell 1 line 2
     22 opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="weighted_random_sampling")
     23 for i in range(60):
---> 24    params = opt.step(cost, params)
     25    print(f"Step {i}: cost = {cost(params):.2f}, shots_used = {opt.total_shots_used}")

File ~.venv/lib/python3.11/site-packages/pennylane/optimize/shot_adaptive.py:421, in ShotAdaptiveOptimizer.step(self, objective_fn, *args, **kwargs)
    417 self.total_shots_used += self.shots_used
    419 # compute the gradient, as well as the variance in the gradient,
    420 # using the number of shots determined by the array s.
--> 421 grads, grad_variances = self.compute_grad(objective_fn, args, kwargs)
    422 new_args = self.apply_grad(grads, args)
    424 if self.xi is None:

File ~.venv/lib/python3.11/site-packages/pennylane/optimize/shot_adaptive.py:358, in ShotAdaptiveOptimizer.compute_grad(self, objective_fn, args, kwargs)
    345 r"""Compute gradient of the objective function, as well as the variance of the gradient,
    346 at the given point.
    347 
   (...)
    355     :math:`\nabla f(x^{(t)})` and the variance of the gradient
    356 """
    357 if isinstance(objective_fn, qml.QNode) or hasattr(objective_fn, "device"):
--> 358     grads = self._single_shot_qnode_gradients(objective_fn, args, kwargs)
    359 else:
    360     raise ValueError(
    361         "The objective function must be encoded as a single QNode object for the shot "
    362         "adaptive optimizer. "
    363     )

File ~.venv/lib/python3.11/site-packages/pennylane/optimize/shot_adaptive.py:323, in ShotAdaptiveOptimizer._single_shot_qnode_gradients(self, qnode, args, kwargs)
    320     self.check_learning_rate(coeffs)
    322 if self.term_sampling == "weighted_random_sampling":
--> 323     return self.qnode_weighted_random_sampling(
    324         qnode, coeffs, observables, self.max_shots, self.trainable_args, *args, **kwargs
    325     )
    326 if self.term_sampling is not None:
    327     raise ValueError(
    328         f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
    329         "Only term_sampling='weighted_random_sampling' and "
    330         "term_sampling=None currently supported."
    331     )

File ~.venv/lib/python3.11/site-packages/pennylane/optimize/shot_adaptive.py:263, in ShotAdaptiveOptimizer.qnode_weighted_random_sampling(qnode, coeffs, observables, shots, argnums, *args, **kwargs)
    260 else:
    261     cost = qnode
--> 263 jacs = qml.jacobian(cost, argnum=argnums)(*args, **kwargs, shots=new_shots)
    265 if s == 1:
    266     jacs = [np.expand_dims(j, 0) for j in jacs]

File ~.venv/lib/python3.11/site-packages/pennylane/_grad.py:455, in jacobian.<locals>._jacobian_function(*args, **kwargs)
    449 if not _argnum:
    450     warnings.warn(
    451         "Attempted to differentiate a function with no trainable parameters. "
    452         "If this is unintended, please add trainable parameters via the "
    453         "'requires_grad' attribute or 'argnum' keyword."
    454     )
--> 455 jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
    457 return jac[0] if unpack else jac

File ~.venv/lib/python3.11/site-packages/pennylane/_grad.py:455, in <genexpr>(.0)
    449 if not _argnum:
    450     warnings.warn(
    451         "Attempted to differentiate a function with no trainable parameters. "
    452         "If this is unintended, please add trainable parameters via the "
    453         "'requires_grad' attribute or 'argnum' keyword."
    454     )
--> 455 jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
    457 return jac[0] if unpack else jac

File ~.venv/lib/python3.11/site-packages/autograd/wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     18 else:
     19     x = tuple(args[i] for i in argnum)
---> 20 return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File ~.venv/lib/python3.11/site-packages/autograd/differential_operators.py:60, in jacobian(fun, x)
     50 @unary_to_nary
     51 def jacobian(fun, x):
     52     """
     53     Returns a function which computes the Jacobian of `fun` with respect to
     54     positional argument number `argnum`, which must be a scalar or array. Unlike
   (...)
     58     (out1, out2, ...) then the Jacobian has shape (out1, out2, ..., in1, in2, ...).
     59     """
---> 60     vjp, ans = _make_vjp(fun, x)
     61     ans_vspace = vspace(ans)
     62     jacobian_shape = ans_vspace.shape + vspace(x).shape

File ~.venv/lib/python3.11/site-packages/autograd/core.py:10, in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

File ~.venv/lib/python3.11/site-packages/autograd/tracer.py:10, in trace(start_node, fun, x)
      8 with trace_stack.new_trace() as t:
      9     start_box = new_box(x, t, start_node)
---> 10     end_box = fun(start_box)
     11     if isbox(end_box) and end_box._trace == start_box._trace:
     12         return end_box._value, end_box._node

File ~.venv/lib/python3.11/site-packages/autograd/wrap_util.py:15, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f.<locals>.unary_f(x)
     13 else:
     14     subargs = subvals(args, zip(argnum, x))
---> 15 return fun(*subargs, **kwargs)

File ~.venv/lib/python3.11/site-packages/pennylane/optimize/shot_adaptive.py:258, in ShotAdaptiveOptimizer.qnode_weighted_random_sampling.<locals>.cost(*args, **kwargs)
    256 def cost(*args, **kwargs):
    257     # pylint: disable=cell-var-from-loop
--> 258     return qml.math.stack(qnode(*args, **kwargs))

File ~.venv/lib/python3.11/site-packages/pennylane/math/multi_dispatch.py:151, in multi_dispatch.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    148 interface = interface or get_interface(*dispatch_args)
    149 kwargs["like"] = interface
--> 151 return fn(*args, **kwargs)

File ~.venv/lib/python3.11/site-packages/pennylane/math/multi_dispatch.py:497, in stack(values, axis, like)
    468 """Stack a sequence of tensors along the specified axis.
    469 
    470 .. warning::
   (...)
    494        [5.00e+00, 8.00e+00, 1.01e+02]], dtype=float32)>
    495 """
    496 values = np.coerce(values, like=like)
--> 497 return np.stack(values, axis=axis, like=like)

File ~.venv/lib/python3.11/site-packages/autoray/autoray.py:81, in do(fn, like, *args, **kwargs)
     79 backend = _choose_backend(fn, args, kwargs, like=like)
     80 func = get_lib_fn(backend, fn)
---> 81 return func(*args, **kwargs)

File ~.venv/lib/python3.11/site-packages/numpy/core/shape_base.py:445, in stack(arrays, axis, out, dtype, casting)
    443 arrays = [asanyarray(arr) for arr in arrays]
    444 if not arrays:
--> 445     raise ValueError('need at least one array to stack')
    447 shapes = {arr.shape for arr in arrays}
    448 if len(shapes) != 1:

ValueError: need at least one array to stack

Expected behavior
qml.ShotAdaptiveOptimizer should work with qiskit.aer or we should have a better error message if we don't support this.

Desktop (please complete the following information):

Name: PennyLane
Version: 0.36.0
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:.venv/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-qiskit, PennyLane_Lightning

Platform info:           macOS-13.2.1-arm64-arm-64bit
Python version:          3.11.6
Numpy version:           1.26.4
Scipy version:           1.13.0
Installed devices:
- default.clifford (PennyLane-0.36.0)
- default.gaussian (PennyLane-0.36.0)
- default.mixed (PennyLane-0.36.0)
- default.qubit (PennyLane-0.36.0)
- default.qubit.autograd (PennyLane-0.36.0)
- default.qubit.jax (PennyLane-0.36.0)
- default.qubit.legacy (PennyLane-0.36.0)
- default.qubit.tf (PennyLane-0.36.0)
- default.qubit.torch (PennyLane-0.36.0)
- default.qutrit (PennyLane-0.36.0)
- default.qutrit.mixed (PennyLane-0.36.0)
- null.qubit (PennyLane-0.36.0)
- qiskit.aer (PennyLane-qiskit-0.36.0)
- qiskit.basicaer (PennyLane-qiskit-0.36.0)
- qiskit.basicsim (PennyLane-qiskit-0.36.0)
- qiskit.ibmq (PennyLane-qiskit-0.36.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.36.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.36.0)
- qiskit.remote (PennyLane-qiskit-0.36.0)
- lightning.qubit (PennyLane_Lightning-0.36.0)

Additional context
First reported here: https://discuss.pennylane.ai/t/shot-adaptive-optimizer-limitations/4839/3
This looks like it's due to qiskit.aer not supporting shot vectors.
We should have a better error message if anyone tries to use shot vectors with it.

dev = qml.device('qiskit.aer', wires=2,shots=(1,1,1))

@qml.qnode(dev)
def cost():
    return qml.expval(qml.Z(0))

cost()