class_docstr: Tests for the mutual information functions
class_docstr_tok: ["Tests","for","the","mutual","information","functions"]
docstr: Test that the gradient of mutual information works for QNodes
  with the JAX-jit interface
docstr_tok: ["Test","that","the","gradient","of","mutual","information","works","for","QNodes","with","the","JAX-jit","interface"]
code: def test_qnode_grad_jax_jit(self, param, diff_method, interface):
      """Test that the gradient of mutual information works for QNodes
      with the JAX-jit interface"""
      import jax
      import jax.numpy as jnp
  
      dev = qml.device("default.qubit", wires=2)
  
      param = jnp.array(param)
  
      @qml.qnode(dev, interface=interface, diff_method=diff_method)
      def circuit(param):
          qml.RY(param, wires=0)
          qml.CNOT(wires=[0, 1])
          return qml.mutual_info(wires0=[0], wires1=[1])
  
      if param == 0:
          # we don't allow gradients to flow through the discontinuity at 0
          expected = 0
      else:
          expected = jnp.sin(param) * (
              jnp.log(jnp.cos(param / 2) ** 2)
              - jnp.log(jnp.sin(param / 2) ** 2)
          )
  
      # higher tolerance for finite-diff method
      tol = 1e-8 if diff_method == "backprop" else 1e-5
  
      actual = jax.jit(jax.grad(circuit))(param)
      assert np.allclose(actual, expected, atol=tol)
  
imports: import copy
  import numpy
  import pytest
  import pennylane
  from pennylane.measurements import MutualInfo, Shots
  from pennylane.measurements.mutual_info import MutualInfoMP
  from pennylane.wires import Wires
label_desc: Use image node_id 12 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_qnode_grad_jax_jit(param, diff_method, interface) and returns: qml