class_docstr: Integration tests to ensure that the single shot gradient is correctly computed
  for a variety of argument types.
docstr: 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
docstr_tok: ["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"]
code: 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
      """
      # pylint: disable=protected-access
      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.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      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),
      )
  
imports: import pytest
  from flaky import flaky
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 3 for calling the TestSingleShotGradientIntegration obj's underlying member method code with example usage: obj.test_single_array_argument_step(cost_fn, mocker, monkeypatch) without return types