class_docstr: Tests to ensure the transform is fully differentiable in all interfaces
  when the circuit contains two-qubit QubitUnitary operations.
  
  Note that we are not testing whether we can differentiate w.r.t. the two-qubit
  unitary itself, but rather that the entire pipeline remains differentiable even
  when we are decomposing two-qubit unitaries within.
docstr: Tests differentiability in jax interface.
docstr_tok: ["Tests","differentiability","in","jax","interface","."]
code: def test_gradient_unitary_to_rot_two_qubit_jax(self, diff_method):
      """Tests differentiability in jax interface."""
      import jax
      from jax import numpy as jnp
  
      U0 = jnp.array(test_two_qubit_unitaries[0], dtype=jnp.complex128)
      U1 = jnp.array(test_two_qubit_unitaries[1], dtype=jnp.complex128)
  
      def two_qubit_decomp_qnode(x):
          qml.RX(x, wires=0)
          qml.QubitUnitary(U0, wires=[0, 1])
          qml.QubitUnitary(U1, wires=[1, 2])
          return qml.expval(
              qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2)
          )
  
      x = jnp.array(0.1, dtype=jnp.float64)
  
      dev = qml.device("default.qubit", wires=3)
  
      original_qnode = qml.QNode(
          two_qubit_decomp_qnode,
          device=dev,
          interface="jax",
          diff_method=diff_method,
      )
  
      transformed_qfunc = unitary_to_rot(two_qubit_decomp_qnode)
      transformed_qnode = qml.QNode(
          transformed_qfunc,
          dev,
          interface="jax",
          diff_method=diff_method,
      )
  
      assert qml.math.allclose(original_qnode(x), transformed_qnode(x))
  
      # 1 normal operations + 18 for the first decomp and 6 for the second
      assert len(transformed_qnode.qtape.operations) == 25
  
      original_grad = jax.grad(original_qnode, argnums=(0))(x)
      transformed_grad = jax.grad(transformed_qnode, argnums=(0))(x)
  
      assert qml.math.allclose(
          original_grad, transformed_grad, atol=1e-6
      )
  
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 TestTwoQubitUnitaryDifferentiability obj's underlying member method code with example usage: obj.test_gradient_unitary_to_rot_two_qubit_jax(diff_method) and returns: qml