class_docstr: Tests for the Trace Distance function between two QNodes.
docstr: Test the gradient of the trace distance between two trainable circuits with JAX.
docstr_tok: ["Test","the","gradient","of","the","trace","distance","between","two","trainable","circuits","with","JAX","."]
code: def test_trace_distance_qnodes_rx_tworx_jax_grad(
      self, param, wire, device, use_jit
  ):
      """Test the gradient of the trace distance between two trainable circuits with JAX."""
      dev = qml.device(device, wires=wire)
  
      @qml.qnode(dev, interface="jax" if not use_jit else "jax-jit")
      def circuit(x):
          qml.RX(x, wires=wire - 1)
          return qml.state()
  
      td_grad = jax.grad(
          qml.qinfo.trace_distance(
              circuit, circuit, wires0=[wire - 1], wires1=[wire - 1]
          ),
          argnums=[0, 1],
      )
  
      if use_jit:
          td_grad = jax.jit(td_grad)
  
      td_grad = td_grad(
          (jax.numpy.array(param)),
          (jax.numpy.array(2 * param)),
      )
      expected = expected_grad_trace_distance_rx_pauliz(param)
      expected_td = [-expected, expected]
      assert qml.math.allclose(
          td_grad, expected_td, rtol=1e-04, atol=1e-03
      )
  
imports: import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 12 for calling the TestTraceDistanceQnode obj's underlying member method code with example usage: obj.test_trace_distance_qnodes_rx_tworx_jax_grad(param, wire, device, use_jit) and returns: qml