Repository: PennyLaneAI/pennylane
Pull request #4896 (2023-11-29 01:02:58)

copy expvalcost code for qnode compatibility in ShotAdaptiveOptimizer
**Context:**
We deprecated `qml.ExpvalCost` in v0.24, but it is the only way to use the multinomial distribution for distributing shots with the `ShotAdaptiveOptimizer`. Since we recommend simply using a QNode with the expectation of a Hamiltonian, this method should have feature parity with ExpvalCost.

**Description of the Change:**
Copied the multinomial distribution code with tweaks in order to support its use with regular (non-ExpvalCost) qnodes.

**Benefits:**
We are closer to removing ExpvalCost

**Possible Drawbacks:**
- Duplicate code for the time-being
- some not-so-pretty code pieces, like changing `QNode.func` in place.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4896?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`9781f11`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9781f1151bd27104efefb1c7ec78ff8d8bb0728c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65% compared to head [(`03243b0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4896?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65%.
> Report is 1 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4896      +/-   ##
==========================================
- Coverage   99.65%   99.65%   -0.01%     
==========================================
  Files         387      387              
  Lines       34957    34739     -218     
==========================================
- Hits        34837    34618     -219     
- Misses        120      121       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4896?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

timmysilv: did 2 things on this today:
1. confirmed that instead of doing `res0, res1 = qml.jacobian(cost, argnum=0), qml.jacobian(cost, argnum=1)`, we can just do `res0, res1 = qml.jacobian(cost, argnum=[0, 1])`. So, I  removed the unnecessary for-loop + additional jacobian call, and un-indented all the code which makes the diff ugly.
2. chatted with Romain for some time about trying to make this a transform. The problem is that we need to apply post-processing to the jacobians, not to the execution results. This makes it quite difficult (Romain had tried with other things like `classical_fisher` but found it was not worth the effort), so I'm abandoning my efforts. It _can_ be done if we apply param-shift manually and make it all work, but it's rather overkill for some private code.
timmysilv: WRS does do it with single-shot expectations. couldn't tell you why, but it appears to be on purpose.
josh146: > So they are both working, and ultimately use the same number of shots, but right now if you switch out the ExpvalCost for a qnode in code that is already working, your optimization is likely to fail because you aren't doing enough reps.

Ah this sounds like a potential breaking change 🤔 @lillian542 do you have a rough idea why this is the case? Is it an implementation detail in the PR, or something more fundamental that we can't do without `ExpvalCost`?
lillian542: >Ah this sounds like a potential breaking change 🤔 @lillian542 do you have a rough idea why this is the case? Is it an implementation detail in the PR, or something more fundamental that we can't do without ExpvalCost?

@josh146 - actually @timmysilv fixed this in the most recent commit (before my comment), I was having a problem when pulling from GitHub and hadn't noticed, so I was trying out an out-of-date version of this PR. I think he's doing a few more sanity checks to be super sure, but I don't believe there's any issue here. Sorry for the confusion!
timmysilv: @josh146 - chatted with @lillian542 offline, her local git didn't have the latest fix I pushed so they should be fine. I re-ran the example (with only 30 steps, but same initial params for both QNode and ExpvalCost) a few times, and the results seem fairly consistent. Here are a few comparisons to highlight the chaos from this Optimizer:
![fig](https://github.com/PennyLaneAI/pennylane/assets/23283972/95fc1365-7a00-4e7f-934c-90531d7238ea)
![fig](https://github.com/PennyLaneAI/pennylane/assets/23283972/b622a1bf-5413-44df-b0a0-eb4ed0557d57)
![fig](https://github.com/PennyLaneAI/pennylane/assets/23283972/49fe1a9e-a199-494e-a1a4-d9068ad2e2cc)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6fa7ae6ff..3438e03c4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -362,6 +362,11 @@
   to `None` if the operator has not defined it (or has no definition in the pauli basis).
   [(#4915)](https://github.com/PennyLaneAI/pennylane/pull/4915)
 
+* `qml.ShotAdaptiveOptimizer` can now use a multinomial distribution for spreading shots across
+  the terms of a Hamiltonian measured in a QNode. Note that this is equivalent to what can be
+  done with `qml.ExpvalCost`, but this is the preferred method because `ExpvalCost` is deprecated.
+  [(#4896)](https://github.com/PennyLaneAI/pennylane/pull/4896)
+
 <h3>Breaking changes 💔</h3>
 
 * The transforms submodule `qml.transforms.qcut` is now its own module: `qml.qcut`.
@@ -394,6 +399,10 @@
 * The `"pennylane"` MPL-drawer style now draws straight lines instead of sketch-style lines.
   [(#4880)](https://github.com/PennyLaneAI/pennylane/pull/4880)
 
+* The default value for the `term_sampling` argument of `ShotAdaptiveOptimizer` is now
+  `None` instead of `"weighted_random_sampling"`.
+  [(#4896)](https://github.com/PennyLaneAI/pennylane/pull/4896)
+
 <h3>Deprecations 👋</h3>
 
 * `single_tape_transform`, `batch_transform`, `qfunc_transform`, and `op_transform` are deprecated.
diff --git a/pennylane/optimize/shot_adaptive.py b/pennylane/optimize/shot_adaptive.py
index 7beb372f5..6ef47e2d2 100644
--- a/pennylane/optimize/shot_adaptive.py
+++ b/pennylane/optimize/shot_adaptive.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Shot adaptive optimizer"""
 # pylint: disable=too-many-instance-attributes,too-many-arguments,too-many-branches
+from copy import copy
 from scipy.stats import multinomial
 
 import pennylane as qml
@@ -29,10 +30,23 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
     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.
+    In addition, weighted random sampling can be used to further distribute the
+    shot budget across the local terms from which the Hamiltonian is constructed.
+
+    .. warning::
+
+        ``ExpvalCost`` is deprecated. Instead, it is recommended to simply
+        pass Hamiltonians to the :func:`~pennylane.expval` function inside QNodes.
+
+        .. code-block:: python
+
+            @qml.qnode(dev)
+            def ansatz(params):
+                some_qfunc(params)
+                return qml.expval(Hamiltonian)
+
+        You should also set ``term_sampling='weighted_random_sampling'`` to continue using weighted
+        random sampling in ``ShotAdaptiveOptimizer``.
 
     .. note::
 
@@ -47,11 +61,10 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
         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.
+        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. The default value is ``None``, which
+            disables the random sampling behaviour.
         stepsize (float): The learning rate :math:`\eta`. The learning rate *must* be such
             that :math:`\eta < 2/L = 2/\sum_i|c_i|`, where:
 
@@ -63,9 +76,10 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
 
     **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`.
+    For VQE/VQE-like problems, the objective function for the optimizer can be realized
+    as a :class:`~.QNode` object measuring the expectation of a :class:`~.Hamiltonian`.
 
+    >>> from pennylane import numpy as np
     >>> coeffs = [2, 4, -1, 5, 2]
     >>> obs = [
     ...   qml.PauliX(1),
@@ -76,7 +90,10 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
     ... ]
     >>> H = qml.Hamiltonian(coeffs, obs)
     >>> dev = qml.device("default.qubit", wires=2, shots=100)
-    >>> cost = qml.ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
+    >>> @qml.qnode(dev)
+    >>> def cost(weights):
+    ...     qml.StronglyEntanglingLayers(weights, wires=range(2))
+    ...     return qml.expval(H)
 
     Once constructed, the cost function can be passed directly to the
     optimizer's ``step`` method. The attributes ``opt.shots_used`` and
@@ -85,7 +102,7 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
 
     >>> shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=2)
     >>> params = np.random.random(shape)
-    >>> opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+    >>> 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}")
@@ -172,9 +189,7 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
           <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
