code: def optimize(program):
      """Optimizes a `Program`'s variable allocation strategy.
  
      The variable allocation strategies determine how much memory the `Program`
      consumes, and how costly its memory access operations are (see
      `instructions.VariableAllocation`).  In general, a variable holding data with
      a longer or more complex lifetime will need a more expensive storage strategy.
      This analysis examines variables' liveness and opportunistically selects
      inexpensive sound allocation strategies.
  
      Specifically, the algorithm is to:
      - Run liveness analysis to determine the lifespan of each variable.
      - Assume optimistically that no variable needs to be stored at all
        (`instructions.VariableAllocation.NULL`).
      - Traverse the instructions and pattern-match conditions that require
        some storage:
        - If a variable is read by an instruction, it must be at least
          `instructions.VariableAllocation.TEMPORARY`.
        - If a variable is live out of some block (i.e., crosses a block boundary),
          it must be at least `instructions.VariableAllocation.REGISTER`.  This is
          because temporaries do not appear in the loop state in `execute`.
        - If a variable is alive across a call to an autobatched `Function`, it must
          be `instructions.VariableAllocation.FULL`, because that `Function` may
          push values to it that must not overwrite the value present at the call
          point.  (This can be improved by examining the call graph to see whether
          the callee really does push values to this variable, but that's future
          work.)
  
      Args:
        program: `Program` to optimize.
  
      Returns:
        program: A newly allocated `Program` with the same semantics but possibly
          different allocation strategies for some (or all) variables.  Each new
          strategy may be more efficient than the input `Program`'s allocation
          strategy for that variable (if the analysis can prove it safe), but will
          not be less efficient.
      """
      alloc = {
          var: inst.VariableAllocation.NULL
          for var in program.var_defs.keys()
      }
      # The program counter is always read
      _variable_is_read(inst.pc_var, alloc)
      _optimize_1(
          program.graph, inst.pattern_flatten(program.vars_out), alloc
      )
      for varname in program.vars_in:
          # Because there is a while_loop iteration between the inputs and the first
          # block.
          _variable_crosses_block_boundary(varname, alloc)
      for func in program.functions:
          _optimize_1(
              func.graph, inst.pattern_flatten(func.vars_out), alloc
          )
          for varname in func.vars_in:
              _variable_crosses_block_boundary(varname, alloc)
      null_vars = [
          k
          for k, v in six.iteritems(alloc)
          if v is inst.VariableAllocation.NULL
      ]
      if null_vars:
          logging.warning(
              "Found variables with NULL allocation. These are written "
              "but never read: %s",
              null_vars,
          )
      return program.replace(var_alloc=alloc)
  
file_path: probability/tensorflow_probability/python/experimental/auto_batching/allocation_strategy.py
imports: from absl import logging
  import six
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import liveness
label_desc: Use image node_id 1 for calling a global function with example usage: optimize(program) and returns: program