class_docstr: Test the decomposition in the non-autograd interfaces.
class_docstr_tok: ["Test","the","decomposition","in","the","non-autograd","interfaces","."]
docstr: Test that a two-qubit operation is correctly decomposed with JAX-JIT .
docstr_tok: ["Test","that","a","two-qubit","operation","is","correctly","decomposed","with","JAX-JIT","."]
code: def test_two_qubit_decomposition_jax_jit(self, U, wires):
      """Test that a two-qubit operation is correctly decomposed with JAX-JIT ."""
  
      import jax
  
      U = jax.numpy.array(U, dtype=jax.numpy.complex128)
  
      def wrapped_decomposition(U):
          # the jitted function cannot return objects like operators,
          # so we cannot jit two_qubit_decomposition directly
          obtained_decomposition = two_qubit_decomposition(
              U, wires=wires
          )
  
          with qml.queuing.AnnotatedQueue() as q:
              for op in obtained_decomposition:
                  qml.apply(op)
  
          tape = qml.tape.QuantumScript.from_queue(q)
          obtained_matrix = qml.matrix(tape, wire_order=wires)
  
          return obtained_matrix
  
      jitted_matrix = jax.jit(wrapped_decomposition)(U)
  
      assert check_matrix_equivalence(U, jitted_matrix, atol=1e-7)
  
imports: from functools import reduce
  import pytest
  from gate_data import I, Z, S, T, H, X, Y, CNOT, SWAP
  import pennylane
  from pennylane import numpy
  from pennylane.wires import Wires
  from pennylane.ops.op_math.decompositions import one_qubit_decomposition
  from pennylane.ops.op_math.decompositions import two_qubit_decomposition
  from pennylane.ops.op_math.decompositions.two_qubit_unitary import _convert_to_su4, _su2su2_to_tensor_products, _compute_num_cnots
label_desc: Use image node_id 7 for calling the TestTwoQubitUnitaryDecompositionInterfaces obj's underlying member method code with example usage: obj.test_two_qubit_decomposition_jax_jit(U, wires) and returns: obtained_matrix