Repository: PennyLaneAI/pennylane
Pull request #1139 (2021-03-16 08:21:31)

Add shot adaptive optimizer
**Context:** The shots refactor allows for shot adaptive optimizers to be supported.

**Description of the Change:**

Adds a shot adaptive optimizer `qml.ShotAdaptiveOptimizer`, a gradient-descent optimizer based on both iCANS1 and Rosalin (see https://quantum-journal.org/papers/q-2020-05-11-263 and https://arxiv.org/abs/2004.06252).
```pycon
>>> coeffs = [2, 4, -1, 5, 2]
>>> obs = [
...   qml.PauliX(1),
...   qml.PauliZ(1),
...   qml.PauliX(0) @ qml.PauliX(1),
...   qml.PauliY(0) @ qml.PauliY(1),
...   qml.PauliZ(0) @ qml.PauliZ(1)
... ]
>>> H = qml.Hamiltonian(coeffs, obs)
>>> dev = qml.device("default.qubit", wires=2, analytic=False)
>>> cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
```

Once constructed, the cost function can be passed directly to the
optimizer's `step` function:

```pycon
>>> params = qml.init.strong_ent_layers_uniform(n_layers=2, n_wires=2)
>>> opt = qml.ShotAdaptiveOptimizer(min_shots=10)
>>> for i in range(60):
...    params = opt.step(cost, params)
...    print(f"Step {i}: cost = {cost(params)}, shots_used = {opt.shots_used}")
Step 0: cost = -5.686, shots_used = 240
Step 1: cost = -2.983999999999999, shots_used = 336
Step 2: cost = -4.974, shots_used = 624
Step 3: cost = -5.534, shots_used = 1054
Step 4: cost = -6.5, shots_used = 1798
Step 5: cost = -6.684, shots_used = 2942
Step 6: cost = -6.992, shots_used = 4350
Step 7: cost = -6.970000000000001, shots_used = 5814
Step 8: cost = -6.998, shots_used = 7230
Step 9: cost = -6.686, shots_used = 9006
Step 10: cost = -6.85, shots_used = 11286
Step 11: cost = -6.628000000000001, shots_used = 14934
Step 12: cost = -6.862, shots_used = 17934
Step 13: cost = -7.188, shots_used = 22950
Step 14: cost = -6.9879999999999995, shots_used = 28302
Step 15: cost = -7.380000000000001, shots_used = 34134
Step 16: cost = -7.656, shots_used = 41022
Step 17: cost = -7.212, shots_used = 48918
Step 18: cost = -7.532, shots_used = 56286
Step 19: cost = -7.462000000000001, shots_used = 63822
Step 20: cost = -7.306000000000001, shots_used = 72534
Step 21: cost = -7.234, shots_used = 82014
Step 22: cost = -7.314, shots_used = 92838
```

**Benefits:**

- Supports distributing the shot budget across both parameters and Hamiltonian terms
- Support for weighted random sampling

**Possible Drawbacks:**

- Until #1131 is merged, each optimization step will be slightly slower due to a large number of redundant Jacobian computations.

- Only permits QNodes and ExpvalCost objective functions, as the optimizer needs access to the underlying device to change the shot number.

- Only works with Autograd, like all other optimizers

**Related GitHub Issues:** n/a


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
josh146: Edit: just realised I forgot to write tests - will do that now!
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=h1) Report
> Merging [#1139](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=desc) (96ce8a6) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c701627d2c5ee24adc752c0d4013e9a5d777e259?el=desc) (c701627) will **decrease** coverage by `0.77%`.
> The diff coverage is `13.51%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #1139      +/-   ##
==========================================
- Coverage   97.78%   97.00%   -0.78%     
==========================================
  Files         156      157       +1     
  Lines       11900    12011     +111     
==========================================
+ Hits        11636    11651      +15     
- Misses        264      360      +96     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pennylane/optimize/shot\_adaptive.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139/diff?src=pr&el=tree#diff-cGVubnlsYW5lL29wdGltaXplL3Nob3RfYWRhcHRpdmUucHk=) | `12.72% <12.72%> (ø)` | |
| [pennylane/optimize/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139/diff?src=pr&el=tree#diff-cGVubnlsYW5lL29wdGltaXplL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=footer). Last update [c701627...96ce8a6](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1139?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

albi3ro: Also, `step_and_cost` also needs to be optimizer specific, instead of just inheriting the `GradientDescentOptimizer` version.  Right now, attempting to use `opt.step_and_cost(cost, params)` gives

```
TypeError: compute_grad() got an unexpected keyword argument 'grad_fn'
```
josh146: Thanks @albi3ro and @thisac! Have incorporated your comments and suggestions, and the PR is better off for it.

> Also, step_and_cost also needs to be optimizer specific, instead of just inheriting the `GradientDescentOptimizer` version. Right now, attempting to use `opt.step_and_cost(cost, params)` gives
>
> ```
> TypeError: compute_grad() got an unexpected keyword argument 'grad_fn'
> ```

Good catch. I didn't implement `step_and_cost` here for two reasons:

- I wondered if it would be confusing for users to track a cost that is being estimated with a different number of shots
- More technically, this optimizer uses `qml.jacobian`, and not `qml.grad`. We need to update `qml.jacobian` to store the forward pass value like we did with `qml.grad`.

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index b7bd88e9f..2ed70fb58 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,51 @@
 
 <h3>New features since last release</h3>
 
+* Adds a new optimizer `qml.ShotAdaptiveOptimizer`, a gradient-descent optimizer where
+  the shot rate is adaptively calculated using the variances of the parameter-shift gradient.
+  [(#1139)](https://github.com/PennyLaneAI/pennylane/pull/1139)
+
+  By keeping a running average of the parameter-shift gradient and the *variance* of the
+  parameter-shift gradient, this optimizer frugally distributes a shot budget across the partial
+  derivatives of each parameter.
+
+  In addition, if computing the expectation value of a Hamiltonian, weighted random sampling can be
+  used to further distribute the shot budget across the local terms from which the Hamiltonian is
+  constructed.
+
+  This optimizer is based on both the [iCANS1](https://quantum-journal.org/papers/q-2020-05-11-263)
+  and [Rosalin](https://arxiv.org/abs/2004.06252) shot adaptive optimizers.
+
+  ```pycon
+  >>> coeffs = [2, 4, -1, 5, 2]
+  >>> obs = [
+  ...   qml.PauliX(1),
+  ...   qml.PauliZ(1),
+  ...   qml.PauliX(0) @ qml.PauliX(1),
+  ...   qml.PauliY(0) @ qml.PauliY(1),
+  ...   qml.PauliZ(0) @ qml.PauliZ(1)
+  ... ]
+  >>> H = qml.Hamiltonian(coeffs, obs)
+  >>> dev = qml.device("default.qubit", wires=2, shots=100)
+  >>> cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+  >>> params = qml.init.strong_ent_layers_uniform(n_layers=2, n_wires=2)
+  ```
+
+  Once constructed, the cost function can be passed directly to the optimizer's `step` method.  The
+  attribute `opt.total_shots_used` can be used to track the number of shots per iteration.
+
+  ```pycon
+  >>> opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+  >>> for i in range(5):
+  ...    params = opt.step(cost, params)
+  ...    print(f"Step {i}: cost = {cost(params):.2f}, shots_used = {opt.total_shots_used}")
+  Step 0: cost = -5.68, shots_used = 240
+  Step 1: cost = -2.98, shots_used = 336
+  Step 2: cost = -4.97, shots_used = 624
+  Step 3: cost = -5.53, shots_used = 1054
+  Step 4: cost = -6.50, shots_used = 1798
+  ```
+
 * Added the `SingleExcitation` two-qubit operation, which is useful for quantum 
   chemistry applications. [(#1121)](https://github.com/PennyLaneAI/pennylane/pull/1121)
   
diff --git a/doc/introduction/optimizers.rst b/doc/introduction/optimizers.rst
index dd7811174..a1f3594cd 100644
--- a/doc/introduction/optimizers.rst
+++ b/doc/introduction/optimizers.rst
@@ -18,7 +18,8 @@ NumPy Interface
 ^^^^^^^^^^^^^^^
 
 When using the standard NumPy interface, PennyLane offers some built-in optimizers.
-Some of these are specific to quantum optimization, such as the :mod:`~.QNGOptimizer`.
+Some of these are specific to quantum optimization, such as the :mod:`~.QNGOptimizer`,
+:mod:`~.RotosolveOptimizer`, :mod:`~.RotoselectOptimizer`, and :mod:`~.ShotAdaptiveOptimizer`.
 
 :html:`<div class="summary-table">`
 
@@ -34,6 +35,7 @@ Some of these are specific to quantum optimization, such as the :mod:`~.QNGOptim
     ~pennylane.RMSPropOptimizer
     ~pennylane.RotosolveOptimizer
     ~pennylane.RotoselectOptimizer
+    ~pennylane.ShotAdaptiveOptimizer
 
 :html:`</div>`
 
diff --git a/pennylane/math/tensorbox.py b/pennylane/math/tensorbox.py
index abf3358ed..9a015ba4c 100644
--- a/pennylane/math/tensorbox.py
+++ b/pennylane/math/tensorbox.py
@@ -105,7 +105,7 @@ class TensorBox(abc.ABC):
 
         namespace = tensor.__class__.__module__.split(".")[0]
         box = None
-        if namespace == "jax":
+        if namespace in ("jax", "jaxlib"):
             from .jax_box import JaxBox
 
             box = JaxBox.__new__(JaxBox, tensor)
diff --git a/pennylane/optimize/__init__.py b/pennylane/optimize/__init__.py
index 98efc779d..d927095c9 100644
--- a/pennylane/optimize/__init__.py
+++ b/pennylane/optimize/__init__.py
@@ -26,6 +26,7 @@ from .rms_prop import RMSPropOptimizer
 from .qng import QNGOptimizer
 from .rotosolve import RotosolveOptimizer
 from .rotoselect import RotoselectOptimizer
+from .shot_adaptive import ShotAdaptiveOptimizer
 
 
 # Optimizers to display in the docs
@@ -39,4 +40,5 @@ __all__ = [
     "QNGOptimizer",
     "RotosolveOptimizer",
     "RotoselectOptimizer",
+    "ShotAdaptiveOptimizer",
 ]
diff --git a/pennylane/optimize/shot_adaptive.py b/pennylane/optimize/shot_adaptive.py
new file mode 100644
index 000000000..39151d314
--- /dev/null
+++ b/pennylane/optimize/shot_adaptive.py
@@ -0,0 +1,518 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Shot adaptive optimizer"""
+# pylint: disable=too-many-instance-attributes,too-many-arguments,too-many-branches
+from scipy.stats import multinomial
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+from .gradient_descent import GradientDescentOptimizer
+
+
+class ShotAdaptiveOptimizer(GradientDescentOptimizer):
+    r"""Optimizer where the shot rate is adaptively calculated using the variances of the parameter-shift
+    gradient.
+
+    By keeping a running average of the parameter-shift gradient and the *variance*
+    of the parameter-shift gradient, this optimizer frugally distributes a shot
+    budget across the partial derivatives of each parameter.
+
+    In addition, if computing the expectation value of a Hamiltonian using
+    :class:`~.ExpvalCost`, weighted random sampling can be used to further
+    distribute the shot budget across the local terms from which the Hamiltonian
+    is constructed.
+
+    .. note::
+
+        The shot adaptive optimizer only supports single QNodes or :class:`~.ExpvalCost` objects as
+        objective functions. The bound device must also be instantiated with a finite number
+        of shots.
+
+    Args:
+        min_shots (int): The minimum number of shots used to estimate the expectations
+            of each term in the Hamiltonian. Note that this must be larger than 2 for the variance
+            of the gradients to be computed.
+        mu (float): The running average constant :math:`\mu \in [0, 1]`. Used to control how quickly the
+            number of shots recommended for each gradient component changes.
+        b (float): Regularization bias. The bias should be kept small, but non-zero.
+        term_sampling (str): The random sampling algorithm to multinomially distribute the shot budget
+            across terms in the Hamiltonian expectation value.
+            Currently, only ``"weighted_random_sampling"`` is supported.
+            Only takes effect if the objective function provided is an instance of :class:`~.ExpvalCost`.
+            Set this argument to ``None`` to turn off random sampling of Hamiltonian terms.
+        stepsize (float): The learning rate :math:`\eta`. The learning rate *must* be such
+            that :math:`\eta < 2/L = 2/\sum_i|c_i|`, where:
+
+            * :math:`L \leq \sum_i|c_i|` is the bound on the `Lipschitz constant
+              <https://en.wikipedia.org/wiki/Lipschitz_continuity>`__ of the variational quantum
+              algorithm objective function, and
+
+            * :math:`c_i` are the coefficients of the Hamiltonian used in the objective function.
+
+    **Example**
+
+    For VQE/VQE-like problems, the objective function for the optimizer can be
+    realized as an :class:`~.ExpvalCost` object, constructed using a :class:`~.Hamiltonian`.
+
+    >>> coeffs = [2, 4, -1, 5, 2]
+    >>> obs = [
+    ...   qml.PauliX(1),
+    ...   qml.PauliZ(1),
+    ...   qml.PauliX(0) @ qml.PauliX(1),
+    ...   qml.PauliY(0) @ qml.PauliY(1),
+    ...   qml.PauliZ(0) @ qml.PauliZ(1)
+    ... ]
+    >>> H = qml.Hamiltonian(coeffs, obs)
+    >>> dev = qml.device("default.qubit", wires=2, shots=100)
+    >>> cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+
+    Once constructed, the cost function can be passed directly to the
+    optimizer's ``step`` method. The attributes ``opt.shots_used`` and
+    ``opt.total_shots_used`` can be used to track the number of shots per
+    iteration, and across the life of the optimizer, respectively.
+
+    >>> params = qml.init.strong_ent_layers_uniform(n_layers=2, n_wires=2)
+    >>> opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+    >>> for i in range(60):
+    ...    params = opt.step(cost, params)
+    ...    print(f"Step {i}: cost = {cost(params):.2f}, shots_used = {opt.total_shots_used}")
+    Step 0: cost = -5.69, shots_used = 240
+    Step 1: cost = -2.98, shots_used = 336
+    Step 2: cost = -4.97, shots_used = 624
+    Step 3: cost = -5.53, shots_used = 1054
+    Step 4: cost = -6.50, shots_used = 1798
+    Step 5: cost = -6.68, shots_used = 2942
+    Step 6: cost = -6.99, shots_used = 4350
+    Step 7: cost = -6.97, shots_used = 5814
+    Step 8: cost = -7.00, shots_used = 7230
+    Step 9: cost = -6.69, shots_used = 9006
+    Step 10: cost = -6.85, shots_used = 11286
+    Step 11: cost = -6.63, shots_used = 14934
+    Step 12: cost = -6.86, shots_used = 17934
+    Step 13: cost = -7.19, shots_used = 22950
+    Step 14: cost = -6.99, shots_used = 28302
+    Step 15: cost = -7.38, shots_used = 34134
+    Step 16: cost = -7.66, shots_used = 41022
+    Step 17: cost = -7.21, shots_used = 48918
+    Step 18: cost = -7.53, shots_used = 56286
+    Step 19: cost = -7.46, shots_used = 63822
+    Step 20: cost = -7.31, shots_used = 72534
+    Step 21: cost = -7.23, shots_used = 82014
+    Step 22: cost = -7.31, shots_used = 92838
+
+    .. UsageDetails::
+
+        The shot adaptive optimizer is based on the iCANS1 optimizer by
+        `Kübler et al. (2020) <https://quantum-journal.org/papers/q-2020-05-11-263/>`__, and works
+        as follows:
+
+        1. The initial step of the optimizer is performed with some specified minimum
+           number of shots, :math:`s_{min}`, for all partial derivatives.
+
+        2. The parameter-shift rule is then used to estimate the gradient :math:`g_i` with :math:`s_i` shots
+           for each parameter :math:`\theta_i`, parameters, as well as the variances
+           :math:`v_i` of the estimated gradients.
+
+        3. Gradient descent is performed for each parameter :math:`\theta_i`, using
+           the pre-defined learning rate :math:`\eta` and the gradient information :math:`g_i`:
+           :math:`\theta_i \rightarrow \theta_i - \eta g_i`.
+
+        4. A maximum shot number is set by maximizing the improvement in the expected gain per shot.
+           For a specific parameter value, the improvement in the expected gain per shot
+           is then calculated via
+
+           .. math::
+               \gamma_i = \frac{1}{s_i} \left[ \left(\eta - \frac{1}{2} L\eta^2\right)
+                           g_i^2 - \frac{L\eta^2}{2s_i}v_i \right],
+
+           where:
+
+           * :math:`L \leq \sum_i|c_i|` is the bound on the `Lipschitz constant
+             <https://en.wikipedia.org/wiki/Lipschitz_continuity>`__ of the variational quantum algorithm objective function,
+
+           * :math:`c_i` are the coefficients of the Hamiltonian, and
+
+           * :math:`\eta` is the learning rate, and *must* be bound such that :math:`\eta < 2/L`
+             for the above expression to hold.
+
+        5. Finally, the new values of :math:`s_{i+1}` (shots for partial derivative of parameter
+           :math:`\theta_i`) is given by:
+
+           .. math::
+
+               s_{i+1} = \frac{2L\eta}{2-L\eta}\left(\frac{v_i}{g_i^2}\right)\propto
+                     \frac{v_i}{g_i^2}.
+
+        In addition to the above, to counteract the presence of noise in the system, a
+        running average of :math:`g_i` and :math:`s_i` (:math:`\chi_i` and :math:`\xi_i` respectively)
+        are used when computing :math:`\gamma_i` and :math:`s_i`.
+
+        For more details, see:
+
+        * Andrew Arrasmith, Lukasz Cincio, Rolando D. Somma, and Patrick J. Coles. "Operator Sampling
+          for Shot-frugal Optimization in Variational Algorithms." `arXiv:2004.06252
+          <https://arxiv.org/abs/2004.06252>`__ (2020).
+
+        * Jonas M. Kübler, Andrew Arrasmith, Lukasz Cincio, and Patrick J. Coles. "An Adaptive Optimizer
+          for Measurement-Frugal Variational Algorithms." `Quantum 4, 263
+          <https://quantum-journal.org/papers/q-2020-05-11-263/>`__ (2020).
+    """
+
+    def __init__(
+        self, min_shots, term_sampling="weighted_random_sampling", mu=0.99, b=1e-6, stepsize=0.07
+    ):
+        self.term_sampling = term_sampling
+        self.trainable_args = set()
+
+        # hyperparameters
+        self.min_shots = min_shots
+        self.mu = mu  # running average constant
+        self.b = b  # regularization bias
+        self.lipschitz = None
+
+        self.shots_used = 0
+        """int: number of shots used on the current iteration"""
+
+        self.total_shots_used = 0
+        """int: total number of shots used across all iterations"""
+
+        # total number of iterations
+        self.k = 0
+        # Number of shots per parameter
+        self.s = None
+        # maximum number of shots required to evaluate across all parameters
+        self.max_shots = None
+
+        # Running average of the parameter gradients
+        self.chi = None
+        # Running average of the variance of the parameter gradients
+        self.xi = None
+
+        super().__init__(stepsize=stepsize)
+
+    @staticmethod
+    def weighted_random_sampling(qnodes, coeffs, shots, argnums, *args, **kwargs):
+        """Returns an array of length ``shots`` containing single-shot estimates
+        of the Hamiltonian gradient. The shots are distributed randomly over
+        the terms in the Hamiltonian, as per a multinomial distribution.
+
+        Args:
+            qnodes (Sequence[.QNode]): Sequence of QNodes, each one when evaluated
+                returning the corresponding expectation value of a term in the Hamiltonian.
+            coeffs (Sequence[float]): Sequences of coefficients corresponding to
+                each term in the Hamiltonian. Must be the same length as ``qnodes``.
+            shots (int): The number of shots used to estimate the Hamiltonian expectation
+                value. These shots are distributed over the terms in the Hamiltonian,
+                as per a Multinomial distribution.
+            argnums (Sequence[int]): the QNode argument indices which are trainable
+            *args: Arguments to the QNodes
+            **kwargs: Keyword arguments to the QNodes
+
+        Returns:
+            array[float]: the single-shot gradients of the Hamiltonian expectation value
+        """
+
+        # determine the shot probability per term
+        prob_shots = np.abs(coeffs) / np.sum(np.abs(coeffs))
+
+        # construct the multinomial distribution, and sample
+        # from it to determine how many shots to apply per term
+        si = multinomial(n=shots, p=prob_shots)
+        shots_per_term = si.rvs()[0]
+
+        grads = []
+
+        for h, c, p, s in zip(qnodes, coeffs, prob_shots, shots_per_term):
+
+            # if the number of shots is 0, do nothing
+            if s == 0:
+                continue
+
+            # set the QNode device shots
+            h.device.shots = [(1, s)]
+
+            jacs = []
+            for i in argnums:
+                j = qml.jacobian(h, argnum=i)(*args, **kwargs)
+
+                if s == 1:
+                    j = np.expand_dims(j, 0)
+
+                # Divide each term by the probability per shot. This is
+                # because we are sampling one at a time.
+                jacs.append(c * j / p)
+
+            grads.append(jacs)
+
+        return [np.concatenate(i) for i in zip(*grads)]
+
+    @staticmethod
+    def check_device(dev):
+        r"""Verifies that the device used by the objective function is non-analytic.
+
+        Args:
+            dev (.Device): the device to verify
+
+        Raises:
+            ValueError: if the device is analytic
+        """
+        if dev.analytic:
+            raise ValueError(
+                "The Rosalin optimizer can only be used with devices "
+                "that estimate expectation values with a finite number of shots."
+            )
+
+    def check_learning_rate(self, coeffs):
+        r"""Verifies that the learning rate is less than 2 over the Lipschitz constant,
+        where the Lipschitz constant is given by :math:`\sum |c_i|` for Hamiltonian
+        coefficients :math:`c_i`.
+
+        Args:
+            coeffs (Sequence[float]): the coefficients of the terms in the Hamiltonian
+
+        Raises:
+            ValueError: if the learning rate is large than :math:`2/\sum |c_i|`
+        """
+        self.lipschitz = np.sum(np.abs(coeffs))
+
+        if self._stepsize > 2 / self.lipschitz:
+            raise ValueError(f"The learning rate must be less than {2 / self.lipschitz}")
+
+    def _single_shot_expval_gradients(self, expval_cost, args, kwargs):
+        """Compute the single shot gradients of an ExpvalCost object"""
+
+        qnodes = expval_cost.qnodes
+        coeffs = expval_cost.hamiltonian.coeffs
+        device = qnodes[0].device
+
+        self.check_device(device)
+        original_shots = device.shots
+
+        if self.lipschitz is None:
+            self.check_learning_rate(coeffs)
+
+        try:
+            if self.term_sampling == "weighted_random_sampling":
+                grads = self.weighted_random_sampling(
+                    qnodes, coeffs, self.max_shots, self.trainable_args, *args, **kwargs
+                )
+            elif self.term_sampling is None:
+                device.shots = [(1, self.max_shots)]
+                # We iterate over each trainable argument, rather than using
+                # qml.jacobian(expval_cost), to take into account the edge case where
+                # different arguments have different shapes and cannot be stacked.
+                grads = [
+                    qml.jacobian(expval_cost, argnum=i)(*args, **kwargs)
+                    for i in self.trainable_args
+                ]
+            else:
+                raise ValueError(
+                    f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
+                    "Only term_sampling='weighted_random_sampling' and "
+                    "term_sampling=None currently supported."
+                )
+        finally:
+            device.shots = original_shots
+
+        return grads
+
+    def _single_shot_qnode_gradients(self, qnode, args, kwargs):
+        """Compute the single shot gradients of a QNode."""
+        device = qnode.device
+
+        self.check_device(qnode.device)
+        original_shots = device.shots
+
+        if self.lipschitz is None:
+            self.check_learning_rate(1)
+
+        try:
+            device.shots = [(1, self.max_shots)]
+            grads = [qml.jacobian(qnode, argnum=i)(*args, **kwargs) for i in self.trainable_args]
+        finally:
+            device.shots = original_shots
+
+        return grads
+
+    def compute_grad(
+        self, objective_fn, args, kwargs
+    ):  # pylint: disable=signature-differs,arguments-differ
+        r"""Compute gradient of the objective function, as well as the variance of the gradient,
+        at the given point.
+
+        Args:
+            objective_fn (function): the objective function for optimization
+            args: arguments to the objective function
+            kwargs: keyword arguments to the objective function
+
+        Returns:
+            tuple[array[float], array[float]]: a tuple of NumPy arrays containing the gradient
+            :math:`\nabla f(x^{(t)})` and the variance of the gradient
+        """
+        if isinstance(objective_fn, qml.ExpvalCost):
+            grads = self._single_shot_expval_gradients(objective_fn, args, kwargs)
+        elif isinstance(objective_fn, qml.QNode) or hasattr(objective_fn, "device"):
+            grads = self._single_shot_qnode_gradients(objective_fn, args, kwargs)
+        else:
+            raise ValueError(
+                "The objective function must either be encoded as a single QNode or "
+                "an ExpvalCost object for the Shot adaptive optimizer. "
+            )
+
+        # grads will have dimension [max(self.s), *params.shape]
+        # For each parameter, we want to truncate the number of shots to self.s[idx],
+        # and take the mean and the variance.
+        gradients = []
+        gradient_variances = []
+
+        for i, grad in enumerate(grads):
+            p_ind = np.ndindex(*grad.shape[1:])
+
+            g = np.zeros_like(grad[0])
+            s = np.zeros_like(grad[0])
+
+            for idx in p_ind:
+                grad_slice = grad[(slice(0, self.s[i][idx]),) + idx]
+                g[idx] = np.mean(grad_slice)
+                s[idx] = np.var(grad_slice, ddof=1)
+
+            gradients.append(g)
+            gradient_variances.append(s)
+
+        return gradients, gradient_variances
+
+    def step(self, objective_fn, *args, **kwargs):
+        """Update trainable arguments with one step of the optimizer.
+
+        Args:
+            objective_fn (function): the objective function for optimization
+            *args: variable length argument list for objective function
+            **kwargs: variable length of keyword arguments for the objective function
+
+        Returns:
+            list[array]: The new variable values :math:`x^{(t+1)}`.
+            If single arg is provided, list[array] is replaced by array.
+        """
+
+        self.trainable_args = set()
+
+        for index, arg in enumerate(args):
+            if getattr(arg, "requires_grad", True):
+                self.trainable_args |= {index}
+
+        if self.s is None:
+            # Number of shots per parameter
+            self.s = [
+                np.zeros_like(a, dtype=np.int64) + self.min_shots
+                for i, a in enumerate(args)
+                if i in self.trainable_args
+            ]
+
+        # keep track of the number of shots run
+        s = np.concatenate([i.flatten() for i in self.s])
+        self.max_shots = max(s)
+        self.shots_used = int(2 * np.sum(s))
+        self.total_shots_used += self.shots_used
+
+        # compute the gradient, as well as the variance in the gradient,
+        # using the number of shots determined by the array s.
+        grads, grad_variances = self.compute_grad(objective_fn, args, kwargs)
+        new_args = self.apply_grad(grads, args)
+
+        if self.xi is None:
+            self.chi = [np.zeros_like(g, dtype=np.float64) for g in grads]
+            self.xi = [np.zeros_like(g, dtype=np.float64) for g in grads]
+
+        # running average of the gradient
+        self.chi = [self.mu * c + (1 - self.mu) * g for c, g in zip(self.chi, grads)]
+
+        # running average of the gradient variance
+        self.xi = [self.mu * x + (1 - self.mu) * v for x, v in zip(self.xi, grad_variances)]
+
+        for idx, (c, x) in enumerate(zip(self.chi, self.xi)):
+            xi = x / (1 - self.mu ** (self.k + 1))
+            chi = c / (1 - self.mu ** (self.k + 1))
+
+            # determine the new optimum shots distribution for the next
+            # iteration of the optimizer
+            s = np.ceil(
+                (2 * self.lipschitz * self._stepsize * xi)
+                / (
+                    (2 - self.lipschitz * self._stepsize)
+                    * (chi ** 2 + self.b * (self.mu ** self.k))
+                )
+            )
+
+            # apply an upper and lower bound on the new shot distributions,
+            # to avoid the number of shots reducing below min(2, min_shots),
+            # or growing too significantly.
+            gamma = (
+                (self._stepsize - self.lipschitz * self._stepsize ** 2 / 2) * chi ** 2
+                - xi * self.lipschitz * self._stepsize ** 2 / (2 * s)
+            ) / s
+
+            argmax_gamma = np.unravel_index(np.argmax(gamma), gamma.shape)
+            smax = max(s[argmax_gamma], 2)
+            self.s[idx] = np.squeeze(np.int64(np.clip(s, min(2, self.min_shots), smax)))
+
+        self.k += 1
+
+        # unwrap from list if one argument, cleaner return
+        if len(new_args) == 1:
+            return new_args[0]
+
+        return new_args
+
+    def step_and_cost(self, objective_fn, *args, **kwargs):
+        """Update trainable arguments with one step of the optimizer and return the corresponding
+        objective function value prior to the step.
+
+        The objective function will be evaluated using the maximum number of shots
+        across all parameters as determined by the optimizer during the
+        optimization step.
+
+        .. warning::
+
+            Unlike other gradient descent optimizers, the objective function will be evaluated
+            **separately** to the gradient computation, and will result in extra
+            device evaluations.
+
+        Args:
+            objective_fn (function): the objective function for optimization
+            *args : variable length argument list for objective function
+            **kwargs : variable length of keyword arguments for the objective function
+
+        Returns:
+            tuple[list [array], float]: the new variable values :math:`x^{(t+1)}` and the objective
+            function output prior to the step.
+            If single arg is provided, list [array] is replaced by array.
+        """
+        new_args = self.step(objective_fn, *args, **kwargs)
+
+        if isinstance(objective_fn, qml.ExpvalCost):
+            device = objective_fn.qnodes[0].device
+        elif isinstance(objective_fn, qml.QNode) or hasattr(objective_fn, "device"):
+            device = objective_fn.device
+
+        original_shots = device.shots
+
+        try:
+            device.shots = int(self.max_shots)
+            forward = objective_fn(*args, **kwargs)
+        finally:
+            device.shots = original_shots
+
+        return new_args, forward


Test patch:
diff --git a/tests/test_optimize_shot_adapative.py b/tests/test_optimize_shot_adapative.py
new file mode 100644
index 000000000..b1a755a2b
--- /dev/null
+++ b/tests/test_optimize_shot_adapative.py
@@ -0,0 +1,571 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the shot adaptive optimizer"""
+import pytest
+from scipy.stats import multinomial
+
+import pennylane as qml
+from pennylane import numpy as np
+
+
+class TestExceptions:
+    """Test exceptions are raised for incorrect usage"""
+
+    def test_analytic_device_error(self):
+        """Test that an exception is raised if an analytic device is used"""
+        H = qml.Hamiltonian([0.3, 0.1], [qml.PauliX(0), qml.PauliZ(0)])
+        dev = qml.device("default.qubit", wires=1, shots=None)
+        expval_cost = qml.ExpvalCost(lambda x, **kwargs: qml.RX(x, wires=0), H, dev)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        # test expval cost
+        with pytest.raises(ValueError, match="can only be used with devices that"):
+            opt.step(expval_cost, 0.5)
+
+        # test qnode cost
+        with pytest.raises(ValueError, match="can only be used with devices that"):
+            opt.step(expval_cost.qnodes[0], 0.5)
+
+    def test_learning_error(self):
+        """Test that an exception is raised if the learning rate is beyond the
+        lipschitz bound"""
+        coeffs = [0.3, 0.1]
+        H = qml.Hamiltonian(coeffs, [qml.PauliX(0), qml.PauliZ(0)])
+        dev = qml.device("default.qubit", wires=1, shots=100)
+        expval_cost = qml.ExpvalCost(lambda x, **kwargs: qml.RX(x, wires=0), H, dev)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, stepsize=100.)
+
+        # lipschitz constant is given by sum(|coeffs|)
+        lipschitz = np.sum(np.abs(coeffs))
+
+        assert opt._stepsize > 2 / lipschitz
+
+        with pytest.raises(ValueError, match=f"The learning rate must be less than {2 / lipschitz}"):
+            opt.step(expval_cost, 0.5)
+
+        # for a single QNode, the lipschitz constant is simply 1
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, stepsize=100.)
+        with pytest.raises(ValueError, match=f"The learning rate must be less than {2 / 1}"):
+            opt.step(expval_cost.qnodes[0], 0.5)
+
+    def test_unknown_objective_function(self):
+        """Test that an exception is raised if an unknown objective function is passed"""
+        dev = qml.device("default.qubit", wires=1, shots=100)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        def cost(x):
+            return np.sin(circuit(x))
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        # test expval cost
+        with pytest.raises(ValueError, match="The objective function must either be encoded as a single QNode"):
+            opt.step(cost, 0.5)
+
+        # defining the device attribute allows it to proceed
+        cost.device = circuit.device
+        new_x = opt.step(cost, 0.5)
+
+        assert isinstance(new_x, float)
+
+
+class TestSingleShotGradientIntegration:
+    """Integration tests to ensure that the single shot gradient is correctly computed
+    for a variety of argument types."""
+
+    dev = qml.device("default.qubit", wires=1, shots=100)
+    H = qml.Hamiltonian([1.0], [qml.PauliZ(0)])
+
+    ansatz = lambda x, **kwargs: qml.RX(x, wires=0)
+    expval_cost = qml.ExpvalCost(ansatz, H, dev)
+
+    @qml.qnode(dev)
+    def qnode(x):
+        qml.RX(x, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    @pytest.mark.parametrize("cost_fn", [qnode, expval_cost])
+    def test_single_argument_step(self, cost_fn, mocker, monkeypatch):
+        """Test that a simple QNode with a single argument correctly performs an optimization step,
+        and that the single-shot gradients generated have the correct shape"""
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        spy_single_shot_expval = mocker.spy(opt, "_single_shot_expval_gradients")
+        spy_single_shot_qnodes = mocker.spy(opt, "_single_shot_qnode_gradients")
+        spy_grad = mocker.spy(opt, "compute_grad")
+
+        x_init = 0.5
+        new_x = opt.step(cost_fn, x_init)
+
+        assert isinstance(new_x, float)
+        assert new_x != x_init
+
+        spy_grad.assert_called_once()
+
+        if isinstance(cost_fn, qml.ExpvalCost):
+            spy_single_shot_expval.assert_called_once()
+            single_shot_grads = opt._single_shot_expval_gradients(cost_fn, [x_init], {})
+        else:
+            spy_single_shot_qnodes.assert_called_once()
+            single_shot_grads = opt._single_shot_qnode_gradients(cost_fn, [x_init], {})
+
+        # assert single shot gradients are computed correctly
+        assert len(single_shot_grads) == 1
+        assert single_shot_grads[0].shape == (10,)
+
+        # monkeypatch the optimizer to use the same single shot gradients
+        # as previously
+        monkeypatch.setattr(opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads)
+        monkeypatch.setattr(opt, "_single_shot_expval_gradients", lambda *args, **kwargs: single_shot_grads)
+
+        # reset the shot budget
+        opt.s = [np.array(10)]
+
+        # check that the gradient and variance are computed correctly
+        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
+        assert len(grad) == 1
+        assert len(grad_variance) == 1
+        assert np.allclose(grad, np.mean(single_shot_grads))
+        assert np.allclose(grad_variance, np.var(single_shot_grads, ddof=1))
+
+        # check that the gradient and variance are computed correctly
+        # with a different shot budget
+        opt.s = [np.array(5)]
+        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
+        assert len(grad) == 1
+        assert len(grad_variance) == 1
+        assert np.allclose(grad, np.mean(single_shot_grads[0][:5]))
+        assert np.allclose(grad_variance, np.var(single_shot_grads[0][:5], ddof=1))
+
+    def ansatz(x, **kwargs):
+        qml.RX(x[0, 0], wires=0)
+        qml.RY(x[0, 1], wires=0)
+        qml.RZ(x[0, 2], wires=0)
+        qml.RX(x[1, 0], wires=0)
+        qml.RY(x[1, 1], wires=0)
+        qml.RZ(x[1, 2], wires=0)
+
+    expval_cost = qml.ExpvalCost(ansatz, H, dev)
+    qnode = expval_cost.qnodes[0]
+
+    @pytest.mark.parametrize("cost_fn", [qnode, expval_cost])
+    def test_single_array_argument_step(self, cost_fn, mocker, monkeypatch):
+        """Test that a simple QNode with a single array argument correctly performs an optimization step,
+        and that the single-shot gradients generated have the correct shape"""
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        spy_single_shot_expval = mocker.spy(opt, "_single_shot_expval_gradients")
+        spy_single_shot_qnodes = mocker.spy(opt, "_single_shot_qnode_gradients")
+        spy_grad = mocker.spy(opt, "compute_grad")
+
+        x_init = np.array([[1., 2., 3.], [4., 5., 6.]])
+        new_x = opt.step(cost_fn, x_init)
+
+        assert isinstance(new_x, np.ndarray)
+        assert not np.allclose(new_x, x_init)
+
+        if isinstance(cost_fn, qml.ExpvalCost):
+            spy_single_shot_expval.assert_called_once()
+            single_shot_grads = opt._single_shot_expval_gradients(cost_fn, [x_init], {})
+        else:
+            spy_single_shot_qnodes.assert_called_once()
+            single_shot_grads = opt._single_shot_qnode_gradients(cost_fn, [x_init], {})
+
+        spy_grad.assert_called_once()
+
+        # assert single shot gradients are computed correctly
+        assert len(single_shot_grads) == 1
+        assert single_shot_grads[0].shape == (10, 2, 3)
+
+        # monkeypatch the optimizer to use the same single shot gradients
+        # as previously
+        monkeypatch.setattr(opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads)
+        monkeypatch.setattr(opt, "_single_shot_expval_gradients", lambda *args, **kwargs: single_shot_grads)
+
+        # reset the shot budget
+        opt.s = [10 * np.ones([2, 3], dtype=np.int64)]
+
+        # check that the gradient and variance are computed correctly
+        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
+        assert len(grad) == 1
+        assert len(grad_variance) == 1
+        assert grad[0].shape == x_init.shape
+        assert grad_variance[0].shape == x_init.shape
+
+        assert np.allclose(grad, np.mean(single_shot_grads, axis=1))
+        assert np.allclose(grad_variance, np.var(single_shot_grads, ddof=1, axis=1))
+
+        # check that the gradient and variance are computed correctly
+        # with a different shot budget
+        opt.s[0] = opt.s[0] // 2  # all array elements have a shot budget of 5
+        opt.s[0][0, 0] = 8    # set the shot budget of the zeroth element to 8
+
+        grad, grad_variance = opt.compute_grad(cost_fn, [x_init], {})
+        assert len(grad) == 1
+        assert len(grad_variance) == 1
+
+        # check zeroth element
+        assert np.allclose(grad[0][0, 0], np.mean(single_shot_grads[0][:8, 0, 0]))
+        assert np.allclose(grad_variance[0][0, 0], np.var(single_shot_grads[0][:8, 0, 0], ddof=1))
+
+        # check other elements
+        assert np.allclose(grad[0][0, 1], np.mean(single_shot_grads[0][:5, 0, 1]))
+        assert np.allclose(grad_variance[0][0, 1], np.var(single_shot_grads[0][:5, 0, 1], ddof=1))
+
+    def test_multiple_argument_step(self, mocker, monkeypatch):
+        """Test that a simple QNode with multiple scalar arguments correctly performs an optimization step,
+        and that the single-shot gradients generated have the correct shape"""
+        dev = qml.device("default.qubit", wires=1, shots=100)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        spy_single_shot = mocker.spy(opt, "_single_shot_qnode_gradients")
+        spy_grad = mocker.spy(opt, "compute_grad")
+
+        args = [0.1, 0.2]
+        new_x = opt.step(circuit, *args)
+
+        assert isinstance(new_x, list)
+        assert len(new_x) == 2
+
+        spy_single_shot.assert_called_once()
+        spy_grad.assert_called_once()
+
+        # assert single shot gradients are computed correctly
+        single_shot_grads = opt._single_shot_qnode_gradients(circuit, args, {})
+        assert len(single_shot_grads) == 2
+        assert single_shot_grads[0].shape == (10,)
+
+        # monkeypatch the optimizer to use the same single shot gradients
+        # as previously
+        monkeypatch.setattr(opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads)
+
+        # reset the shot budget
+        opt.s = [np.array(10), np.array(10)]
+
+        # check that the gradient and variance are computed correctly
+        grad, grad_variance = opt.compute_grad(circuit, args, {})
+        assert len(grad) == 2
+        assert len(grad_variance) == 2
+        assert np.allclose(grad, np.mean(single_shot_grads, axis=1))
+        assert np.allclose(grad_variance, np.var(single_shot_grads, ddof=1, axis=1))
+
+        # check that the gradient and variance are computed correctly
+        # with a different shot budget
+        opt.s = [np.array(5), np.array(7)]
+        grad, grad_variance = opt.compute_grad(circuit, args, {})
+        assert len(grad) == 2
+        assert len(grad_variance) == 2
+
+        for p, s in zip(range(2), opt.s):
+            assert np.allclose(grad[p], np.mean(single_shot_grads[p][:s]))
+            assert np.allclose(grad_variance[p], np.var(single_shot_grads[p][:s], ddof=1))
+
+    def test_multiple_array_argument_step(self, mocker, monkeypatch):
+        """Test that a simple QNode with multiple array arguments correctly performs an optimization step,
+        and that the single-shot gradients generated have the correct shape"""
+        dev = qml.device("default.qubit", wires=1, shots=100)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(x[0, 1], wires=0)
+            qml.RZ(x[0, 2], wires=0)
+            qml.RX(x[1, 0], wires=0)
+            qml.RY(x[1, 1], wires=0)
+            qml.RZ(x[1, 2], wires=0)
+            qml.RX(y[0], wires=0)
+            qml.RY(y[1], wires=0)
+            qml.RZ(y[2], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        spy_single_shot = mocker.spy(opt, "_single_shot_qnode_gradients")
+        spy_grad = mocker.spy(opt, "compute_grad")
+
+        args = [np.array([[1., 2., 3.], [4., 5., 6.]]), np.array([1., 2., 3.])]
+        new_x = opt.step(circuit, *args)
+
+        assert isinstance(new_x, list)
+        assert len(new_x) == 2
+
+        spy_single_shot.assert_called_once()
+        spy_grad.assert_called_once()
+
+        # assert single shot gradients are computed correctly
+        single_shot_grads = opt._single_shot_qnode_gradients(circuit, args, {})
+        assert len(single_shot_grads) == 2
+        assert single_shot_grads[0].shape == (10, 2, 3)
+        assert single_shot_grads[1].shape == (10, 3)
+
+        # monkeypatch the optimizer to use the same single shot gradients
+        # as previously
+        monkeypatch.setattr(opt, "_single_shot_qnode_gradients", lambda *args, **kwargs: single_shot_grads)
+
+        # reset the shot budget
+        opt.s = [10 * np.ones([2, 3], dtype=np.int64), 10 * np.ones([3], dtype=np.int64)]
+
+        # check that the gradient and variance are computed correctly
+        grad, grad_variance = opt.compute_grad(circuit, args, {})
+        assert len(grad) == 2
+        assert len(grad_variance) == 2
+
+        for i in range(2):
+            assert grad[i].shape == args[i].shape
+            assert grad_variance[i].shape == args[i].shape
+
+        assert np.allclose(grad[0], np.mean(single_shot_grads[0], axis=0))
+        assert np.allclose(grad_variance[0], np.var(single_shot_grads[0], ddof=1, axis=0))
+
+        assert np.allclose(grad[1], np.mean(single_shot_grads[1], axis=0))
+        assert np.allclose(grad_variance[1], np.var(single_shot_grads[1], ddof=1, axis=0))
+
+        # check that the gradient and variance are computed correctly
+        # with a different shot budget
+        opt.s[0] = opt.s[0] // 2  # all array elements have a shot budget of 5
+        opt.s[0][0, 0] = 8    # set the shot budget of the zeroth element to 8
+
+        opt.s[1] = opt.s[1] // 5  # all array elements have a shot budget of 2
+        opt.s[1][0] = 7    # set the shot budget of the zeroth element to 7
+
+        grad, grad_variance = opt.compute_grad(circuit, args, {})
+        assert len(grad) == 2
+        assert len(grad_variance) == 2
+
+        # check zeroth element of arg 0
+        assert np.allclose(grad[0][0, 0], np.mean(single_shot_grads[0][:8, 0, 0]))
+        assert np.allclose(grad_variance[0][0, 0], np.var(single_shot_grads[0][:8, 0, 0], ddof=1))
+
+        # check other elements of arg 0
+        assert np.allclose(grad[0][0, 1], np.mean(single_shot_grads[0][:5, 0, 1]))
+        assert np.allclose(grad_variance[0][0, 1], np.var(single_shot_grads[0][:5, 0, 1], ddof=1))
+
+        # check zeroth element of arg 1
+        assert np.allclose(grad[1][0], np.mean(single_shot_grads[1][:7, 0]))
+        assert np.allclose(grad_variance[1][0], np.var(single_shot_grads[1][:7, 0], ddof=1))
+
+        # check other elements of arg 1
+        assert np.allclose(grad[1][1], np.mean(single_shot_grads[1][:2, 1]))
+        assert np.allclose(grad_variance[1][1], np.var(single_shot_grads[1][:2, 1], ddof=1))
+
+
+class TestWeightedRandomSampling:
+    """Tests for weighted random Hamiltonian term sampling"""
+
+    def test_wrs_expval_cost(self, mocker):
+        """Checks that cost functions that are expval costs can
+        make use of weighted random sampling"""
+        coeffs = [0.2, 0.1]
+        dev = qml.device("default.qubit", wires=2, shots=100)
+        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)])
+
+        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+        weights = qml.init.strong_ent_layers_normal(n_layers=3, n_wires=2)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        spy = mocker.spy(opt, "weighted_random_sampling")
+
+        new_weights = opt.step(expval_cost, weights)
+        spy.assert_called_once()
+
+        grads = opt.weighted_random_sampling(expval_cost.qnodes, coeffs, 10, [0], weights)
+        assert len(grads) == 1
+        assert grads[0].shape == (10, *weights.shape)
+
+    def test_wrs_disabled(self, mocker):
+        """Checks that cost functions that are expval costs can
+        disable use of weighted random sampling"""
+        coeffs = [0.2, 0.1]
+        dev = qml.device("default.qubit", wires=2, shots=100)
+        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)])
+
+        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+        weights = qml.init.strong_ent_layers_normal(n_layers=3, n_wires=2)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling=None)
+        spy = mocker.spy(opt, "weighted_random_sampling")
+
+        new_weights = opt.step(expval_cost, weights)
+        spy.assert_not_called()
+
+    def test_unknown_term_sampling_method(self):
+        """Checks that an exception is raised if the term sampling method is unknown"""
+        coeffs = [0.2, 0.1]
+        dev = qml.device("default.qubit", wires=2, shots=100)
+        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1)])
+
+        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+        weights = qml.init.strong_ent_layers_normal(n_layers=3, n_wires=2)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="uniform_random_sampling")
+
+        with pytest.raises(ValueError, match="Unknown Hamiltonian term sampling method"):
+            opt.step(expval_cost, weights)
+
+    def test_zero_shots(self, mocker, monkeypatch):
+        """Test that, if the shot budget for a single term is 0,
+        that the jacobian computation is skipped"""
+        coeffs = [0.2, 0.1, 0.1]
+        dev = qml.device("default.qubit", wires=2, shots=100)
+        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliX(1), qml.PauliZ(0) @ qml.PauliZ(1)])
+
+        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+        weights = qml.init.strong_ent_layers_normal(n_layers=3, n_wires=2)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        spy = mocker.spy(qml, "jacobian")
+        mocker.patch("scipy.stats._multivariate.multinomial_gen.rvs", return_value=np.array([[4, 0, 6]]))
+        grads = opt.weighted_random_sampling(expval_cost.qnodes, coeffs, 10, [0], weights)
+
+        assert len(spy.call_args_list) == 2
+        assert len(grads) == 1
+        assert grads[0].shape == (10, *weights.shape)
+
+    def test_single_shots(self, mocker, monkeypatch):
+        """Test that, if the shot budget for a single term is 1,
+        that the number of dimensions for the returned Jacobian is expanded"""
+        coeffs = [0.2, 0.1, 0.1]
+        dev = qml.device("default.qubit", wires=2, shots=100)
+        H = qml.Hamiltonian(coeffs, [qml.PauliZ(0), qml.PauliX(1), qml.PauliZ(0) @ qml.PauliZ(1)])
+
+        expval_cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+        weights = qml.init.strong_ent_layers_normal(n_layers=3, n_wires=2)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        spy = mocker.spy(qml, "jacobian")
+        spy_dims = mocker.spy(np, "expand_dims")
+        mocker.patch("scipy.stats._multivariate.multinomial_gen.rvs", return_value=np.array([[4, 1, 5]]))
+        grads = opt.weighted_random_sampling(expval_cost.qnodes, coeffs, 10, [0], weights)
+
+        spy_dims.assert_called_once()
+        assert len(spy.call_args_list) == 3
+        assert len(grads) == 1
+        assert grads[0].shape == (10, *weights.shape)
+
+
+class TestOptimization:
+    """Integration test to ensure that the optimizer
+    minimizes simple examples"""
+
+    def test_multi_qubit_rotation(self):
+        """Test that multiple qubit rotation can be optimized"""
+        dev = qml.device("default.qubit", wires=2, shots=100)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=1)
+            qml.RZ(x[2], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(0) @ qml.PauliY(1))
+
+        params = np.array([0.1, 0.3, 0.4], requires_grad=True)
+        initial_loss = circuit(params)
+
+        min_shots = 10
+        loss = initial_loss
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        for i in range(100):
+            params = opt.step(circuit, params)
+            loss = circuit(params)
+
+        assert loss < initial_loss
+        assert np.allclose(circuit(params), -1, atol=0.1, rtol=0.2)
+        assert opt.shots_used > min_shots
+
+    def test_vqe_optimization(self):
+        """Test that a simple VQE circuit can be optimized"""
+        dev = qml.device("default.qubit", wires=2, shots=100)
+        coeffs = [0.1, 0.2]
+        obs = [qml.PauliZ(0), qml.PauliX(0)]
+        H = qml.Hamiltonian(coeffs, obs)
+
+        def ansatz(x, **kwargs):
+            qml.Rot(*x[0], wires=0)
+            qml.Rot(*x[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.Rot(*x[2], wires=0)
+            qml.Rot(*x[3], wires=1)
+            qml.CNOT(wires=[0, 1])
+
+        cost = qml.ExpvalCost(ansatz, H, dev)
+        params = np.random.random((4, 3), requires_grad=True)
+        initial_loss = cost(params)
+
+        min_shots = 10
+        loss = initial_loss
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        for i in range(100):
+            params = opt.step(cost, params)
+            loss = cost(params)
+
+        assert loss < initial_loss
+        assert np.allclose(loss, -1 / (2 * np.sqrt(5)), atol=0.1, rtol=0.2)
+        assert opt.shots_used > min_shots
+
+
+class TestStepAndCost:
+    """Tests for the step_and_cost method"""
+
+    def test_qnode_cost(self, tol):
+        """Test that the cost is correctly returned
+        when using a QNode as the cost function"""
+        dev = qml.device("default.qubit", wires=1, shots=10)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        params = np.array([0.1, 0.3], requires_grad=True)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        for i in range(100):
+            params, res = opt.step_and_cost(circuit, params)
+
+        assert np.allclose(res, -1, atol=tol, rtol=0)
+
+    def test_expval_cost(self, tol):
+        """Test that the cost is correctly returned
+        when using a QNode as the cost function"""
+        dev = qml.device("default.qubit", wires=1, shots=10)
+
+        def ansatz(x, **kwargs):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=0)
+
+        H = qml.Hamiltonian([1.0], [qml.PauliZ(0)])
+        circuit = qml.ExpvalCost(ansatz, H, dev)
+        params = np.array([0.1, 0.3], requires_grad=True)
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+
+        for i in range(100):
+            params, res = opt.step_and_cost(circuit, params)
+
+        assert np.allclose(res, -1, atol=tol, rtol=0)