code: def execute(
      program, args, max_stack_depth, backend, block_code_cache=None
  ):
      """Executes or stages a complete auto-batching VM program.
  
      Whether this executes or stages computation depends on whether the backend has
      an eager or deferred computation model.
  
      The dimensions of the inputs and internal variables are split into
      one top batch dimension and an arbitrary number (here `E`) event
      dimensions.  The event rank may be different for different inputs,
      outputs, and internal variables.
  
      Args:
        program: A `instructions.Program` to execute or stage.
        args: Input values, a list of arrays, each of shape `[batch_size,
          e1, ..., eE]`.  The batch size must be the same for all inputs.
          The other dimensions must agree with the declared shapes of the
          variables they will be stored in, but need not in general be the
          same as one another.
        max_stack_depth: Python `int`. Maximum depth of stack to allocate.
        backend: Object implementing required backend operations.
        block_code_cache: Dict (allows cache to live across calls to `vm.execute`,
          or `None` (in which case a dict is created and used per call).
  
      Returns:
        results: A list of the output values. Each returned value is an
          array of shape `[batch_size, e1, ..., eE]`.  The results are
          returned in the same order as the variables appear in
          `program.out_vars`.
      """
      program = select_block_priority(program)
      halt_index = program.graph.exit_index()
      logging.vlog(
          1, "Staging computation: %d blocks to stage", halt_index
      )
      valid_indices = range(halt_index)
      assert len(program.vars_in) == len(args)
      init_vals = dict(zip(program.vars_in, args))
      environment = _initialize_environment(
          program, init_vals, max_stack_depth, backend
      )
      next_block_index = _choose_next_op(environment, backend)
  
      if block_code_cache is None:
          block_code_cache = {}
  
      def make_run_block_callable(env):
          """Makes a i->next_env callable using cached, backend-wrapped _run_block."""
  
          def run_block_callable(i):
              if i not in block_code_cache:
                  logging.vlog(1, "Fill block code cache: block %d", i)
                  block_code_cache[
                      i
                  ] = backend.wrap_straightline_callable(
                      lambda env_arg: _run_block(
                          program.graph, i, env_arg, backend
                      )
                  )
              else:
                  logging.vlog(1, "Use cached block code: block %d", i)
              return block_code_cache[i](env)
  
          return run_block_callable
  
      def cond(_, next_block_index):
          return backend.not_equal(next_block_index, halt_index)
  
      def body(
          env_dict, next_block_index
      ):  # pylint:disable=missing-docstring
          # This `_staged_apply` turns into the block dispatch tree (see
          # docstring of `_staged_apply`).
          # Briefly, this will build a graph snippet for each basic block
          # in the control flow graph, and glue them together with a switch
          # on the runtime value of `next_block_index`.
          env_dict = backend.prepare_for_cond(env_dict)
          f = make_run_block_callable(env_dict)
          env_dict = backend.switch_case(
              next_block_index,
              [functools.partial(f, i) for i in valid_indices],
          )
          next_block_index = _choose_next_op(
              inst.Environment(env_dict, backend), backend
          )
          return env_dict, next_block_index
  
      env_dict, _ = backend.while_loop(
          cond, body, [environment.env_dict, next_block_index]
      )
      final_env = inst.Environment(env_dict, backend)
      return inst.pattern_map(final_env.read, program.vars_out)
  
file_path: probability/tensorflow_probability/python/experimental/auto_batching/virtual_machine.py
imports: import contextlib
  import functools
  from absl import logging
  import six
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 1 for calling a global function with example usage: execute(program, args, max_stack_depth, backend, block_code_cache) and returns: inst, run_block_callable, backend, env_dict,  next_block_index, block_code_cache