class_docstr: Tests that the ControlledSequence is executable and differentiable in a QNode context
class_docstr_tok: ["Tests","that","the","ControlledSequence","is","executable","and","differentiable","in","a","QNode","context"]
docstr: Test that the QNode executes and is differentiable with JAX. The shots
  argument controls whether autodiff or parameter-shift gradients are used.
docstr_tok: ["Test","that","the","QNode","executes","and","is","differentiable","with","JAX",".","The","shots","argument","controls","whether","autodiff","or","parameter-shift","gradients","are","used","."]
code: def test_qnode_jax(self, shots, use_jit):
      """Test that the QNode executes and is differentiable with JAX. The shots
      argument controls whether autodiff or parameter-shift gradients are used.
      """
      import jax
  
      jax.config.update("jax_enable_x64", True)
  
      dev = qml.device("default.qubit", shots=shots, seed=10)
      diff_method = "backprop" if shots is None else "parameter-shift"
      qnode = qml.QNode(
          self.circuit, dev, interface="jax", diff_method=diff_method
      )
      if use_jit:
          qnode = jax.jit(qnode)
  
      x = jax.numpy.array(self.x)
      res = qnode(x)
      assert qml.math.shape(res) == (16,)
      assert np.allclose(res, self.exp_result, atol=0.005)
  
      jac_fn = jax.jacobian(qnode)
      if use_jit:
          jac_fn = jax.jit(jac_fn)
  
      jac = jac_fn(x)
      assert jac.shape == (16,)
      assert np.allclose(jac, self.exp_jac, atol=0.006)
  
imports: import pytest
  import numpy
  from pennylane import numpy
  import pennylane
  from pennylane.wires import Wires
label_desc: Use image node_id 4 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_qnode_jax(shots, use_jit) without return types