class_docstr: Test the jax interface execute function
  integrates well for both forward and backward execution
class_docstr_tok: ["Test","the","jax","interface","execute","function","integrates","well","for","both","forward","and","backward","execution"]
docstr: Test that operation and nested tapes expansion
  is differentiable
docstr_tok: ["Test","that","operation","and","nested","tapes","expansion","is","differentiable"]
code: def test_differentiable_expand(self, execute_kwargs, tol):
      """Test that operation and nested tapes expansion
      is differentiable"""
  
      class U3(qml.U3):
          def expand(self):
              theta, phi, lam = self.data
              wires = self.wires
              return [
                  qml.Rot(lam, theta, -lam, wires=wires),
                  qml.PhaseShift(phi + lam, wires=wires),
              ]
  
      def cost_fn(a, p, device):
          qscript = qml.tape.QuantumScript(
              [qml.RX(a, wires=0), U3(*p, wires=0)],
              [qml.expval(qml.PauliX(0))],
          )
          qscript = qscript.expand(
              stop_at=lambda obj: device.supports_operation(obj.name)
          )
          return execute([qscript], device, **execute_kwargs)[0]
  
      a = jax.numpy.array(0.1)
      p = jax.numpy.array([0.1, 0.2, 0.3])
  
      dev = qml.device("default.qubit.legacy", wires=1)
      res = jax.jit(cost_fn, static_argnums=2)(a, p, device=dev)
      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)
  
      jac_fn = jax.jit(jax.grad(cost_fn, argnums=1), static_argnums=2)
      res = jac_fn(a, p, device=dev)
      expected = jax.numpy.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 numpy
  import pennylane
  from pennylane.gradients import param_shift
  from pennylane import execute
label_desc: Use image node_id 9 for calling the TestJaxExecuteIntegration obj's underlying member method code with example usage: obj.test_differentiable_expand(execute_kwargs, tol) and returns: execute, list