class_docstr: Tests that the template is compatible with all interfaces, including the computation
  of gradients.
class_docstr_tok: ["Tests","that","the","template","is","compatible","with","all","interfaces",",","including","the","computation","of","gradients","."]
docstr: Tests the jax interface.
docstr_tok: ["Tests","the","jax","interface","."]
code: def test_jax(self):
      """Tests the jax interface."""
      import jax
      import jax.numpy as jnp
  
      dev = qml.device("default.qubit", wires=2)
  
      input_default = [0.5, 0.2]
      input_jax = jnp.array(input_default)
  
      circuit_default = qml.QNode(manual_rx_circuit, dev)
      circuit_jax = qml.QNode(select_rx_circuit, dev)
  
      assert qml.math.allclose(
          qml.matrix(circuit_default)(input_default),
          qml.matrix(circuit_jax)(input_jax),
      )
      assert (
          qml.math.get_interface(qml.matrix(circuit_jax)(input_jax))
          == "jax"
      )
  
      grad_fn = jax.grad(circuit_default)
      grads = grad_fn(input_jax)
  
      grad_fn2 = jax.grad(circuit_jax)
      grads2 = grad_fn2(input_jax)
  
      assert qml.math.allclose(grads, grads2)
  
imports: import copy
  import pytest
  import numpy
  from pennylane import numpy
  import pennylane
label_desc: Use image node_id 5 for calling the TestInterfaces obj's underlying member method code with example usage: obj.test_jax() without return types