class_docstr: Test gradients for a single measurement
class_docstr_tok: ["Test","gradients","for","a","single","measurement"]
docstr: Test derivatives when using JAX
docstr_tok: ["Test","derivatives","when","using","JAX"]
code: def test_jax_jit(self, diff_method, tol, interface):
      """Test derivatives when using JAX"""
      import jax
      import jax.numpy as jnp
  
      jax.config.update("jax_enable_x64", True)
  
      dev = qml.device("default.qubit", wires=2)
  
      @jax.jit
      @partial(qml.batch_input, argnum=0)
      @qml.qnode(dev, diff_method=diff_method, interface=interface)
      def circuit(input, x):
          qml.RY(input, wires=1)
          qml.CNOT(wires=[0, 1])
          qml.RX(x, wires=0)
          return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
  
      batch_size = 3
  
      def cost(input, x):
          return jnp.sum(circuit(input, x))
  
      input = jnp.linspace(0.1, 0.5, batch_size)
      x = jnp.array(0.1)
  
      res = jax.grad(cost, argnums=1)(input, x)
      expected = -np.sin(0.1) * sum(np.sin(input))
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
imports: from functools import partial
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 3 for calling the TestDiffSingle obj's underlying member method code with example usage: obj.test_jax_jit(diff_method, tol, interface) and returns: qml, jnp