class_docstr: Tests for the mutual information functions
class_docstr_tok: ["Tests","for","the","mutual","information","functions"]
docstr: Test that the mutual information transform works for QNodes by comparing
  against analytic values, for the JAX-jit interface
docstr_tok: ["Test","that","the","mutual","information","transform","works","for","QNodes","by","comparing","against","analytic","values",",","for","the","JAX-jit","interface"]
code: def test_qnode_state_jax_jit(self, params):
      """Test that the mutual information transform works for QNodes by comparing
      against analytic values, for the JAX-jit interface"""
      import jax
      import jax.numpy as jnp
  
      dev = qml.device("default.qubit", wires=2)
  
      params = jnp.array(params)
  
      @qml.qnode(dev, interface="jax-jit")
      def circuit(params):
          qml.RY(params, wires=0)
          qml.CNOT(wires=[0, 1])
          return qml.state()
  
      actual = jax.jit(
          qml.qinfo.mutual_info(circuit, wires0=[0], wires1=[1])
      )(params)
  
      # compare transform results with analytic values
      expected = -2 * jnp.cos(params / 2) ** 2 * jnp.log(
          jnp.cos(params / 2) ** 2 + 1e-10
      ) - 2 * jnp.sin(params / 2) ** 2 * jnp.log(
          jnp.sin(params / 2) ** 2 + 1e-10
      )
  
      assert np.allclose(actual, expected)
  
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 8 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_qnode_state_jax_jit(params) and returns: qml