class_docstr: Test that using the QNode with JAX integrates with the PennyLane
  stack
class_docstr_tok: ["Test","that","using","the","QNode","with","JAX","integrates","with","the","PennyLane","stack"]
docstr: Test that operation and nested tape expansion
  is differentiable
docstr_tok: ["Test","that","operation","and","nested","tape","expansion","is","differentiable"]
code: def test_differentiable_expand(
      self, dev_name, diff_method, grad_on_execution, interface, tol
  ):
      """Test that operation and nested tape expansion
      is differentiable"""
  
      kwargs = dict(
          diff_method=diff_method,
          interface=interface,
          grad_on_execution=grad_on_execution,
      )
      if diff_method == "spsa":
          spsa_kwargs = dict(
              sampler_rng=np.random.default_rng(SEED_FOR_SPSA),
              num_directions=10,
          )
          kwargs = {**kwargs, **spsa_kwargs}
          tol = TOL_FOR_SPSA
  
      class U3(qml.U3):
          def decomposition(self):
              theta, phi, lam = self.data
              wires = self.wires
  
              with qml.queuing.AnnotatedQueue() as q_tape:
                  qml.Rot(lam, theta, -lam, wires=wires)
                  qml.PhaseShift(phi + lam, wires=wires)
  
              tape = QuantumScript.from_queue(q_tape)
              return tape.operations
  
      num_wires = 1
  
      if diff_method == "hadamard":
          num_wires = 2
  
      dev = qml.device(dev_name, wires=num_wires)
      a = jax.numpy.array(0.1)
      p = jax.numpy.array([0.1, 0.2, 0.3])
  
      @qnode(dev, **kwargs)
      def circuit(a, p):
          qml.RX(a, wires=0)
          U3(p[0], p[1], p[2], wires=0)
          return qml.expval(qml.PauliX(0))
  
      res = circuit(a, p)
      expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
          np.cos(p[2]) * np.sin(p[1])
          + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
      )
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
      res = jax.grad(circuit, argnums=1)(a, p)
      expected = np.array(
          [
              np.cos(p[1])
              * (
                  np.cos(a) * np.cos(p[0])
                  - np.sin(a) * np.sin(p[0]) * np.sin(p[2])
              ),
              np.cos(p[1]) * np.cos(p[2]) * np.sin(a)
              - np.sin(p[1])
              * (
                  np.cos(a) * np.sin(p[0])
                  + np.cos(p[0]) * np.sin(a) * np.sin(p[2])
              ),
              np.sin(a)
              * (
                  np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2])
                  - np.sin(p[1]) * np.sin(p[2])
              ),
          ]
      )
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
imports: import pytest
  import pennylane
  from pennylane import numpy
  from pennylane import qnode
  from pennylane.tape import QuantumScript
label_desc: Use image node_id 5 for calling the TestQNode obj's underlying member method code with example usage: obj.test_differentiable_expand(dev_name, diff_method, grad_on_execution, interface, tol) and returns: qml, tape