class_docstr: Test the jax interface execute function
  integrates well for both forward and backward execution
class_docstr_tok: ["Test","the","jax","interface","execute","function","integrates","well","for","both","forward","and","backward","execution"]
docstr: Test jacobian calculation
docstr_tok: ["Test","jacobian","calculation"]
code: def test_jacobian(self, execute_kwargs, shots, device):
      """Test jacobian calculation"""
  
      a = jnp.array(0.1)
      b = jnp.array(0.2)
  
      def cost(a, b):
          ops = [
              qml.RY(a, wires=0),
              qml.RX(b, wires=1),
              qml.CNOT(wires=[0, 1]),
          ]
          m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
          tape = qml.tape.QuantumScript(ops, m, shots=shots)
          return execute([tape], device, **execute_kwargs)[0]
  
      res = cost(a, b)
      expected = [jnp.cos(a), -jnp.cos(a) * jnp.sin(b)]
      if shots.has_partitioned_shots:
          assert np.allclose(
              res[0], expected, atol=atol_for_shots(shots), rtol=0
          )
          assert np.allclose(
              res[1], expected, atol=2 * atol_for_shots(shots), rtol=0
          )
      else:
          assert np.allclose(
              res, expected, atol=atol_for_shots(shots), rtol=0
          )
  
      g = jax.jacobian(cost, argnums=[0, 1])(a, b)
      assert isinstance(g, tuple) and len(g) == 2
  
      expected = (
          [-jnp.sin(a), jnp.sin(a) * jnp.sin(b)],
          [0, -jnp.cos(a) * jnp.cos(b)],
      )
  
      if shots.has_partitioned_shots:
          for i in (0, 1):
              assert np.allclose(
                  g[i][0][0],
                  expected[0][0],
                  atol=atol_for_shots(shots),
                  rtol=0,
              )
              assert np.allclose(
                  g[i][1][0],
                  expected[0][1],
                  atol=atol_for_shots(shots),
                  rtol=0,
              )
              assert np.allclose(
                  g[i][0][1],
                  expected[1][0],
                  atol=atol_for_shots(shots),
                  rtol=0,
              )
              assert np.allclose(
                  g[i][1][1],
                  expected[1][1],
                  atol=atol_for_shots(shots),
                  rtol=0,
              )
      else:
          assert np.allclose(
              g[0][0],
              expected[0][0],
              atol=atol_for_shots(shots),
              rtol=0,
          )
          assert np.allclose(
              g[1][0],
              expected[0][1],
              atol=atol_for_shots(shots),
              rtol=0,
          )
          assert np.allclose(
              g[0][1],
              expected[1][0],
              atol=atol_for_shots(shots),
              rtol=0,
          )
          assert np.allclose(
              g[1][1],
              expected[1][1],
              atol=atol_for_shots(shots),
              rtol=0,
          )
  
imports: import pytest
  import numpy
  from param_shift_dev import ParamShiftDerivativesDevice
  import pennylane
  from pennylane.devices import DefaultQubit
  from pennylane.gradients import param_shift
  from pennylane import execute
  from pennylane.measurements import Shots
label_desc: Use image node_id 3 for calling the TestJaxExecuteIntegration obj's underlying member method code with example usage: obj.test_jacobian(execute_kwargs, shots, device) and returns: execute