code: def _interpret(program, mask, backend, block_code_cache, *inputs):
      """Worker function for `execute`; operates under a mask."""
      environment = inst.Environment.initialize(
          backend,
          program.var_alloc,
          program.var_defs,
          0,
          backend.batch_size(mask),
      )
      for var, inp in inst.pattern_zip(program.vars_in, inputs):
          environment[var] = environment.push(var, inp, mask)
      program_counter = 0  # Index of initial block
      queue = ExecutionQueue(backend)
      while program_counter != program.graph.exit_index():
          block = program.graph.block(program_counter)
          for split_idx, split in enumerate(
              _split_fn_calls(block.instructions)
          ):
              if isinstance(split, inst.FunctionCallOp):
                  op = split
                  if inst.pc_var in inst.pattern_flatten(
                      op.vars_in
                  ) or inst.pc_var in inst.pattern_flatten(op.vars_out):
                      raise ValueError(
                          "FunctionCallOp reads or writes program counter: {}".format(
                              op
                          )
                      )
                  inputs = [
                      inst.pattern_map(environment.read, var_pat)
                      for var_pat in op.vars_in
                  ]
                  # Option: Could gather+scatter at function boundaries.  Then the inner
                  # interpreter would not need to accept the mask, but would need to
                  # recompute the batch size and make a new mask of all ones.
                  outputs = _invoke_fun(
                      program,
                      mask,
                      backend,
                      block_code_cache,
                      op.function,
                      inputs,
                  )
                  new_vars = [
                      (varname, environment.push(varname, output, mask))
                      for varname, output in inst.pattern_zip(
                          op.vars_out, outputs
                      )
                  ]
                  environment = inst.Environment(
                      environment.env_dict,
                      environment.backend,
                      update=new_vars,
                  )
              else:  # This split is not a FunctionCallOp.
                  block_code_key = (
                      id(program.graph),
                      program_counter,
                      split_idx,
                  )
                  if block_code_key not in block_code_cache:
                      logging.vlog(
                          1,
                          "Fill block cache for block %s",
                          block_code_key,
                      )
                      varnames = inst.extract_referenced_variables(
                          split
                      )
                      code = backend.wrap_straightline_callable(
                          functools.partial(
                              _run_straightline,
                              split,
                              environment.backend,
                          )
                      )
                      block_code_cache[block_code_key] = (
                          varnames,
                          code,
                      )
                  else:
                      logging.vlog(
                          1,
                          "Use cached code for block %s",
                          block_code_key,
                      )
                  varnames, code = block_code_cache[block_code_key]
                  filtered_env = dict(
                      {  # Only pass variables relevant to these ops
                          k: v
                          for k, v in six.iteritems(
                              environment.env_dict
                          )
                          if k in varnames
                      }
                  )
                  environment = inst.Environment(
                      environment.env_dict,
                      environment.backend,
                      update=code(filtered_env, mask),
                  )
          op = block.terminator
          if isinstance(op, inst.BranchOp):
              if inst.pc_var == op.cond_var:
                  raise ValueError(
                      "Branching on program counter: {}".format(op)
                  )
              condition = environment.read(op.cond_var)
              true_index = program.graph.block_index(op.true_block)
              false_index = program.graph.block_index(op.false_block)
              queue.enqueue(true_index, mask & condition)
              queue.enqueue(false_index, mask & ~condition)
          elif isinstance(op, inst.GotoOp):
              next_index = program.graph.block_index(op.block)
              queue.enqueue(next_index, mask)
          else:
              raise TypeError("Unexpected op type: {}".format(type(op)))
          program_counter, mask = queue.dequeue()
      # assert the queue is now empty
      return inst.pattern_map(environment.read, program.vars_out)
  
file_path: probability/tensorflow_probability/python/experimental/auto_batching/stackless.py
imports: import contextlib
  import functools
  import heapq
  from absl import logging
  import six
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 4 for calling a global function with example usage: _interpret(program, mask, backend, block_code_cache) and returns: inst