class_docstr: Tests to ensure the transform itself works in all interfaces.
docstr: Test that the transform works in the JAX interface with JIT.
docstr_tok: ["Test","that","the","transform","works","in","the","JAX","interface","with","JIT","."]
code: def test_unitary_to_rot_jax_jit(
      self, U, expected_gates, expected_params
  ):
      """Test that the transform works in the JAX interface with JIT."""
      # pylint: disable=unused-argument
      import jax
  
      U = jax.numpy.array(U, dtype=jax.numpy.complex128)
  
      dev = qml.device("default.qubit", wires=["a", "b"])
  
      def test_qfunc(U):
          qfunc(U)
          return qml.expval(qml.PauliX("a") @ qml.PauliZ("b"))
  
      original_qnode = qml.QNode(test_qfunc, dev)
      transformed_qnode = qml.QNode(unitary_to_rot(test_qfunc), dev)
      jitted_qnode = jax.jit(transformed_qnode)
  
      original_result = original_qnode(U)
      transformed_result = transformed_qnode(U)
      jitted_result = jitted_qnode(U)
      assert np.allclose(transformed_result, original_result)
      assert np.allclose(jitted_result, original_result)
  
imports: from itertools import product
  import pytest
  from gate_data import I, Z, S, T, H, X
  from test_optimization.utils import check_matrix_equivalence
  import pennylane
  from pennylane import numpy
  from pennylane.wires import Wires
  from pennylane.transforms import unitary_to_rot
label_desc: Use image node_id 7 for calling the TestDecomposeSingleQubitUnitaryTransform obj's underlying member method code with example usage: obj.test_unitary_to_rot_jax_jit(U, expected_gates, expected_params) and returns: qml