code: def _execute_wrapper(
      params, tapes, execute_fn, _, device
  ) -> ResultBatch:
      """
      Execute tapes using a pure-callback.
      """
      shape_dtype_structs = tuple(
          _result_shape_dtype_struct(t, device) for t in tapes.vals
      )
  
      def pure_callback_wrapper(p):
          new_tapes = _set_all_parameters_on_copy(tapes.vals, p)
          res = tuple(_to_jax(execute_fn(new_tapes)))
          # When executed under `jax.vmap` the `result_shapes_dtypes` will contain
          # the shape without the vmap dimensions, while the function here will be
          # executed with objects containing the vmap dimensions. So res[i].ndim
          # will have an extra dimension for every `jax.vmap` wrapping this execution.
          #
          # The execute_fn will return an object with shape `(n_observables, batches)`
          # but the default behaviour for `jax.pure_callback` is to add the extra
          # dimension at the beginning, so `(batches, n_observables)`. So in here
          # we detect with the heuristic above if we are executing under vmap and we
          # swap the order in that case.
          return jax.tree_map(
              lambda r, s: r.T if r.ndim > s.ndim else r,
              res,
              shape_dtype_structs,
          )
  
      out = jax.pure_callback(
          pure_callback_wrapper,
          shape_dtype_structs,
          params,
          vectorized=True,
      )
      return out
  
file_path: pennylane/pennylane/workflow/interfaces/jax_jit.py
imports: import jax
  import jax.numpy
  import pennylane
  from pennylane.typing import ResultBatch
  from ..jacobian_products import _compute_jvps
  from .jax import _NonPytreeWrapper
label_desc: Use image node_id 8 for calling a global function with example usage: _execute_wrapper(params, tapes, execute_fn, _, device) and returns: out, jax