class_docstr: Test that the transform is differentiable
class_docstr_tok: ["Test","that","the","transform","is","differentiable"]
docstr: Tests that the output of the finite-difference transform
  can be differentiated using JAX, yielding second derivatives.
docstr_tok: ["Tests","that","the","output","of","the","finite-difference","transform","can","be","differentiated","using","JAX",",","yielding","second","derivatives","."]
code: def test_jax(self, dev_name, approx_order, strategy):
      """Tests that the output of the finite-difference transform
      can be differentiated using JAX, yielding second derivatives."""
      import jax
      from jax import numpy as jnp
  
      dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
      execute_fn = (
          dev.execute
          if dev_name == "default.qubit"
          else dev.batch_execute
      )
      params = jnp.array([0.543, -0.654])
  
      def cost_fn(x):
          with qml.queuing.AnnotatedQueue() as q:
              qml.RX(x[0], wires=[0])
              qml.RY(x[1], wires=[1])
              qml.CNOT(wires=[0, 1])
              qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
  
          tape = qml.tape.QuantumScript.from_queue(
              q, shots=many_shots_shot_vector
          )
          tape.trainable_params = {0, 1}
          tapes, fn = finite_diff(
              tape,
              n=1,
              approx_order=approx_order,
              strategy=strategy,
              h=h_val,
          )
          jac = fn(execute_fn(tapes))
          return jac
  
      all_res = jax.jacobian(cost_fn)(params)
  
      assert isinstance(all_res, tuple)
      assert len(all_res) == len(many_shots_shot_vector)
  
      x, y = params
      expected = np.array(
          [
              [-np.cos(x) * np.sin(y), -np.cos(y) * np.sin(x)],
              [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)],
          ]
      )
      for res in all_res:
          assert isinstance(res, tuple)
          assert np.allclose(res, expected, atol=0.3, rtol=0)
  
imports: import numpy
  import pytest
  from pennylane import numpy
  import pennylane
  from pennylane.gradients import finite_diff
  from pennylane.devices import DefaultQubitLegacy
  from pennylane.measurements import Shots
  from pennylane.operation import Observable, AnyWires
label_desc: Use image node_id 6 for calling the TestFiniteDiffGradients obj's underlying member method code with example usage: obj.test_jax(dev_name, approx_order, strategy) and returns: jac