class_docstr: Gradient tests for the jvp function
class_docstr_tok: ["Gradient","tests","for","the","jvp","function"]
docstr: Tests that the output of the JVP transform
  can be differentiated using JAX.
docstr_tok: ["Tests","that","the","output","of","the","JVP","transform","can","be","differentiated","using","JAX","."]
code: def test_jax(self, tol, dev_name, batch_dim):
      """Tests that the output of the JVP transform
      can be differentiated using JAX."""
      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_np = np.array([0.543, -0.654], requires_grad=True)
      if batch_dim is not None:
          params_np = np.outer(
              np.arange(1, 1 + batch_dim), params_np, requires_grad=True
          )
      tangent_np = np.array([1.2, -0.3], requires_grad=False)
      params = jnp.array(params_np)
      tangent = jnp.array(tangent_np)
  
      def cost_fn(params, tangent):
          with qml.queuing.AnnotatedQueue() as q:
              ansatz(params[..., 0], params[..., 1])
  
          tape = qml.tape.QuantumScript.from_queue(q)
          tape.trainable_params = {0, 1}
          tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
          jvp = fn(execute_fn(tapes))
          return jvp
  
      res = cost_fn(params, tangent)
      exp = expected_jvp(params_np, tangent_np)
      assert np.allclose(res, exp, atol=tol, rtol=0)
  
      res = jax.jacobian(cost_fn)(params, tangent)
      exp = qml.jacobian(expected_jvp)(params_np, tangent_np)
      assert np.allclose(res, exp, atol=tol, rtol=0)
  
imports: import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.gradients import param_shift
label_desc: Use image node_id 4 for calling the TestJVPGradients obj's underlying member method code with example usage: obj.test_jax(tol, dev_name, batch_dim) and returns: jvp