class_docstr: Tests for purity transform
docstr: Test purity for a QNode gradient with jax-jit.
docstr_tok: ["Test","purity","for","a","QNode","gradient","with","jax-jit","."]
code: def test_IsingXX_qnode_purity_grad_jax_jit(
      self, device, param, wires, is_partial, diff_method, interface
  ):
      """Test purity for a QNode gradient with jax-jit."""
  
      import jax
  
      dev = qml.device(device, wires=2)
  
      @qml.qnode(dev, interface=interface, diff_method=diff_method)
      def circuit_state(x):
          qml.IsingXX(x, wires=[0, 1])
          return qml.state()
  
      grad_purity = jax.jit(
          jax.grad(qml.qinfo.purity(circuit_state, wires=wires))
      )(jax.numpy.array(param))
      grad_expected_purity = (
          expected_purity_grad_ising_xx(param) if is_partial else 0
      )
  
      assert qml.math.allclose(
          grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05
      )
  
imports: import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 11 for calling the TestPurity obj's underlying member method code with example usage: obj.test_IsingXX_qnode_purity_grad_jax_jit(device, param, wires, is_partial, diff_method, interface) and returns: qml