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 multiple array arguments correctly performs an optimization step,
  and that the single-shot gradients generated have the correct shape
docstr_tok: ["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"]
code: 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
      """
      # pylint: disable=protected-access
      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.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
          np.array([1.0, 2.0, 3.0]),
      ]
      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),
      )
  
imports: import pytest
  from flaky import flaky
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 6 for calling the TestSingleShotGradientIntegration obj's underlying member method code with example usage: obj.test_multiple_array_argument_step(mocker, monkeypatch) and returns: qml