class_docstr: Tests for the Trace Distance function between two QNodes.
docstr: Test the gradient of the trace distance between Rx and PauliZ circuits with Jax and
  two params.
docstr_tok: ["Test","the","gradient","of","the","trace","distance","between","Rx","and","PauliZ","circuits","with","Jax","and","two","params","."]
code: def test_trace_distance_qnodes_rx_pauliz_jax_grad_two_params(
      self, param, wire, use_jit
  ):
      """Test the gradient of the trace distance between Rx and PauliZ circuits with Jax and
      two params."""
      dev = qml.device("default.qubit", wires=wire)
  
      @qml.qnode(dev, interface="jax" if not use_jit else "jax-jit")
      def circuit0(x):
          qml.RX(x, wires=wire - 1)
          return qml.state()
  
      @qml.qnode(dev, interface="jax" if not use_jit else "jax-jit")
      def circuit1(x):
          qml.RX(x, wires=wire - 1)
          qml.RX(-x, wires=wire - 1)
          qml.PauliZ(wires=wire - 1)
          return qml.state()
  
      td_grad = jax.grad(
          qml.qinfo.trace_distance(
              circuit0, circuit1, 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.0))
      expected_td_grad = expected_grad_trace_distance_rx_pauliz(param)
      assert qml.math.allclose(
          td_grad, (expected_td_grad, 0.0), rtol=1e-03, atol=1e-04
      )
  
      td_grad = jax.grad(
          qml.qinfo.trace_distance(
              circuit1, circuit0, 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(2.0), jax.numpy.array(param))
      assert qml.math.allclose(
          td_grad, (0.0, expected_td_grad), rtol=1e-03, atol=1e-04
      )
  
imports: import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 16 for calling the TestTraceDistanceQnode obj's underlying member method code with example usage: obj.test_trace_distance_qnodes_rx_pauliz_jax_grad_two_params(param, wire, use_jit) and returns: qml, qml