class_docstr: Test that the top-level compile function works across all interfaces.
class_docstr_tok: ["Test","that","the","top-level","compile","function","works","across","all","interfaces","."]
docstr: Test that compilation pipelines work with jax.jit, unitary_to_rot, and fusion.
docstr_tok: ["Test","that","compilation","pipelines","work","with","jax.jit",",","unitary_to_rot",",","and","fusion","."]
code: def test_compile_jax_jit(self, diff_method):
      """Test that compilation pipelines work with jax.jit, unitary_to_rot, and fusion."""
      import jax
      from jax import numpy as jnp
  
      dev = qml.device("default.qubit", wires=2)
  
      def test_qfunc(x):
          qml.Rot(x, x + 1, x + 2, wires=0)
          qml.RX(2 * x, wires=0)
          qml.CNOT(wires=[0, 1])
          qml.QubitUnitary(qml.RX.compute_matrix(x), wires=1)
          qml.QubitUnitary(qml.RZ.compute_matrix(x + 1), wires=1)
          return qml.expval(qml.PauliX(0) @ qml.PauliZ(1))
  
      original_qnode = qml.QNode(
          test_qfunc, dev, diff_method=diff_method
      )
  
      pipeline = [cancel_inverses, unitary_to_rot, single_qubit_fusion]
  
      compiled_qfunc = qml.compile(test_qfunc, pipeline=pipeline)
      compiled_qnode = qml.QNode(
          compiled_qfunc, dev, diff_method=diff_method
      )
  
      jitted_compiled_qnode = jax.jit(compiled_qnode)
  
      x = jnp.array(0.1, dtype=jnp.float64)
  
      # Check that the numerical output is the same
      assert qml.math.allclose(
          original_qnode(x), jitted_compiled_qnode(x)
      )
  
      # Check that the gradient is the same
      assert qml.math.allclose(
          jax.grad(original_qnode)(x),
          jax.grad(jitted_compiled_qnode)(x),
          atol=1e-7,
      )
  
imports: from functools import partial
  import pytest
  from test_optimization.utils import compare_operation_lists
  import pennylane
  from pennylane import numpy
  from pennylane.wires import Wires
  from pennylane.transforms.compile import compile
  from pennylane.transforms import unitary_to_rot
  from pennylane.transforms.optimization import cancel_inverses, commute_controlled, merge_rotations, single_qubit_fusion
  from pennylane.transforms.optimization.optimization_utils import _fuse_global_phases
label_desc: Use image node_id 5 for calling the TestCompileInterfaces obj's underlying member method code with example usage: obj.test_compile_jax_jit(diff_method) and returns: qml