class_docstr: Test that the transform is differentiable
class_docstr_tok: ["Test","that","the","transform","is","differentiable"]
docstr: Tests that the output of the SPSA gradient transform
  can be differentiated using JAX, yielding second derivatives.
docstr_tok: ["Tests","that","the","output","of","the","SPSA","gradient","transform","can","be","differentiated","using","JAX",",","yielding","second","derivatives","."]
code: def test_jax(self, dev_name, sampler, num_directions, atol):
      """Tests that the output of the SPSA gradient transform
      can be differentiated using JAX, yielding second derivatives."""
      import jax
      from jax import numpy as jnp
  
      dev = qml.device(dev_name, wires=2)
      execute_fn = (
          dev.execute
          if dev_name == "default.qubit"
          else dev.batch_execute
      )
      params = jnp.array([0.543, -0.654])
      rng = np.random.default_rng(42)
  
      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)
          tape.trainable_params = {0, 1}
          tapes, fn = spsa_grad(
              tape,
              n=1,
              num_directions=num_directions,
              sampler=sampler,
              sampler_rng=rng,
          )
          jac = fn(execute_fn(tapes))
          if sampler is coordinate_sampler:
              jac = tuple(2 * _jac for _jac in jac)
          return jac
  
      res = jax.jacobian(cost_fn)(params)
      assert isinstance(res, tuple)
      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)],
          ]
      )
      assert np.allclose(res, expected, atol=atol, rtol=0)
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.devices import DefaultQubitLegacy
  from pennylane.gradients import spsa_grad
  from pennylane.gradients.spsa_gradient import _rademacher_sampler
  from pennylane.operation import AnyWires, Observable
label_desc: Use image node_id 6 for calling the TestSpsaGradientDifferentiation obj's underlying member method code with example usage: obj.test_jax(dev_name, sampler, num_directions, atol) and returns: jac