class_docstr: Gradient tests for the vjp function
class_docstr_tok: ["Gradient","tests","for","the","vjp","function"]
docstr: Tests that the output of the VJP transform
  can be differentiated using JAX.
docstr_tok: ["Tests","that","the","output","of","the","VJP","transform","can","be","differentiated","using","JAX","."]
code: def test_jax(self, dev_name, tol):
      """Tests that the output of the VJP 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)
      params = jnp.array(params_np)
  
      @partial(jax.jit)
      def cost_fn(x):
          with qml.queuing.AnnotatedQueue() as q:
              ansatz(x[0], x[1])
          tape = qml.tape.QuantumScript.from_queue(q)
          dy = jax.numpy.array([-1.0, 0.0, 0.0, 1.0])
          tape.trainable_params = {0, 1}
          tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
          vjp = fn(execute_fn(tapes))
          return vjp
  
      res = cost_fn(params)
      assert np.allclose(res, expected(params), atol=tol, rtol=0)
  
      res = jax.jacobian(cost_fn, argnums=0)(params)
      exp = qml.jacobian(expected)(params_np)
      assert np.allclose(res, exp, atol=tol, rtol=0)
  
imports: from functools import partial
  import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.gradients import param_shift
label_desc: Use image node_id 4 for calling the TestVJPGradients obj's underlying member method code with example usage: obj.test_jax(dev_name, tol) and returns: vjp