class_docstr: Tests to ensure the transform is fully differentiable in all interfaces for
  single-qubit unitary decompositions.
docstr: Tests differentiability in jax interface.
docstr_tok: ["Tests","differentiability","in","jax","interface","."]
code: def test_gradient_unitary_to_rot_jax(self, rot_angles, diff_method):
      """Tests differentiability in jax interface."""
      import jax
      from jax import numpy as jnp
  
      def qfunc_with_qubit_unitary(angles):
          z = angles[0]
          x = angles[1]
  
          Z_mat = jnp.array(
              [[jnp.exp(-1j * z / 2), 0.0], [0.0, jnp.exp(1j * z / 2)]]
          )
  
          c = jnp.cos(x / 2)
          s = jnp.sin(x / 2) * 1j
          X_mat = jnp.array([[c, -s], [-s, c]])
  
          qml.Hadamard(wires="a")
          qml.QubitUnitary(Z_mat, wires="a")
          qml.QubitUnitary(X_mat, wires="b")
          qml.CNOT(wires=["b", "a"])
          return qml.expval(qml.PauliX(wires="a"))
  
      # Setting the dtype to complex64 causes the gradients to be complex...
      angles = jnp.array(rot_angles, dtype=jnp.float64)
  
      dev = qml.device("default.qubit", wires=["a", "b"])
  
      original_qnode = qml.QNode(
          original_qfunc_for_grad, dev, diff_method=diff_method
      )
      original_result = original_qnode(angles)
  
      transformed_qfunc = unitary_to_rot(qfunc_with_qubit_unitary)
      transformed_qnode = qml.QNode(
          transformed_qfunc, dev, diff_method=diff_method
      )
      transformed_result = transformed_qnode(angles)
      assert qml.math.allclose(original_result, transformed_result)
  
      original_grad = jax.grad(original_qnode)(angles)
      transformed_grad = jax.grad(transformed_qnode)(angles)
      assert qml.math.allclose(
          original_grad, transformed_grad, atol=1e-7
      )
  
      # Check that we can also JIT
      grad_of_jit = jax.grad(jax.jit(transformed_qnode))(angles)
      jit_of_grad = jax.jit(jax.grad(transformed_qnode))(angles)
      assert np.allclose(original_grad, jit_of_grad, atol=1e-7)
      assert np.allclose(original_grad, grad_of_jit, atol=1e-7)
  
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 4 for calling the TestQubitUnitaryDifferentiability obj's underlying member method code with example usage: obj.test_gradient_unitary_to_rot_jax(rot_angles, diff_method) and returns: qml