code: def _vm_staging():
      global _staging
      old_staging = _staging
      try:
          _staging = True
          yield
      finally:
          _staging = old_staging
  
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 4 for calling a global function with example usage: _vm_staging() without return types