code: def _initialize_pc(environment, program, batch_size, mask, backend):
      """Write the requisite initial values to the program counter."""
      # Initializing the program counter is a bit subtle.  The initial function call
      # to main is a tail call, so doesn't change the PC's stack depth.  However,
      # main will try to "return" by executing IndirectGotoOp when it's done.  This
      # should have the effect of setting that thread's (top) PC to the halt index,
      # so the VM knows not to keep executing it.
      #
      # We arrange this effect here by initializing the PC to depth 2: the bottom
      # row is the halt index, and the top row is the initial PC, i.e., 0.  This
      # way, when main pops the PC at the end, the current value will become the
      # halt index, as desired.
      #
      # This solution may appear somewhat unfortunate, because it forces the PC to
      # have a stack, i.e., be allocated as a FULL variable.  However, if the
      # program has any IndirectGotoOp in it, we need a stack for the PC anyway; and
      # if it doesn't, then it's ok to allocate the PC as a REGISTER and drop the
      # halt index on the floor, because there won't be any IndirectGotoOp to read
      # it.
      #
      # Other alternatives considered:
      # - Define a REGISTER_WITH_DEFAULT variable allocation strategy, which is
      #   a register that writes a default value into itself when popped.
      # - Separate the PC into an explicit register holding the current PC, together
      #   with a Variable for the PC stack.  Then the latter could itself be a
      #   register (holding the halt index) if no other PC pushes are needed in the
      #   program.
      var_defs = program.var_defs
      halt_pc = backend.fill(
          program.graph.exit_index(),
          batch_size,
          dtype=var_defs[inst.pc_var].tensors.dtype,
          shape=[],
      )
      environment[inst.pc_var] = environment.push(
          inst.pc_var, halt_pc, mask
      )
      start_pc = backend.fill(
          0,
          batch_size,
          dtype=var_defs[inst.pc_var].tensors.dtype,
          shape=[],
      )
      environment[inst.pc_var] = environment.push(
          inst.pc_var, start_pc, mask
      )
      return environment
  
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 7 for calling a global function with example usage: _initialize_pc(environment, program, batch_size, mask, backend) and returns: environment