docstr: Test qinfo meausrements work with jax and jitting.
docstr_tok: ["Test","qinfo","meausrements","work","with","jax","and","jitting","."]
code: def test_qinfo_jax(self, use_jit):
      """Test qinfo meausrements work with jax and jitting."""
  
      import jax
  
      def f(x):
          qs = qml.tape.QuantumScript(
              [qml.IsingXX(x, wires=(0, 1))], self.measurements
          )
          return simulate(qs)
  
      if use_jit:
          f = jax.jit(f)
  
      phi = jax.numpy.array(-0.792)
  
      results = f(phi)
      for val1, val2 in zip(results, self.expected_results(phi)):
          assert qml.math.allclose(val1, val2)
  
      def real_out(phi):
          return tuple(jax.numpy.real(r) for r in f(phi))
  
      grad_real = jax.jacobian(real_out)(phi)
  
      def imag_out(phi):
          return tuple(jax.numpy.imag(r) for r in f(phi))
  
      grad_imag = jax.jacobian(imag_out)(phi)
      grads = tuple(r + 1j * i for r, i in zip(grad_real, grad_imag))
      expected_grads = self.expected_grad(phi)
  
      # Writing this way makes it easier to figure out which is failing
      # density 0
      assert qml.math.allclose(grads[0], expected_grads[0])
      # density 1
      assert qml.math.allclose(grads[1], expected_grads[1])
      # density both
      assert qml.math.allclose(grads[2], expected_grads[2])
      # entropy 0
      assert qml.math.allclose(grads[3], expected_grads[3])
      # entropy 1
      assert qml.math.allclose(grads[4], expected_grads[4])
      # mutual info
      assert qml.math.allclose(grads[5], expected_grads[5])
  
imports: import pytest
  import numpy
  import pennylane
  from pennylane.devices.qubit import simulate, get_final_state, measure_final_state
label_desc: Use image node_id 7 for calling the TestQInfoMeasurements obj's underlying member method code with example usage: obj.test_qinfo_jax(use_jit) and returns: simulate, tuple, tuple