-    ):
+    def __init__(self, min_shots, term_sampling=None, mu=0.99, b=1e-6, stepsize=0.07):
         self.term_sampling = term_sampling
         self.trainable_args = set()
 
@@ -244,26 +259,91 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
             # set the QNode device shots
             new_shots = 1 if s == 1 else [(1, int(s))]
 
-            jacs = []
-            for i in argnums:
-                if s > 1:
+            if s > 1:
+
+                def cost(*args, **kwargs):
+                    # pylint: disable=cell-var-from-loop
+                    return qml.math.stack(h(*args, **kwargs))
+
+            else:
+                cost = h
+
+            jacs = qml.jacobian(cost, argnum=argnums)(*args, **kwargs, shots=new_shots)
+
+            if s == 1:
+                jacs = [np.expand_dims(j, 0) for j in jacs]
+
+            # Divide each term by the probability per shot. This is
+            # because we are sampling one at a time.
+            grads.append([(c * j / p) for j in jacs])
+
+        return [np.concatenate(i) for i in zip(*grads)]
+
+    @staticmethod
+    def qnode_weighted_random_sampling(qnode, coeffs, observables, 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:
+            qnode (.QNode): A QNode that returns the expectation value of a Hamiltonian.
+            coeffs (List[float]): The coefficients of the Hamiltonian being measured
+            observables (List[Observable]): The terms of the Hamiltonian being measured
+            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 QNode
+            **kwargs: Keyword arguments to the QNode
+
+        Returns:
+            array[float]: the single-shot gradients of the Hamiltonian expectation value
+        """
+        qnode = copy(qnode)
+        base_func = qnode.func
+
+        # 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 o, c, p, s in zip(observables, coeffs, prob_shots, shots_per_term):
+            # if the number of shots is 0, do nothing
+            if s == 0:
+                continue
+
+            def func(*qnode_args, **qnode_kwargs):
+                qs = qml.tape.make_qscript(base_func)(*qnode_args, **qnode_kwargs)
+                for op in qs.operations:
+                    qml.apply(op)
+                return qml.expval(o)  # pylint:disable=cell-var-from-loop
+
+            qnode.func = func
+
+            new_shots = 1 if s == 1 else [(1, int(s))]
+
+            if s > 1:
 
-                    def cost(*args, **kwargs):
-                        # pylint: disable=cell-var-from-loop
-                        return qml.math.stack(h(*args, **kwargs))
+                def cost(*args, **kwargs):
+                    # pylint: disable=cell-var-from-loop
+                    return qml.math.stack(qnode(*args, **kwargs))
 
-                else:
-                    cost = h
+            else:
+                cost = qnode
 
-                j = qml.jacobian(cost, argnum=i)(*args, **kwargs, shots=new_shots)
+            jacs = qml.jacobian(cost, argnum=argnums)(*args, **kwargs, shots=new_shots)
 
-                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)
+            if s == 1:
+                jacs = [np.expand_dims(j, 0) for j in jacs]
 
-            grads.append(jacs)
+            # Divide each term by the probability per shot. This is
+            # because we are sampling one at a time.
+            grads.append([(c * j / p) for j in jacs])
 
         return [np.concatenate(i) for i in zip(*grads)]
 
@@ -335,8 +415,26 @@ class ShotAdaptiveOptimizer(GradientDescentOptimizer):
         """Compute the single shot gradients of a QNode."""
         self.check_device(qnode.device)
 
+        qnode.construct(args, kwargs)
+        tape = qnode.tape
+        [expval] = tape.measurements
+        coeffs, observables = (
+            expval.obs.terms() if isinstance(expval.obs, qml.Hamiltonian) else ([1.0], [expval.obs])
+        )
+
         if self.lipschitz is None:
-            self.check_learning_rate(1)
+            self.check_learning_rate(coeffs)
+
+        if self.term_sampling == "weighted_random_sampling":
+            return self.qnode_weighted_random_sampling(
+                qnode, coeffs, observables, self.max_shots, self.trainable_args, *args, **kwargs
+            )
+        if self.term_sampling is not None:
+            raise ValueError(
+                f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
+                "Only term_sampling='weighted_random_sampling' and "
+                "term_sampling=None currently supported."
+            )
 
         new_shots = [(1, int(self.max_shots))]
 


Test patch:
diff --git a/tests/optimize/test_optimize_shot_adaptive.py b/tests/optimize/test_optimize_shot_adaptive.py
index 862fff19e..06d327459 100644
--- a/tests/optimize/test_optimize_shot_adaptive.py
+++ b/tests/optimize/test_optimize_shot_adaptive.py
@@ -490,7 +490,7 @@ class TestWeightedRandomSampling:
         expval_cost = catch_warn_ExpvalCost(qml.templates.StronglyEntanglingLayers, H, dev)
         weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
 
-        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="weighted_random_sampling")
         spy = mocker.spy(opt, "weighted_random_sampling")
 
         _ = opt.step(expval_cost, weights)
@@ -574,6 +574,154 @@ class TestWeightedRandomSampling:
         assert grads[0].shape == (10, *weights.shape)
 
 
+class TestQNodeWeightedRandomSampling:
+    """Tests for weighted random Hamiltonian term sampling"""
+
+    def test_wrs_qnode(self, mocker):
+        """Checks that cost functions that are qnodes 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)])
+
+        @qml.qnode(dev)
+        def circuit(weights, x):
+            qml.StronglyEntanglingLayers(weights, wires=range(2))
+            qml.RX(x, 0)
+            qml.CNOT([0, 1])
+            return qml.expval(H)
+
+        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
+        x = np.array(1.1)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="weighted_random_sampling")
+        spy = mocker.spy(opt, "qnode_weighted_random_sampling")
+
+        _ = opt.step(circuit, weights, x)
+        spy.assert_called_once()
+
+        grads = opt.qnode_weighted_random_sampling(circuit, coeffs, H.ops, 10, [0], weights, x)
+        assert len(grads) == 1
+        assert grads[0].shape == (10, *weights.shape)
+
+    def test_wrs_qnode_multiple_args(self, mocker):
+        """Checks that cost functions that are qnodes works with multiple args"""
+        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)])
+
+        @qml.qnode(dev)
+        def circuit(weights, x):
+            qml.StronglyEntanglingLayers(weights, wires=range(2))
+            qml.RX(x, 0)
+            qml.CNOT([0, 1])
+            return qml.expval(H)
+
+        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
+        x = np.array(1.1)
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="weighted_random_sampling")
+        spy = mocker.spy(opt, "qnode_weighted_random_sampling")
+
+        _ = opt.step(circuit, weights, x)
+        spy.assert_called_once()
+
+        weight_grad, x_grad = opt.qnode_weighted_random_sampling(
+            circuit, coeffs, H.ops, 10, [0, 1], weights, x
+        )
+        assert weight_grad.shape == (10, *weights.shape)
+        assert x_grad.shape == (10,)
+
+    def test_wrs_disabled(self, mocker):
+        """Checks that cost functions that are qnodes 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)])
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.StronglyEntanglingLayers(weights, wires=range(2))
+            return qml.expval(H)
+
+        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling=None)
+        spy = mocker.spy(opt, "qnode_weighted_random_sampling")
+
+        opt.step(circuit, 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)])
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.StronglyEntanglingLayers(weights, wires=range(2))
+            return qml.expval(H)
+
+        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
+
+        opt = qml.ShotAdaptiveOptimizer(min_shots=10, term_sampling="uniform_random_sampling")
+
+        with pytest.raises(ValueError, match="Unknown Hamiltonian term sampling method"):
+            opt.step(circuit, weights)
+
+    def test_zero_shots(self, mocker):
+        """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)])
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.StronglyEntanglingLayers(weights, wires=range(2))
+            return qml.expval(H)
+
+        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 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.qnode_weighted_random_sampling(circuit, coeffs, H.ops, 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):
+        """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)])
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            qml.StronglyEntanglingLayers(weights, wires=range(2))
+            return qml.expval(H)
+
+        weights = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 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, 1, 5]])
+        )
+        grads = opt.qnode_weighted_random_sampling(circuit, coeffs, H.ops, 10, [0], weights)
+
+        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"""