class_docstr: Unit tests for jax execution
class_docstr_tok: ["Unit","tests","for","jax","execution"]
docstr: Test that grad on execution uses the `device.execute_and_gradients` pathway
docstr_tok: ["Test","that","grad","on","execution","uses","the","`","device.execute_and_gradients","`","pathway"]
code: def test_grad_on_execution(self, mocker):
      """Test that grad on execution uses the `device.execute_and_gradients` pathway"""
      dev = qml.device("default.qubit.legacy", wires=1)
      spy = mocker.spy(dev, "execute_and_gradients")
  
      def cost(a):
          with qml.queuing.AnnotatedQueue() as q:
              qml.RY(a[0], wires=0)
              qml.RX(a[1], wires=0)
              qml.expval(qml.PauliZ(0))
  
          tape = qml.tape.QuantumScript.from_queue(q)
  
          return execute(
              [tape],
              dev,
              gradient_fn="device",
              gradient_kwargs={
                  "method": "adjoint_jacobian",
                  "use_device_state": True,
              },
          )[0]
  
      a = jax.numpy.array([0.1, 0.2])
      jax.jit(cost)(a)
  
      # adjoint method only performs a single device execution
      # gradients are not requested when we only want the results
      assert dev.num_executions == 1
      spy.assert_not_called()
  
      # when the jacobian is requested, we always calculate it at the same time as the results
      jax.grad(jax.jit(cost))(a)
      spy.assert_called()
  
imports: import pytest
  import numpy
  import pennylane
  from pennylane.gradients import param_shift
  from pennylane import execute
label_desc: Use image node_id 4 for calling the TestJaxExecuteUnitTests obj's underlying member method code with example usage: obj.test_grad_on_execution(mocker) and returns: execute