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=2)
      spy = mocker.spy(dev, "execute_and_gradients")
  
      def cost(params):
          tape1 = qml.tape.QuantumScript(
              [qml.RY(params[0], 0), qml.RX(params[1], 0)],
              [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))],
          )
  
          tape2 = qml.tape.QuantumScript(
              [qml.RY(np.array(0.5), 0)],
              [qml.expval(qml.PauliZ(0))],
          )
  
          tape3 = qml.tape.QuantumScript(
              [qml.RY(params[0], 0), qml.RX(params[1], 0)],
              [qml.expval(qml.PauliZ(0))],
          )
          return execute(
              [tape1, tape2, tape3],
              dev,
              gradient_fn="device",
              gradient_kwargs={
                  "method": "adjoint_jacobian",
                  "use_device_state": True,
              },
          )
  
      a = jax.numpy.array([0.1, 0.2])
      res = cost(a)
  
      x, y = a
      assert np.allclose(res[0][0], np.cos(x) * np.cos(y))
      assert np.allclose(res[0][1], 1)
      assert np.allclose(res[1], np.cos(0.5))
      assert np.allclose(res[2], np.cos(x) * np.cos(y))
  
      # adjoint method only performs a single device execution per tape, but gets both result and gradient
      assert dev.num_executions == 3
      spy.assert_not_called()
  
      g = jax.jacobian(cost)(a)
      spy.assert_called()
      expected_g = (-np.sin(x) * np.cos(y), -np.cos(x) * np.sin(y))
      assert qml.math.allclose(g[0][0], expected_g)
      assert qml.math.allclose(g[0][1], np.zeros(2))
      assert qml.math.allclose(g[1], np.zeros(2))
      assert qml.math.allclose(g[2], expected_g)
  
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