code: def _directly_writes(graph):
      """Set of variables directly written by the given graph."""
      answer = set()
      for i in range(graph.exit_index()):
          block = graph.block(i)
          for op in block.instructions:
              if isinstance(op, inst.PrimOp):
                  answer = answer.union(
                      set(inst.pattern_flatten(op.vars_out))
                  )
              elif isinstance(op, inst.FunctionCallOp):
                  answer = answer.union(
                      set(inst.pattern_flatten(op.vars_out))
                  )
                  # These because the caller writes them before the goto.
                  answer = answer.union(
                      set(inst.pattern_flatten(op.function.vars_in))
                  )
              elif isinstance(op, inst.PopOp):
                  # If the stack discipline is followed, any local variable will be
                  # written by something before it is ever popped.  Formal parameters are
                  # written by the caller and popped before returning.
                  # Pops should not prevent a variable from being allocated as a register
                  # instead of a full variable, because pops as such do not cause
                  # registers to lose and data that a full variable would have kept.
                  pass
      return answer
  
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 7 for calling a global function with example usage: _directly_writes(graph) and returns: answer